7 papers
FlexGrad: a root-free approach to adaptive step sizes
Pablo Barros, Aaron Defazio, Vincent Guigues
We introduce the FlexGrad optimizer, an adaptive gradient method that addresses a key limitation of classical AdaGrad-type methods: their inability to exploit correlations between…
ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models
Aaron Defazio
Schedule-Free Learning has shown promise as a practical anytime training method for machine learning, showing success across dozens of standard benchmark problems. However, strong…
Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs
Aaron Defazio, Konstantin Mishchenko, Parameswaran Raman +2
We propose Generalized Primal Averaging (GPA), an extension of Nesterov's method that unifies and generalizes recent averaging-based optimizers like single-worker DiLoCo and Schedu…
Why Gradients Rapidly Increase Near the End of Training
Aaron Defazio
During long-duration Large Language Model (LLM) training runs the gradient norm increases rapidly near the end of training. In this short note, we show that this increase is due to…
Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
Runa Eschenhagen, Aaron Defazio, Tsung-Hsien Lee +2
The recent success of Shampoo in the AlgoPerf contest has sparked renewed interest in Kronecker-factorization-based optimization algorithms for training neural networks. Despite it…
MuLoCo: Muon is a practical inner optimizer for DiLoCo
Benjamin Thérien, Xiaolong Huang, Aaron Defazio +2
DiLoCo is a powerful framework for training large language models (LLMs), enabling larger optimal batch sizes and increased accelerator utilization under networking constraints. Ho…
PARQ: Piecewise-Affine Regularized Quantization
Lisa Jin, Jianhao Ma, Zechun Liu +3
We develop a principled method for quantization-aware training (QAT) of large-scale machine learning models. Specifically, we show that convex, piecewise-affine regularization (PAR…