Adaptive Orthogonal Momentum: A Novel Optimizer for Transformer Language Models

View PDF

Abstract:We present Adaptive Orthogonal Momentum (AOM), a novel optimizer for transformer language models that combines selective orthogonalization with adaptive learning rates. AOM achieves a validation loss of 3.808 on the FineWeb benchmark, outperforming the AdamW baseline (4.927) and approaching the state-of-the-art Muon optimizer (3.537). Our key innovation is the integration of layer-specific orthogonal gradient processing with momentum-based adaptation, enabling more stable training and faster convergence. Extensive ablations demonstrate the effectiveness of our approach, particularly in attention layers where orthogonalization provides the most benefit.