Loss Landscape Geometry of Partial Differential Equation Emulators: Or, Symmetry Learning via Gradient Alignment

AI summary

Overview

Research area: Scientific machine learning / neural PDE emulation, specifically the geometry of loss landscapes and the mechanisms by which models acquire physical symmetries (equivariance) from data.

Technical level: Advanced. The paper is readable in its motivation and results, but the core diagnostic is built from differential geometry, influence functions, and Gauss-Newton curvature, and assumes familiarity with equivariance, group orbits, and training dynamics.

Scope: A study of how UNet and ViT fluid-flow emulators learn (or fail to learn) translation and dihedral symmetries, measured through a new gradient-alignment diagnostic that quantifies whether training updates propagate coherently across symmetry orbits.

What This Paper Is About

Neural emulators of partial differential equation solvers can achieve strong accuracy on their training distribution while still violating the symmetries of the underlying physics, such as translation and rotation invariance in fluid flow. Standard tests only check whether the trained model's forward pass is equivariant, which says nothing about how or whether the model learned the symmetry during training. This paper introduces a diagnostic that measures whether the gradients produced by symmetry-related training examples align with one another, and uses it to explain when symmetry learning succeeds and when it silently fails.

Key Contributions

1. A hat-matrix gradient-alignment diagnostic. The authors define influence between symmetry-related examples as a metric-weighted overlap of loss gradients, structured through the (λ-regularized) hat-matrix J_μ χ^{μν} J_ν. Because this quantity is the Lie derivative of the cost along the gradient flow induced by an individual example, it is formulated as a dynamical quantity observable during training rather than only a static post-training measurement.

2. An empirical correspondence between equivariance error and cross-orbit influence. Across translations and the dihedral group, the authors show that forward-pass equivariance error and gradient coherence across group orbits move together: group elements with large equivariance error are precisely those with suppressed cross-influence.

3. Architecture-specific gradient transport profiles. Both the UNet (13M parameters, 4 down-sampling blocks, 24 embedding channels) and the ViT (5M parameters, 6 layers, 256 channels) attain approximate translation equivariance, but the UNet distributes influence nearly uniformly over the translation orbit while the ViT concentrates it periodically — linked to the ViT's patch-lattice token ordering.

4. A diagnosis of dihedral failure as decoupled learning. On Navier-Stokes data, where models achieve excellent pointwise accuracy, the failure to learn the full dihedral group D₄ is traced to a loss-landscape basin in which the two symmetry sectors never receive comparable cross-influence.

Main Findings

• Translation equivariance is learned without explicit enforcement. Neither model was trained on translated data augmentation nor designed to enforce exact translation equivariance, and the initial-condition distributions violate translation stationarity (four-quadrant initial conditions for CE-RP and CE-RPUI). Both architectures nonetheless achieve median relative equivariance error on the order of one percent or less, with nontrivial cross-influence between translated states.

• UNet and ViT distribute influence differently over translations. The ViT shows an axis-dependent resonance structure with periodicity of wavelength 16 pixels under horizontal translations and 64 pixels under vertical translations, attributed to flattening the 2D patch grid into a 1D token sequence. The UNet's coarsest representation has spatial extent 8×8, producing a smoother influence dependence with only low-amplitude variation at shorter 8-pixel wavelengths and their subharmonics. Across an 18-pixel-scale lattice on joint horizontal/vertical translation grids, the UNet shows an approximately eight-pixel lattice.

• Complementarity between coherence and equivariance error. On CE data, elevated third-quartile relative equivariance error coincides with reduced gradient coherence; the same pattern holds for the UNet on NS data. The ViT–NS grid shows the weakest correspondence, with error phase-banded but coherence decaying approximately isotropically away from the identity translation.

• A trade-off between inductive bias and optimization flexibility. The UNet's near-uniform gradient coherence contrasts with the ViT's non-uniform but locally stronger influence on privileged group elements. The authors suggest that enforcing uniform coupling may constrain admissible update directions, while weaker inductive bias lets a model concentrate learning signal on a subset of symmetry transformations.

