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

TL;DR

POET-X is introduced, a scalable and memory-efficient variant that performs orthogonal equivalence transformations with significantly reduced computational cost and maintains the generalization and stability benefits of POET while achieving substantial improvements in throughput and memory efficiency.

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.

Figures (7)

• Figure 1: Fully-stochastic POET (with $b=1/8$) vs. block-stochastic POET (with $b=8$) for the weight matrix update coverage. In the toy experiment, we use a $64\times64$ weight matrix and run both POET variants for 100-step update, so 200 are the largest possible update steps (multiplication by $\bm{R}$ and $\bm{P}$ counts as two updates). Block-stochastic POET ensures balanced update for the weight matrix while fully-stocahstic POET does not.
• Figure 2: Illustration of efficient Cayley-Neumann parameterization (batch-wise implementation).
• Figure 3: Latency breakdown of POET, POET-X, and PyTorch Linear Layers with sequence length 2048 and block size $b = 256$.
• Figure 4: Memory breakdown for training Llama-8B on a single GPU across different methods with batch size 1, sequence lengths 1024 and block size $b = 256$. Since POET qiu2025reparameterized runs OOM under this setting, we estimate its memory footprint by profiling memory usage across different numbers of decoder layers (i.e., parameter sizes) and applying scaling.
• Figure 5: Validation perplexity dynamics with respect to GPU hours for training Llama-8B with $\bm{L}_{\max}=256$ (5B tokens) and $\bm{L}_{\max}=1024$ (10B tokens), respectively.