Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization
2026-06-15 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study a way to make training language models faster using special optimizers called matrix based optimizers. They find that as models get bigger, the usual improvements become smaller with standard weight decay techniques. To fix this, they propose Hyperball, a method that keeps the size of weight matrices and their updates fixed during training. Their tests show that Hyperball can speed up training by 20-30% compared to traditional methods and helps with adjusting learning rates across different model sizes. This idea is based on theory that weight decay controls the size of weights in a way that affects how quickly their direction changes.
matrix based optimizersMuon optimizerAdamWweight decayFrobenius normlanguage model pretraininglearning ratetoken speedupangular learning rate
Authors
Kaiyue Wen, Xingyu Dang, Kaifeng Lyu, Tengyu Ma, Percy Liang
Abstract
Matrix based optimizers such as Muon can substantially speed up language model pretraining, but their gains over AdamW are observed to shrink as model size and data scale grow when using standard constant decoupled weight decay. We propose Hyperball, a simple optimizer wrapper that addresses this issue. Given a base optimizer such as Adam or Muon, Hyperball sets the Frobenius norms of weight matrices and their corresponding optimizer updates to fixed constants. On Qwen3 style models up to 1.2B parameters, Muon Hyperball achieves 20--30% token equivalent speedup over weight decay baselines. Hyperball also improves learning rate transfer across widths and depths compared to decoupled weight decay. This method is motivated by prior theory showing that training with weight decay leads to an equilibrium weight norm that only depends on the training hyperparameters. Through this mechanism, the weight decay then decides the angular learning rate, i.e. how fast the direction of the weight matrix changes.