GWT: Scalable Optimizer State Compression for Large Language Model Training

[2501.07237] GWT: Scalable Optimizer State Compression for Large Language Model Training Skip to main content Search Submit Donate Log in Computer Science > Machine Learning arXiv:2501.07237 (cs) [Submitted on 13 Jan 2025 ( v1 ), last revised 25 Aug 2026 (this version, v6)] Title: GWT: Scalable Optimizer State Compression for Large Language Model Training Authors: Ziqing Wen , Ping Luo , Jiahuan Wang , Junlin Zeng , Kun Yuan , Dongsheng Li , Tao Sun View PDF HTML (experimental) 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. Subjects: Machine Learning (cs.LG) ; Artificial Intelligence (cs.AI) Cite as: arXiv:2501.07237 [cs.LG] (or arXiv:2501.07237v6 [cs.LG] for this version) https://doi.org/10.48550/arXiv.2501.07237 Focus to learn more arXiv-issued DOI via DataCite Submission history From: Ziqing Wen [ view email ] [v1] Mon, 13 Jan 2025 11:35:09 UTC (322 KB) [v2] Tue, 29 Jul 2025 09:24:44 UTC (216 KB) [v3] Wed, 30 Jul 2025 01:07:39 UTC (216 KB) [v4] Mon, 30 Mar 2026 13:25:46 UTC (210 KB) [v5] Mon, 27 Apr 2026 13:03:55 UTC (208 KB) [v6] Tue, 25 Aug 2026 05:32:44 UTC (314 KB) Full-text links: Access Paper: View PDF HTML (experimental) TeX Source view license Additional Features Audio Summary Current browse context: cs.LG < prev | next > new | recent | 2025-01 Change to browse by: cs cs.AI References & Citations SciX Google Scholar Semantic Scholar export BibTeX citation Bookmark Bibliographic Tools Bibliographic and Citation Tools Bibliographic Explorer Toggle Bibliographic Explorer ( What is the Explorer? ) Connected Papers Toggle Connected Papers ( What is Connected Papers? ) Litmaps Toggle Litmaps ( What is Litmaps? ) scite.ai Toggle scite Smart Citations ( What are Smart Citations? ) Code, Data, Media Code, Data and Media Associated with this Article alphaXiv Toggle alphaXiv ( What is alphaXiv? ) Links to Code Toggle CatalyzeX Code Finder for Papers ( What is CatalyzeX? ) DagsHub Toggle DagsHub ( What is DagsHub? ) GotitPub Toggle Gotit.pub ( What is GotitPub? ) Huggingface Toggle Hugging Face ( What is Huggingface? ) ScienceCast Toggle ScienceCast ( What is ScienceCast? ) Demos Demos Replicate Toggle Replicate ( What is Replicate? ) Spaces Toggle Hugging Face Spaces ( What is Spaces? ) Spaces Toggle TXYZ.AI ( What is TXYZ.AI? ) Related Papers Recommenders and Search Tools Link to Influence Flower Influence Flower ( What are Influence Flowers? ) Core recommender toggle CORE Recommender ( What is CORE? ) IArxiv recommender toggle IArxiv Recommender ( What is IArxiv? ) Author Venue Institution Topic About arXivLabs arXivLabs: experimental projects with community collaborators arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website. Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them. Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs . Which authors of this paper are endorsers? | Disable MathJax ( What is MathJax? ) We gratefully acknowledge support from our major funders , member institutions , and all contributors. About · Help · Contact · Subscribe · Copyright · Privacy · Accessibility · Operational Status (opens in new tab) Major funding support from (opens in new tab) (opens in new tab) (opens in new tab)

Related documents

• GWT: Scalable Optimizer State Compression for Large Language Model Training#670313
• JeanKaddour/minipile#786962
• h2oGPT: Democratizing Large Language Models#446540
• CLASP: Class-Adaptive Layer Fusion and Dual-Stage Pruning for Multimodal Large Language Models#13147
• WhiteMatter: All-to-All Cross-Layer Connections via KV Source Mixing#1118690