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

James Amarel, Robyn Miller, Nicolas Hengartner, Benjamin Migliori, Emily Taylor, Alexei Skurikhin, Earl Lawrence, Gerd J. Kunde
Proceedings of the 2nd Conference on Topology, Algebra, and Geometry in Data Science(TAG-DS 2026), PMLR 334(2):364-386, 2026.

Abstract

We study how neural emulators of partial differential equation solution operators learn physical symmetries from data by introducing a hat-matrix diagnostic that quantifies the alignment of parameter updates between symmetry related training examples. The diagnostic is a metric-weighted overlap of loss gradients evaluated across group orbits, giving a proximal influence function for symmetry-related examples. Our measurements of gradient alignment across both translations and rotations for models trained as autoregressive fluid-flow emulators suggest that equivariance arises when training dynamics propagate gradients coherently throughout symmetry orbits. This finding is based on an empirical correspondence between equivariance error and cross-orbit influence. Both our UNet and ViT architectures exhibit approximate translation equivariance, yet their gradient alignment profiles differ by uniform versus periodically concentrated influence over the orbit. On a Navier-Stokes dataset, pronounced dihedral equivariance error coincides with suppressed cross-influence, identifying the failure as due to decoupled learning of symmetry group elements. Our diagnostic is architecture-agnostic and isolates the mechanism of learning symmetries from data, extending beyond forward-pass equivariance tests by directly assessing whether learning dynamics share information across physically equivalent configurations. We identify symmetry generalization performance with symmetry-compatible gradient transport, implying that evaluation of scientific machine learning models requires dynamical probes of loss landscape geometry in addition to predictive accuracy.

Cite this Paper

BibTeX
@InProceedings{pmlr-v334-amarel26a, title = {Loss Landscape Geometry of Partial Differential Equation Emulators: Or, Symmetry Learning via Gradient Alignment}, author = {Amarel, James and Miller, Robyn and Hengartner, Nicolas and Migliori, Benjamin and Taylor, Emily and Skurikhin, Alexei and Lawrence, Earl and Kunde, Gerd J.}, booktitle = {Proceedings of the 2nd Conference on Topology, Algebra, and Geometry in Data Science(TAG-DS 2026)}, pages = {364--386}, year = {2026}, editor = {Berman, Eddie and Bernárdez, Guillermo and Chen, Samantha and Cloninger, Alex and Doster, Timothy and Emerson, Tegan and Grigsby, J. Elisenda and Kvinge, Henry and Lawrence, Hannah and Marrinan, Tim and Myers, Audun and Papillon, Mathilde and Tahmasebi, Behrooz and Telyatnikov, Lev and Walters, Robin and Weber, Melanie and Xie, YuQing and Yeats, Eric}, volume = {334}, number = {2}, series = {Proceedings of Machine Learning Research}, month = {18--20 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v334/main/assets/amarel26a/amarel26a.pdf}, url = {https://proceedings.mlr.press/v334/amarel26a.html}, abstract = {We study how neural emulators of partial differential equation solution operators learn physical symmetries from data by introducing a hat-matrix diagnostic that quantifies the alignment of parameter updates between symmetry related training examples. The diagnostic is a metric-weighted overlap of loss gradients evaluated across group orbits, giving a proximal influence function for symmetry-related examples. Our measurements of gradient alignment across both translations and rotations for models trained as autoregressive fluid-flow emulators suggest that equivariance arises when training dynamics propagate gradients coherently throughout symmetry orbits. This finding is based on an empirical correspondence between equivariance error and cross-orbit influence. Both our UNet and ViT architectures exhibit approximate translation equivariance, yet their gradient alignment profiles differ by uniform versus periodically concentrated influence over the orbit. On a Navier-Stokes dataset, pronounced dihedral equivariance error coincides with suppressed cross-influence, identifying the failure as due to decoupled learning of symmetry group elements. Our diagnostic is architecture-agnostic and isolates the mechanism of learning symmetries from data, extending beyond forward-pass equivariance tests by directly assessing whether learning dynamics share information across physically equivalent configurations. We identify symmetry generalization performance with symmetry-compatible gradient transport, implying that evaluation of scientific machine learning models requires dynamical probes of loss landscape geometry in addition to predictive accuracy.} }
Endnote
%0 Conference Paper %T Loss Landscape Geometry of Partial Differential Equation Emulators: Or, Symmetry Learning via Gradient Alignment %A James Amarel %A Robyn Miller %A Nicolas Hengartner %A Benjamin Migliori %A Emily Taylor %A Alexei Skurikhin %A Earl Lawrence %A Gerd J. Kunde %B Proceedings of the 2nd Conference on Topology, Algebra, and Geometry in Data Science(TAG-DS 2026) %C Proceedings of Machine Learning Research %D 2026 %E Eddie Berman %E Guillermo Bernárdez %E Samantha Chen %E Alex Cloninger %E Timothy Doster %E Tegan Emerson %E J. Elisenda Grigsby %E Henry Kvinge %E Hannah Lawrence %E Tim Marrinan %E Audun Myers %E Mathilde Papillon %E Behrooz Tahmasebi %E Lev Telyatnikov %E Robin Walters %E Melanie Weber %E YuQing Xie %E Eric Yeats %F pmlr-v334-amarel26a %I PMLR %P 364--386 %U https://proceedings.mlr.press/v334/amarel26a.html %V 334 %N 2 %X We study how neural emulators of partial differential equation solution operators learn physical symmetries from data by introducing a hat-matrix diagnostic that quantifies the alignment of parameter updates between symmetry related training examples. The diagnostic is a metric-weighted overlap of loss gradients evaluated across group orbits, giving a proximal influence function for symmetry-related examples. Our measurements of gradient alignment across both translations and rotations for models trained as autoregressive fluid-flow emulators suggest that equivariance arises when training dynamics propagate gradients coherently throughout symmetry orbits. This finding is based on an empirical correspondence between equivariance error and cross-orbit influence. Both our UNet and ViT architectures exhibit approximate translation equivariance, yet their gradient alignment profiles differ by uniform versus periodically concentrated influence over the orbit. On a Navier-Stokes dataset, pronounced dihedral equivariance error coincides with suppressed cross-influence, identifying the failure as due to decoupled learning of symmetry group elements. Our diagnostic is architecture-agnostic and isolates the mechanism of learning symmetries from data, extending beyond forward-pass equivariance tests by directly assessing whether learning dynamics share information across physically equivalent configurations. We identify symmetry generalization performance with symmetry-compatible gradient transport, implying that evaluation of scientific machine learning models requires dynamical probes of loss landscape geometry in addition to predictive accuracy.
APA
Amarel, J., Miller, R., Hengartner, N., Migliori, B., Taylor, E., Skurikhin, A., Lawrence, E. & Kunde, G.J.. (2026). Loss Landscape Geometry of Partial Differential Equation Emulators: Or, Symmetry Learning via Gradient Alignment. Proceedings of the 2nd Conference on Topology, Algebra, and Geometry in Data Science(TAG-DS 2026), in Proceedings of Machine Learning Research 334(2):364-386 Available from https://proceedings.mlr.press/v334/amarel26a.html.

Related Material