Transformers for In-Context PDE Solving and Inverse Problems

Part 5: Generalization and Optimization Error

This post is a map of failure modes. When the model is wrong, the error can come from several places: the prompt may identify the wrong PDE, the decoder may not run enough solver steps, pretraining may not have seen enough tasks, or the architecture may not represent the right update.

The bounds below are deliberately narrower than the full replica program. They record the exact certificates we can already verify for encoder-only, decoder-only, and composed encoder-decoder systems. The plots are diagnostics for these certificates, not evidence that the global nonlinear replica closure is solved.

The focused bounds report separates the three objects that matter:

  1. the encoder;
  2. the decoder;
  3. the composition of encoder and decoder.

Encoder-decoder optimization and generalization certificates

Compact diagnostic plots (simple experiments) ![Encoder-decoder certificate decomposition animation](/assets/img/blog/transformer-icl-pde/visuals/optgen_encoder_decoder_decomposition.gif) This animation tracks the three quantities that matter during pretraining: encoder task error, decoder solver error, and their composed certificate. The point is diagnostic. If the decoder bar shrinks but the composition remains large, the encoder is the bottleneck. If the encoder bar is small but the decoder bar remains large, the solver is the bottleneck. ![Measured and predicted risk curves](/assets/img/blog/transformer-icl-pde/visuals/optgen_risk_curves.png) The risk plot overlays measured curves with the corresponding predicted or certified curves. Encoder-only can have an exact population prediction, decoder-only has an exact frozen-checkpoint Richardson identity, and the composition inherits both pieces plus interaction error. ![Prompt and task scaling laws](/assets/img/blog/transformer-icl-pde/visuals/optgen_scaling_laws.png) The scaling-law plot displays the two statistical axes: prompt size \(m\) and number of training tasks \(N\). The former controls task identification at test time. The latter controls how well pretraining learns the shared mechanism.

For the encoder-only low-rank recovery experiment,

\[z\sim\mathcal N(0,\Sigma_z),\qquad b=G_\star z+\xi.\]

Population gradient descent has the exact trajectory

\[G_s-G_\star=(G_0-G_\star)(I-\eta\Sigma_z)^s.\]

The final held-out risk matched the exact population curve:

\[\widehat R_{\rm enc}(800)=4.0020\times 10^{-4}, \qquad R_{\rm enc}(800)=4.0000\times 10^{-4}.\]

For the decoder,

\[e_L=(I-P_\theta H)^Le_0.\]

The decoder-only learned-preconditioner run had zero optimization-bound violations and final held-out depth risk

\[7.6747\times 10^{-3}.\]

The clipped empirical Bernstein upper was

\[1.7710\times 10^{-1}.\]

For encoder–decoder, the important decomposition is

\[\widehat x_L-x_\star = (\widehat x_L-x^\dagger(\widehat z)) + (x^\dagger(\widehat z)-x_\star).\]

The first term is decoder optimization. The second is encoder task recovery. In the PDE K=4 recast, the encoder was essentially exact:

\[\mathrm{MSE}(z_{\rm enc})=6.1356\times 10^{-8}, \qquad \frac{\|A(\widehat z)-A(z)\|}{\|A(z)\|} =6.0764\times 10^{-5}.\]

The best end-to-end PDE method reached

\[5.9760\times 10^{-6}\]

with zero optimization-bound violations.

The result is not just that the model works. The point is that each stage has a separate certificate: exact encoder trajectory where available, exact decoder optimization identity, and clipped held-out generalization bounds.

Why the certificates are different

The encoder-only experiment is special because its population dynamics are linear. The loss is quadratic in (G), and the data covariance is fixed. That is why the optimization trajectory can be written in closed form.

The decoder certificate is different. It does not describe how the preconditioner was trained. It describes what any frozen preconditioner does after training. Once (P_\theta) is fixed, the decoder is a deterministic linear iteration. This is why the certificate applies equally to free linear preconditioners, no-softmax linear attention, softmax attention, and Flexformer-style kernels.

The encoder–decoder certificate combines both. It first asks how accurately the prompt identifies the task, then asks how well the decoder solves the query system induced by that estimate.

What the numbers say

The encoder-only run reached the noise floor. The final held-out risk was essentially the exact population prediction:

\[4.0020\times 10^{-4} \quad\text{versus}\quad 4.0000\times 10^{-4}.\]

The decoder-only learned preconditioner had a nontrivial raw risk but a clean optimization certificate. The bound was conservative, as expected, because it uses the taskwise spectral norm:

\[\|e_L\|\leq \|I-P_\theta H\|^L\|e_0\|.\]

The encoder–decoder Flexformer run decomposed the held-out risk into an encoded floor and a decoder component:

\[R_{\rm floor}=2.4169\times 10^{-3}, \qquad R_{\rm dec}=1.7448\times 10^{-3}.\]

That is useful diagnostically. It says the remaining error is not a single opaque number. Part of it comes from task representation, part from the finite depth of the decoder.

The PDE recast is the cleanest composed result: the encoder almost exactly recovers the task coefficient, and the decoder certificate has zero violations. That is the setting where the mathematical story and the numerical evidence are most aligned.

How to use the report

The attached TeX/PDF is deliberately narrower than the full Flexformer report. It leaves out architecture speculation and focuses only on statements that can be written as bounds:

  • exact encoder trajectory in the low-rank Gaussian model;
  • exact decoder Richardson identity;
  • exact encoder–decoder decomposition;
  • clipped empirical-Bernstein generalization certificates;
  • numerical tables for encoder-only, decoder-only, low-rank enc+dec, and PDE enc+dec.

That narrower scope is useful. It gives a clean reference for the mathematical claims without mixing them with broader replica or Flexformer training questions.

The main caveat

The empirical-Bernstein bound is intentionally conservative. It is designed to be a finite-sample certificate for clipped held-out losses, not a sharp asymptotic prediction. The sharp-looking curves come from exact optimization identities or exact population dynamics; the clipped bound is there to make the held-out risk statement legally finite-sample.

So the takeaway is not “the Bernstein upper is tight.” It is: the raw held-out risk is low, the pointwise optimization certificate has zero violations, and the reported finite-sample upper is valid for the clipped risk.

Downloads