Layer-Adaptive Orthogonal Momentum: A Novel Optimizer for Transformer Training

View PDF

Abstract:We present Layer-Adaptive Orthogonal Momentum (LAOM), a novel optimization method for training transformer-based language models. LAOM combines layer-specific learning rate adaptation with orthogonal momentum updates, particularly benefiting attention layers. Through extensive experiments on the FineWeb benchmark using a 134M parameter Qwen 3 architecture, we demonstrate that LAOM achieves a validation loss of 4.63, outperforming the AdamW baseline (4.9266) and ranking second on the AardXiv optimizer leaderboard. Our method introduces three key innovations: (1) layer-specific learning rate scaling based on component type, (2) Newton-Schulz orthogonalization for attention layer gradients, and (3) dynamic variance stabilization techniques. The paper includes complete implementation details, ablation studies, and analysis of training dynamics to facilitate reproducibility and future research.