Warm-starting PDE solvers with any-dimensional machine learning

Wilson G. Gregory; George A. Kevrekidis; Ben Blum-Smith; Soledad Villar

stat.MLSep 30, 2026

Warm-starting PDE solvers with any-dimensional machine learning

Authors: Wilson G. Gregory, George A. Kevrekidis, Ben Blum-Smith, Soledad Villar

Organizations: Department of Statistics and Irving Institute for Cancer Dynamics, Columbia University · Division of Applied Mathematics, Brown University · Department of Applied Mathematics and Statistics, Johns Hopkins University

Abstract

Any-dimensional machine learning models, such as graph neural networks (GNNs), can be naturally trained and evaluated on inputs of different sizes and dimensions. Inspired by the GNN transferability literature, we show mathematical conditions under which a partial differential equation (PDE) learning-based solver can be trained in small dimensions and directly applied to solve a higher dimensional PDE in a zero-shot fashion. These conditions are based on symmetries in both the partial differential equation and the initial data. When the equations satisfy the symmetries but the data does not, which is the case for many PDEs arising from physics, we show that our theory gives a principled way of warm-starting low-dimensional PDE solvers for higher dimensional PDEs. We apply this method on the heat equation, Burgers' equation, and the compressible Navier--Stokes equations, improving the performance in both zero-shot and typical training regimes on high dimensional data. For example, we train a surrogate model on 2D Navier--Stokes data and achieve better results on 3D test data than a baseline surrogate model trained on 3D data, while only using 12 $\%$ of the flops and 20 $\%$ of the total data size.

Figures & tables

Panel of Figure 1 — Figure 1 : (Left) The 5x5 scalar, pseudoscalar, vector, and pseudovector symmetric filters for 2-D images. (Right) Convolution of a scalar image with different geometric filters. The convolution with a scalar filter can be seen as the discretization of a diffusion operator, and the convolution with the vector filter approximates a discretization of the gradient. In general any symmetric discretization of any coordinate-free operator (gradient, divergence, curl) could be expressed in terms of these symmetric filters.

Panel of Figure 2 — Figure 2 : Relative error performance of (a) the heat equation, (b) Burgers’ equation, and (c) the compressible Navier–Stokes equations over varying numbers of fine tuning points, indexed by the total number of gigaflops to pre-train and fine-tune the models. Note the order of magnitude of the x-tick values, and the log values of the y-axis. We ran 5 trials, with the lines representing the median value, the shaded region representing the interquartile range, and a dot for each value.

PDE	Models	0	1	4	8	32	128
Heat	Baseline	105.3%	21.0%	20.8%	19.8%	3.3%	0.5%
	Nepo	5.7%	2.2%	2.2%	2.1%	0.7%	0.5 %
Burgers	Baseline	105.3%	2.5%	1.8%	1.6%	Time out	Time out
	Nepo	12.8%	1.8%	1.6%	1.4%	Time out	Time out
N-S	Baseline	1,565.7%	97.0%	56.3%	27.2%	5.2%	Time out
	Nepo	22.8%	28.9%	17.8%	14.3%	3.8%	Time out

Table 1 : The median relative error performance over 5 trials of each PDE over varying numbers of fine tuning points. N-S refers to compressible Navier–Stokes. For the heat equation, the column labels represent the number of fine-tuning points, while for Burgers and Navier–Stokes, they refer to the number of trajectories of 17 and 49 time steps each, respectively. The positions in the table labeled “Time out” were training loops that would have taken longer than 15 hours per trial. With higher precision, the baseline method does better than Nepo for the heat equation by

0.001

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Type	Dimension	Weights
Scalar filters 1D	$\alpha_{0}$	$\alpha_{1}$
2D Nepo	$\alpha_{0}-\frac{2}{3}\alpha_{1}$	$\frac{1}{3}\alpha_{1}$	$\frac{1}{3}\alpha_{1}$
2D Copy	$\alpha_{0}+2\alpha_{1}$	$-\alpha_{1}$	$\alpha_{1}$
2D Zeros	$\alpha_{0}-2\alpha_{1}$	$\alpha_{1}$	0
Scalar filters 2D	$\alpha_{0}$	$\alpha_{1}$	$\alpha_{2}$
3D	$\alpha_{0}-\frac{10}{9}\alpha_{1}+\frac{4}{9}\alpha_{2}$	$\frac{5}{9}\alpha_{1}-\frac{2}{9}\alpha_{2}$	$\frac{2}{9}\alpha_{1}+\frac{1}{9}\alpha_{2}$	$-\frac{1}{9}\alpha_{1}+\frac{4}{9}\alpha_{2}$

