Rich Representations, Small Heads: Fast PDE Solves and Flows in Weight Space

The study of embeddings of partial differential equations (PDEs) using Physics-Informed Neural Networks (PINNs) involves learning a latent space basis that represents a family of solutions. For linear PDEs, a sufficiently expressive learned basis yields approximate solutions for unseen equation parameters (EPs) or initial conditions (ICs) through a direct solve for the output-head weights, without retraining the shared network. This procedure is referred to as one-shot transfer learning (OTL). For nonlinear equations, however, a direct closed-form solution for the head weights is generally unavailable, and existing extensions of OTL instead use approaches that rely on linearizing the nonlinearity, such as perturbative decompositions or Chebyshev surrogates. In this work, we investigate second-order methods that adapt the head weights to the nonlinear problem directly while keeping the learned basis fixed. We train multi-head PINNs on families of linear and nonlinear PDEs with varying EPs and ICs, learning a shared basis for each equation with a separate linear head for each training instance. This low-dimensional learned basis restricts adaptation to a small space of solution coefficients, making second-order optimization tractable. With the learned basis fixed, we investigate two second-order methods for obtaining new solutions at unseen EPs or ICs. The first applies Gauss–Newton to the least-squares problem for the head weights. The second uses continuation in head-weight space, for which we derive an analytic expression, and returns the path of heads between a trained instance and the target. We compare both methods against training a new head in terms of accuracy and computational cost. On the heat, wave, Allen–Cahn and Burgers equations, both methods improve the accuracy of a trained head for the same computational cost, improving the efficiency of transfer learning in PINNs.