Transformers for In-Context PDE Solving and Inverse Problems

Part 1: The Decoder as a Reusable PDE Solver

This series studies one question: can a transformer read a few examples from a PDE task, understand the task, and solve a new query from the same task?

Read this first if the topic is new. A PDE is an equation whose unknown is a function. A numerical method turns that PDE into a large linear system. In this post the PDE itself does not change, so the transformer is only asked to learn a fast reusable solver for that one system.

The plots below are simple diagnostic simulations, not the full benchmark. They make the solver mechanism visible before we add changing PDE coefficients and replica-theory questions later in the series.

The cleanest case is the one where the PDE coefficient is not task-dependent. Take a linear elliptic equation on a bounded domain,

\[-\nabla\cdot(a\nabla u)+Vu=f,\qquad u|_{\partial\Omega}=0,\]

and suppose for this first note that (a) and (V) are fixed across tasks. After Galerkin discretization, every prompt pair ((f_i,u_i)) satisfies the same finite-dimensional system

\[A u_i=f_i .\]

The in-context problem is therefore not to infer a new operator. The operator is fixed. The problem is to learn how to apply (A^{-1}) quickly to a new query source (f_\star).

That makes the decoder the central object. A depth-(L) decoder can be viewed as a recurrent solver

\[u_{\ell+1}=u_\ell+P_\theta(f_\star-Au_\ell),\]

where (P_\theta) is the preconditioner represented by the trained attention and feed-forward layers. The exact error identity is

\[e_L=(I-P_\theta A)^L e_0,\qquad e_\ell=u_\ell-u_\star.\]

So the role of depth is transparent: each layer is one solver step. The generalization question is not whether the model memorized a particular source, but whether the learned preconditioner works on new right-hand sides drawn from the same source distribution.

Decoder-only learned preconditioner diagnostics

Compact diagnostic plots (simple experiments) ![Constant coefficient PDE reconstruction animation](/assets/img/blog/transformer-icl-pde/visuals/constant_solution_reconstruction.gif) The animation shows the decoder as an iterative solver. The black curve is the Galerkin solution \(u_\star=A^{-1}f_\star\). The blue curve is the depth-\(L\) decoder iterate. In the constant-\(a\) case, this is the cleanest possible visual test: the model is not changing the operator, it is only improving the reconstruction of the solution for a fixed linear system. ![Constant coefficient residual and energy decay](/assets/img/blog/transformer-icl-pde/visuals/constant_residual_and_energy.png) The residual, solution error, and energy error all fall together. This is what one expects if the decoder has really learned a preconditioned Richardson step rather than only fitting the pointwise solution curve. Residual decay means the iterate is satisfying the Galerkin equation; energy decay means it is improving in the PDE norm. ![Constant coefficient spectrum and learned preconditioner contraction](/assets/img/blog/transformer-icl-pde/visuals/constant_spectrum_preconditioner.png) The spectrum plot is the preconditioner story in one picture. Without preconditioning, the wide spectrum forces a slow worst-mode contraction. A Jacobi-like or learned inverse scaling compresses the effective spectrum, so the same decoder depth gives a much smaller \(\rho^{2L}\) solver term.

For fixed (P_\theta), the optimization certificate is exact:

\[\|e_L\|_2\le \|I-P_\theta A\|_2^L\|e_0\|_2.\]

This explains why preconditioning matters. If (P_\theta=\eta I), convergence slows down when (A) is ill-conditioned. If (P_\theta\approx A^{-1}), the contraction (I-P_\theta A) becomes small and the same number of layers gives a much lower error.

In transformer language, attention is useful when it routes residual information into directions where the system is hard. The values carry residuals. The query-key scores decide which equations or spectral directions matter. The feed-forward block turns the routed residual information into a preconditioned correction.

For the constant-coefficient case, this is already enough: the whole problem collapses to learning a reusable solver for one Galerkin matrix.

What the prompt contributes

In this setting, the prompt does not need to identify a new coefficient field. It is better to think of it as calibration data for the solver behavior. Each pair ((f_i,u_i)) tells the model what the inverse map (f\mapsto u) looks like on another right-hand side. If the source distribution concentrates on a low-dimensional subspace, the prompt reveals which directions are common. If the source distribution is broad, the prompt pushes the model toward a more uniform approximation of (A^{-1}).

This is also the regime where the distinction between approximation and optimization is cleanest. The approximation question is whether the network class can represent a good (P_\theta). The optimization question is what happens after a particular (P_\theta) is fixed. The latter is completely answered by

\[e_L=(I-P_\theta A)^L e_0.\]

There is no hidden transformer mystery in the frozen decoder. Once the model is interpreted as a recurrent solver, its depth-(L) error is exactly the error of (L) preconditioned Richardson steps.

Generalization in the constant case

The held-out risk is measured on fresh right-hand sides. With clipped loss (\ell_c=\min(\ell,c)), the finite-sample certificate is

\[R_c\leq \widehat R_c+ \sqrt{\frac{2\widehat V_c\log(2/\delta)}{n}} + \frac{7c\log(2/\delta)}{3(n-1)}.\]

The clipping is not cosmetic. Raw squared Gaussian losses are unbounded, so a distribution-free statement needs either clipping or an explicit tail assumption. The experiments report both the raw risk and the clipped certificate.

In this constant-(A) regime, the expected scaling is simple. More decoder depth reduces the solver term like (\rho^{2L}). More pretraining tasks reduce uncertainty about the reusable preconditioner. More prompt examples help only to the extent that the architecture uses them to infer the active source subspace; they do not need to identify a new operator.

This is why the constant-coefficient case is the correct first sanity check. If the decoder cannot learn a useful preconditioner for one fixed Galerkin matrix, then the variable-coefficient task is premature. If it can, then the next question is how to encode the operator before applying the same solver logic.

What this says about attention

In the constant case, the most honest description of attention is not that it “solves the PDE” directly. Rather, it organizes residual information. If the tokens contain weak equations, a query can ask which equations are informative for a direction of the coefficient space. A value can carry the residual in that equation. The output projection and feed-forward block then turn the routed residuals into a correction.

For no-softmax linear attention, this story is especially literal: with enough head dimension, the attention block can reconstruct signed moments of the form

\[G^\top r.\]

For softmax attention, positive normalization changes the algebra. Softmax is still useful as a routing mechanism, but the signed information must be carried carefully by the value stream or by paired positive/negative queries. This is why the experiments distinguish low-rank values, projected values, and full vector values.

The constant-coefficient setting is therefore a microscope. It isolates whether attention and the feed-forward block can implement a preconditioned residual update before any variable-coefficient inference is added. When the learned preconditioner decreases risk across held-out right-hand sides, it is evidence that the architecture has learned a solver primitive rather than a collection of task-specific answers.

Takeaway

The constant-(a) story is: one Galerkin matrix, many right-hand sides, one reusable learned preconditioner. The exact frozen optimization law is Richardson. The generalization test is held-out source prediction. This is the base case on which the variable-(a) encoder is built.

Downloads