🤖 machine learning

How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks

This paper demonstrates that the Muon optimizer consistently outperforms Adam in training equivariant neural networks by navigating more regular loss surfaces and yielding higher-rank representations, highlighting the critical yet underexplored role of optimizer design in shaping solutions for geometric deep learning.

Original authors: Teodor-Mihai Stupariu, Andrei Manolache

Published 2026-05-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Teodor-Mihai Stupariu, Andrei Manolache

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to recognize 3D objects, like a chair or a cat, just by looking at a cloud of dots (a point cloud). To make the robot smart, you build it with special "rules" baked into its brain. These rules say, "It doesn't matter if the chair is rotated or flipped; it's still a chair." In the scientific world, these are called equivariant neural networks. They are designed to respect the geometry of the world.

However, there's a problem. Because these rules are so strict, the robot's brain is hard to train. It's like trying to walk through a maze where the walls are made of glass; you can see the exit, but you keep bumping into invisible barriers. Usually, scientists try to fix this by making the rules slightly less strict (relaxing the constraints).

This paper asks a different question: What if we don't change the rules, but instead change the "coach" that guides the robot through the maze?

The Two Coaches: Adam vs. Muon

The researchers compared two different "coaches" (optimizers) that tell the robot how to adjust its brain:

Adam: The standard, popular coach used by almost everyone.
Muon: A newer, specialized coach that uses a clever trick to keep its guidance steps "orthogonal" (at right angles to each other), ensuring it explores the maze more evenly.

The Race Results

The researchers put these coaches to work on three different types of robot brains (architectures) using two main test sets:

ModelNet40: A dataset of 3D shapes (like furniture).
QM9: A dataset of molecules (chemical structures).

The Outcome:
In the 3D shape tests, Muon consistently beat Adam, no matter which robot brain architecture was used. Even when the data was "corrupted" (noisy or messy), Muon still performed better. In the molecular tests, Muon also improved results on most tasks.

The paper claims that simply swapping the coach from Adam to Muon made the robots smarter, without changing the robot's design or the rules it follows.

Why Did Muon Win? (The Secret Sauce)

The researchers didn't just look at the scores; they looked inside the robot's brain to see what was different. They found three key differences:

1. The Terrain of the Solution
Imagine the training process as walking down a mountain to find the lowest valley (the best solution).

Adam tended to find valleys that looked a bit jagged and irregular on the surface.
Muon found valleys that looked smoother and more regular, even though the ground was actually steeper (higher curvature).
Analogy: It's like Adam found a rocky, uneven path down the mountain, while Muon found a smooth, paved slide. Even though the slide was steeper, the smoothness helped the robot settle into a better spot.

2. The "Rank" of the Brain
The researchers looked at how the robot's brain used its "muscle memory" (weights and internal representations).

Adam tended to rely heavily on a few specific directions, like a musician who only plays three notes over and over. This is called having a "low rank."
Muon used a much wider variety of directions, like a full orchestra playing a complex symphony. This is called having a "higher stable and effective rank."
Analogy: Adam was like a painter using only a few colors, while Muon used the entire rainbow. The paper suggests that for these geometric robots, using the full range of colors (directions) leads to a better picture.

3. The Path Taken
The paper suggests that Muon's special "orthogonal" trick prevents the robot from getting stuck in narrow, dead-end paths that Adam might fall into. It forces the robot to explore the landscape more broadly, finding a better solution that Adam missed.

The Bottom Line

The paper concludes that the choice of optimizer (the coach) is a powerful tool that has been overlooked. You don't always need to redesign the robot or loosen its rules to make it better. Sometimes, you just need a better coach who knows how to navigate the specific geometry of the problem.

What the paper does NOT claim:

It does not claim this works for every type of AI problem (it worked best on 3D shapes and molecules, but was less clear on graph-based data).
It does not claim to have solved the problem of training these networks forever; it just shows a promising new direction.
It does not discuss medical or clinical applications; the focus is strictly on the mechanics of training these specific types of neural networks.

In short: If you are training a geometric AI, don't just tweak the design. Try changing the optimizer, and you might find a smoother, smarter path to success.

