Alternative Unit-norm momentum-updates by Normalized nonlinear scaling
A linear-time optimizer that achieves remarkable performance without producing semi-orthogonal matrices while preserving structure to guide better-aligned progress and recondition ill-posed updates.
📄 Paper: A Survey For Linear-time Orthogonal Optimizer
- Linear-time complexity: No expensive SVD or matrix inverse operations
- Cosh-based normalization: Stable gradient scaling through hyperbolic cosine
- Drop-in replacement: Compatible with PyTorch's optimizer API
- H100 optimized: Includes FP8 operations for modern GPU acceleration
# Clone the repository
git clone https://github.com/ryyzn9/AuON.git
cd AuON
# Install dependencies
pip install -r requirements.txt
# Or install as package
pip install -e .from auon import AuON, Adam
# Create model
model = YourModel()
# Separate parameters: Adam for small params, AuON for matrix params
small_params = [p for p in model.parameters() if p.dim() < 2]
matrix_params = [p for p in model.parameters() if p.dim() >= 2]
# Create optimizers
adam_optimizer = Adam(small_params, lr=0.008, betas=(0.8, 0.95))
auon_optimizer = AuON(matrix_params, lr=0.08, momentum=0.95, weight_decay=0.01)
# Training loop
for batch in dataloader:
loss = model(batch)
loss.backward()
adam_optimizer.step()
auon_optimizer.step()
adam_optimizer.zero_grad()
auon_optimizer.zero_grad()python scripts/train.py --optimizer auon --num-chunks 10 --iterations 5000python scripts/compare_optimizers.py --num-chunks 10 --iterations 10000AuON/
├── auon/ # Core optimizer package
│ ├── optimizer.py # AuON optimizer
│ ├── muon.py # Muon optimizer (baseline)
│ ├── adam.py # Custom Adam
│ └── utils.py # Newton-Schulz, helpers
│
├── models/ # Model implementations
│ ├── gpt.py # GPT model
│ ├── components.py # Attention, MLP, etc.
│ └── custom_ops.py # FP8 operations
│
├── data/ # Data loading
│ └── fineweb.py # FineWebEDU loader
│
├── training/ # Training infrastructure
│ ├── config.py # Hyperparameters
│ └── trainer.py # Training loop
│
├── scripts/ # Executable scripts
│ ├── train.py # Main training
│ └── compare_optimizers.py
│
└── examples/ # Usage examples
└── basic_usage.py
The AuON optimizer applies:
- Momentum: Exponential moving average of gradients
- Decoupled weight decay: Applied before the update step
- Cosh normalization: Scales gradients using
sqrt(mean(cosh(g)²))
# Core AuON update (simplified)
m = momentum * m + (1 - momentum) * grad
g_normalized = m / (norm(m) + eps)
g_scaled = g_normalized * gamma
r_scaled = sqrt(mean(cosh(g_scaled) ** 2))
update = m / (r_scaled + eps)
param -= lr * update- Python >= 3.10
- PyTorch >= 2.4.0
- CUDA-capable GPU (H100 recommended for FP8 operations)
@misc{auon2024,
title={AuON: A Linear-time Alternative to Orthogonal Momentum Updates
},
author={Dipan Maity},
year={2025},
url={https://arxiv.org/abs/2509.24320/}
}Apache-2.0 License. See LICENSE for details.