POET-X: Memory-efficient LLM Training by Scaling Orthogonal Transformation

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.

Training large language models (LLMs) is difficult and expensive. It requires enormous computing power, can take a long time, and may fail if the training process becomes unstable. Previous work introduced POET, a method that improves training stability by enforcing the principle of singular spectrum preservation. However, while POET is parameter-efficient, it does not automatically reduce memory use or training time. To address this limitation, we introduce POET-X. POET-X follows the same training principle as POET, but uses optimized implementation strategies to carry out the same algorithmic operations at much lower computational cost. In practice, this gives researchers the benefits of POET’s stable training while making the method much faster and more memory-efficient. The impact of POET-X is practical and significant. It enables the training of an 8 billion-parameter LLM on a single GPU, such as an NVIDIA H100. By reducing the hardware needed for large-scale training, POET-X takes a step toward making frontier LLM research more accessible to a wider community.