Technical Summary: How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks

Problem Statement
Equivariant neural networks encode geometric symmetries directly into their architecture, providing a natural inductive bias for data with known geometric symmetries. However, hard equivariance constraints often make optimization difficult, leading to loss landscapes containing critical points and spurious minima. Consequently, these models can underperform less constrained alternatives at scale. The dominant response in recent literature has been to soften architectural constraints through approximate or relaxed equivariance methods. In contrast, the role of the optimizer in shaping the solutions of equivariant models remains underexplored. This paper investigates whether the choice of optimizer, independent of architectural modifications, can significantly alter the solutions reached during training and improve performance.

Methodology
The authors conduct a comparative study between Adam (Kingma & Ba, 2015) and Muon (Jordan et al., 2024), a recently proposed optimizer that orthogonalizes its momentum buffer via Newton-Schulz iterations. The evaluation spans three distinct settings:

3D Point Cloud Classification: Using the ModelNet40 dataset (and its corrupted variant, ModelNet40-C) with three architectures possessing different geometric inductive biases:
- EGNN: Explicit $E(3)$ -equivariant graph neural network.
- DGCNN: Permutation-equivariant graph CNN with dynamic $k$ -NN connectivity.
- PointNet: Permutation-invariant architecture via symmetric pooling.
Molecular Learning: Using the QM9 dataset with GotenNet ( $E(3)$ -equivariant) and GINE (permutation-equivariant) on peptide and ZINC datasets.
Geometric Analysis: To understand why performance differs, the authors analyze the trained checkpoints using:
- Local Loss Geometry: Visualizing 2D loss slices and 3D surfaces around checkpoints, and estimating Hessian summaries (top eigenvalue and trace) via power iteration and Hutchinson estimators.
- Spectral Structure: Computing the stable rank and effective rank (based on Shannon entropy of singular values) for both weight matrices and intermediate layer representations.

Key Results

Performance Gains: On ModelNet40, Muon consistently outperforms Adam across all architectures (EGNN, PointNet, DGCNN) in both clean and corrupted settings. For example, on clean ModelNet40, Muon improved EGNN accuracy by +5.17%, PointNet by +2.67%, and DGCNN by +1.96%. Similar trends were observed on the QM9 molecular dataset, where Muon improved most target properties.
Loss Landscape Characteristics: Visualizations of the loss surface around trained checkpoints reveal that Muon solutions exhibit smoother, more regular local geometry compared to the irregular contours often seen with Adam. However, Hessian estimates contradict the intuition that smoother geometry implies lower curvature; Muon checkpoints actually possess larger top eigenvalues and trace estimates (indicating higher curvature) than Adam checkpoints.
Spectral Properties: Muon consistently yields higher stable and effective ranks for both weights and intermediate representations across all layers and architectures. This suggests that Muon distributes capacity across singular directions more evenly, avoiding the concentration of spectra into a few dominant directions often associated with gradient-based methods.

Significance and Claims
The paper posits that the optimizer is a meaningful source of inductive bias that interacts with geometric constraints. The authors claim that:

Optimizer Choice Matters: Changing the optimizer from Adam to Muon can improve performance in equivariant models without modifying the architecture or relaxing constraints.
Distinct Solution Spaces: Muon and Adam converge to qualitatively different solutions, characterized by different local loss geometries and spectral structures.
Mechanism Hypothesis: The orthogonalized momentum updates in Muon appear to provide a trajectory-level bias toward more favorable regions of the loss landscape, treating singular directions of updates more evenly.
Design Principle: The results suggest a potential design principle for equivariant optimizers: update rules that prevent the singular spectra of weights and representations from concentrating in a few directions may improve training.

The authors maintain a modest stance, noting that while the correlation between spectral properties and accuracy gains is consistent, a causal link has not been established. They conclude that the interaction between optimizer design and geometric inductive bias warrants further investigation, particularly regarding how optimizers can exploit the structure of symmetry groups rather than treating equivariance as a black-box constraint.

The Two Coaches: Adam vs. Muon

The Race Results

Why Did Muon Win? (The Secret Sauce)

The Bottom Line

Technical Summary: How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks

More like this