Equivariance Breaks the Learning Rate
cs.LG, cs.AI
Submitted: 2026-09-08
Updated: 2026-09-30
License: http://creativecommons.org/licenses/by/4.0/
The gist: Equivariant networks are commonly trained with Adam, yet recent work reports that matrix-structured optimizers such as Muon can perform better on these architectures without explaining why.
Terminology
Abstract
Equivariant networks are commonly trained with Adam, yet recent work reports that matrix-structured optimizers such as Muon can perform better on these architectures without explaining why. We identify one source of this difference inside equivariant linear layers. Each irrep block learns a channel-mixing matrix W l shared across its 2l+1 components, giving the expanded map W l I 2l+1. For a single application of the layer, the gradient of W l sums 2l+1 outer product contributions and has rank at most 2l+1. Adam rescales stored weights individually without using the irrep boundaries, so one learning rate can produce different spectral step sizes across blocks within a layer. We address this mismatch by normalizing each block update separately, without introducing a new hyperparameter. This changes only the scale of the update, leaving Adam's moment estimates and its direction within each block unchanged. We evaluate the mechanism in a controlled SO(3) -equivariant model with a matched dense control and in an e3nn interatomic potential model trained on rMD17 and MD22. The toy setup isolates a mismatch that grows with width while the dense control shows no corresponding growth. In the interatomic potential model, block normalization and tuning Adam's momentum coefficients independently improve performance, but neither alone matches Muon. Combined, they make Adam competitive with Muon on all datasets, indicating that blockwise step control and momentum accumulation account for much of Muon's advantage.
Sources
- Does equivariance matter at scale?
- Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges
- Gaussian Error Linear Units (GELUs)
- Adam: A Method for Stochastic Optimization
- DPA4: Pushing the Accuracy-Cost Frontier of Interatomic Potentials with EMFA SO(2) Convolution
- Optimizers Qualitatively Alter Solutions And We Should Leverage This
- Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
- A Spectral Condition for Feature Learning
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