GWT: Scalable Optimizer State Compression for Large Language Model Training

Abstract

Training large language models (LLMs) requires substantial memory, a significant fraction of which is consumed by the moment states maintained by adaptive optimizers such as Adam. Existing memory-efficient approaches commonly compress parameters or gradients through low-rank projections, which may discard information outside the selected subspace or incur additional costs from subspace construction and singular value decomposition (SVD). We introduce \textbf{Gradient Wavelet Transform (GWT)}, an SVD-free framework that applies a multilevel Haar transform to matrix gradients. GWT maintains first- and second-moment states only for the compact approximation coefficients, while retaining all detail coefficients transiently in each update. We show that the orthogonal transform preserves the complete current gradient and establish gradient alignment and a stationary-point guarantee for smooth objectives under bounded preconditioning and a directly verifiable coarse-momentum coherence condition. Experiments on language-model pre-training and downstream fine-tuning show that GWT reduces estimated model-and-optimizer-state memory by up to 52\% while maintaining competitive or improved model quality. On LLaMA-3B pre-training, GWT achieves $1.9\times$ the throughput of 8-bit Adam and is compatible with multiple optimizer families, including Adam, Adam-mini, and Muon.