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

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 37 of its 55 references that are in the map.
- First posted on arXiv 28 Jan 2026.

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

- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [A Closer Look at Memorization in Deep Networks](https://mlsanity.com/paper/1706.05394/) (2017)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [Making Convolutional Networks Shift-Invariant Again](https://mlsanity.com/paper/1904.11486/) (2019)
- [3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data](https://mlsanity.com/paper/1807.02547/) (2018)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [Steerable CNNs](https://mlsanity.com/paper/1612.08498/) (2016)
- [PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers](https://mlsanity.com/paper/2308.05732/) (2023)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [Why do deep convolutional networks generalize so poorly to small image transformations?](https://mlsanity.com/paper/1805.12177/) (2018)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [On Translation Invariance in CNNs: Convolutional Layers can Exploit Absolute Spatial Location](https://mlsanity.com/paper/2003.07064/) (2020)
- [Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior](https://mlsanity.com/paper/2306.00258/) (2023)
- [Approximately Equivariant Networks for Imperfectly Symmetric Dynamics](https://mlsanity.com/paper/2201.11969/) (2022)
- [On Second-Order Group Influence Functions for Black-Box Predictions](https://mlsanity.com/paper/1911.00418/) (2019)
- [Lie Point Symmetry Data Augmentation for Neural PDE Solvers](https://mlsanity.com/paper/2202.07643/) (2022)
- [Residual Pathway Priors for Soft Equivariance Constraints](https://mlsanity.com/paper/2112.01388/) (2021)
- [On the Symmetries of Deep Learning Models and their Internal Representations](https://mlsanity.com/paper/2205.14258/) (2022)
- [Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspective](https://mlsanity.com/paper/2502.00604/) (2025)
- [The Lie Derivative for Measuring Learned Equivariance](https://mlsanity.com/paper/2210.02984/) (2022)
- [Emergent properties of the local geometry of neural loss landscapes](https://mlsanity.com/paper/1910.05929/) (2019)
- [Don't Unroll Adjoint: Differentiating SSA-Form Programs](https://mlsanity.com/paper/1810.07951/) (2018)

## Abstract (arXiv)

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.