Loss Landscape Geometry and the Learning of Symmetries: Or, What Influence Functions Reveal About Robust Generalization

This paper introduces an influence-based diagnostic that analyzes the local geometry of the loss landscape to determine whether neural emulators of partial differential equations have successfully internalized physical symmetries by measuring the coherence of gradient updates along symmetry-related orbits, thereby offering a novel method to evaluate robust generalization beyond standard forward-pass tests.

The Gist

The Big Picture: Are AI Models "Learning" Physics or Just "Memorizing" Patterns?

Imagine you are teaching a student to predict how water flows in a river. You show them thousands of pictures of water moving.

• The Good Student (True Learning): If you show them a picture of the river flowing left, and then you show them the exact same river but flipped to flow right, they understand the physics. They know, "Oh, if I flip the scene, the water just flows the other way, but the rules are the same."
• The Bad Student (Memorization): This student memorizes the specific pictures you showed them. If you flip the picture, they get confused. They might say, "I've never seen water flow that way before, so I don't know what to do." They got a perfect score on the test, but they didn't actually learn the rules of water.

This paper asks: How can we tell if an AI is the "Good Student" or the "Bad Student"?

Most AI models for science (like predicting weather or fluid flow) are great at getting the right answer for the data they've seen. But often, they fail when the situation changes slightly (like rotating an image or moving it to a different spot). This paper introduces a new "diagnostic tool" to peek inside the AI's brain to see if it truly understands the symmetries of physics.

The New Tool: The "Echo Chamber" Test

The authors invented a way to measure something called Influence Functions. Here is a simple analogy:

Imagine the AI is a large group of people in a room, and the "Loss" is a measure of how confused they are.

1. The Standard Test (Forward Pass): You ask the group, "What happens if I rotate this image?" They give an answer. If the answer is wrong, you know they failed. But this doesn't tell you why.
2. The New Test (Influence Functions): Instead of just asking for an answer, you whisper a correction to the group based on one specific image. Then, you check: Does that whisper help them understand a different image that is just a rotated version of the first one?

• If the AI is learning physics: The whisper travels easily. If you correct them on a "North-facing" river, that correction instantly helps them understand a "South-facing" river. The "echo" is loud and clear. This means the AI has connected these two states in its brain.
• If the AI is just memorizing: The whisper dies out. Correcting the "North" image does nothing for the "South" image. The AI treats them as totally unrelated strangers.

The paper calls this "Orbit-wise Gradient Coherence." In plain English: Do the AI's learning signals travel smoothly between physically equivalent situations?

What They Found: Two Types of AI Students

The researchers tested two popular types of AI architectures (UNets and Vision Transformers) on fluid flow problems.

1. The Vision Transformers (The "Flexible" Students)

• How they act: These models are very flexible. They can learn quickly and get very high scores on standard tests.
• The Problem: When the researchers used their new "Echo Chamber" test, they found that the learning signals were uneven. The AI would learn the "North" river perfectly, but the "South" river got almost no help from that learning.
• The Result: They got good answers for the specific data they saw, but they failed to generalize. They were essentially memorizing specific patterns rather than learning the universal rules of fluid dynamics. They converged into a "basin" (a state of learning) that broke the rules of symmetry.

2. The UNets (The "Structured" Students)

• How they act: These models are built with more rigid rules (like a grid). They are less flexible but more structured.
• The Result: Their "Echo Chamber" test showed uniform coherence. When they learned about one direction, that learning spread evenly to all other directions.
• The Trade-off: They might learn a tiny bit slower or be less flexible, but when they do learn, they truly understand the symmetry. They treat all physically equivalent situations as the same.

The "Anisotropy" Surprise

The paper also found something interesting about how these models handle rotation.

• Imagine a grid of tiles. If you rotate a picture by 90 degrees, a "Good Student" should see no difference in difficulty.
• The researchers found that for some models, rotating the image by 90 degrees made the AI suddenly much worse at predicting, even though the physics hadn't changed.
• Why? The AI had learned to rely on the specific "grid" of the data. It was like a student who only knows how to read a book held upright. If you turn the book sideways, they can't read it, even though the words are the same. The AI's internal "map" of the world was distorted by the data it was fed.

The Main Takeaway

The paper concludes that getting a low error rate on a test isn't enough. You can have an AI that looks perfect on paper but fails to understand the underlying physics.

To trust an AI for scientific predictions (like climate change or fluid dynamics), you need to check how it learns, not just what it predicts.

• If the AI's learning signals (the "whispers") travel coherently between symmetrical states, it is likely learning real physics.
• If the signals get stuck or die out, the AI is just memorizing correlations and will likely fail when the real world presents a new, rotated, or shifted scenario.

In short: The authors built a "symmetry detector" that checks if an AI's brain is wired to understand the laws of physics, rather than just memorizing a photo album.

Technical Summary

Technical Summary: Loss Landscape Geometry and the Learning of Symmetries

Problem Statement

Deep learning emulators for partial differential equation (PDE) solvers frequently achieve high in-distribution accuracy but often fail to respect the fundamental physical symmetries (e.g., translations, rotations, reflections) of the governing equations. This limitation compromises their ability to extrapolate and generalize, raising the question of whether these models are learning underlying physical processes or merely fitting correlations within the training data. Existing diagnostic methods primarily rely on forward-pass equivariance tests, which measure output consistency under symmetry transformations but do not probe the learning dynamics or the internal geometry of the loss landscape that governs generalization.

Methodology

The authors introduce a geometry-aware, symmetry-conditioned diagnostic based on influence functions to probe how training updates propagate between symmetry-related states.

• Core Metric: The study defines a metric-weighted overlap of loss gradients evaluated along group orbits. Specifically, the influence of a parameter update induced by an input $x$ on the loss of a transformed input $gx$ is calculated as the Lie derivative of the cost along the gradient directions:
  $$L_V C_{gx} = (\partial_\mu C_{gx}) \chi^{\mu\nu} (-\partial_\nu C_x)$$
  Here, $\chi^{\mu\nu}$ represents the regularized neural tangent kernel metric, acting as a Fisher-information analog on the parameter space.
• Interpretation: This quantity measures whether learning signals propagate coherently across symmetry orbits. High coherence implies that the model couples physically equivalent configurations, suggesting the learning dynamics have selected a symmetry-compatible basin in the loss landscape. Low coherence indicates that the model is memorizing localized patterns or that the loss geometry decouples symmetry-related states.
• Experimental Setup: The diagnostic is applied to autoregressive emulators of two-dimensional compressible Euler flows and Navier-Stokes flows. Two architectures are compared: a UNet (13M parameters) and a Vision Transformer (ViT, 5M parameters). The models are trained on Riemann-type initial conditions (CE-RP, CE-RPUI, CE-CRP) and Navier-Stokes datasets (NS-BB, NS-Gauss, NS-Sines).
• Evaluation: The authors pair the influence analysis with standard forward-pass equivariance error tests. They evaluate performance under the dihedral group $D_4$ (rotations and reflections) and the translation group, analyzing both median errors and upper-tail (Q3) errors to capture symmetry violations.

Key Results

1. Dihedral Group ($D_4$) Learning

• Navier-NS Failure: Models trained on Navier-Stokes data exhibited catastrophic failure in equivariance for specific group elements (e.g., 90-degree rotations followed by flips), with relative errors increasing by orders of magnitude ($10^4$).
• Gradient Decoupling: Crucially, the group elements with high equivariance error corresponded precisely to those with suppressed cross-influence. The training dynamics drove the models into loss basins where gradient signals did not accumulate coherently across the orbit.
• Architecture Differences: UNets assigned near-zero cross-influence to challenging rotations, indicating a symmetry-incompatible geometry. ViTs showed a consistent but weak response. In both cases, data-induced anisotropies were absorbed into the local loss geometry, reinforcing symmetry breaking despite high pointwise accuracy on training-distribution data.
• Compressible Euler Success: Conversely, models trained on Compressible Euler data showed low equivariance error and a uniformly distributed influence profile across the $D_4$ orbit, suggesting that the training distribution adequately represented the symmetries to induce orbit-wise coupling.

2. Translation Group Learning

• Generalization without Hard Constraints: Both architectures demonstrated nontrivial cross-influence across translated states, even without explicit data augmentation or hard symmetry constraints.
• Architectural Trade-offs:
  • UNets: Exhibited nearly uniform, constructive gradient coherence across translations, consistent with their convolutional inductive bias.
  • ViTs: Distributed influence non-uniformly, showing axis-dependent resonance structures (e.g., periodicity of 16 vs. 32 pixels). This suggests ViTs concentrate learning signals on specific subsets of translation phases, allowing for rapid convergence but resulting in heterogeneous orbit-wise coupling.
• Error Correlation: Regions of elevated forward-pass error (Q3) aligned with regions of weak parameter-update coupling in the influence landscape, confirming that the local geometry of the loss surface dictates generalization capabilities.

Key Contributions

1. Novel Diagnostic Framework: The paper introduces a method to assess symmetry learning by measuring the propagation of parameter updates between symmetry-related states, moving beyond static forward-pass checks to analyze the dynamics of learning.
2. Loss Landscape Geometry: It frames symmetry learning as a problem of basin selection in the loss landscape, governed by orbit-wise gradient coherence. The work demonstrates that a model can achieve low test error while converging to a basin with a local geometry that explicitly breaks physical symmetries.
3. Architectural Insights: The study highlights a trade-off between inductive bias and optimization flexibility. Rigid architectures (UNets) promote principled symmetry learning but may constrain update directions, while flexible architectures (ViTs) optimize efficiently but may only partially internalize symmetry structures, leading to "interpolators" rather than true physics emulators.

Significance and Claims

The authors claim that their influence-based diagnostic provides a principled tool for evaluating whether surrogate models have genuinely learned the symmetries of the underlying solution operator. The paper argues that:

• Robustness Indicator: Apparent accuracy in the absence of gradient coherence is an indicator of reduced robustness under symmetry transformations.
• Mechanism of Failure: The failure to generalize is often rooted in the local geometry of the loss landscape, where training dynamics fail to couple physically equivalent states, rather than just in the representation space.
• Practical Utility: This approach allows researchers to distinguish between models that learn shared physical structures and those that assemble collections of local estimators. It suggests that for data-driven symmetry learning, exhaustive data augmentation may be unnecessary if the influence landscape confirms that unsampled translations lie in the same response-equivalence classes.

The work concludes that while symmetry-agnostic architectures can achieve low test error, true robust generalization requires training dynamics that propagate information coherently along symmetry orbits, a property that can be directly measured and diagnosed using the proposed influence functions.