Beyond Adam: SOAP and Muon Optimizers for Faster MLIP Training

Paper: 2607.02499 Authors: Gil Harari, Yoel Zimmermann, Ola Tangen Kulseng, Laura Zichi, Chuin Wei Tan, Marc L. Descoteaux, Boris Kozinsky Categories: cs.LG, cs.AI, physics.chem-ph, physics.comp-ph

The Gap

The MLIP community has poured enormous energy into architecture design (NequIP, Allegro, MACE, etc.) and dataset curation (MD17, OC20, etc.), but the optimizer has been treated as a settled question: everyone uses Adam or AdamW. The assumption is that these adaptive gradient methods are “good enough” and the real gains come from better equivariant architectures or larger datasets. Recently, a new class of matrix-structured optimizers — SOAP, Muon, and hybrids — have shown dramatic improvements in vision and language tasks, yet nobody had systematically tested them in the MLIP setting. The gap is clear: optimizer choice is an unexplored design axis in a field where every percent of accuracy or reduction in training cost matters for practical deployment.

[Optimizer choice ignored in MLIP field]
          |
          v
[Assumption: Adam is good enough]
          |
          v
[New matrix optimizers excel in vision/NLP]
          |
          v
[This paper: systematic comparison for MLIPs]
          |
          v
[SOAP/SOAP-Muon >> Adam in speed AND accuracy]

The Increment

One sentence: Before this paper, every MLIP practitioner defaulted to Adam; after this paper, they have strong evidence that switching to SOAP or SOAP-Muon gives meaningfully faster convergence and better final accuracy — especially when force labels are incomplete.

Core Mechanism

The paper evaluates three optimizer families against Adam baselines on two canonical MLIP architectures (NequIP and Allegro). The key players are SOAP (which decomposes gradient updates into spectral components using a structured matrix approach), Muon (which uses a momentum-based method with matrix-aware updates), and SOAP-Muon (a hybrid that combines both ideas). Each optimizer replaces Adam’s element-wise adaptive scaling with matrix-structured updates that respect the geometry of the parameter space more faithfully.

The experimental pipeline is straightforward: train identical model architectures with identical hyperparameter sweeps for each optimizer, measure convergence curves (loss vs. wall-clock time) and final energy/force MAE on held-out test sets. The authors test under both full supervision (all forces available) and partial supervision (only a fraction of forces labeled), which is the realistic setting where expensive DFT calculations limit data availability.

   [Gradient from loss]
           |
           v
   +------------------+
   | Matrix-structured |
   | spectral decomp.  |
   +------------------+
           |
     +-----+-----+
     |           |
     v           v
  [SOAP]     [Muon]
  spectral   momentum +
  factors    matrix norm
     |           |
     +-----+-----+
           |
           v
   [SOAP-Muon hybrid]
           |
           v
   [Parameter update]

Think of it like tuning a piano versus tightening a guitar string. Adam is like the guitar tuner: it adjusts each string (parameter) independently, based on how that individual string sounds (its own gradient history). SOAP is like a piano technician: it listens to the harmonic relationships between all the strings at once, understanding that tightening one string changes the resonance of the whole instrument. The “spectral decomposition” is the technician’s ability to hear overtones — it decomposes the gradient information into frequency-like components that capture how parameters interact with each other. Muon adds a memory of past tuning patterns (momentum) but applied at the matrix level. The hybrid SOAP-Muon is the master technician who uses both the harmonic ear and the muscle memory of years of tuning. The result: the piano is in tune faster and holds its tuning better than a guitar tuned string-by-string.

Key Concepts

• Matrix-structured optimizers: Traditional optimizers like Adam treat each parameter as an independent scalar with its own learning rate. Matrix-structured optimizers instead treat groups of parameters as a matrix and perform operations (like spectral decomposition or matrix normalization) on that matrix as a whole. The intuition: parameters in neural networks are not independent — a weight in layer 3 interacts with weights in layer 2. By respecting these interactions during the update step, you converge faster. Think of it as the difference between adjusting 1000 independent thermostat dials versus turning a single steering wheel that controls all of them in a coordinated way.

