Loss Landscape Geometry of Partial Differential Equation Emulators: Or, Symmetry Learning via Gradient Alignment
cs.LG, physics.comp-ph
Submitted: 2026-01-28
Updated: 2026-08-25
Journal ref: Proceedings of the 2nd Conference on Topology, Algebra, and Geometry in Data Science(TAG-DS 2026), PMLR 334(2):364-386, 2026
Code: https://github.com/aai-institute/pyDVL
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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
Terminology
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.
Sources
- A Closer Look at Memorization in Deep Networks
- Coherent Gradients: An Approach to Understanding Generalization in Gradient Descent-based Optimization
- Lie Point Symmetry and Physics Informed Networks
- Steerable CNNs
- Don't Unroll Adjoint: Differentiating SSA-Form Programs
- Neural Tangent Kernel: Convergence and Generalization in Neural Networks
- Residual Pathway Priors for Soft Equivariance Constraints
- Emergent properties of the local geometry of neural loss landscapes
- Understanding Black-box Predictions via Influence Functions
- Stiffness: A New Perspective on Generalization in Neural Networks
- PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers
- Equivariant geometric convolutions for emulation of dynamical systems
- The Lie Derivative for Measuring Learned Equivariance
- Towards Multi-spatiotemporal-scale Generalized PDE Modeling
- Optimizing Neural Networks with Kronecker-factored Approximate Curvature
- ClimaX: A foundation model for weather and climate
- The Local Elasticity of Neural Networks
- Poseidon: Efficient Foundation Models for PDEs
- Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
- Approximately Equivariant Networks for Imperfectly Symmetric Dynamics
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks