Equivariance Breaks the Learning Rate
Andrei Manolache, Mathias Niepert
cs.LG · cs.AI
Quick verdict
Argues Adam's irrep-block step mismatch partly explains Muon's edge on equivariant nets; fixing it makes Adam match Muon — small-scale but well-controlled.
Summary
The paper argues that part of Muon's reported advantage over Adam on equivariant networks originates inside equivariant linear layers. In e3nn-style layers, each irrep block stores a channel-mixing matrix W_l shared across its 2l+1 components (expanded map W_l ⊗ I_{2l+1}); a single layer application therefore contributes up to 2l+1 rank-one terms to the block gradient, and Adam's weight-wise rescaling, blind to irrep boundaries, yields unequal spectral step sizes across blocks. The authors propose block normalization: each block of Adam's update is rescaled to spectral norm η, preserving the update direction within blocks and adding no hyperparameter. In a controlled SO(3)-invariant toy task with a matched dense control, the largest/smallest spectral-step ratio grows from 1.5 to 2.4 with width for equivariant Adam while staying near 1.7 for the dense model; combining block normalization with tuned momentum beats Muon at widths 64 and 128. On rMD17 (aspirin, ethanol) and MD22 (Ac-Ala3-NHMe) with a small NequIP-style e3nn potential, each modification improves Adam alone, and together Adam reaches force MAE 0.827 vs Muon 0.851 (aspirin), 0.961 vs 0.962 (peptide), 0.372 vs 0.348 (ethanol). Limitations: small models and datasets, three seeds for molecular runs, no significance tests, no SOAP baseline, and the promised code package is not yet released.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 2
● low 3
5 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- In e3nn-style equivariant layers, an irrep block's gradient accumulates up to 2l+1 rank-one terms, and Adam's weight-wise rescaling yields unequal spectral step sizes across blocks within a layer.
- Block normalization — rescaling each block of Adam's update to spectral norm η — removes this mismatch without new hyperparameters while preserving Adam's within-block direction and moment estimates.
- In a controlled SO(3) experiment, the spectral step mismatch grows with width (1.5→2.4) for the equivariant model but not a matched dense control (~1.7), and block normalization plus tuned momentum beats Muon at widths 64 and 128.
- On rMD17/MD22 interatomic potentials, block normalization and momentum tuning each improve Adam, and combined they make Adam competitive with Muon: better on aspirin (0.827 vs 0.851), matched on Ac-Ala3-NHMe (0.961 vs 0.962), slightly behind on ethanol (0.372 vs 0.348).
- Blockwise step control plus momentum accumulation accounts for much, though not all, of Muon's reported advantage on equivariant networks.
Strengths
Precise mechanistic diagnosis
Section 2 and Appendix A derive the gradient rank bound rank(G_l) ≤ min(m_in, m_out, 2l+1) and prove ||∆W_l ⊗ I_{2l+1}||₂ = ||∆W_l||₂, giving a clean account of why a single Adam learning rate produces heterogeneous spectral steps across irrep blocks.
Matched dense control isolates the effect
Figure 3's feature-matched dense MLP keeps the spectral-step ratio near 1.7 across widths while the equivariant model grows from 1.5 to 2.4, supporting an irrep-specific rather than width-generic origin of the mismatch.
Disciplined ablation and confound checks
Calibrated claims
Practical drop-in rule with implementation detail
Weaknesses
Attribution gap between mechanism and gains
Small scale and seed count
Missing baselines
Momentum effect is inconsistent across settings
Code not yet released
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
The paper targets a narrow, well-posed question: why matrix-structured optimizers such as Muon outperform Adam on equivariant networks. The authors localize one cause inside e3nn-style equivariant linear layers. Each irrep block stores a channel-mixing matrix W_l tied across its 2l+1 components (expanded map W_l ⊗ I_{2l+1}), so a single layer application contributes up to 2l+1 rank-one terms to the block gradient; Appendix A derives rank(G_l) ≤ min(m_in, m_out, 2l+1) and shows ||∆W_l ⊗ I_{2l+1}||₂ = ||∆W_l||₂. Because Adam rescales stored weights coordinate-wise without regard to irrep boundaries, one global learning rate produces unequal spectral step sizes ||∆W_l||₂ across blocks. The proposed fix, block normalization (Section 2), rescales each block of Adam's update to spectral norm η, changing only scale — within-block singular value ratios and moment estimates are preserved. Claims are calibrated throughout: the abstract and conclusion say blockwise step control plus momentum accumulation accounts for *much of* Muon's advantage, not all of it.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 5 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
Why does Muon outperform Adam on equivariant neural networks?
According to this paper, one reason is structural: equivariant linear layers store one channel-mixing matrix per irrep block, shared across its 2l+1 components. Adam rescales weights individually and ignores these boundaries, so blocks within the same layer receive different spectral step sizes. Muon normalizes its update per matrix block and avoids the mismatch. The authors show that adding per-block normalization plus tuned momentum makes Adam competitive with Muon on molecular benchmarks.
What is block normalization in this paper?
After Adam computes its usual update, the update is sliced into irrep blocks and each block is rescaled so its spectral norm (largest singular value) equals the learning rate η. Only the scale changes: the update direction within each block and Adam's moment estimates are untouched, no new hyperparameter is introduced, and Algorithm 1 shows a drop-in e3nn implementation.
Does Adam beat Muon after the proposed fix?
Not universally. On rMD17 aspirin and MD22 Ac-Ala3-NHMe, Adam with block normalization plus tuned momentum slightly beats or matches Muon (force MAE 0.827 vs 0.851 and 0.961 vs 0.962); on ethanol it remains slightly behind (0.372 vs 0.348). In the toy experiment the combination wins at widths 64 and 128 but not at 16 and 32. The paper's claim is competitiveness, not dominance.
What data and models were used in the experiments?
Two settings: a synthetic SO(3)-invariant regression task on random 16-point clouds (predicting the largest eigenvalue of the second-moment matrix) with a 3-layer equivariant network and a matched dense control, and a small NequIP-style e3nn interatomic potential (2 interaction layers, 32 channels, l_max=2) trained on rMD17 aspirin and ethanol plus the MD22 peptide Ac-Ala3-NHMe, using 950 training samples per dataset and force MAE as the molecular metric.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.