A Transformer Can Learn the PDE, Then Solve It

A PDE transformer becomes surprisingly legible when its two jobs are separated: infer the operator from the prompt, then solve the query equation.

The result

The encoder recovered the latent PDE operator to 6.1 × 10^-5 relative error; the best end-to-end run reached 6.0 × 10^-6 error with no violations of its solver certificate.

The whole construction

After Galerkin discretization, a family of elliptic PDEs becomes

\[A(z)u=f,\qquad A(z)=A_0+\sum_{k=1}^K z_kA_k.\]

The prompt contains source-solution pairs generated by one hidden task $z$. Each pair contributes weak equations, which stack into a small inverse problem $G_mz\simeq b_m$. The encoder estimates $z$; the decoder then acts on the query residual with a learned preconditioner.

A prompt encoder inferring a PDE coefficient followed by a recurrent decoder solving the query equation — The prompt identifies the task. The recurrent decoder solves the resulting PDE. Keeping these jobs separate makes both failure modes measurable.

Prompt as experiment design

The prompt is not a bag of demonstrations. It is a sensing matrix for the unknown operator. Its small singular directions are precisely the coefficient combinations the examples do not identify.

The encoder therefore has a familiar closed form,

\[\widehat z=(G_m^\top G_m+\lambda I)^{-1}G_m^\top b_m,\]

and a simple scientific interpretation: good examples excite complementary weak-form directions. As the prompt grows, the coefficient estimate improves and the query solution follows it.

The inferred variable PDE coefficient and reconstructed query solution improving together — Left: the operator estimate forms from prompt equations. Right: the query solution improves once the recovered operator is used.

Coefficient and PDE solution errors decreasing with prompt size — More prompt equations reduce operator uncertainty. Decoder depth cannot remove this floor; only a more informative prompt can.

Depth as computation

Once $\widehat z$ is fixed, every decoder layer is one preconditioned residual step,

\[u_{\ell+1}=u_\ell+P_\theta\bigl(f-A(\widehat z)u_\ell\bigr), \qquad e_L=(I-P_\theta A(\widehat z))^Le_0.\]

This identity is the cleanest part of the story. Depth is not mysterious capacity here; it is an iteration budget. Pretraining matters when attention learns a $P_\theta$ that compresses the effective spectrum, so every layer removes error in both stiff and easy directions.

Encoder, decoder, and end-to-end PDE errors with their finite solver certificates — The end-to-end error separates into task identification and finite-depth solution error. The reported solver trajectories remain inside their Richardson certificates.

What is actually learned

In the $K=4$ variable-coefficient experiment, the latent coefficient MSE is $6.14\times10^{-8}$ and the recovered operator has relative error $6.08\times10^{-5}$. The best composed PDE method reaches $5.98\times10^{-6}$ end-to-end error. For the frozen decoder, the Richardson bound has zero observed violations.

The attractive conclusion is not that a transformer has replaced numerical analysis. It is that attention has learned two classical numerical objects from tasks: a low-dimensional inverse map for the operator and a reusable preconditioner for the solve. Once written this way, prompt length controls information, depth controls computation, and conditioning controls both.