PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariance Breaks the Learning Rate

Andrei Manolache, Mathias Niepert

arXiv:2609.08381v1cs.LGcs.AI

TL;DR

Equivariant layers create an optimization mismatch because Adam applies one global, elementwise-scaled learning rate across irrep blocks with different shared-component structures. The paper removes this mismatch with parameter-free block normalization and evaluates it alongside momentum tuning, finding that their combination makes Adam competitive with Muon across the tested molecular datasets.

  • Problem

    Prior work reported that Muon can outperform Adam on equivariant networks, but the reason for this advantage was unidentified.

  • Method

    The paper characterizes irrep-block spectral-step mismatch, normalizes each block update without a new hyperparameter, and evaluates the correction with tuned Adam momentum in controlled and molecular models.

  • Results

    Combining block normalization with tuned momentum makes Adam competitive with Muon across the tested molecular datasets, outperforming it on aspirin, matching it on Ac-Ala3-NHMe, and remaining close on ethanol.

  • Takeaways & Limitations

    Blockwise step control and momentum accumulation account for much of Muon’s advantage in the evaluated equivariant settings, while Muon remains useful without additional moment tuning.

  • Takeaways & Limitations

    The evaluation is limited to Adam and Muon, SO(3)-equivariant networks, and three molecular datasets; other models, optimizers, and symmetry groups remain open questions.

Abstract

