POET-X: Memory-efficient LLM Training by Scaling Orthogonal Transformation
Zeju Qiu, Lixin LIU, Adrian Weller, Han Shi, Weiyang Liu
Abstract
Efficient and stable training of large language models (LLMs) remains a core challenge in modern machine learning systems. To address this challenge, Reparameterized Orthogonal Equivalence Training (POET), a spectrum-preserving framework that optimizes each weight matrix through orthogonal equivalence transformation, has been proposed. Although POET provides strong training stability, its original implementation incurs high memory consumption and computational overhead due to intensive matrix multiplications. To overcome these limitations, we introduce POET-X, a scalable and memory-efficient variant that performs orthogonal equivalence transformations with significantly reduced computational cost. POET-X maintains the generalization and stability benefits of POET while achieving substantial improvements in throughput and memory efficiency. In our experiments, POET-X enables the pretraining of billion-parameter LLMs on a single Nvidia H100 GPU, and in contrast, standard optimizers such as AdamW run out of memory under the same settings.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 3a9d2794-e4fb-4ef8-bdd4-e733d4fda833Builds on15
- LoRA: Low-Rank Adaptation of Large Language ModelsEdward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu et al.ICLR 2022 · 18,833 citations
- GaLore: Memory-Efficient LLM Training by Gradient Low-Rank ProjectionJiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang et al.ICML 2024 · 433 citations
- Controlling Text-to-Image Diffusion by Orthogonal FinetuningZeju Qiu, Weiyang Liu, Haiwen Feng, Yuxuan Xue et al.NeurIPS 2023 · 277 citations
- Scatterbrain: Unifying Sparse and Low-rank AttentionBeidi Chen, Tri Dao, Eric Winsor, Zhao Song et al.NeurIPS 2021 · 165 citations
- Monarch: Expressive Structured Matrices for Efficient and Accurate TrainingTri Dao, Beidi Chen, Nimit Sharad Sohoni, Arjun D. Desai et al.ICML 2022 · 125 citations
Related papers
- Reparameterized LLM Training via Orthogonal Equivalence TransformationZeju Qiu, Simon Buchholz, Tim Z. Xiao, Maximilian Dax et al.NeurIPS 2025 · 11 citations
- Gradient Weight-normalized Low-rank Projection for Efficient LLM TrainingJia-Hong Huang, Yixian Shen, Hongyi Zhu, Stevan Rudinac et al.AAAI 2025 · 1 citation
- ReLoRA: High-Rank Training Through Low-Rank UpdatesVladislav Lialin, Sherin Muckatira, Namrata Shivagunde, Anna RumshiskyICLR 2024 · 214 citations
- Train Small, Infer Large: Memory-Efficient LoRA Training for Large Language ModelsJun Zhang, Jue Wang, Huan Li, Lidan Shou et al.ICLR 2025
- Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai et al.NeurIPS 2025 · 48 citations
