---
title: 'POET-X: Scalable LLM Pretraining'
url: https://www.emergentmind.com/topics/poet-x
type: topic
---

# POET-X: Scalable LLM Pretraining

Searching arXiv for POET-X and closely related POET work to ground the article in current papers.
POET-X is a scalable, memory-efficient variant of POET for large language model pretraining. It retains POET’s spectrum-preserving orthogonal equivalence formulation—training each weight matrix through left and right orthogonal transformations of a fixed base matrix—while replacing the original implementation’s weight-centric matrix-multiplication pattern with an input-centric design, blockwise orthogonal structure, permutation-by-index kernels, and more efficient orthogonal parameterization routines. In the reported experiments, this redesign yields a **3× GPU memory reduction** and **8× runtime speed-up** relative to original POET, and enables pretraining of **billion-parameter LLMs on a single NVIDIA H100 GPU**, whereas **AdamW** runs out of memory under the same setting [2603.05500].

## 1. Conceptual basis and relation to POET

POET-X inherits its optimization principle from POET, described as **reParameterized Orthogonal Equivalence Training**. The core parameterization is
\[
\bm{W}_{RP} = \bm{R}\bm{W}_0\bm{P},
\]
where \(\bm{W}_0\in\mathbb{R}^{m\times n}\) is fixed and \(\bm{R}\in\mathbb{R}^{m\times m}\), \(\bm{P}\in\mathbb{R}^{n\times n}\) are trainable orthogonal matrices. The orthogonal equivalence transformation is
\[
\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},
\]
and the forward pass is written as
\[
\begin{aligned}
&~~~~~~~~~~~~~~~\bm{y}=\bm{W}_{RP}^\top\bm{x}=(\bm{R}\bm{W}_0\bm{P})^\top\bm{x},\\
&\text{s.t.}~\big\{\bm{R}^\top\bm{R}=\bm{R}\bm{R}^\top=\bm{I},~\bm{P}^\top\bm{P}=\bm{P}\bm{P}^\top=\bm{I}\big\}.
\end{aligned}
\]

The defining theoretical property is spectrum preservation. If
\[
\bm{W}_0 = \bm{U}\bm{\Sigma}_0\bm{V}^\top,
\]
then
\[
\bm{W}_{RP} = \bm{R}\bm{U}\bm{\Sigma}_0\bm{V}^\top\bm{P}.
\]
Because \(\bm{R}\bm{U}\) and \(\bm{V}^\top\bm{P}\) remain orthogonal, the singular values are unchanged. The paper explicitly attributes POET’s training stability to **spectrum preservation** together with **provably small hyperspherical energy** under Gaussian initialization [2603.05500].

A central misconception is that POET-X introduces a new optimization objective. The paper instead presents it as a systems-level reformulation of POET: the optimization principle is preserved, and the contribution is to make orthogonal equivalence training practical at LLM scale. As in POET, \(\bm{R}\) and \(\bm{P}\) can be merged into the weight after training, so the method incurs **no inference overhead**.

## 2. Orthogonal parameterization and input-centric execution

