Equivariance Breaks the Learning Rate
Not on the map. Why?
Abstract (arXiv)
Equivariant networks are commonly trained with Adam, yet recent work reports that matrix structured optimizers such as Muon can perform better, with the reasons for these gains only partly understood. We identify one source of this difference inside equivariant layers. An equivariant layer learns one channel mixing matrix Wl per degree l, which we call an irrep block, and shares it across the 2l+1 components, giving the expanded map Wl otimes I2l+1. This sharing sums gradient contributions across components and can produce different update scales under SGD. Adam's entrywise normalization reduces sensitivity to gradient scale, but neither optimizer directly controls the effective step size of each block. A single learning rate can therefore produce different effective step sizes across blocks. Muon instead controls the effective step size by approximately equalizing the singular values of each momentum matrix. We normalize each irrep block update by a single scalar, preserving its singular value ratios while letting the learning rate control its size. We implement this with spectral normalization or a simpler root-mean-square normalization. We evaluate spectral normalization in a controlled SO(3)-equivariant model with a matched non-equivariant model. In this setting, the step size mismatch grows with width in the equivariant model but not in the non-equivariant model. We evaluate both variants across molecular force prediction on the rMD17 and MD22 datasets, QM9 molecular property prediction, and charged particle dynamics. Across these applications, block normalization generally improves Adam and closes part of its gap to Muon. These results highlight an overlooked interaction between equivariant architectures and their optimizers. Studying and designing the two together may help explain and address training difficulties often attributed to equivariance itself.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 30 of its 42 references that are in the map. First posted on arXiv 8 Sep 2026.
Builds on
30 of 42 references.
References · 30
- 2016Gaussian Error Linear Units (GELUs)cited by 3,431
- 2021Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gaugescited by 1,131
- 2021E(n) Equivariant Graph Neural Networkscited by 951
- 2018Adafactor: Adaptive Learning Rates with Sublinear Memory Costcited by 896
- 2018Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point cloudscited by 735
- 2021E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentialscited by 701
- 2019Large Batch Optimization for Deep Learning: Training BERT in 76 minutescited by 671
- 2017Large Batch Training of Convolutional Networkscited by 529
- 2022MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fieldscited by 467
- 20183D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Datacited by 348
- 2018Neural Relational Inference for Interacting Systemscited by 339
- 2023A Spectral Condition for Feature Learningcited by 101
18 more
- 2022Approximately Equivariant Networks for Imperfectly Symmetric Dynamicscited by 95
- 2021Residual Pathway Priors for Soft Equivariance Constraintscited by 74
- 2024Scalable Optimization in the Modular Normcited by 64
- 2024Does equivariance matter at scale?cited by 49
- 2023Approximation-Generalization Trade-offs under (Approximate) Group Equivariancecited by 45
- 2022Relaxing Equivariance Constraints with Non-stationary Continuous Filterscited by 41
- 2023Symmetry Breaking and Equivariant Neural Networkscited by 24
- 2024Improving Equivariant Model Training via Constraint Relaxationcited by 22
- 2023Optimization Dynamics of Equivariant and Augmented Neural Networkscited by 17
- 2025Optimizers Qualitatively Alter Solutions And We Should Leverage Thiscited by 17
- 2024Relaxed Equivariance via Multitask Learningcited by 13
- 2025Learning (Approximately) Equivariant Networks via Constrained Optimizationcited by 10
- 2025A Tale of Two Symmetries: Exploring the Loss Landscape of Equivariant Modelscited by 8
- 2024Relaxed Equivariant Graph Neural Networkscited by 7
- 2026Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentialscited by 3
- 2026PolarAdamW: Disentangling Spectral Control and Schur Gauge-Equivariance in Matrix Optimisationcited by 2
- 2026Protein Fold Classification at Scale: Benchmarking and Pretrainingcited by 2
- 2026How the Optimizer Shapes Learned Solutions in Equivariant Neural Networkscited by 2
Built on by
Not cited yet.