Technical Summary: Optimizer-Induced Mode Connectivity: From AdamW to Muon

Problem Statement

Mode connectivity, the phenomenon where independently trained neural network models can be connected by a low-loss path, has been extensively studied. However, existing literature largely treats the solution set as a monolithic entity, often assuming connectivity independent of the specific optimization dynamics used to reach those solutions. This paper posits that different optimizers impose distinct implicit biases (regularization), effectively constraining solutions to different regions of the loss landscape. The central question addressed is: How does mode connectivity behave when restricted to solutions consistent with a specific optimizer's implicit bias? Specifically, the authors investigate whether solutions found by the same optimizer form a connected set (intra-optimizer connectivity) and whether solutions found by different optimizers can be connected without incurring a loss barrier (inter-optimizer connectivity).

Methodology

Theoretical Framework

The authors formalize the concept of optimizer-regularized solution sets. Drawing on recent theoretical characterizations of adaptive optimizers (e.g., AdamW and the Lion-K family), they define the limit points of these optimizers as solutions to constrained optimization problems:

• AdamW: Converges to solutions satisfying an $\ell_\infty$-norm constraint ($\|\theta\|_{\max} \leq 1/\lambda$).
• Muon (and Lion-K with nuclear norm): Converges to solutions satisfying an operator-norm constraint ($\|\theta\|_{\text{op}} \leq 1/\lambda$).

The paper introduces two definitions:

1. Intra-optimizer connectivity: The set of solutions $O_A(\lambda)$ for a specific optimizer $A$ is connected.
2. Inter-optimizer connectivity: The union of solution sets $O_{A_1}(\lambda_1) \cup O_{A_2}(\lambda_2)$ is connected.

Theoretical analysis is conducted on two-layer ReLU networks in the interpolation regime (where zero training loss is achievable). The authors analyze the geometry of these regularized sets under varying network widths ($m$).

Experimental Setup

Empirical validation is performed on GPT-2 pretraining tasks across five datasets (enwik8, WikiText-103, Stories, BookCorpus, One Billion Word) using two model scales (~30M and ~50M parameters).

• Optimizers: AdamW and Muon.
• Path Construction: The authors employ a modified polychain interpolation algorithm combined with symmetric alignment (permutation and orthogonal alignment of attention heads and MLP neurons). This method constructs a piecewise linear path with a learnable bend-point to minimize loss barriers between models.
• Metrics: The study tracks validation loss barriers and the evolution of singular value spectra (specifically stable rank) along the interpolation paths.

Key Contributions

1. Theoretical Proofs of Intra-Optimizer Connectivity

The paper proves that for sufficiently wide two-layer ReLU networks, the solution sets induced by specific optimizers are connected:

• Theorem 3.1: Intra-optimizer connectivity holds for Muon (and Normalized GD with momentum) when the network width $m \geq 4P$, where $P$ is the number of distinct hyperplane arrangement patterns.
• Theorem 3.2: Intra-optimizer connectivity holds for AdamW (and Signum) when $m \geq m^*$, where $m^*$ is a critical width determined by the data and regularization parameters.
• Mechanism: The proofs involve reducing arbitrary solutions to "sparse" or "equalized" canonical forms within the regularized set and demonstrating that these forms can be connected via continuous paths that preserve the implicit norm constraints.

2. Characterization of Inter-Optimizer Connectivity

The authors establish that the connectivity between different optimizer-induced regions depends critically on network width and regularization strength:

• Large Width Regime: The union of two regularized sets is connected if and only if the sets overlap (i.e., the zero-loss set passes through the intersection of their respective norm balls). The number of connected components can transition from one to two as regularization strength increases.
• Finite-Width Construction: The paper constructs a specific dataset and width ($m=2$) where the zero-loss solution set decomposes into multiple disconnected components. Crucially, they prove that AdamW and Muon can be forced into different components by choosing appropriate regularization parameters. In this scenario, no zero-loss path exists between an AdamW solution and a Muon solution; any connecting path must incur a provable loss barrier (at least $1/2$ for squared loss).

3. Empirical Spectral Transitions

Experiments on GPT-2 reveal distinct spectral behaviors:

• Same-Optimizer Paths: Interpolating between two models trained with the same optimizer (e.g., AdamW-AdamW) preserves the singular value spectrum (e.g., AdamW retains outliers, Muon remains isotropic) and maintains low loss.
• Cross-Optimizer Paths: Interpolating between different optimizers (e.g., AdamW-Muon) results in a smooth spectral phase transition. The spectrum evolves from the AdamW signature to the Muon signature along the path. While the loss remains relatively low, the intermediate models exhibit spectral profiles distinct from either endpoint.

4. Out-of-Distribution (OOD) Generalization

The paper demonstrates that cross-optimizer interpolation paths can improve OOD generalization. Models sampled along the AdamW-Muon path exhibit negative loss barriers on OOD datasets (performing better than the linear interpolation of the endpoints), suggesting that traversing the space between different optimizer biases yields structurally diverse models with superior generalization properties.

Results and Significance

Summary of Findings

• Optimizer-Dependent Structure: The loss landscape is not a single connected blob of low-loss solutions; rather, it contains optimizer-specific "basins" or regions defined by implicit regularization.
• Connectivity is Conditional: While same-optimizer solutions are connected in wide networks, different optimizers may converge to disconnected components in finite-width settings, separated by loss barriers.
• Spectral Fingerprinting: The singular value spectrum serves as a robust proxy for the optimizer-induced region. Same-optimizer paths preserve this fingerprint; cross-optimizer paths traverse a transition.

Claimed Significance

The paper claims to reveal optimizer-dependent structure that is invisible to classical mode connectivity literature. By refining the definition of mode connectivity to account for implicit bias, the authors show that:

1. The choice of optimizer fundamentally shapes the topology of the reachable solution space.
2. Cross-optimizer interpolation is a viable mechanism for model merging that accesses diverse regions of the solution space, potentially improving downstream properties like quantization and generalization.
3. Theoretical guarantees for connectivity must now consider the specific regularizer induced by the optimizer, not just the network architecture.

The work bridges the gap between theoretical implicit bias analysis and practical model fusion, suggesting that the "best" path between models may depend on traversing the boundaries of different optimizer-induced regularizations.