PaperlayerView exact revision on arXiv

Source-linked AI summary

Warm-starting PDE solvers with any-dimensional machine learning

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

arXiv:2609.38916v1stat.MLcs.LG

TL;DR

PDE solvers have limited theory explaining when they generalize across spatial domains, despite the computational burden of high-dimensional data. This paper derives symmetry-based conditions for exact transfer and proposes low-to-high-dimensional warm-starting; across three PDE families, its 2D-trained Navier–Stokes model outperforms a 3D-trained baseline with less compute and data.

  • Problem

    There is limited theory explaining when PDE solvers generalize across spatial domains, while high-dimensional data can be computationally burdensome.

    Source: §1 · p. 1

  • Method

    The paper derives symmetry-based conditions for exact cross-dimensional PDE transfer and uses low-dimensional solvers to warm-start higher-dimensional solvers when those conditions do not hold.

    Source: §1 · p. 2

  • Results

    Across the heat, Burgers’, and compressible Navier–Stokes equations, a 2D-trained Navier–Stokes model outperforms a 3D-trained baseline while using 12% of the flops and 20% of the total data size.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The approach provides computational and memory benefits for high-dimensional PDE learning, including memory benefits from processing spatial domains in sections.

    Source: §1 · p. 2, §5 · p. 10

    Exact generalization applies on a subspace where the higher-dimensional problem is essentially an n-dimensional problem embedded in n + 1 dimensions.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

The paper develops conditions for PDE solvers to transfer across spatial dimensions and uses low-dimensional solvers to warm-start higher-dimensional ones when exact transfer conditions fail. It demonstrates the approach on three PDE families, including a 2D-trained Navier–Stokes model that outperforms a 3D-trained baseline with less compute and data.

1 Introduction

  • Any-dimensional models reuse fixed parameters across variable-sized inputs, and prior work studies their transferability and representation stability.The paper addresses limited theory on when PDE solvers generalize across spatial domains and the computational burden of high-dimensional data scaling as O(N^n).Source: §1 · p. 1
  • Warm-starting requires symmetry of the equation, but not of the initial data, boundary conditions, or domain.The paper assumes Euclidean symmetry, while noting that other symmetry assumptions could be developed with additional work.Source: §1 · p. 2
  • The framework expresses PDEs through orthogonally invariant differential operators on jet space, including the Laplacian and Dirichlet energy density.It gives conditions for lifting an n-dimensional PDE solution to a corresponding solution in R^(n+1).Source: §1 · p. 2
  • When exact-transfer assumptions fail, Nepo warm-starting uses an n-dimensional solver to initialize an (n + 1)-dimensional solver.Under the stated assumptions, the lift maps PDE solutions in n dimensions to solutions in n + 1 dimensions.Source: §1 · p. 2
  • A 2D-trained Navier–Stokes model outperforms a 3D-trained baseline while using 12% of the flops and 20% of the total data size.The approach is demonstrated on the heat equation, Burgers’ equation, and compressible Navier–Stokes equations.Source: §1 · p. 2

2 Any-dimensional PDEs

The section defines dimension-indexed PDEs using differential operators on jet spaces, then specifies embeddings and a symmetry-based compatibility condition under which the equations agree across dimensions.

2 Any-dimensional PDEs

  • The PDE is represented by dimension-indexed functions on kth-order jet spaces, whose coordinates encode the function and its partial derivatives.This formalism is motivated by equations, such as the heat equation, that can be written without fixing the spatial dimension.Source: §2 · pp. 2–3
  • The duplication embedding extends a function to R^(n+1) by making it constant in the new coordinate, so derivatives involving that coordinate vanish.The induced embedding of jet sections zero-pads coordinates for derivatives involving the added coordinate.Source: §2 · pp. 3–4
  • Compatibility requires the dimension-indexed equations to commute with the embeddings, with coordinate-independent, O(n)-invariant operators ensuring this for alternative projections.The condition constrains the equation itself, not the initial data or boundary conditions.Source: §2 · p. 4
  • For fixed PDE degree and order, the differential operators can be expressed in a finite basis independent of dimension.The Laplacian and Dirichlet energy density are examples of operators in this framework.Source: §2 · pp. 4–5

3 Any-dimensional solution operators

For compatible scalar PDEs with unique solutions, symmetry and dimension-embedding results relate solution operators across dimensions. The relation holds for embedded initial conditions, suggesting a low-dimensional warm start rather than full operator equivalence.

