# OmniOpt: Taxonomy, Geometry, and Benchmarking of Modern Optimizers

> OmniOpt unifies over 100 optimizers under a five-stage meta-pipeline and

- **Source:** [arXiv](https://arxiv.org/abs/2607.04033)
- **Published:** 2026-07-08
- **Permalink:** https://picx.dev/p/IU1PkT
- **Whiteboard:** https://picx.dev/p/IU1PkT/image

## Summary

## 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

---

_Markdown view of https://picx.dev/p/IU1PkT, served by PicX — AI-generated visual whiteboard summaries of research papers._