from arXiv · show

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 \otimes 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 $\mathrm{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.

1 Introduction and Related Work

Equivariant networks can be difficult to optimize, and prior work reported Muon gains without identifying their cause. This work attributes part of the gap to an intra-layer learning-rate mismatch and proposes parameter-free block normalization, validated in controlled and molecular settings.

  • Equivariant architectures encode geometric symmetries but can produce unfavorable loss geometry and underperform less constrained alternatives at scale.
  • Prior studies report that Muon can outperform Adam on several equivariant and geometric models, but do not explain why.
  • The paper identifies an intra-layer learning-rate mismatch in which different irrep blocks receive unequal spectral step sizes under one global learning rate.
  • A parameter-free blockwise normalization rule removes this mismatch while preserving the optimizer framework.
  • In controlled SO(3) models and molecular interatomic potentials, block normalization and tuned momentum each improve Adam, while combining them makes Adam competitive with Muon.

2 Irrep Blocks and Spectral Step Size

Equivariant linear layers share each channel-mixing matrix across an irrep’s components, creating block-specific gradient structure. Adam ignores these boundaries, so equal stored-weight updates can yield unequal spectral steps; blockwise normalization corrects the scale without changing update geometry within a block.

  • An irrep of degree l has 2l+1 components, and its multiplicity counts the number of channels sharing that representation.Scalars and vectors correspond to l = 0 and l = 1, respectively.
  • Each degree-l block uses a channel-mixing matrix W_l shared across all 2l+1 components, forming W_l⊗I_{2l+1}.The same stored parameter can therefore be repeated one, three, or five times for l = 0, 1, or 2.
  • An irrep block gradient has rank at most 2l+1 because its shared matrix receives one rank-one contribution from each component.Adam instead rescales stored weights independently from their gradient histories.
  • Equal-sized stored-weight updates can produce different spectral steps after reconstruction because Adam does not distinguish the irrep matrices.Changing the global learning rate scales all blocks equally and cannot remove this difference; Muon normalizes momentum separately per block.
  • Block normalization gives every nonzero block spectral step η without adding a hyperparameter.Unlike Muon, it uniformly rescales Adam’s update, preserving singular-value ratios while Adam’s β1 and β2 still govern accumulation.

3 Empirical Evaluation

The evaluation tests block normalization and momentum tuning in controlled SO(3)-equivariant and realistic interatomic-potential models. Blockwise correction becomes more useful with width, while combining both modifications makes Adam comparable to or better than Muon across molecular datasets.

  • Experimental settings: The evaluation uses a toy SO(3)-equivariant model with a matched dense control and an e3nn interatomic-potential model.The experiments use shared tuning and fixed training budgets, with validation-based selection.
  • Controlled setting: The equivariant model’s spectral-step ratio grows from 1.5 to 2.4 with width, while the dense control remains around 1.7.Block normalization and Muon eliminate this spread by assigning every nonzero matrix a spectral step of η.
  • Controlled setting: At widths 64 and 128, combining block normalization with tuned momentum reaches losses of 0.376 and 0.225, outperforming Muon at both widths.Block normalization alone also outperforms Muon at width 128, whereas tuned momentum alone is not consistently beneficial.
  • Interatomic potentials: On aspirin, ethanol, and Ac-Ala3-NHMe, block normalization reduces Adam’s force MAE from 1.054 to 0.972, 0.537 to 0.464, and 1.521 to 1.145.The realistic model uses e3nn and trains on rMD17 aspirin and ethanol plus the MD22 peptide Ac-Ala3-NHMe.
  • Interatomic potentials: Combining block normalization with tuned momentum reaches force MAEs of 0.827, 0.372, and 0.961, outperforming Muon on aspirin, matching it on Ac-Ala3-NHMe, and remaining slightly behind on ethanol.Neither modification alone matches Muon across the three datasets.

4 Conclusion and Future Work

The paper attributes part of Muon’s advantage to blockwise step control and momentum accumulation, while showing that Adam can become competitive when both are addressed. The evaluation remains limited to Adam and Muon in SO(3)-equivariant networks and three molecular datasets.

  • Conclusion: Separately normalizing irrep blocks removes Adam’s intra-layer learning-rate mismatch without adding a hyperparameter.In the controlled setting, combining block normalization with tuned momentum outperforms Muon at widths 64 and 128.
  • Conclusion: On molecular datasets, the combined method outperforms Muon on aspirin, matches it on Ac-Ala3-NHMe, and remains close on ethanol.Muon remains attractive because it controls matrix steps by construction and needs no additional moment tuning.
  • Future work: The evaluation is limited to Adam and Muon in SO(3)-equivariant networks and three molecular datasets.Future work includes larger models, more realistic datasets, other optimizers, other symmetry groups, and other parameter-sharing forms.

A Additional Mathematical Details

Equivariant linear maps mix repeated irrep channels with the same matrix across components, producing structured gradients and block-dependent spectral updates. Block normalization equalizes each block’s largest update singular value while preserving Adam’s within-block update geometry.

  • Irrep features and equivariant maps: A degree-l irrep has 2l+1 components, and repeated copies are mixed across channels by a matrix W_l shared across every component.The component axis transforms under rotations, while the channel axis is mixed by W_l.
  • Irrep features and equivariant maps: Under channel-first flattening, the expanded equivariant map is W_l ⊗ I_{2l+1}, which preserves the component axis while mixing channels.Schur’s lemma gives this form for repeated copies of the same irrep.
  • Gradient structure: The degree-l gradient sums one outer product for each of its 2l+1 tied component applications, so its rank is bounded by 2l+1 and block dimensions.Dense layers produce one outer product per application, whereas batching and repeated use add further terms.
  • Spectral step of the expanded layer: Sharing W_l across more components repeats its singular values in the expanded layer without increasing the spectral norm, while changing gradient accumulation.The expanded update is ΔW_l ⊗ I_{d_l}, with each singular value repeated d_l times.
  • Block normalization and Muon: Block normalization rescales each Adam update so every nonzero block has spectral step η, preserving Adam’s singular vectors and singular-value ratios.It changes the update scale uniformly within each block rather than changing the update direction.
  • Block normalization and Muon: Block normalization matches Muon on the largest update singular value, but Muon additionally flattens the momentum spectrum while block normalization preserves Adam’s spectrum.This separates blockwise step control from Muon’s additional change in update geometry.

B PyTorch Implementation

The implementation identifies e3nn linear-weight irrep blocks, applies Adam first, then reshapes and normalizes each block separately before writing the updated flat weights.

  • Update procedure: Algorithm 1 first lets Adam compute its usual update and update its moment estimates, then processes the resulting flat weights.The implementation stores old weights before applying blockwise updates.
  • Block discovery: The implementation scans model modules for e3nn linear layers and derives block ranges from each instruction’s path shape.It yields each layer weight together with its block specifications.
  • Block normalization: Each selected update slice is reshaped into a matrix, its matrix 2-norm is computed, and the normalized result is flattened back into the parameter layout.The blockwise operation uses the target scale η when the norm is nonzero.
  • Update procedure: The normalized update is added to the saved weight, changing only the selected block update while retaining the model’s flat storage representation.The operation is performed without gradient tracking and writes the updated parameter in place.

C.1 Model Architectures

The experiments use matched toy and interatomic-potential architectures with validation-based optimizer selection and controlled Muon comparisons. The reported protocol includes specified datasets, force-weighted training, cosine decay, and no weight decay.

  • Toy architecture: The toy equivariant network uses m scalar and m vector channels, three gated layers, and a scalar readout, while a dense MLP matches its hidden feature dimension.The toy model is designed as a matched architecture comparison.
  • Optimization protocol: Hyperparameters are selected using validation data only, with learning-rate and moment-coefficient tuning performed before reporting test results.The interatomic experiments re-sweep the learning rate after changing Adam’s moment coefficients.
  • Optimization protocol: All runs use cosine learning-rate decay to zero and no weight decay, while MD22 transfers tuned moment coefficients from rMD17 and selects its learning rate again.This transfer protocol is specified for the main results.
  • Muon comparison: Muon uses Nesterov momentum 0.95, and exact SVD-based and five-step Newton–Schulz updates have similar force errors on both rMD17 datasets.Table 5 reports force MAE as mean ± standard deviation over three seeds, with lower values better.

D AI Usage Statement

The authors report using language models as assistive tools for coding, experimental design, manuscript editing, and checking, while retaining responsibility for the research decisions and conclusions.

  • AI usage: Language models assisted with code, experimental design, manuscript editing, proofreading, and checks of implementation and mathematical derivations.The authors state that all generated suggestions were reviewed.
  • AI usage: The authors retain responsibility for the research questions, central ideas, experimental decisions, interpretations, and conclusions.
Loading 2609.08381v1…