Warm-starting PDE solvers with any-dimensional machine learning

2026-09-30

AI Summary

This study addresses the prohibitive training costs and extensive data requirements associated with high-dimensional partial differential equation (PDE) solvers by proposing a novel cross-dimensional transfer paradigm for graph neural networks grounded in symmetry theory. By establishing rigorous mathematical conditions governing the symmetries of both equations and data, the proposed method enables zero-shot transfer or warm-start acceleration from low-dimensional models to arbitrary higher dimensions. Extensive evaluations on benchmark equations, including the Navier-Stokes equations, demonstrate that this approach surpasses fully trained baselines while requiring only 12% of the computational resources and 20% of the training data. Consequently, this work provides a generalizable framework for the efficient solution of high-dimensional PDEs.

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.

Problem

• Partial Differential Equations
• Graph Neural Networks
• Warm-starting
• Zero-shot transfer
• Dimensionality

Innovation

• Graph Neural Networks
• Partial Differential Equations
• Zero-shot Transfer
• Warm-starting
• Dimensional Symmetry