OLion: Approaching the Hadamard Ideal by Intersecting Spectral and $\ell_{\infty}$ Implicit Biases

TL;DR

OLion tackles the problem of merging complementary implicit biases in large-scale optimization by intersecting spectral control from orthogonalized updates with $\ell_\infty$-style coordinate control via sign updates. The method operates as sign-after-orthogonalization, implemented with momentum and RMS alignment to approximate an intersection over a Hadamard-like set for matrix parameters. The authors prove convergence under a diagonal-isotropy assumption and validate OLion across GPT-2, Llama-2, SiT, and Llama-3.1-8B, showing faster convergence, robust learning-rate tolerance, and reduced pretrain–finetune mismatch compared to AdamW and Muon. Practically, OLion offers memory efficiency, potential communication benefits from sign updates, and improved implicit bias alignment for large-scale training.

Abstract

Many optimizers can be interpreted as steepest-descent methods under norm-induced geometries, and thus inherit corresponding implicit biases. We introduce \nameA{} (\fullname{}), which combines spectral control from orthogonalized update directions with $\ell_\infty$-style coordinate control from sign updates. \nameA{} forms a Lion-style momentum direction, approximately orthogonalizes it via a few Newton--Schulz iterations, and then applies an entrywise sign, providing an efficient approximation to taking a maximal step over the intersection of the spectral and $\ell_\infty$ constraint sets (a scaled Hadamard-like set for matrix parameters). Despite the strong nonlinearity of orthogonalization and sign, we prove convergence under a mild, empirically verified diagonal-isotropy assumption. Across large-scale language and vision training, including GPT-2 and Llama pretraining, SiT image pretraining, and supervised fine-tuning, \nameA{} matches or outperforms AdamW and Muon under comparable tuning while using only momentum-level optimizer state, and it mitigates optimizer mismatch when fine-tuning AdamW-pretrained checkpoints.

Figures (13)

• Figure 1: The geometry motivation: We view Muon and Lion as maximal-update methods under two norm-induced geometries: a spectral geometry (orthogonalization / polar factor) and an $\ell_\infty$ geometry (sign-based coordinate normalization). Their intersection suggests a scaled Hadamard set as an idealized target for matrix-shaped updates, motivating an intersection-seeking design.
• Figure 2: Evolution of spectral and $\ell_\infty$ norms during GPT-2 small pretraining for weight matrices of different shapes. Adam favors a small $\ell_\infty$ norm, Muon favors a small spectral norm, while OLion simultaneously exhibits both implicit biases.
• Figure 3: Validation losses for GPT-2 pretraining with AdamW, Lion, Muon, and OLion. OLion converges faster across all model sizes.
• Figure 4: Llama-2-7B pretraining curves with different optimizers: AdamW, Lion, Muon, and OLion.
• Figure 5: Validation loss at step 10,000 for GPT-2 small trained with different optimizers under four learning rates: $3\times 10^{-4}$, $1\times 10^{-3}$, $2\times 10^{-3}$, and $5\times 10^{-3}$. OLion consistently outperforms both baselines across all learning rates, demonstrating robustness to step-size selection.

Theorems & Definitions (6)

• Lemma 4.3: Cancellation-aware upper bound
• proof
• Theorem 4.4: Descent and OLion convergence
• Theorem B.1: Diagonal-isotropy for random Gaussian singular vectors
• proof
• proof