Appendix

Table 2 : Coefficients for positive parity filters with side length 3. For each set of two, the first row is the weights for the lower dimensional filters, in order from the center of the filter outwards. The second row is the weights for the higher dimensional filters in terms of the weights of the lower dimensional filters. The 2D Copy and 2D Zeros coefficient transfer strategies are explained in Section G.3 .

PDE	Models	Params	1	4	8	32	128
	1D Pre-train	3,926,737					1e-3
Heat	2D Baseline	5,691,697	5e-4	5e-4	5e-4	5e-4	5e-4
	2D Fine-tune	5,691,697	5e-4	5e-4	5e-4	5e-4	5e-4
	2D Pre-train	11,576,305			1e-4
Burgers	3D Baseline	16,871,185	1e-4	1e-4	1e-4
	3D Fine-tune	16,871,185	1e-4	1e-4	1e-4

Appendix

Table 3 : Pre-training and fine-tuning learning rates for the number of data points or trajectories.

Figure 3 : The median relative error performance of each warm-start strategy with the heat equation over varying numbers of fine tuning points, indexed by the total number of gigaflops to pre-train and fine-tune the models. Note the order of magnitude of the x-tick values, and the log values of the y-axis. Each experiment was run 5 times, with the lines representing the median value, the shaded region representing the interquartile range, and a dot for each value.

PDE	Models	0	1	4	8	32	128
	Baseline	105.3%	21.0%	20.8%	19.8%	3.3%	0.5%
Heat	Copy	12,797.3%	131.6%	138.5%	152.5%	52.7%	15.0%
	Zeros	138.3%	8.1%	8.1%	8.4%	3.1%	1.0%
	Nepo	5.7%	2.2%	2.2%	2.1%	0.7%	0.5 %

Appendix

Table 4 : The median relative error performance over 5 trials of each warm-start strategy with the heat equation over varying numbers of fine tuning points. The column labels represent the number of fine-tuning points. With higher precision, the baseline method does better than Nepo by

0.001

Panel of Figure 4 — Figure 4 : $\ell_{2}$ error performance of each PDE over varying numbers of fine tuning points, indexed by the total number of gigaflops to pre-train and fine-tune the models. Note the order of magnitude of the x-tick values, and the log values of the y-axis. Each experiment was run 5 times, with the lines representing the median value, the shaded region representing the interquartile range, and a dot for each value. (a) The heat equation trained on 0, 1, 4, 8, 32, and 128 fine-tuning points. (b) Burgers’ equation trained on 0, 1, 4, and 8 trajectories. (c) The compressible Navier–Stokes equations trained on 0, 1, 4, 8, and 32 trajectories.

PDE	Models	0	1	4	8	32	128
	Baseline	4.4e-02	1.7e-03	1.7e-03	1.5e-03	4.5e-05	1.0e-06
Heat	Copy	6.5e+02	6.9e-02	7.6e-02	9.3e-02	1.1e-02	9.0e-04
	Zeros	7.6e-02	2.6e-04	2.9e-04	2.8e-04	3.8e-05	4.8e-06
	Nepo	1.3e-04	2.0e-05	1.9e-05	1.9e-05	2.2e-06	1.0e-06
Burgers	Baseline	3.3e-02	6.3e-05	4.1e-05	3.7e-05	Time out	Time out
	Nepo	5.5e-04	3.8e-05	3.5e-05	3.2e-05	Time out	Time out

Appendix

Table 5 : The median

\ell_{2}

