SymbolicPullback build time grows multiplicatively with n_integrators

[michakraus]

Reported by @benedict-96 in review of #207 ("massive performance issues" with the SymbolicPullback). Measured; the summary is that the per-call performance is excellent and the build is what explodes, multiplicatively in n_integrators.

Per call, the symbolic pullback is what you want

dim = 4, nhidden = 1, batch 32, comparing SymbolicPullback(nn, ParametricLoss(), μ) with ZygotePullback(ParametricLoss()):

network	params	ZygotePullback	SymbolicPullback
dim=2, width=2, n_int=1	32	10.4 ms	0.1 ms
dim=2, width=4, n_int=1	80	23.6 ms	0.1 ms
dim=4, width=4, n_int=1	88	33.6 ms	0.1 ms
dim=4, width=4, n_int=2	176	140.1 ms	— (never builds)

The symbolic call is flat at 0.1 ms while Zygote degrades from 10 ms to 140 ms — a 100–1000× advantage that grows with the network. So this is worth fixing rather than abandoning.

The build explodes with n_integrators

Phase-by-phase, dim = 4, width = 4, nhidden = 1:

phase	n_integrators = 1	n_integrators = 2
symbolic variables	0.08 s	0.01 s
build the symbolic loss	0.87 s, 3.4 ⋅ 10⁵ chars	0.03 s, 1.4 ⋅ 10⁹ chars
symbolic_differentials	0.03 s	0.01 s
symbolic_derivative	0.66 s, 1.5 ⋅ 10⁸ chars	never finishes
build_nn_function	1.34 s	—

Two observations:

1. The loss expression grows ~4,100× when n_integrators goes 1 → 2, before any differentiation. At two integrators it is already 1.4 billion characters; the process passes 8 GB and does not return.
2. Differentiating multiplies by a further ~440× (3.4 ⋅ 10⁵ → 1.5 ⋅ 10⁸ at one integrator).

Cause

Each SymplecticEuler layer's forward pass calls the executable gradient that build_gradient produced for its energy network. Traced on symbolic Nums, that call inlines the whole gradient expression. Stacking integrators inlines that expression inside itself, so the loss grows multiplicatively in the depth of the chain rather than additively. cse = true in build_nn_function cannot help: it runs at code generation, long after the expression has been materialised.

Suggested fix

Do not trace the inlined chain. Compose the pullback layer by layer — each SymplecticEuler has already built the gradient of its own energy network, so the chain rule can reuse those instead of re-expanding them. That keeps the expression linear in depth.

This is upstream-shaped work: it needs either a layerwise pullback protocol in SymbolicNeuralNetworks, or a build_nn_function that can take a composition without inlining it. Related: #243, which tracks the type piracy in the same area.

For now

SymbolicPullback(nn, ::ParametricLoss, μ) throws an ArgumentError for n_integrators > 1 rather than appearing to hang (GML #207). n_integrators = 1 builds in ≈1.4 s and is fully supported. The default training path uses ZygotePullback and is unaffected.

🤖 Generated with Claude Code

[michakraus]

The upstream half of this is now SymbolicNeuralNetworks #49, with the design worked out and measured on plain Dense chains — no GML involved, so it stands on SNN's own terms.

Two findings from that write-up that sharpen the diagnosis here:

The blow-up is not specific to the SymplecticEuler layer. A four-layer, width-16 Chain of ordinary Dense layers — 626 parameters — already produces a gradient expression of 209 million nodes. GML's generalized architectures are the extreme case, not the cause.

It factorises into two independent multiplicative effects:

• the forward pass, because applychain substitutes each layer's output expression into the next, so the loss expression alone is already w^depth (113 → 417 → 1 633 → 6 497 → 25 953 nodes as a width-4 chain goes from 2 to 6 layers);
• symbolic_derivative, which walks and expands that entire expression once per scalar parameter with nothing shared between them, contributing a further factor of nparams × ~4.

Multiplying: nparams × w^depth. Here n_integrators = 2 adds a whole inlined gradient expression to the depth, which is why it goes from 1.4 s to never.

What layerwise composition buys, measured as the sum of the per-layer Jacobian and parameter-gradient expressions: 2 983× smaller for a six-layer network, 2 743× for four layers at width 16 — and it grows linearly in depth (+880 nodes per identical layer) rather than by ~5.5× per layer.

So the fix for this issue is: adopt whatever SNN #49 lands. Nothing GML-side is required beyond dropping the n_integrators > 1 guard once it does.

🤖 Generated with Claude Code

[michakraus]

Status against SymbolicNeuralNetworks 0.6: the upstream construction exists but does not reach this case. The guard stays.

SNN #49 shipped in 0.6.0, and its changelog names this issue:

This also fixes the second-derivative case of GML #245 […] with a seam at each layer the layer is called rather than traced, so stacking two costs a function call instead of inlining a gradient expression inside a gradient expression.

Measured, it does not — not yet, and not for want of the idea being right.

What happens

A GHNN chain does decompose, so :auto commits to the layerwise path:

arch = GeneralizedHamiltonianArchitecture(4; width = 4, nhidden = 1, n_integrators = 2)
snn  = SymbolicNeuralNetwork(arch)
symbolic_steps(snn)       # 4 steps
composes_layerwise(snn)   # true

Seeding each layer individually shows where it stops:

layer 1 (L1, SymplecticEuler): layer_seed FAILED -- MethodError: scalar_expressions(::Tuple{Vector{Num}, NullParameters})
layer 2 (L2, SymplecticEuler): layer_seed FAILED -- …
layer 3 (L3, SymplecticEuler): layer_seed FAILED -- …
layer 4 (L4, SymplecticEuler): layer_seed OK

The three that fail are the ones built with return_parameters = true, which pass the system parameters to the next layer and so return a Tuple. SNN's seam is Symbolics.variables(:x, 1:input_dimension(layer)) — a plain vector — so the construction assumes every layer maps an array to an array.

And it raises rather than declining

At n_integrators = 1, where the monolithic construction builds in 3.2 s:

	result
layerwise = false	builds, 3.2 s
layerwise = :auto	throws

which is what SNN documents :auto must never do. Filed as SNN #54.

So

Two things have to happen before this issue can close, and both are upstream:

1. SNN AttentionLayer tests #54 — :auto declines instead of raising when a layer cannot be seeded. That is the correctness fix, and on its own it only restores the fallback.
2. The seam has to model a layer that carries something alongside the state. That is what would actually let a GHNN chain compose layer by layer, and it is the part this issue is waiting for.

Nothing on the GML side is blocked or wrong: the guard refuses n_integrators > 1 with an explanation, one integrator builds and evaluates about 100× faster than the Zygote pullback, and the default training path uses ZygotePullback regardless. The guard's docstring now records the above.

🤖 Generated with Claude Code