[Project POET-X] Scaling Orthogonal Equivalence: Training 13B LLMs with LoRA-level Memory Efficiency

TL;DR

Training Large Language Models (LLMs) is notoriously resource-heavy and often unstable. While Reparameterized Orthogonal Equivalence Training (POET) was proposed to solve stability issues via spectrum preservation, it was historically too slow and memory-hungry for the "Billion-parameter" era. POET-X changes the game by redesigning the underlying math into an input-centric form and leveraging custom Triton kernels, allowing a single H100 GPU to pretrain models up to 13B parameters—surpassing AdamW in performance while consuming significantly less memory.

The "Weight-Centric" Wall

The original POET algorithm transforms a weight matrix $W_0$ using two trainable orthogonal matrices, $R$ and $P$, such that $W_{RP}=R{W}_{0}P$. This "spectrum-preserving" property ensures that the singular values of the weights remain stable during training, preventing gradient explosions.

However, the implementation was Weight-Centric: it explicitly computed $R{W}_{0}P$ and stored it. For a model like Llama-8B, this approach causes immediate Out-of-Memory (OOM) errors because:

1. Storing the transformed weight matrix doubles memory overhead.
2. Matrix-matrix multiplications at this scale are computationally prohibitive.
3. Permutation operations for block-sparsity were unoptimized in standard PyTorch.

POET-X: Methodology & The Quest for Efficiency

The authors solve these bottlenecks through three key technical pillars:

1. Input-Centric Reformulation

Instead of transforming the weights, POET-X transforms the input activations. By rewriting the operation as $\mathbf{z}={\mathbf{P}}^{o}p({\mathbf{W}}_0^{o}p({\mathbf{R}}^{o}p\mathbf{x}))$, the complexity shifts from Matrix-Matrix multiplication to a sequence of efficient Matrix-Vector multiplications.

2. Optimized Cayley-Neumann Parameterization (CNP)

To keep $R$ and $P$ orthogonal, the authors use CNP, which approximates a matrix inverse via a Neumann series. POET-X optimizes this by:

• Half-Storage: Storing only the upper-triangular part of skew-symmetric matrices.
• Kernel Fusion: Using Triton to load tensors once into shared memory and compute high-order terms ($Q^2,Q^3,Q^4$) in a single pass.

Figure: Comparison of update coverage between fully-stochastic and block-stochastic POET.

3. Parallel Batch-wise Computation

Instead of constructing sparse block-diagonal matrices explicitly, POET-X treats each block as an independent matrix in a batch, drastically reducing the memory footprint of the operator itself.

Results: Breaking the Memory Barrier

In empirical tests, POET-X shows a massive leap in efficiency:

• Throughput: 8x faster than the original POET implementation.
• Memory: A 3x reduction in GPU footprint.
• Scaling: Successfully pretrains a 13B model on a single H100—a feat impossible for the de facto standard AdamW due to its optimizer state overhead.

Performance vs. State-of-the-Art

When compared to AdamW, Muon, and GaLore on the C4 dataset, POET-X ($b=512$) consistently achieves lower validation perplexity than AdamW.

Figure: Breakdown of per-operation time. POET-X (fast/mem) drastically reduces the backward-pass latency compared to original POET.

Quantized Training (POET-XQ)

Because POET-X uses custom CUDA kernels to handle dequantization on-the-fly, it supports 8-bit quantized pretraining (POET-XQ) with even lower memory (51.6G for a 3B model). This makes it one of the most versatile frameworks for resource-constrained training.

Critical Insight: Why This Matters

The real breakthrough of POET-X isn't just that it's faster; it's that it enables Full-Rank Training with PEFT-like Memory. Unlike LoRA which restricts updates to a low-rank subspace, POET-X allows the model to explore a richer parameter space more effectively, leading to better convergence stability.

Limitations

While POET-X is highly efficient, there is still a small computational overhead compared to a raw cuBLAS Linear layer. Additionally, the block size ($b$) introduces a new hyperparameter that needs tuning—larger blocks improve performance but increase memory.

Future Outlook

POET-X opens the door for training massive models on consumer-grade or mid-tier enterprise hardware. By scaling orthogonal transformations, we move closer to a future where "stability" and "efficiency" are no longer a trade-off in LLM development.