Source-linked AI summary
Lie Point Symmetry Data Augmentation for Neural PDE Solvers
Johannes Brandstetter, Max Welling, Daniel E. Worrall
TL;DR
Neural PDE solvers need abundant training data generated by the classical solvers they are intended to replace. The paper introduces Lie point symmetry data augmentation (LPSDA), which uses solution-preserving transformations derived from PDE symmetries, and reports sample-complexity improvements of over an order of magnitude across multiple evolution PDEs and training modalities.
Problem
Neural PDE solvers require abundant training data that must be generated with classical solvers.
Source: §1 · p. 1Method
Results
Over an order of magnitude: LPSDA improves neural PDE solver sample complexity across multiple evolution PDEs and training modalities.
Source: §1 · p. 1Takeaways & Limitations
The authors report clear sample-complexity decreases, measured by generalization error, across PDEs, models, and training regimes.
Source: §7 · p. 8The experiments implemented only global symmetries; purely local symmetries were not addressed.
Source: §6 · p. 8
Abstract
from arXiv · showhide
Neural networks are increasingly being used to solve partial differential equations (PDEs), replacing slower numerical solvers. However, a critical issue is that neural PDE solvers require high-quality ground truth data, which usually must come from the very solvers they are designed to replace. Thus, we are presented with a proverbial chicken-and-egg problem. In this paper, we present a method, which can partially alleviate this problem, by improving neural PDE solver sample complexity -- Lie point symmetry data augmentation (LPSDA). In the context of PDEs, it turns out that we are able to quantitatively derive an exhaustive list of data transformations, based on the Lie point symmetry group of the PDEs in question, something not possible in other application areas. We present this framework and demonstrate how it can easily be deployed to improve neural PDE solver sample complexity by an order of magnitude.
1. Introduction
The introduction frames neural PDE solver training as a data bottleneck and presents LPSDA as a mathematically grounded way to augment training data using solution-preserving transformations. It highlights improved sample complexity and additional evaluations of solver speed, stability, and equivariance.
1. Introduction
- Classical PDE solvers are expensive and custom-designed per PDE class, while training neural solvers still requires abundant data generated by classical solvers.Source: §1 · p. 1
- LPSDA comprises the full set of solution-preserving, pointwise, continuous data transformations for any analytic PDE.The paper identifies these permissible augmentations as Lie point symmetries, providing a mathematical basis for choosing augmentation transformations.Source: §1 · p. 1
- LPSDA improves neural PDE solver sample complexity by over an order of magnitude across multiple evolution-type PDEs and training modalities.The demonstration uses a state-of-the-art neural PDE solver.Source: §1 · p. 1
- The paper includes ablations and interpolation analyses of how individual symmetries affect neural PDE solvers.Source: §1 · p. 1
- The paper reports that LPSDA is faster and stabler than a classical solver producing the same solution trajectories and measures equivariance in classical and neural solvers.Source: §1 · p. 1
2. Background
PDEs can be represented through their prolonged derivatives, while Lie point symmetries preserve their solution sets. These symmetries can therefore generate new solutions from an existing one, motivating their use for data augmentation.
Partial Differential Equations
- A PDE relates a solution to its derivatives; its prolongation collects derivatives in jet space, where the equation defines an algebraic subvariety of prolonged solutions (Olver, 1986, p. 79).The equation is expressed as an algebraic system that must hold for every point in the domain.Source: §2 · p. 2
Lie Point Symmetries of PDEs
- Lie point symmetries are smooth, invertible transformations acting pointwise on (x,u), independent of derivatives; examples include rotations and translations, while blur is excluded.Transformations that depend on derivatives are Lie-Bäcklund transformations, not considered in this paper.Source: §2 · p. 2
- A procured solution can generate a collection of new solutions through PDE symmetries, at relatively low computational cost.The paper notes this is useful when computing transformed solutions is relatively cheap.Source: §2 · p. 2
One-parameter Lie groups
- A complete simply-connected Lie point symmetry group can be decomposed into one-parameter transformations, each corresponding to a data augmentation; distinct transformations may not commute, and reflections are excluded.The simply-connected condition limits consideration to transformations that can be continuously parameterized.Source: §2 · p. 2
- Complete sets of simply-connected Lie point symmetries can be derived for each PDE, using published tables for common equations or symbolic software for novel ones.The passages cite Ibragimov (1993) for tabulated symmetries and Baumann (1998) for symbolic derivation software.Source: §2 · p. 2
3. Lie Point Symmetry Data Augmentation
Lie point symmetry data augmentation (LPSDA) augments PDE training examples by composing symmetry transformations with sampled parameters. The section illustrates the method with KdV and describes the conditions and boundary limitations on its symmetry list.
3. Lie Point Symmetry Data Augmentation
- LPSDA augments a sampled training solution by composing one-parameter symmetries, with each parameter drawn independently from a density chosen using standard techniques such as cross-validation.The symmetry list is exhaustive under the nondegeneracy conditions discussed in Section 3.1.Source: §3 · p. 3
- With noisy simulation or real-world data, transformed examples need not be exact solutions; LPSDA assumes the noise is small enough that they remain close to the solution set.The data may also consist of known analytical solutions.Source: §3 · p. 3
- For KdV, the listed augmentation symmetries are time shift, space shift, and Galilean boost.The Galilean boost maps (x, t, u) to (x + ϵt, t, u + ϵ).Source: §3 · p. 3
- Because transformations can move grid points off a fixed Eulerian mesh, the augmented solutions are resampled on the original grid using trigonometric interpolation in space.The example describes this for a spatial translation.Source: §3 · p. 3, Figure 2 · p. 4
3.1. How complete is our symmetry list?
- The symmetry list is exhaustive when the PDE is nondegenerate—maximal rank and local solvability—and analyticity implies local solvability.The text says nondegeneracy need not be demonstrated in practice.Source: §3.1 · p. 3
3.2. A note on boundary conditions
- Symmetries are defined locally and can break at boundaries, so the method uses wrapped domains; augmentation on local patches is left for future work.The section says the restricted setting still provides enough symmetries to improve generalization in experiments.Source: §3.2 · p. 4
4. Experiments
The experiments test LPSDA across equations, model classes, training methods, and symmetry choices, assessing sample complexity, rollout performance, and equivariance. They report gains of up to ×16 in effective training-set size on the tested datasets.
4.1. Experimental setup
- Experiments tested ResNet and FNO under neural-operator and autoregressive training, reporting rollout-averaged NMSE with bootstrapped 95% confidence intervals.The neural-operator setup predicts future times directly, while the autoregressive setup predicts them recurrently in temporally bundled steps.Source: §4.1 · pp. 4–5
4.2. Different PDEs
- ×4–×16 sample-complexity gains appeared across KdV, KS, and Burgers, with the strongest improvement on Burgers attributed to its solution-combining gα subalgebra.The authors suggest this symmetry effectively doubles the dataset size, accounting for the near-doubling in slope observed for Burgers.Source: §4.2 · p. 5, Figure 4 · p. 6
4.3. Different models
- ×4–×8 effective training-set gains appeared for FNO and ResNet on KdV; ResNet was evaluated at T = 20 s rather than T = 40 s to make the task easier.The improvement differed by model and appeared to grow with dataset size for FNO.Source: §4.3 · p. 5, Figure 4 · p. 6
4.5. Incremental effect of symmetries
- Time translation (g1) and Galilean boosts (g3) improved KdV-FNO sample complexity, while space translation (g2) and scaling (g4) did not.The authors attribute g2’s lack of effect to the models’ existing translation equivariance and suggest that scaling was too simple a transformation for their parameterization.Source: Figure 4 · p. 6, §4.5 · p. 6
4.4. Different Training Methods
- LPSDA improved both neural-operator and autoregressive FNO training, although autoregressive generalization did not follow a power law.The neural-operator model predicted 100 future timesteps at once, while the autoregressive model predicted 20 at a time.Source: §4.4 · p. 6
4.6. Very long rollouts
- LPSDA improved long-rollout performance for autoregressive FNO on KdV and KS, including rollouts of 400 timesteps.The experiments compared single-step prediction, temporal bundling, temporal bundling with the pushforward trick, and the latter combined with LPSDA; cumulative error grew subexponentially over time.Source: §4.6 · p. 7
4.7. How equivariant are classical/neural solvers
- LPSDA encouraged equivariance for Galilean boosts within ϵ3 ∈ [−0.4, 0.4], while classical solvers had lower equivariance error but were two orders slower.Data augmentation took 0.76 ms; the authors propose transforming solutions of simpler problems to generate data in challenging regions. Without LPSDA, models mapped boosted inputs close to unboosted outputs for small ϵ3.Source: §4.7 · p. 7
5. Related Work
The section contrasts LPSDA with symmetry-equivariant neural solvers and situates it among classical and neural PDE-solving approaches. Figure 7 reports equivariance behavior with and without LPSDA.
5. Related Work
- Unlike Wang et al. (2021), who incorporate only one non-translation/time symmetry at a time, LPSDA learns multiple symmetries and is simpler than equivariant networks.Equivariant-network computational budgets grow exponentially with the number of symmetries, limiting cases with six or seven symmetries, which the passage says are unseen in that literature (Bronstein et al., 2021).Source: §5 · p. 7
- Classical-solver research studied Burgers’ finite-difference stability under individual symmetries (Hoarau & David (2007); Hoarau et al. (2007)), solver equivariance (Dorodnitsyn (1994)), and symmetry-preserving discretization (Valiquette & Winternitz (2005)).Source: §5 · p. 7
- Machine-learning work uses symplectic integrators for Hamiltonian Monte Carlo, learns Hamiltonians directly or indirectly, and conserves volume in conservation-form evolution equations.Symplectic integrators are limited to symplectic systems; volume conservation is implemented through finite-volume projection on solver outputs.Source: §5 · p. 7
- Related neural PDE solvers include finite-dimensional operators (Raissi, 2018; Sirignano & Spiliopoulos, 2018; Zhu & Zabaras, 2018), infinite-dimensional operators (Li et al., 2020; Bhattacharya et al., 2020), autoregressive methods (Sanchez-Gonzalez et al., 2020; Brandstetter et al., 2022), and neural augmentation of finite-element, multigrid, and eikonal solvers (Hsieh et al., 2019; Greenfeld et al., 2019; Lichtenstein et al., 2019).Source: §5 · p. 7
- For neural models tested across Galilean boost parameters ϵ3 ∈ [−0.4, 0.4], LPSDA is clearly equivariant, whereas the model trained without LPSDA (∅) is not.For small ϵ3, the unaugmented model maps f(gu0) close to f(u0), while LPSDA does not; the authors suspect it maps unseen ϵ3 toward ϵ3 = 0, often encountered in training.Source: §5 · p. 8
6. Limitations and Restrictions
The experiments are restricted to Eulerian solvers for 1D evolution equations on periodic domains, while methodological limits include global symmetries, unknown quantitative efficacy, and user-implemented augmentation functions.
6. Limitations and Restrictions
- The experiments cover Eulerian solvers for 1D evolution equations on periodic domains, a setup restriction rather than a limitation of the core idea.The authors expect LPSDA to generalize easily to other settings.Source: §6 · p. 8
- Only globally acting symmetries were implemented; local symmetries would require augmentation on local patches, which remains undeveloped.Although Lie point symmetries can be exhaustively listed for a PDE, their individual efficacy cannot be predicted quantitatively.Source: §6 · p. 8
- Users must implement augmentation functions from the mathematical symmetry descriptions, though the authors view LPSDA as a significant expansion of neural PDE solver training techniques.Mechanizing this implementation would make the method more attractive.Source: §6 · p. 8
7. Conclusion and Future work
Lie point symmetry data augmentation derives its transformations exhaustively from a PDE’s definition, grounding data augmentation mathematically. The authors report reduced sample complexity across three fluid-dynamics PDEs and identify broader symmetry and equivariance directions for future work.
Conclusion
- The authors show that Lie point symmetry transformations for data augmentation can be exhaustively derived from the PDE definition, replacing intuition with a mathematical basis.They present this as applying commonplace data augmentation to neural PDE solver training.Source: §7 · p. 8
- The method clearly decreased sample complexity, measured by generalization error, across the Korteweg-de Vries, Kuramoto-Shivashinsky, and Burgers’ equations, models, and training regimes.The authors also investigated individual symmetries and equivariance properties of classical and learned solvers.Source: §7 · p. 8
Future work
- Future extensions include generalized symmetries depending on solution derivatives, discrete symmetries such as reflections, and symmetries of variational forms that imply conserved quantities.The KdV equation has infinitely many generalized symmetries; discrete symmetries can be integrated into the framework, while Noether’s theorem connects variational symmetries to conserved quantities.Source: §7 · p. 8
- Future work could incorporate equivariance into solver architectures and combine neural PDE solvers with variable-tolerance classical solvers to build equivariance to complicated groups.The authors note that work beyond translation or scale symmetries is limited, and speculate that variable-tolerance solvers may be fruitful.Source: §7 · p. 8
A. Training data
The training data were generated numerically for KdV and Kuramoto–Shivashinsky equations, with Burgers’ equation data derived through the Cole–Hopf transformation. A finite-volume method provided a crosscheck for KdV and Kuramoto–Shivashinsky data generation.
A. Training data
- KdV and Kuramoto–Shivashinsky groundtruth data were computed on a periodic domain using method-of-lines spatial discretization and pseudospectral derivatives.The domain was represented on a regular finite mesh; spatial derivatives were computed in the frequency domain using FFTs and inverse FFTs, and time integration used fifth-order Radau IIA.Source: A. Training data · p. 11
- A finite-volume method crosschecked the KdV and Kuramoto–Shivashinsky data generation by estimating boundary fluxes for these conservation-form equations.Fluxes at the left and right cell boundaries were estimated using the pseudospectral method.Source: A. Training data · p. 11
- Mean numerical mean-squared errors for pseudospectral and FVM pseudospectral data generation were on the order of 10^-5 for KdV and Kuramoto–Shivashinsky.Figure 8 compares the two data-generation methods for both equations.Source: A. Training data · pp. 11–12, Figure 8 · p. 12
- Burgers’ equation data were obtained by solving the heat equation at viscosity ν = 0.01 and applying the inverse Cole–Hopf transformation.Heat-equation initial conditions came from applying the inverse transformation to Burgers’ initial conditions; the solved trajectories were then transformed back.Source: A. Training data · pp. 11–12
- Burgers’ trajectories 1+2 were formed by nonlinear combination of trajectories 1 and 2, matching the generator gα in Table 1.The paper identifies these trajectories as examples of the generator’s operation.Source: A. Training data · p. 12
B. Interpolation scheme
The method shifts or rotates solutions in the x-direction using trigonometric interpolation based on the Fourier shift theorem. A benchmark found no statistically significant difference in test NMSE between trigonometric and linear interpolation.
B. Interpolation scheme
- Trigonometric interpolation shifts or rotates solutions in x by applying the Fourier shift theorem on a periodic domain of length L.The method Fourier-transforms the function, shifts its phase, and applies the inverse Fourier transform to translate it by ϵ.Source: B. Interpolation scheme · p. 12
- The authors initially chose Fourier interpolation based on signal-processing considerations and its fit with the pseudospectral derivatives used in the classical solver.Source: B. Interpolation scheme · p. 12
- Test NMSE did not differ statistically significantly between trigonometric and linear interpolation; the authors suggest models may average out linear-interpolation artifacts across the Nin-step input trajectory.The models receive a trajectory rather than only a single-time initial condition.Source: B. Interpolation scheme · p. 12
C. Experimental Architectures
The experiments compare a 1D ResNet-like model with a Fourier Neural Operator (FNO), both using K history timesteps, and adapt FNO temporal processing by bundling those timesteps as channels. The section also specifies architecture and training configurations.
C. Experimental Architectures
- The study compares a baseline 1D ResNet-like model and an FNO, each taking K history timesteps and trained autoregressively or as a neural operator.The ResNet-like baseline is attributed to He et al. (2016), and the FNO to Li et al. (2020).Source: C. Experimental Architectures · p. 12
- The 1D ResNet uses 8 residual convolutional layers, while the 1D FNO uses 5 residual Fourier layers with 32 modes; the tested 2D FNO performed worse.The ResNet maps channels to 128 dimensions and uses kernel sizes 3–15 with ELU; the FNO maps to 256 dimensions and uses GeLU. Outputs map to K channels autoregressively or to the desired timestep count in neural-operator mode.Source: C. Experimental Architectures · p. 13
- Bundling K input timesteps into FNO channels enhances temporal processing, yielding K-times faster rollouts, improved stability, and decreased errors.This combines the temporal bundling idea of Brandstetter et al. (2022) with the FNO framework described by Li et al. (2020).Source: C. Experimental Architectures · p. 13
- Models were trained for 20 epochs with AdamW at learning rate 1e-4 and weight decay 1e-8, minimizing normalized mean squared error.FNO architectures used roughly 1 million parameters; experiments took 12–24 hours on average on a GeForce RTX 2080 Ti GPU.Source: C. Experimental Architectures · p. 13
D. Results
Across KdV, Kuramoto–Shivashinsky, and Burgers experiments, LPSDA performance consistently improves with additional symmetry generators. Long-rollout evaluations also test LPSDA with autoregressive FNO models and training tricks.
D. Results
- Across KdV and Kuramoto–Shivashinsky tasks, performance consistently improves with additional symmetry generators across solvers and generator combinations.The experiments compare autoregressive and neural-operator errors using a ResNet-like solver and FNO.Source: D. Results · p. 14, Table 3 · p. 14
- Burgers experiments also show consistently better performance with additional symmetry generators using an FNO solver.The comparison tests different combinations of Lie point symmetries.Source: D. Results · p. 14, Table 4 · p. 14
- Long-rollout evaluations test autoregressive FNO models on KdV and Kuramoto–Shivashinsky, with temporal bundling and pushforward training tricks ablated.The tables report cumulative error normalized over seconds for 400 output timesteps; rollouts span 200 s for KdV and 100 s for KS.Source: Table 5 · p. 15, Table 6 · p. 15