POET-X is built on **block-stochastic POET**. At iteration \(i\), the left orthogonal factor is parameterized as
\[
\bm{R}_i=\bm{\Psi}_i^\top\cdot \text{Diag}(\tilde{\bm{G}^1_i,\tilde{\bm{G}^2_i,\cdots,\tilde{\bm{G}^{\lceil\frac{m}{b}\rceil}_i)}) \cdot\bm{\Psi}_i,
\]
with analogous structure for \(\bm{P}_i\). The weight update is
\[
\bm{W}_{i}=\bm{R}_i\bm{W}_{i-1}\bm{P}_i.
\]
The block-diagonal orthogonal matrices are denoted
\[
\bm{G}_P=\text{Diag}(\tilde{\bm{G}^1_P,\cdots,\tilde{\bm{G}^{\lceil\frac{n}{b}\rceil}_P),~~ \bm{G}_R=\text{Diag}(\tilde{\bm{G}^1_R,\cdots,\tilde{\bm{G}^{\lceil\frac{m}{b}\rceil}_R).
\]

The paper identifies the original POET implementation as **weight-centric**: it explicitly forms transformed weights through updates of the form
\[
\bm{W}\leftarrow\bm{R}_i\bm{W}\bm{P}_i,
\]
and states that this implementation has complexity
\[
\mathcal{O}(nm^2)
\]
for input \(\bm{x}\in\mathbb{R}^m\). POET-X replaces this with an **input-centric** formulation. Rather than computing \((\bm{P}_i^\top\bm{W}^\top\bm{R}_i^\top)\bm{x}\) by materializing transformed weights, it evaluates
\[
\bm{P}_i^\top\big(\bm{W}^\top(\bm{R}_i^\top\bm{x})\big).
\]
The paper emphasizes the operational difference: the weight-centric form requires **two matrix-matrix multiplications and one matrix-vector multiplication**, whereas the input-centric form requires **three matrix-vector multiplications**.

In simplified notation, the forward is written as
\[
\bm{z}=\bm{G}_P^\top\bm{W}\bm{G}_R^\top\bm{x},
\]
with execution decomposed into
\[
\text{mm1:~}\bm{a} = \bm{G}_R^\top \bm{x},\qquad
\text{mm2:~}\bm{b} = \bm{W} \bm{a},\qquad
\text{mm3:~}\bm{z} = \bm{G}_P^\top \bm{b}.
\]
This avoids explicit transformed-weight construction and reduces saved activations.

## 3. Systems-level optimizations

A large fraction of POET-X’s contribution is in the implementation of orthogonal blocks, permutations, and blockwise multiplication. The orthogonal blocks use the **Cayley-Neumann Parameterization (CNP)**. A skew-symmetric matrix \(\bm{Q}\) is constructed with
\[
\bm{Q}=-\bm{Q}^\top,
\]
and, for \(k=3\),
\[
\bm{G} \approx (\bm{I} + \bm{Q}) (\bm{I}+\sum_{i=1}^3\bm{Q}^i) = \bm{I}+2\bm{Q}+2\bm{Q}^2+2\bm{Q}^3+\bm{Q}^4.
\]
POET-X stores only the **upper-triangular part** of \(\bm{Q}\), reducing trainable degrees of freedom to
\[
\frac{b(b-1)}{2}.
\]
The paper states that this cuts the POET-related memory footprint roughly in half because gradients and optimizer states are computed for the compact representation rather than the full matrix.

For efficient kernel execution, the CNP expansion is rewritten as
\[
\bm{G} \approx 2 (\bm{Q} + \bm{Q}^2 + \bm{Q}^2 \cdot \bm{Q}) + \bm{Q}^2\cdot \bm{Q}^2 + \bm{I}.
\]
The paper’s key observation is that the higher-order terms depend only on \(\bm{Q}\) and \(\bm{Q}^2\), enabling fused Triton kernels with shared-memory reuse. It reports **2–3×** speedups for orthogonalization relative to PyTorch-native CNP [2603.05500].

Permutation handling is also redesigned. Instead of explicit permutation matrices, POET-X uses **index mapping**. For row and column permutations, the paper gives
\[
\begin{aligned}
\bm{\Psi}_m \bm{W}      \equiv \bm{W}'~~&\Leftrightarrow~~(\bm{W}')_{i, :} = \bm{W}_{\pi_p(i), :} \\
\bm{\Psi}_m^T \bm{W}    \equiv \bm{W}'~~&\Leftrightarrow~~(\bm{W}')_{i, :} = \bm{W}_{\pi^{-1}_p(i), :} \\
\bm{W} \bm{\Psi}_n      \equiv \bm{W}'~~&\Leftrightarrow~~(\bm{W}')_{:, j} = \bm{W}_{:, \pi^{-1}_q(j)} \\
\bm{W} \bm{\Psi}_n^T    \equiv \bm{W}'~~&\Leftrightarrow~~(\bm{W}')_{:, j} = \bm{W}_{:, \pi_q(j)}.
\end{aligned}
\]
Custom CUDA kernels yield up to **20× speedup** over PyTorch-native permutation, with reported examples of **18.75×**, **19.57×**, **16.76×**, and **14.00×** at hidden sizes 2048, 4096, 8192, and 16384 respectively.

A further refinement, termed **permutation reduction**, pre-merges two of the four permutations into the fixed weight matrix:
\[
\bm{z}=\bm{\Phi}_{n}\bm{G}_P^\top\underbrace{\bm{\Phi}_{n}^\top\bm{W}\bm{\Phi}_{m}}_{\text{Pre-computed by permuting }\bm{W}}
\bm{G}_R^\top\bm{\Phi}_{m}^\top\bm{x}.
\]
The reported speedup from this reduction is approximately **1.1×–1.8×**.

Finally, POET-X avoids explicit block-diagonal sparse matrix construction by treating the blocks as a batched collection of small matrices. The paper reports **2.3×–2.4×** runtime speedups and **9.55%–31.43%** memory reductions from this batch-parallel multiplication strategy, depending on sequence length [2603.05500].

## 4. Variants, training procedure, and implementation constraints

The simplified forward
\[
\bm{z}=\bm{G}_P^\top\bm{W}\bm{G}_R^\top\bm{x}
\]
induces a specific activation-memory trade-off. In standard autograd, the intermediate \(\bm{b}\) from `mm3` must be saved to compute \(\nabla_{\bm{G}_P}\), while \(\bm{a}\) from `mm2` need not be saved because \(\bm{W}\) has no gradient. POET-X accordingly defines two variants.

| Variant | Core idea | Emphasis |
|---|---|---|
| **POET-X\(_{\text{fast}}\)** | Standard autograd, saves \(\bm{b}\) | Faster |
| **POET-X\(_{\text{mem}}\)** | Gradient checkpointing, recomputes instead of saving \(\bm{b}\) | Lower memory |
| **POET-XQ** | Quantized base weights, dequantized on the fly | Quantized memory efficiency |

The paper states that **POET-XQ** can be implemented efficiently only on **POET-X\(_{\text{mem}}\)**, because activation recomputation eliminates the need to keep high-precision weights for saved activations.

In the reported LLM experiments, POET-X is applied to **Llama-style** models—**Llama-3B**, **Llama-8B**, and **Llama-13B**—trained on **C4**. Main experiments use block sizes \(b=256\) and \(b=512\). A practical constraint is that linear dimensions must be divisible by \(b\); the paper states that FFN intermediate dimensions were slightly modified for this reason [2603.05500].

POET-related parameters are updated with **AdamW**. The learning-rate scaling rule is
\[
\eta_{\text{POET}} = \gamma \cdot \eta_{\text{AdamW}},
\qquad
\gamma = 0.5.
\]
The paper also specifies a transient gradient clipping schedule after merge-then-reinitialize steps: for the first **2000 global steps**, the clipping norm is set to **0.01** immediately after a merge step and linearly increased to **1.0** over the next **10** steps. For the reported Llama-3B POET-X and POET-XQ setting, the base learning rate is **\(1\text{e-}3\)** and the reset gap is **400**.

## 5. Empirical performance

At the single-layer level, the paper reports the headline systems gains relative to original POET: **3× GPU memory reduction** and **8× runtime speed-up**. Reported forward-plus-backward times are **10.59 ms** for POET, **1.38 ms** for POET-X\(_{\text{fast}}\), and **1.89 ms** for POET-X\(_{\text{mem}}\). The paper further notes that POET-X\(_{\text{fast}}\) attains backward-pass latency comparable to a standard linear layer because of higher parameter efficiency [2603.05500].

For Llama-3B pretraining on **C4** at \(L_{\max}=256\) over **60B tokens**, the reported validation perplexity and memory figures are:

| Method | Memory (G) | Val PPL |
|---|---:|---:|
| **AdamW** | 81.03 | 12.69 |
| **POET-X\(_{b=256}\)** | 60.58 | 12.76 |
| **POET-X\(_{b=512}\)** | 68.52 | 12.05 |

This comparison is the core quality-efficiency result. POET-X with \(b=512\) beats **AdamW** on validation perplexity (**12.05** vs **12.69**) while using less memory. The paper also reports **Muon** at **11.45** validation perplexity and **70.94 G** memory, positioning POET-X as a favorable trade-off rather than a universal optimum across all axes.

The strongest scaling result concerns peak memory on a single H100. For **Llama-13B** at sequence length **2048**, the paper reports **OOM** for **AdamW** and **Muon**, while POET-X fits:

| Configuration | Peak memory |
|---|---:|
| **AdamW** | OOM |
| **POET-X\(_{\text{fast},b=256}\)** | 61.87 GB |
| **POET-X\(_{\text{mem},b=256}\)** | 47.21 GB |

For **Llama-8B** at sequence length **2048**, peak memory is **78.69 GB** for AdamW, **43.08 GB** for POET-X\(_{\text{fast},b=256}\), and **31.74 GB** for POET-X\(_{\text{mem},b=256}\). These figures underpin the paper’s claim that POET-X enables **13B pretraining on a single H100** and is **LoRA-like in memory**.

Throughput is more nuanced. The paper states that **AdamW** can be faster on a **single GPU** in some less memory-constrained settings, but often fails to fit larger configurations. POET-X retains competitive throughput while fitting larger models and longer sequences, and scales well in distributed settings because it can use **DDP** rather than **FSDP**. Reported scaling from **1 GPU to 64 GPUs** is often around **50–64×**, which the paper describes as close to linear. A plausible implication is that POET-X’s memory efficiency improves not only fit but also systems simplicity, since DDP avoids some of the collective-communication overhead imposed by FSDP.

The quantized extension, **POET-XQ**, is evaluated on Llama-3B at \(L_{\max}=256\). Reported validation perplexities are **20.49** for quantized 8-bit APOLLO, **17.74** for quantized 8-bit GaLore, **16.21** for POET-XQ\(_{b=256}\), and **14.78** for POET-XQ\(_{b=512}\). Reported memory values are **66.37 G**, **66.28 G**, **51.66 G**, and **60.65 G** respectively. The paper’s interpretation is that POET-XQ is both more memory-efficient and higher-quality than those quantized baselines [2603.05500].

## 6. Limitations and technical significance

The paper presents POET-X as a scalable realization of spectrum-preserving orthogonal equivalence training rather than as a new optimization theorem. It explicitly states that the stability rationale is inherited from POET, and that POET-X mainly contributes a practical implementation. Several limitations follow directly from the design.

First, many of the gains depend on specialized **CUDA** and **Triton** kernels. This improves performance but raises implementation complexity and portability concerns. Second, POET-X’s single-GPU throughput is not uniformly superior; **AdamW** can be faster when memory is not the dominant bottleneck. Third, the method imposes block-size divisibility constraints through \(b\), and the reported experiments modify FFN dimensions accordingly. Fourth, CNP is only approximate: the paper states that it improves numerical efficiency “at the cost of a slight loss of orthogonality,” although it also states that this does not hurt empirical performance. Fifth, the strongest advantages appear in regimes where memory pressure is severe; where dense optimizers already fit comfortably, the trade-off is less decisive [2603.05500].

Within those bounds, POET-X is significant because it converts an elegant but previously expensive orthogonal-training framework into a practical LLM pretraining method. Its main contribution is the conjunction of three properties: **spectrum preservation**, **LoRA-like memory efficiency**, and **LLM-scale training feasibility**. In that sense, POET-X is best understood not as a generic memory-saving heuristic, but as a systems architecture for scaling orthogonal equivalence training to modern transformer workloads.

Source: https://www.emergentmind.com/topics/poet-x