Summary (Overview)

• Unified Framework: OmniOpt introduces a universal five-stage meta-pipeline (S0–S5) that decomposes every optimizer update into a structured transformation of the training signal, and a Linear Minimization Oracle (LMO)-driven four-axis decomposition that geometrically unifies sign, spectral, Kronecker, low-rank, and adaptive-moment directions.
• Dual-Dimension Taxonomy: The paper organizes over 100 optimizers along two complementary dimensions: a methodological dimension (T1–T5 families: element-wise adaptive moments, matrix-level structural methods, discretization/quantization, state compression, curvature-aware regularization) and an effect-objective dimension (O1–O6: convergence efficiency, per-step cost, memory, stability, hyperparameter robustness, generalization).
• Large-Scale Controlled Benchmark: The empirical study spans 24 representative optimizers, LLM pretraining from 60M to 1B parameters across four architectures (Transformer++ and three linear-attention variants) at context lengths from 256 to 32k tokens, plus image classification on CIFAR100 (ResNet50, DeiT-S, CAFormer-S12).
• Key Empirical Findings: No single optimizer dominates all objectives. Aggressive state compression (T4) excels at short context but degrades sharply under long context. Matrix-structured methods (T2) provide the strongest quality ceiling but at substantial per-step cost. Element-wise adaptive methods (T1) remain the most stable baseline. Rankings exhibit systematic crossings with scale, context length, and architecture.
• Mechanistic Ablation: A controlled ablation of Muon shows that Newton–Schulz orthogonalization is the core mechanism, while learning-rate scaling and Nesterov provide secondary gains that stack on standard Transformers but not on linear-attention architectures.

Introduction and Theoretical Foundation

The paper addresses the fragmentation of the deep learning optimizer literature, where over one hundred methods are described in incompatible vocabularies and supported by protocol-sensitive evidence. Modern LLM training makes optimizer selection a system-level design decision constrained by compute, memory, tuning budget, and task diversity. Existing surveys either lack geometric unification, do not cover LLM-specific optimizers, or do not provide a mechanism-aware benchmark.

The theoretical foundation rests on two key ideas:

1. Universal meta-pipeline: Every optimizer can be expressed as a five-stage process:

   • S0: Training signal acquisition (first-order gradient, variance-reduced, or curvature-augmented)
   • S1: Parameter scoping and routing (partition parameters by tensor topology)
   • S2: Gradient transformation (identity, orthogonalization, projection, sign map)
   • S3: State evolution (moment EMAs, Kronecker factors, quantized states)
   • S4: Update reconstruction (inverse of S2 if representation changed)
   • S5: Update finalization (learning rate, weight decay, clipping, trust ratios) Most optimizers act nontrivially in only one or two stages, leaving the rest as identity mappings.

2. LMO-driven four-axis decomposition: Using norm-constrained linear minimization oracles, the update direction of any optimizer can be characterized by four axes:

   • Axis I (Update domain): full space $R^d$, matrix space $R^{m\times n}$, low-rank subspace
   • Axis II (State estimator): how the effective signal and curvature proxy are formed (e.g., momentum, second-moment EMA, variance reduction)
   • Axis III (Geometry and precondition operator): the direction operator $\Phi_t$, equivalently an LMO over a norm ball or a preconditioner matrix $H_t$ (e.g., adaptive $\ell_\infty$ box, spectral-norm ball, Kronecker metric)
   • Axis IV (Finalization wrapper): learning rate schedule, weight decay, projection-back, routing, refresh schedule

The key insight is that AdamW corresponds to a diagonal second-moment metric (Axis III) with identity analysis basis (Axis I) and full state (Axis II). Muon corresponds to a spectral-norm LMO (Axis III) in matrix space (Axis I). Shampoo uses a Kronecker-factored metric, and GaLore uses a projected adaptive box in a low-rank subspace.

Methodology

Dual-Dimension Taxonomy

• Dimension A (Methodological): Five mutually exclusive families based on the primary mechanism:

  • T1: Element-wise adaptive moment and scalar control (AdamW, MARS-AdamW, etc.)
  • T2: Matrix-level structural methods (Muon, Shampoo, SOAP, GaLore, RMNP)
  • T3: Discretization and directional quantization (Lion, SignSGD, MARS-Lion)
  • T4: State compression and structural aggregation (AdaFactor, 8-bit Adam, APOLLO, Adam-mini)
  • T5: Curvature-aware and geometric regularization (SAM, Sophia, LAMB, Cautious optimizers)

• Dimension B (Effect-oriented): Six measurable training objectives:

  • O1: Convergence efficiency (loss reduction)
  • O2: Per-step computational cost
  • O3: Memory overhead (optimizer state)
  • O4: Training stability (gradient-norm dynamics)
  • O5: Hyperparameter robustness (LR perturbation)
  • O6: Generalization (cross-architecture, cross-context transfer)

Benchmark Design

• Stage 1: Broad screening on C4 dataset (sequence length 256) under LLaMA architecture at 60M, 130M, 350M, 1B parameters. All 24 optimizers evaluated on validation perplexity, optimizer-state memory, and per-step runtime. Weight decay and gradient clipping were disabled to isolate S2/S3 mechanisms.

• Stage 2: Transfer to FineWeb-Edu dataset (32k context) at 340M and 1B across four architectures: Transformer++ (standard attention), Gated DeltaNet, DeltaNet, GLA. Weight decay and gradient clipping enabled identically. Evaluated on WikiText test PPL and downstream commonsense reasoning accuracy (10 tasks).

• Auxiliary analyses: Gradient-norm stability (GNormCV), learning-rate perturbation robustness (0.2×, 1×, 5× tuned LR), sequence-length sensitivity (256 vs 32k), and image classification on CIFAR100 across CNN, ViT, and MetaFormer backbones.