• No conflicting updates in the translation experiments. Median influence values were exclusively positive across the translation-group experiments, indicating optimization is not frustrated by conflicting updates; the difference lies in how the signal is distributed.

• Catastrophic dihedral failure on Navier-Stokes. For D₄ group actions, relative SMSE increases by nearly 10⁴ — the models are not even approximately equivariant across all group actions. Both architectures perform competitively on the identity e, 180-degree rotations r², 90-degree rotation followed by vertical-axis reflection sr, and 270-degree rotation followed by reflection sr³, but fail on the complementary actions r, r³, s, and sr². The data distribution (and the models) are approximately equivariant only under the Klein four-subgroup V₄ = {e, r², sr, sr³}.

• Orbit-wise separation develops early and persists. Checkpointed influence profiles show that although the two sectors begin with comparable cross-influence, both the UNet and ViT rapidly establish and maintain stronger transfer within V₄∖e than into D₄∖V₄. Because influence into the unlearned actions is near zero, symmetry-related gradient signals enter only as subleading perturbations and never accumulate coherently.

• The invariance group alone does not predict generalization. The initial-condition distribution is not invariant under the translation group in either the Euler or Navier-Stokes settings, yet the models learn approximate translation equivariance while failing on D₄. The distinguishing observable between the two scenarios is the cross-influence profile.

Methodology in Plain English

The authors train two standard emulator backbones as autoregressive surrogates for fluid flow. The UNet has 13M parameters with four down-sampling blocks and 24 embedding channels; the ViT has 5M parameters, six layers, and 256 channels. Training data comes from PDEGym: three classes of Riemann-type initial conditions (CE-RP, CE-RPUI, CE-CRP), each with 6,500 trajectories of 16 time steps, where each snapshot is a 128×128 grid of mass density, Cartesian momentum density, and energy density. They also use Navier-Stokes velocity fields from NS-BB, NS-Gauss, and NS-Sines initial conditions, which have smoother, viscosity-regularized, vorticity-dominated dynamics. Training uses Adam with learning rate 5×10⁻⁴, weight decay λ = 10⁻⁴, and mini-batches of N = 48 transitions, optimizing a scaled mean-squared error (SMSE) that normalizes errors by channel RMS so shocks and wavefronts are captured without losing sensitivity to quiescent flow. Models run in distributed mode on two 40GB A100 GPUs using Lux.jl and Zygote.jl, with three seeds controlling initialization and dataset splits; results are reported with quantile range bars.

The measurement apparatus is the influence function. For each example, the gradient flow defines a vector field V^μ = −χ^{μν} ∂_ν C_x, where χ_μν = η_μν + λ δ_μν combines an isotropic regularization term with η, the pullback of the Euclidean metric on function space to parameter space. The influence on the group-transformed input gx is then the metric-weighted gradient overlap (∂_μ C_gx) χ^{μν} (∂_ν C_x). This asks a simple question: if the model takes a training step driven by one example, does the loss on a symmetry-equivalent example go down too? Because χ is the regularized Gauss-Newton metric, sandwiching it between Jacobians yields the hat-matrix, which reduces to the neural tangent kernel in the Euclidean limit. Values are cosine-normalized by both gradient magnitudes, which sets each self-response to unity and isolates cross-response magnitude — appropriate when influence effects accumulate over many optimization steps. Influence is evaluated over 18 test mini-batches, each containing 3 distinct initial conditions over 16 time-steps, for each seed and architecture. Rather than solving for the metric iteratively or approximating curvature, the authors QR-decompose the mini-batch gradient matrix to obtain an orthonormal basis, so inversion is both exact and efficient because the spanned subspace dimension is bounded by the mini-batch size rather than the parameter count. Results were not sensitive to choosing λ = 10⁻⁶ or λ = 10⁻⁴.

For translation, the group action is implemented as cyclic permutation of array indices; for the square group, standard image rotation and flipping commands are used, taking care to mix vector indices appropriately for the velocity field alongside rotating the coordinate basis. Training was halted in the late-stage learning regime, so the hat-matrix diagnostics probe local gradient geometry after the models entered low-test-loss basins but before learning saturated. The paper primarily measures influence at a single parameter point due to computational expense; checkpointed influence evolution appears in the supplementary figures.

