ARO: A New Lens On Matrix Optimization For Large Models

TL;DR

This paper introduces Adaptively Rotated Optimization (ARO), a novel matrix optimization framework that enhances large language model training efficiency by leveraging gradient rotation and a new benchmarking protocol.

Contribution

ARO offers a new paradigm beyond orthogonalization for matrix optimization in LLM training, improving efficiency and sample performance with a rigorous evaluation protocol.

Findings

ARO outperforms AdamW by 1.3-1.35x in LLM pretraining.

ARO surpasses orthogonalization methods by 1.1-1.15x.

ARO maintains efficiency up to 8B parameters and 8x overtrain budget.

Abstract

Matrix-based optimizers have attracted growing interest for improving LLM training efficiency, with significant progress centered on orthogonalization/whitening based methods. While yielding substantial performance gains, a fundamental question arises: can we develop new paradigms beyond orthogonalization, pushing the efficiency frontier further? We present \textbf{Adaptively Rotated Optimization (ARO}, a new matrix optimization framework that treats gradient rotation as a first class design principle. ARO accelerates LLM training by performing normed steepest descent in a rotated coordinate system, where the rotation is determined by a novel norm-informed policy. This perspective yields update rules that go beyond existing orthogonalization and whitening optimizers, improving sample efficiency in practice. To make comparisons reliable, we propose a rigorously controlled benchmarking protocol that reduces confounding and bias. Under this protocol, ARO consistently outperforms AdamW (by 1.3 $\sim$1.35$\times$) and orthogonalization methods (by 1.1$\sim$1.15$\times$) in LLM pretraining at up to 8B activated parameters, and up to $8\times$ overtrain budget, without evidence of diminishing returns. Finally, we discuss how ARO can be reformulated as a symmetry-aware optimizer grounded in rotational symmetries of residual streams, motivating advanced designs that enable computationally efficient exploitation of cross-layer/cross module couplings.