3 Any-dimensional solution operators

  • For a compatible scalar PDE with a uniqueness theorem, its solution operator is E(n)-equivariant.The compatibility condition makes the equation independent of a chosen origin and orthonormal coordinates.Source: §3 · p. 5
  • Embedding a solution of a compatible scalar PDE in n dimensions yields a solution of the corresponding equation in n + 1 dimensions.The result also holds for the alternative embeddings described in the lemma.Source: §3 · p. 5
  • For compatible PDEs with unique solutions, projecting the higher-dimensional solution operator recovers the lower-dimensional operator on embedded initial conditions.This is a restricted agreement, not strong equivalence between operators across dimensions.Source: §3 · pp. 5–6
  • That restricted agreement suggests transferring a lower-dimensional operator's parametrization to warm-start learning in higher dimensions.The higher-dimensional operator need only agree on the subset of initial data produced by the embedding.Source: §3 · p. 6

4 Lifting PDE solvers to higher dimensions

The paper builds dimension-transferable PDE models from geometric convolutional networks on tensor images, whose layers can be compatible with tensor-image embeddings. It then uses a soft compatibility objective to transfer filter coefficients for generic higher-dimensional inputs and warm-start further training.

4 Lifting PDE solvers to higher dimensions

  • The lifting framework uses equivariant convolutional networks whose grid pixels carry scalar, vector, and tensor features representing discretized fields.The models extend equivariant convolutional networks to vector and tensor images.Source: §4 · p. 6

4.1 Geometric convolutional neural networks

  • Geometric convolutions use symmetric filters and tensor contractions to preserve tensor-image outputs and characterize linear G_d-equivariant functions.Scalar and vector filters illustrate diffusion-like and gradient-like discretizations, respectively.Source: §4.1 · p. 6, Figure 1 · p. 7

4.2 Embedding of tensor images

  • Tensor-image embeddings extend the lower-dimensional image across a new coordinate, and compatibility conditions specify when layers commute with that embedding.Compatible geometric convolutions and common pointwise or pooling layers can be composed into compatible networks.Source: §4.2 · p. 7

4.3 Algorithm for warm-starting PDEs

  • Because generic higher-dimensional images are not simply replicated lower-dimensional images, the warm-start transfers coefficients by seeking filters whose random-slice convolutions are close.The coefficient objective imposes soft compatibility, while the embedding compatibility condition is hard; the derivation uses iid, zero-mean scalar pixels.Source: §4.3 · p. 8
  • The Nepo warm-start analytically converts lower-dimensional filter coefficients, replaces the filters in the higher-dimensional model, and optionally fine-tunes it.The coefficient conversion can be computed in advance from the lower-dimensional filter coefficients.Source: §4.3 · p. 8

5 Numerical experiments

Across the heat, Burgers, and compressible Navier–Stokes experiments, low-dimensional pretraining generally improves relative error over training only on higher-dimensional data. The largest reported advantage is for Navier–Stokes, alongside lower compute and data use; the method also offers memory benefits in higher dimensions.

5 Numerical experiments

  • The experiments train on n-dimensional data, convert the model to n + 1 dimensions, and optionally fine-tune before comparison with a baseline trained only on fine-tuning data.The study covers heat (1D to 2D), Burgers (2D to 3D), and Navier–Stokes (2D to 3D).Source: §5 · pp. 8–9
  • Pretraining gives better relative error than the baseline in 12/13 cases; for Navier–Stokes, the untuned warm-start beats the 8-trajectory baseline using 12% of the flops and 20% of the data.The authors report mild overfitting when fine-tuning the Navier–Stokes warm-start with one trajectory.Source: Figure 2 · p. 9, §5 · p. 9, Table 1 · p. 10
  • Relative error is evaluated across varying fine-tuning data amounts, with five trials summarized by medians, interquartile ranges, and individual values.Figure 2 indexes pretraining-plus-fine-tuning compute in gigaflops and displays relative error on a log scale.Source: Figure 2 · p. 9, Table 1 · p. 10
  • For heat, each n-dimensional tensor image requires O(N^n) memory; for Burgers and Navier–Stokes, it requires O(n N^n), yielding additional memory benefits from the any-dimensional method.The paper notes that oversized data may otherwise require splitting the spatial domain into sections.Source: §5 · p. 10

