Orthogonal Witness Control for Muon Optimization via Sigmoid Spectral Reshaping
Organizations: Center for AI Research, VinUniversity, Hanoi, Vietnam · Indiana University, Bloomington, IN, USA · Monash University, Clayton, VIC, Australia
Abstract
Matrix-valued optimizers such as Muon exploit the spectral structure of neural network updates through Newton--Schulz orthogonalization, but their near-flattening of the singular spectrum discards relative magnitude information across gradient modes. We introduce Soren (Spectral Orthogonal Reshaping), a matrix-valued optimizer that preserves the singular subspaces of the gradient while applying a bounded, monotone sigmoid transformation to its singular values. This smoothly compresses dominant modes without fully flattening the spectrum. We interpret Soren as a positive-definite preconditioned gradient method and establish convergence guarantees under relative smoothness and metric Polyak--Łojasiewicz geometry. To avoid explicit singular value decomposition, we further develop a finite-depth Soft Newton--Schulz (SNS) polynomial realization of the sigmoid spectral map and characterize how its spectral approximation affects the induced convergence geometry. Experiments across LLM pre-training, supervised fine-tuning, and direct preference optimization demonstrate the effectiveness and robustness of Soren against established optimizers.
Figures & tables
| LLaMA-60M | LLaMA-135M | LLaMA-350M | LLaMA-1B | |
|---|---|---|---|---|
| AdamW | 44.54 | 36.96 | 23.49 | 20.82 |
| Muon | 36.86 | 27.42 | 21.73 | 18.41 |
| HTMuon | 35.32 | 26.15 | 22.58 | 18.44 |
| Soren | 35.14 | 26.15 | 21.53 | 18.29 |
| Optimizer | ARC-c | ARC-e | PIQA | HellaSwag | OBQA | WinoGrande | Avg. |
|---|---|---|---|---|---|---|---|
| AdamW | 23.13 | 25.05 | 50.16 | 25.49 | 24.00 | 48.80 | 32.77 |
| Muon | 23.46 | 26.89 | 50.11 | 26.97 | 24.40 | 50.51 | 33.72 |
| HTMuon | 25.77 | 26.81 | 50.71 | 27.25 | 24.60 | 49.88 | 34.17 |
| Soren | 25.51 | 27.36 | 51.25 | 28.47 | 25.80 | 51.99 | 35.06 |
| Method | HellaSwag | ARC | MMLU | TruthfulQA | Winogrande | GSM8K | Average |
|---|---|---|---|---|---|---|---|
| AdamW | 73.08 ±0.44 | 55.97 ±1.45 | 61.79 ±0.39 | 46.75 ±1.48 | 68.90 ±1.30 | 64.44 ±1.32 | 61.82 |
| Muon | 73.72 ±0.44 | 56.14 ±1.45 | 62.83 ±0.39 | 47.44 ±1.49 | 68.90 ±1.30 | 69.83 ±1.26 | 63.14 |
| HTMuon | 73.75 ±0.44 | 56.31 ±1.45 | 62.93 ±0.39 | 47.50 ±1.48 | 70.24 ±1.30 | 69.22 ±1.27 | 63.33 |
| Soren (Ours) | 74.31 ±0.44 | 57.17 ±1.45 | 65.62 ±0.38 | 48.14 ±1.50 | 71.11 ±1.27 | 72.55 ±1.23 | 64.82 |
| Method | HellaSwag | ARC | MMLU | TruthfulQA | Winogrande | GSM8K | Average |
|---|---|---|---|---|---|---|---|
| AdamW | 80.69 ±0.36 | 61.09 ±1.42 | 51.09 ±0.39 | 45.32 ±1.58 | 76.40 ±1.14 | 30.93 ±1.27 | 57.59 |
| Muon | 81.89 ±0.38 | 61.95 ±1.42 | 55.07 ±0.40 | 49.56 ±1.57 | 77.35 ±1.18 | 31.69 ±1.28 | 59.59 |
| HTMuon | 81.85 ±0.38 | 61.86 ±1.42 | 55.13 ±0.40 | 49.52 ±1.57 | 76.95 ±1.18 | 31.31 ±1.28 | 59.44 |
| Soren (Ours) | 83.25 ±0.37 | 62.46 ±1.42 | 57.21 ±0.39 | 49.54 ±1.56 | 77.74 ±1.17 | 36.32 ±1.32 | 61.09 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| LLaMA-60M | Qwen2.5-3B | Mistral-7B-v0.1 | ||
|---|---|---|---|---|
| Optimizer | Time | Relative to AdamW | Time | Time |
| AdamW | 0h23m | 1.00x | 5h23m | 10h3m |
| Muon | 0h28m | 1.22x | 5h33m | 10h11m |
| HTMuon | 2h04m | 5.39x | 5h38m | 11h19m |
| Soren | 0h39m | 1.69x | 5h33m | 10h29m |
| Params | Hidden | Intermediate | Heads | Blocks | Steps | Data amount | Batch Size |
|---|---|---|---|---|---|---|---|
| 60M | 512 | 1376 | 8 | 8 | 5K | 0.5B | 512 |
| 135M | 768 | 2048 | 12 | 12 | 10K | 1B | 512 |
| 350M | 1024 | 2736 | 16 | 24 | 20K | 2B | 512 |
| 1B | 2048 | 5461 | 32 | 24 | 60K | 6B | 512 |
| LLaMA-60M | LLaMA-135M | |||
|---|---|---|---|---|
| Optimizer | LR | WD | LR | WD |
| AdamW | 1e-3 | 0.1 | 1e-3 | 0.1 |
| Muon | { 0.01 , 0.02, 0.03, 0.04} | 0.1 | { 0.01 , 0.02, 0.03, 0.04} | 0.1 |
| HTMuon | {0.01, 0.02, 0.03 , 0.04} | 0.1 | {0.01, 0.02, 0.03 , 0.04} | 0.1 |
| Soren | {0.01, 0.02, 0.03 , 0.04} | 0.1 | {0.01, 0.02, 0.03 , 0.04} | 0.1 |
| LLaMA-350M | LLaMA-1B | |||
|---|---|---|---|---|
| Optimizer | LR | WD | LR | WD |
| AdamW | 1e-3 | 0.1 | 6e-4 | 0.1 |
| Muon | {0.0025, 0.005 , 0.01, 0.015} | 0.1 | { 0.005 , 0.01} | 0.1 |
| HTMuon | {0.0025, 0.005 , 0.01, 0.015} | 0.1 | { 0.005 , 0.01} | 0.1 |
| Soren | {0.0025, 0.005, 0.01 , 0.015} | 0.1 | {0.005, 0.01 } | 0.1 |
| Qwen2.5-3B | Mistral-7B-v0.1 | |||
|---|---|---|---|---|
| Optimizer | LR | WD | LR | WD |
| AdamW | {1e-6, 5e-6, 2e-5, 3e-5 } | 0.1 | { 1e-7 , 4e-7, 5e-7} | 0.1 |
| Muon | {1e-6, 5e-6, 2e-5 , 3e-5} | 0.1 | {1e-7, 4e-7, 5e-7 } | 0.1 |
| HTMuon | {1e-6, 5e-6, 2e-5 , 3e-5} | 0.1 | {1e-7, 4e-7, 5e-7 } | 0.1 |
| Soren | { 1e-6 , 5e-6, 2e-5, 3e-5} | 0.1 | {1e-7, 4e-7 , 5e-7} | 0.1 |
| Metric | Epoch 2 | Epoch 4 | Epoch 6 | Epoch 8 | Epoch 10 |
|---|---|---|---|---|---|
| MAE | |||||
| MaxErr |
Neural Network Optimization
Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.
See all 321 papers