• Mechanistic ablation: Muon decomposed into its sub-operations (NS orthogonalization, LR scaling, Nesterov, operator ordering) on C4-LLaMA 350M, with cross-scale and cross-architecture validation.

Empirical Validation / Results

Stage 1: Short-Context Screening

• Table 13 shows PPL, memory, and runtime at four scales. Key observations:

  • Best 1B PPL: APOLLO (13.53), MARS-Shampoo (13.72), Muon (13.72), RMNP (13.87)
  • Fastest per-step: Lion (12.48 ms), AdamW (18.62 ms), GaLore (15.29 ms)
  • Lowest memory: AdaFactor (0.004 GB), APOLLO (0.79 GB), GaLore (0.79 GB)

• Pareto frontiers (Figure 14): RMNP occupies a favorable quality–runtime compromise; APOLLO is best in quality–memory at short context; heavy matrix methods (SOAP, Shampoo) dominate quality but at extreme cost.

Stage 2: Long-Context Generalization

• Table 14 reports WikiText PPL and commonsense reasoning accuracy across eight scenario-architecture combinations.

  • SOAP is the most stable: top-2 PPL in 7 of 8 scenarios.
  • MARS-AdamW consistently improves over AdamW within T1.
  • Muon shows architecture sensitivity: strong on GLA, weaker on standard Transformer++.
  • APOLLO collapses: from best (13.53) at 256 tokens to worst (35.40) at 32k tokens, degradation 3× that of AdamW.

• Sequence-length sensitivity (Table 15): APOLLO degrades by +21.87 PPL (vs. AdamW +7.39). Framework attributes this to rank-bounded compression: as context grows, gradient effective rank rises, making fixed low-rank projection lossy.

Stability and Robustness

• Gradient-norm stability (Figure 17): Muon has the best aggregate stability rank (lowest GNormCV) across architectures. T2 methods generally smoother, but some (SOAP, Conda) exhibit rare single-step spikes on certain architectures (GLA). GNormCV reveals "soft instability" not captured by NaN/Inf counts.

• Learning-rate robustness (Figure 18): Lion and MARS-Lion have flattest LR response (sensitivity <10%), but at weaker tuned PPL. AdamW, MARS-Shampoo, APOLLO show high sensitivity (≥25%). T3 methods are locally tolerant but not strong.

Vision Backbones (CIFAR100, Table 16)

• Best accuracy varies: AdaBelief on ResNet50 (80.53%), Muon on DeiT-S (77.38%), Adan on CAFormer-S12 (84.89%).
• Muon shows clear advantage on Transformer-style models (+5% on DeiT-S over AdamW).
• MARS-Lion collapses on DeiT-S (33.70%), showing dangerous architecture sensitivity.

Mechanistic Ablation of Muon (Figure 20, Table 18)

• Removing AdamW's second moment is catastrophic (PPL 17.78→70.74); adding NS orthogonalization recovers to 16.86 (better than AdamW).
• Symmetric LR scaling and post-NS Nesterov provide secondary gains (final 16.51 at 350M).
• Gains stack on standard Transformer but not on Gated DeltaNet (best single gain vs. combined shown in Table 18).
• Operator order matters: momentum in orthogonalized space hurts (23.01 vs 16.60).

Theoretical and Practical Implications

• No universal best optimizer: Optimizer selection is a multi-objective problem. The binding constraint (quality, runtime, memory, stability, robustness, transfer) determines the appropriate family.

• Benefit carriers: Geometry-sensitive direction maps (Axis III) and the structured state that feeds them (Axis II) carry the strongest gains. Scalar control tweaks alone contribute little.

• Composability governed by locality: Mechanisms on different pipeline stages or axes tend to compose (e.g., variance reduction + any base direction; low-rank projection + state quantization). Mechanisms sharing a slot require explicit ordering and may conflict (e.g., spectral orthogonalization + low-rank projection; streaming updates + global gradient operations).

• Two key empirical constraints:

  1. Aggressive state compression is rank-bounded: short-context wins do not certify long-context reliability.
  2. Spectral matrix geometry is architecture-conditional: Muon-style gains transfer across scale but not across attention mechanisms.

• Practical guidance (Table 17):

  • Tier I (primary candidates): AdamW (default reference), RMNP (balanced matrix optimizer), Muon (interpretable, architecture-sensitive)
  • Tier II (scenario-dependent): SOAP (quality ceiling for long context), MARS-AdamW (stable AdamW enhancement), APOLLO (short-context memory champion, long-context risk)
  • Tier III (diagnostic failure cases): Most T1 variants, GaLore, Shampoo, Sophia, etc.

• Design principle: Optimizer geometry should be matched to model architecture and training dynamics. RMNP is cited as an example of structure-aware preconditioning leveraging row-wise matrix properties.

Conclusion

The paper concludes that optimizer selection is a constraint-driven decision. The framework provides an operational coordinate system: any optimizer can be located by its pipeline stages and four-axis coordinates, its composability predicted from that location, and its trade-offs evaluated under explicit mechanism and objective assumptions. The key recommendation is to start from AdamW and move to RMNP, SOAP, or memory-efficient methods only when a specific constraint (quality ceiling, memory pressure, cross-architecture transfer) demands it.

Future directions include:

• Developing quantitative interpretability metrics (effective rank, basis staleness)
• Adaptive compression schemes that respond to gradient rank growth
• Architecture-aware geometry design
• Automated compositional search over compatible four-axis combinations
• Cost-effective curvature estimation
• Standardized matched-budget evaluation protocols