6 Discussion and limitations

Exact cross-dimensional generalization is limited to particular symmetry and data conditions, while low-dimensional solvers can still provide practical warm starts. The discussion also states the geometric scope and situates the work among related transferability and equivariant-learning research.

6 Discussion and limitations

  • The framework assumes orthogonal symmetry, requires some symmetry constraint, and is developed only for flat geometries; adaptation to general manifolds is unknown.The authors note that other symmetries could be considered with further theoretical development.Source: §6 · p. 10
  • Exact generalization applies only when the higher-dimensional problem lies in a subspace that is essentially the lower-dimensional problem; the demonstrated practical use is warm-starting one dimension higher.The authors note that they do not use an infinite-dimensional limit in their PDE application.Source: §6 · p. 10

AI use statement

Reproducibility statement

A Related work

  • Related work connects representation stability and any-dimensional learning to PDE identification [40], while this paper applies the GNN transferability principle to PDE learning across dimensions.The discussion also situates the work alongside equivariant models for PDEs and warm-starting smaller models [37].Source: A Related work · p. 15

B Proof of Proposition 2.3

The proof establishes compatibility across dimensions by translating a commuting diagram into conditions on the differential operators. Orthogonal invariance then allows the result to hold for other affine-orthogonal projections.

B Proof of Proposition 2.3

  • The result assumes the diagram commutes for standard projections, the operators are independent of the space coordinate, and they are O(n)-invariant; it extends commutativity to affine-orthogonal projections.The proof first translates the diagram into an equality quantified over dimensions, functions, and spatial points.Source: B Proof of Proposition 2.3 · pp. 15–16
  • The proof unwinds the diagram using the explicit projection and induced-map definitions, then applies operator independence and orthogonal invariance to establish the transformed equality.The passage describes the remaining proof as short once these definitions and hypotheses are in place.Source: B Proof of Proposition 2.3 · pp. 15–16

C Representation stability

Representation stability supplies dimension-independent families of orthogonal invariants that yield compatible differential operators. The accompanying background describes tensor-image features and geometric convolutions used in equivariant networks.

C Representation stability

  • Representation stability means that, beyond sufficiently high dimensions, the same irreducible representations occur with the same multiplicities under the natural identifications.The canonical O(n) action on R^n gives a representation-stable sequence that extends to jet spaces and their polynomial functions.Source: C Representation stability · p. 16
  • A basis of O(n)-invariants in polynomial functions on jet spaces can be described independently of n, and its elements generate compatible operator families.The Laplacian, the function itself, and squared gradient norm are examples; the Hessian determinant does not satisfy this compatibility condition.Source: C Representation stability · p. 17

D Tensor Image and Geometric Convolution Background

  • The background defines tensors through O(n) actions and describes Vector Neuron nonlinearities and tensor-norm max pooling as network operations.Tensor products and contractions provide the operations used in the geometric convolution construction.Source: D Tensor Image and Geometric… · pp. 17–18
  • Tensor images assign a tensor-valued feature to each lattice site, while geometric convolutions combine input and filter tensors through tensor products and contractions.The described 2D filters respect 90-degree rotations and are symmetric or antisymmetric under reflections.Source: D Tensor Image and Geometric… · p. 18

E.1 Compatibility of Solution Operators

For compatible scalar PDEs, the paper relates solution operators across dimensions and spatial symmetries. Under uniqueness, an embedded lower-dimensional solution is the higher-dimensional solution for embedded initial data; perturbation bounds qualify this result.

E.1 Compatibility of Solution Operators

  • When the PDE is E(n)-equivariant, its solution operator respects the same spatial transformations, provided solutions are unique.The proof compares the transformed solution with the solution initialized by the transformed data.Source: §E.1 · p. 19
  • For compatible scalar PDEs, embedding an initial condition embeds its solution into a solution of the higher-dimensional PDE, and uniqueness identifies that solution with the higher-dimensional evolution.The proposition also permits the corresponding alternative projections described in Section 2.Source: §E.1 · p. 19
  • Perturbation bounds make the lifted compatible part a good approximation to the perturbed higher-dimensional solution, but their constants may depend on dimension and time and yield vacuous bounds.The constants can also be more forgiving and force small perturbations to become smaller over time.Source: §E.1 · p. 20

E.2 Compatibility of Tensor Image Functions