Why This Matters

Impact on research. The paper argues that forward-pass equivariance checks are insufficient on their own. A model can be accurate and approximately equivariant on some group elements while its local loss geometry actively breaks symmetry on others. The gradient-alignment diagnostic is architecture-agnostic and probes a mechanism — whether learning updates share information across physically equivalent configurations — rather than only an outcome. This reframes symmetry generalization as symmetry-compatible gradient transport, and suggests that evaluating scientific ML models requires dynamical probes of loss landscape geometry in addition to predictive accuracy.

Real-world applications:

• Neural surrogates for fluid simulation, where a model that quietly fails on rotations or reflections may be silently unreliable on configurations outside the narrow training distribution.
• Weather and climate forecasting models, which increasingly use transformer and convolutional architectures that are not equivariant by design.
• Models trained on data with broken symmetry — boundary conditions, material inhomogeneities, grid discretization, measurement noise — where approximate rather than exact equivariance is the realistic target.
• Deployment settings where symmetry-transformed inputs are physically admissible but out of distribution statistically, making the gap between training and operational distributions a genuine correctness concern.

Industry relevance. The diagnostic gives practitioners a way to tell whether apparent performance reflects genuine physics learning or the assembly of collections of local estimators. That distinction matters for trust in scientific ML systems, particularly where robustness under symmetry transformations is required. The code is released at https://github.com/lanl/PDEHats. The work was supported by the Los Alamos National Laboratory Laboratory Directed Research and Development program under project numbers 20250637DI, 20250638DI, and 20250639DI, and is published under LA-UR-25-29466.

Future Directions

• Galilean boosts. Exact equivariance to boosts can only be achieved by affine transformations, so generic nonlinear architectures can at best satisfy this symmetry approximately. The authors propose extending the influence-based analysis to probe responses under small Galilean boosts.
• Scaling symmetry. The continuum Navier-Stokes scaling symmetry was not examined, because spatial discretization, numerical regularization, and coarse graining generically break it, making faithful assessment ambiguous in the present framework.
• Beyond UNets and ViTs. The study is restricted to two backbones dominant in PDE emulation. Whether the measured gradient coherence properties depend on the backbone or on the autoregressive training paradigm itself remains only partially resolved. The appendix notes that Fourier neural operators were noncompetitive on these flows due to their plane-wave basis being poorly suited to localized discontinuities and sharp shock fronts, and that DeepONets currently lack competitive models on PDEGym.
• Approximate or relaxed symmetry mechanisms. The conclusion motivates developing symmetry mechanisms that retain enough structure to guide generalization while preserving the flexibility needed for efficient optimization.
• Richer influence measurement. The reported profiles aggregate over multiple trajectory times and initial-condition classes, so their variability reflects both estimator noise and the intrinsic heterogeneity of the operator-learning problem: shocks, smooth advective structures, and vorticity-dominated states need not express symmetry-compatible gradient transport equally.

Target Audience

Researchers in scientific machine learning and neural PDE emulation, particularly those working on equivariant and geometric deep learning, operator learning, and PDE foundation models. The paper is also relevant to machine learning theorists interested in loss landscape geometry, influence functions, and training dynamics, and to practitioners who deploy surrogate models and need diagnostics that go beyond held-out accuracy. Readers should be comfortable with equivariance, group orbits, and basic differential geometry; the applied findings in the results sections are accessible without that background.

Authors’ abstract

We study how neural emulators of partial differential equation solution operators internalize physical symmetries by introducing an influence-based diagnostic that measures the propagation of parameter updates between symmetry-related states, defined as the metric-weighted overlap of loss gradients evaluated along group orbits. This quantity probes the local geometry of the learned loss landscape and goes beyond forward-pass equivariance tests by directly assessing whether learning dynamics couple physically equivalent configurations. Applying our diagnostic to autoregressive fluid flow emulators, we show that orbit-wise gradient coherence provides the mechanism for learning to generalize over symmetry transformations and indicates when training selects a symmetry compatible basin. The result is a novel technique for evaluating if surrogate models have internalized symmetry properties of the known solution operator.