Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials

Beyond Adam: SOAP and Muon for MLIP Training§

Core Methodology§

The paper systematically evaluates three matrix-structured optimizers—Muon, SOAP, and a hybrid SOAP-Muon—against Adam for training equivariant MLIPs (NequIP and Allegro). The key insight is that leveraging second-order (or pseudo-second-order) information via structured preconditioning accelerates convergence and improves final accuracy, particularly under partial force supervision.

SOAP (Second-order Optimization with Approximate Preconditioning) is a variant of the Shampoo algorithm. For a parameter matrix $W \in \mathbb{R}^{m \times n}$, it maintains preconditioners $L$ and $R$ as running averages of the gradient's Gram matrices: $$L_{t+1} = \beta L_t + (1-\beta) G_t G_t^\top, \quad R_{t+1} = \beta R_t + (1-\beta) G_t^\top G_t$$ where $G_t$ is the gradient. The update is then: $$W_{t+1} = W_t - \eta \cdot L_t^{-1/4} G_t R_t^{-1/4}$$

Muon uses the Newton-Schulz iteration to compute the matrix orthogonalization of the gradient. For a gradient matrix $G$, it iterates:

X = G
for _ in range(K):
    X = 0.5 * X @ (3 * I - X @ X.T)

Then $W_{t+1} = W_t - \eta \cdot X$ (after scaling). Muon effectively enforces a near-orthogonal update direction.

SOAP-Muon Hybrid: Apply Muon to weight matrices (typically square or near-square) and SOAP to biases, embeddings, and other non-matrix parameters. This combines the benefits of both: Muon's fast convergence on large weight matrices and SOAP's robustness on remaining parameters.

Implementation Details§

The optimizers are implemented using the torch.optim API with optional Newton-Schulz iterations (5 iterations for Muon). The code snippet below illustrates training with SOAP-Muon:

import torch
from torch.optim.optimizer import Optimizer

# Assuming soap_muon optimizer class is defined as per paper
optimizer = soap_muon(model.parameters(), lr=0.001, muon_modules=[model.layers]) 

for epoch in range(epochs):
    for batch in dataloader:
        pred = model(batch['positions'], batch['atomic_numbers'])
        loss = energy_force_loss(pred['energy'], pred['forces'], batch['energy'], batch['forces'])
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

Key Results§

• SOAP and SOAP-Muon consistently outperform Adam in both convergence speed (fewer steps to reach target accuracy) and final validation error.
• Muon alone provides partial gains but is less robust than SOAP or hybrid.
• Under partial force supervision (e.g., only 10% of atoms have force labels), the improvement is even more pronounced: SOAP/Muon achieve ~30% lower force errors than Adam.
• Experiments on NequIP and Allegro with the 3BPA dataset (small organic molecules) and a water dataset show consistent benefits.

Theoretical Insight§

The advantage stems from better conditioning of the optimization landscape. Matrix-structured preconditioning captures gradient correlations across parameters, which is particularly beneficial for equivariant architectures where weight matrices exhibit structured symmetries. The orthogonalization enforced by Muon prevents greedy directions, while SOAP's full covariance approximation reduces ill-conditioning.

Impact§

This work highlights that optimizer choice is an underappreciated design axis for MLIPs. The recommended default is SOAP-Muon (or SOAP alone) for faster training and improved accuracy, reducing computational costs and label requirements.