This section gives compatibility conditions for tensor-image operations across dimensions, then introduces zero-padding for lifting filters. It shows how convolutional and pointwise components can respect the embedding between dimensions.

E.2 Compatibility of Tensor Image Functions

  • Geometric convolution filters are compatible across dimensions when the higher-dimensional filter coefficients satisfy the stated relation to the lower-dimensional coefficients.The appendix extends the result to linear combinations of filters.Source: §E.2 · pp. 20–21
  • Point-wise functions are compatible when their single-pixel maps commute with the tensor embedding; the Vector Neuron nonlinearity meets this condition.Max pooling is also shown to be compatible with the embedding.Source: §E.2 · pp. 21–23

F Lifting filters

  • The zero-pad embedding places an n-dimensional filter in the center slice of an (n + 1)-dimensional filter and sets other entries to zero.Table 2 lists positive-parity filter coefficients and the corresponding higher-dimensional coefficients.Source: F Lifting filters · p. 23, Table 2 · p. 25
  • The lifting-filter derivation decomposes the squared tensor norm into matching-vector norms and cross terms, using the assumption that the image is scalar.The cross terms use the full k-contraction.Source: F Lifting filters · p. 24

G.1 Data

The experiments use heat, Burgers’, and compressible Navier–Stokes data, with all fields discretized as images of side length 64. The data-generation procedures differ by PDE and dimension.

G.1 Data

  • Heat-equation data use a uniform initial field, with values at t = 1 computed using the analytic heat kernel and torus distance.The initial values are sampled in each pixel; the passage’s range is truncated.Source: §G.1 · p. 26
  • Burgers’ trajectories come from Apebench’s difficulty scenario, with dimension-adjusted coefficients chosen to keep parameters consistent across dimensions.Each 50-step trajectory is split into 49 one-step input-output pairs.Source: §G.1 · p. 26
  • Compressible Navier–Stokes data use PDEBench code with matching generation hyperparameters across dimensions, including shear and bulk viscosity 1e-08 and Mach number 0.1.Each 21-step trajectory is split into 17 examples with four input steps and one output step.Source: §G.1 · p. 26

G.2 Model and Training

The experiments use a UNet-based model, train it with relative error, and report relative error and SMSE. Training uses tuned learning rates and a common 50-epoch AdamW schedule.

G.2 Model and Training

  • The models use a UNet with max-pooling downsampling, two convolution layers between downsampling blocks, and Vector Neuron nonlinearities with GeLU.The default width is 48 channels, doubled at each downsampling step.Source: §G.2 · p. 26
  • Training uses relative error, while evaluation reports relative error and sum of mean squared error (SMSE) across fields.The relative-error denominator includes ϵ = 1e-5 to prevent division-by-zero issues.Source: §G.2 · p. 26
  • Models train for 50 epochs with AdamW, weight decay 1e-5, cosine decay to 1e-7, and five warmup epochs.Peak learning rates are selected from six candidates using problem- and dimension-specific validation sets.Source: §G.2 · p. 26
  • Alternative coefficient-transfer strategies use the learning rates tuned for Nepo, although individually tuning those strategies could improve their performance.The authors note that zero-tuning-point results suggest Nepo would likely remain best.Source: §G.2 · pp. 26–27

G.3 Warm-start Ablation

The ablation finds that Nepo’s coefficient transfer outperforms the tested alternatives, even though those alternatives also satisfy compatibility. Additional figures and tables organize error results by fine-tuning data and compute.

G.3 Warm-start Ablation

  • Nepo performs better than the Copy and Zeros strategies, though at higher precision the baseline method does better than Nepo by 0.001%.The ablation compares median relative error over five trials across fine-tuning-point counts.Source: §G.3 · p. 27, Table 4 · p. 27

G.4 Additional Results

  • Figure 3 compares heat-equation median relative error across fine-tuning counts against total pre-training and fine-tuning gigaflops.It shows five-run medians, interquartile ranges, and individual values on a logarithmic error axis.Source: Figure 3 · p. 28
  • Figure 4 and Table 5 report ℓ2 error across heat, Burgers’, and compressible Navier–Stokes experiments at varying fine-tuning counts.Table 5 marks runs exceeding 15 hours per trial as “Time out”; its caption notes that the baseline does better than Nepo on heat by 1e-09 at higher precision.Source: Figure 4 · p. 28, Table 5 · p. 29

Related papers

Loading 2609.38916v1…