error performance over 5 trials of each PDE over varying numbers of fine tuning points. N-S refers to compressible Navier–Stokes. For the heat equation, the column labels represent the number of fine-tuning points, while for Burgers and Navier–Stokes, they refer to the number of trajectories of 17 and 49 time steps each, respectively. The positions in the table labeled “Time out” were training loops that would have taken longer than 15 hours per trial. With higher precision, the baseline method does better than Nepo for the heat equation by 1e-09.

PDE	Models	n=0	n=1	n=4	n=8	n=32	n=128
Heat	Baseline	0	8.732e+04	3.493e+05	6.986e+05	2.794e+06	1.117e+07
	Nepo	3.594e+05	4.467e+05	7.087e+05	1.058e+06	3.153e+06	1.153e+07
Burgers	Baseline	0	2.965e+09	1.186e+10	2.372e+10
	Nepo	1.396e+08	3.105e+09	1.200e+10	2.386e+10
N-S	Baseline	0	1.846e+09	7.386e+09	1.477e+10	5.908e+10
	Nepo	1.724e+09	3.571e+09	9.110e+09	1.649e+10	6.081e+10

Appendix

Table 6 : Estimated gigaflops of training process. Nepo warm-start includes the time to pre-train on lower dimensional data and convert the model.

Explore similar work

Error-Conditioned Neural Solvers
Jun 25, 2026Haina Jiang, Liam Wang, Peng-Chen Chen +4Neural Partial Differential Equation Solvers Neural Solvers
University of Michigan · KAIST AI · Los Alamos National Laboratory
Physics-Informed Neural Embeddings of PDE Solution Families
Jul 7, 2026Raul Jimenez, Svitlana Mayboroda, Pavlos Protopapas +3Partial Differential Equations Physics-Informed Learning
Institute of Cosmos Sciences (ICC), University of Barcelona, Martí i Franquès 1, ES-08028, Barcelona, Spain. · Department of Mathematics, ETH Zurich, Rämistrasse 101, 8092, Zürich, Switzerland. · Institute for Applied Computational Science, Harvard University, Cambridge, MA, USA. +1
INEUS: Iterative Neural Solver for High-Dimensional PIDEs
May 7, 2026Jean-Loup Dupret, Davide Gallon, Patrick CheriditoNeural Partial Differential Equation Solvers Partial Differential Equations
Department of Mathematics and RiskLab, ETH Zurich, Switzerland · Applied Mathematics: Institute for Analysis and Numerics, University of Münster, Germany

Jun 25, 2026, cs.LGJ/K move · Enter open · S save

Error-Conditioned Neural Solvers

Haina Jiang, Liam Wang, Peng-Chen Chen, Min Seop Kwak+3

University of Michigan · KAIST AI · Los Alamos National Laboratory

Neural Partial Differential Equation Solvers Neural Solvers Partial Differential Equations

Open paper AskPDF arXiv

Neural surrogate models offer fast approximate mappings from PDE parameters to solutions, but they typically treat solving as a purely statistical task: once trained, they struggle to correct their own constraint violations and extrapolate beyond the training distribution. Recent hybrid methods promote physical correctness by targeting the PDE residual via gradient descent or Gauss--Newton steps, but inherit the compute cost and instability of the underlying classical optimizers. We show, theoretically and empirically, that numerically minimizing the PDE residual can be an unreliable proxy for reconstruction accuracy in ill-conditioned systems, explaining why these methods often do not make accurate predictions despite achieving low residuals. We propose error-conditioned Neural Solvers (ENS), built on a different principle: rather than an optimization target, the PDE residual field is passed as a direct input to the network at each iteration, enabling it to read the spatial structure of its own errors and learn an update policy to iteratively correct its predictions. Across four PDE families, ENS attains the highest prediction accuracy in the large majority of settings, with gains reaching

10\times

on turbulent Kolmogorov flow, while avoiding the expensive compute cost of hybrid methods. ENS's learned correction policy generalizes under distribution shift, including zero-shot parameter changes and cross-equation transfer, where its relative advantage is largest in the ill-conditioned regimes where residual minimization is least reliable. Project website: https://neuralsolver.github.io/.

Abstract

Explore similar work

Error-Conditioned Neural Solvers

Physics-Informed Neural Embeddings of PDE Solution Families

INEUS: Iterative Neural Solver for High-Dimensional PIDEs