POET-X: Scalable LLM Pretraining

• The paper introduces POET-X, which replaces weight-centric matrix multiplications with an input-centric design to achieve 3× memory reduction and 8× runtime speed-up.
• It employs blockwise orthogonal transformations and permutation-by-index kernels to preserve the spectrum while maintaining efficient, scalable LLM pretraining.
• Custom CUDA and Triton kernels, along with variants like POET-X_mem and POET-XQ, enable practical pretraining of billion-parameter models without inference overhead.

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 LLM 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 (Qiu et al., 5 Mar 2026).

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 (Qiu et al., 5 Mar 2026).

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{W}_0\in\mathbb{R}^{m\times n}$0 and $\bm{W}_0\in\mathbb{R}^{m\times n}$1 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 $\bm{W}_0\in\mathbb{R}^{m\times n}$2, the left orthogonal factor is parameterized as

$\bm{W}_0\in\mathbb{R}^{m\times n}$3

with analogous structure for $\bm{W}_0\in\mathbb{R}^{m\times n}$4. The weight update is

$\bm{W}_0\in\mathbb{R}^{m\times n}$5

The block-diagonal orthogonal matrices are denoted

$\bm{W}_0\in\mathbb{R}^{m\times n}$6

The paper identifies the original POET implementation as weight-centric: it explicitly forms transformed weights through updates of the form

$\bm{W}_0\in\mathbb{R}^{m\times n}$7

and states that this implementation has complexity

$\bm{W}_0\in\mathbb{R}^{m\times n}$8

for input $\bm{W}_0\in\mathbb{R}^{m\times n}$9. POET-X replaces this with an input-centric formulation. Rather than computing $\bm{R}\in\mathbb{R}^{m\times m}$0 by materializing transformed weights, it evaluates

$\bm{R}\in\mathbb{R}^{m\times m}$1

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{R}\in\mathbb{R}^{m\times m}$2

with execution decomposed into

$\bm{R}\in\mathbb{R}^{m\times m}$3

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{R}\in\mathbb{R}^{m\times m}$4 is constructed with

$\bm{R}\in\mathbb{R}^{m\times m}$5

and, for $\bm{R}\in\mathbb{R}^{m\times m}$6,

$\bm{R}\in\mathbb{R}^{m\times m}$7

POET-X stores only the upper-triangular part of $\bm{R}\in\mathbb{R}^{m\times m}$8, reducing trainable degrees of freedom to

$\bm{R}\in\mathbb{R}^{m\times m}$9

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{P}\in\mathbb{R}^{n\times n}$0

The paper’s key observation is that the higher-order terms depend only on $\bm{P}\in\mathbb{R}^{n\times n}$1 and $\bm{P}\in\mathbb{R}^{n\times n}$2, enabling fused Triton kernels with shared-memory reuse. It reports 2–3× speedups for orthogonalization relative to PyTorch-native CNP (Qiu et al., 5 Mar 2026).

Permutation handling is also redesigned. Instead of explicit permutation matrices, POET-X uses index mapping. For row and column permutations, the paper gives

$\bm{P}\in\mathbb{R}^{n\times n}$3

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{P}\in\mathbb{R}^{n\times n}$4 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 (Qiu et al., 5 Mar 2026).

4. Variants, training procedure, and implementation constraints

The simplified forward

$\bm{P}\in\mathbb{R}^{n\times n}$5

induces a specific activation-memory trade-off. In standard autograd, the intermediate $\bm{P}\in\mathbb{R}^{n\times n}$6 from mm3 must be saved to compute $\bm{P}\in\mathbb{R}^{n\times n}$7, while $\bm{P}\in\mathbb{R}^{n\times n}$8 from mm2 need not be saved because $\bm{P}\in\mathbb{R}^{n\times n}$9 has no gradient. POET-X accordingly defines two variants.

Variant	Core idea	Emphasis
POET-X$$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$0	Standard autograd, saves $$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$1	Faster
POET-X$$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$2	Gradient checkpointing, recomputes instead of saving $$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$3	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{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$4, 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 $$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$5 and $$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$6. A practical constraint is that linear dimensions must be divisible by $$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$7; the paper states that FFN intermediate dimensions were slightly modified for this reason (Qiu et al., 5 Mar 2026).

POET-related parameters are updated with AdamW. The learning-rate scaling rule is

$$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$8

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 $$\text{OET}(\bm{W};\bm{R},\bm{P})=\bm{R}\bm{W}\bm{P},$$9 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$$\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}$$0, and 1.89 ms for POET-X$$\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}$$1. The paper further notes that POET-X$$\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}$$2 attains backward-pass latency comparable to a standard linear layer because of higher parameter efficiency (Qiu et al., 5 Mar 2026).

For Llama-3B pretraining on C4 at $$\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}$$3 over 60B tokens, the reported validation perplexity and memory figures are:

Method	Memory (G)	Val PPL
AdamW	81.03	12.69
POET-X$$\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}$$4	60.58	12.76
POET-X$$\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}$$5	68.52	12.05

This comparison is the core quality-efficiency result. POET-X with $$\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}$$6 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$$\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}$$7	61.87 GB
POET-X$$\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}$$8	47.21 GB

For Llama-8B at sequence length 2048, peak memory is 78.69 GB for AdamW, 43.08 GB for POET-X$$\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}$$9, and 31.74 GB for POET-X$\bm{W}_0 = \bm{U}\bm{\Sigma}_0\bm{V}^\top,$0. 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 $\bm{W}_0 = \bm{U}\bm{\Sigma}_0\bm{V}^\top,$1. Reported validation perplexities are 20.49 for quantized 8-bit APOLLO, 17.74 for quantized 8-bit GaLore, 16.21 for POET-XQ$\bm{W}_0 = \bm{U}\bm{\Sigma}_0\bm{V}^\top,$2, and 14.78 for POET-XQ$\bm{W}_0 = \bm{U}\bm{\Sigma}_0\bm{V}^\top,$3. 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 (Qiu et al., 5 Mar 2026).

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 $\bm{W}_0 = \bm{U}\bm{\Sigma}_0\bm{V}^\top,$4, 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 (Qiu et al., 5 Mar 2026).

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.