Warm-starting PDE solvers with any-dimensional machine learning
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.
Related
- Transformer-Based Warm-Starting for Feasible and Optimal Terminal Approach to Tumbling Objects with Space Manipulators
- PTED: A multi-dimensional two-sample test for scientific inference and generative machine learning
- Deep-Learning Solvers and Surrogates for Infinity and p-Laplace Problems
- Towards learning-enhanced solvers for multi-parametric optimization
- BandPC: Learning Residual-Band Preconditioner Combinations for Flexible Conjugate Gradient Solvers
BibTeX
@misc{gregory2026warmstarting,
title = {{Warm-starting PDE solvers with any-dimensional machine learning}},
author = {Wilson G. Gregory and George A. Kevrekidis and Ben Blum-Smith and Soledad Villar},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.38916},
eprint = {2609.38916},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.38916},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Warm-starting PDE solvers with any-dimensional machine learning
AU - Wilson G. Gregory
AU - George A. Kevrekidis
AU - Ben Blum-Smith
AU - Soledad Villar
PY - 2026
DA - 2026/09/30
PB - arXiv
DO - 10.48550/arxiv.2609.38916
UR - https://arxiv.org/abs/2609.38916
AB - 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.
ER -