# Warm-starting PDE solvers with any-dimensional machine learning

arXiv: https://arxiv.org/abs/2609.38916 · HTML: https://mlsanity.com/paper/2609.38916/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 53 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Builds on (most cited first, 30 of 32)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [Vector Neurons: A General Framework for SO(3)-Equivariant Networks](https://mlsanity.com/paper/2104.12229/) (2021)
- [On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups](https://mlsanity.com/paper/1802.03690/) (2018)
- [e3nn: Euclidean Neural Networks](https://mlsanity.com/paper/2207.09453/) (2022)
- [Scalars are universal: Equivariant machine learning, structured like classical physics](https://mlsanity.com/paper/2106.06610/) (2021)
- [Equivariance with Learned Canonicalization Functions](https://mlsanity.com/paper/2211.06489/) (2022)
- [Transferability of Spectral Graph Convolutional Neural Networks](https://mlsanity.com/paper/1907.12972/) (2019)
- [The Logic of Graph Neural Networks](https://mlsanity.com/paper/2104.14624/) (2021)
- [Graph Neural Networks: Architectures, Stability and Transferability](https://mlsanity.com/paper/2008.01767/) (2020)
- [APEBench: A Benchmark for Autoregressive Neural Emulators of PDEs](https://mlsanity.com/paper/2411.00180/) (2024)
- [Transferability of Graph Neural Networks: an Extended Graphon Approach](https://mlsanity.com/paper/2109.10096/) (2021)
- [N-body Networks: a Covariant Hierarchical Neural Network Architecture for Learning Atomic Potentials](https://mlsanity.com/paper/1803.01588/) (2018)
- [A graphon-signal analysis of graph neural networks](https://mlsanity.com/paper/2305.15987/) (2023)
- [Fine-grained Expressivity of Graph Neural Networks](https://mlsanity.com/paper/2306.03698/) (2023)
- [Some Might Say All You Need Is Sum](https://mlsanity.com/paper/2302.11603/) (2023)
- [Dimensionless machine learning: Imposing exact units equivariance](https://mlsanity.com/paper/2204.00887/) (2022)
- [Machine learning and invariant theory](https://mlsanity.com/paper/2209.14991/) (2022)
- [On Transferring Transferability: Towards a Theory for Size Generalization](https://mlsanity.com/paper/2505.23599/) (2025)
- [Any-dimensional equivariant neural networks](https://mlsanity.com/paper/2306.06327/) (2023)
- [A Galois theorem for machine learning: Functions on symmetric matrices and point clouds via lightweight invariant features](https://mlsanity.com/paper/2405.08097/) (2024)
- [Invariant Kernels: Rank Stabilization and Generalization Across Dimensions](https://mlsanity.com/paper/2502.01886/) (2025)
- [Graph neural networks and non-commuting operators](https://mlsanity.com/paper/2411.04265/) (2024)
- [Any-Dimensional Invariant Universality](https://mlsanity.com/paper/2605.23156/) (2026)
- [Towards Coordinate- and Dimension-Agnostic Machine Learning for Partial Differential Equations](https://mlsanity.com/paper/2505.16549/) (2025)
- [Any-Dimensional Polynomial Optimization via de Finetti Theorems](https://mlsanity.com/paper/2507.15632/) (2025)

## Abstract (arXiv)

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.