• Force supervision ratio: In MLIP training, you have two types of labels: energies (one number per atomic configuration) and forces (a 3D vector per atom, derived from the energy gradient). Forces are expensive to compute via DFT but contain rich local information. “Partial force supervision” means only some configurations have force labels. The paper shows optimizer gains are largest in this regime — when data is scarce or partially labeled, the optimizer’s efficiency matters most. This is the practically important setting, because in real-world molecular simulation, you rarely have complete force labels.

• Equivariant neural network architectures: NequIP and Allegro are not ordinary neural networks — they enforce physical symmetries (rotation, translation, permutation) by construction. The parameters live in structured spaces, not arbitrary Euclidean space. This is exactly why matrix-structured optimizers make sense: they can exploit the geometric structure that these architectures bake in, whereas Adam’s element-wise approach ignores it entirely.

Framework Shift

Before (mainstream approach):              After (this paper):

[MLIP Architecture design]                [MLIP Architecture design]
        |                                          |
        v                                          v
[Dataset curation]                         [Dataset curation]
        |                                          |
        v                                          v
[Adam/AdamW -- fixed]                      [Optimizer choice -- active axis]
        |                                   /    |    \
        v                                 v     v     v
[Final accuracy]                        [SOAP] [Muon] [SOAP-Muon]
                                               |
                                               v
                                        [Faster, more accurate]

From treating the optimizer as a fixed constant to treating it as an active design axis, the core shift is recognizing that optimizer geometry and architecture geometry must be aligned for efficient training.

Expert Assessment

Problem choice: This is a genuine and well-identified gap. The MLIP community has been so focused on architecture innovation that the optimizer has become invisible infrastructure — nobody questions it, nobody benchmarks it. This paper does the obvious experiment that should have been done years ago, and the results are non-trivial. It sits at a productive intersection: the optimizer community has been shipping new methods, but they are only tested on vision and language benchmarks; the MLIP community needs every edge it can get for practical deployment. Someone had to bridge this gap.

Method maturity: The approach is comparative benchmarking rather than algorithmic novelty — the authors did not invent SOAP or Muon, they applied them to MLIPs. This is valuable but not technically deep. The hyperparameter sweeps appear thorough, though the paper would benefit from ablations on *why these optimizers help in this specific geometry. Is it the equivariant structure? The loss landscape shape? The force/energy multi-task setup? These mechanistic questions remain unanswered.

Experimental integrity: The baselines are fair — same architectures, same data, same compute budget, with careful hyperparameter sweeps for each optimizer. The partial-force-supervision experiments are a nice touch because they probe the practically relevant regime. One concern: the improvements, while consistent, vary in magnitude across tasks, and the paper could be more transparent about variance and statistical significance. The wall-clock comparisons assume a fair computational cost per step, which is reasonable but worth verifying independently.

Writing quality: The paper is clearly written and well-structured. The main weakness is the discussion section — it largely restates results without offering mechanistic hypotheses for *why SOAP excels on MLIPs specifically. A deeper analysis connecting optimizer geometry to equivariant network structure would elevate this from “here’s a useful benchmark” to “here’s a principled insight.” The related work section could also better position the contribution against the broader optimizer literature.

Verdict: weak accept — The contribution is empirically solid and practically useful, but it lacks the mechanistic depth that would make it a landmark paper. It is a strong recommendation for practitioners to try SOAP, but it does not advance our theoretical understanding of optimization in equivariant architectures.

Takeaways

Three concrete things a practitioner can steal:

1. Switch your optimizer today: If you are training NequIP or Allegro (or likely any equivariant MLIP), replace Adam with SOAP or SOAP-Muon. The paper provides enough detail on hyperparameters to reproduce. This is low-risk, high-reward.

2. Partial supervision amplifies optimizer choice: If your dataset has incomplete force labels (which is almost always the case in practice), optimizer selection matters *more than in the full-supervision regime. Invest in optimizer tuning precisely when data is scarce.

3. The “optimizer as design axis” framing transfers: In any domain where architectures have strong geometric priors (equivariant networks, physics-informed networks, etc.), Adam may be leaving performance on the table. The general principle — match your optimizer’s inductive bias to your architecture’s structure — is worth testing in your own work.