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

Zeju Qiu; Lixin Liu; Adrian Weller; Han Shi; Weiyang Liu

cs.LGMar 5, 2026

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

Authors: Zeju Qiu, Lixin Liu, Adrian Weller, Han Shi, Weiyang Liu

Organizations: The Chinese University of Hong Kong · MPI for Intelligent Systems · University of Cambridge · Huawei Technologies

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 & tables

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\times 64$ weight matrix and run both POET variants for a 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.

Hidden Dim.	PyTorch (ms)	Ours (ms)	Speedup
2048	1.621	0.086	18.75 $\times$
4096	3.269	0.167	19.57 $\times$
8192	6.584	0.393	16.76 $\times$
16384	13.246	0.946	14.00 $\times$

Table 1: Comparison between PyTorch-native and our customized permutation under different hidden dimensions.

Hidden Dim.	4 permute (ms)	2 permute (ms)	Speedup
Block size = 256, Compiled = False
4096	15.858	12.017	1.32 $\times$
8192	38.583	31.028	1.24 $\times$
16384	120.869	94.280	1.28 $\times$
Block size = 256, Compiled = True
4096	8.705	7.496	1.16 $\times$

Table 2: Runtime improvement of permutation reduction.

Seq. Length	PyTorch (ms)	Ours (ms)	Speedup
2048	0.790	0.334	2.37 $\times$
4096	1.525	0.642	2.38 $\times$
8192	2.918	1.266	2.30 $\times$
16384	6.120	2.661	2.30 $\times$

Table 3: Runtime comparison between PyTorch-native block-diagonal matrix construction/multiplication and our batch-wise matrix multiplication strategy under different sequence lengths.

Seq. Length	PyTorch (MB)	Ours (MB)	Reduction
2048	280	192	31.43%
4096	360	272	24.44%
8192	520	432	16.92%
16384	840	752	9.55%

Table 4: GPU memory comparison between PyTorch-native block-diagonal matrix construction/multiplication and our batch-wise matrix multiplication strategy under different sequence lengths.

Figure 2: Illustration of efficient Cayley-Neumann parameterization (batch-wise implementation).

N	PyTorch (ms)	Triton (ms)	Speedup
Block Size = 256
64	0.316	0.107	2.96 $\times$
128	0.610	0.204	2.99 $\times$
192	0.881	0.297	2.97 $\times$
256	1.156	0.388	2.98 $\times$
320	1.404	0.479	2.93 $\times$

Table 5: Runtime comparison between PyTorch-native implementation of the Cayley-Neumann parameterization and our optimized Triton kernel under different a number of blocks in POET-X.

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 ( Qiu et al., 2025a ) 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.

Method	Params (M)	Mem (G)	Val PPL
AdamW	2764.47	81.03	12.69
Muon (Kimi) ( Liu et al., 2025a )	2764.47	70.94	11.45
APOLLO ( Zhu et al., 2025 )	2764.47	80.60	12.97
GaLore ( Zhao et al., 2024 )	2764.47	74.50	14.88
$\text{POET-X}_{b=256}$	366.64	60.58	12.76
$\text{POET-X}_{b=512}$	570.06	68.52	12.05

Table 6: Validation Perplexity (PPL) comparison. The column labeled Params (M) reports the total number of trainable parameters for each optimizer, where

b

is the block size and

\bm{L}_{\max}\!=\!256

Method	Params (M)	Mem (G)	Val PPL
Quantized 8-bit APOLLO	2764.47	66.37	20.49
Quantized 8-bit GaLore	2764.47	66.28	17.74
$\text{POET-XQ}_{b=256}$	366.64	51.66	16.21
$\text{POET-XQ}_{b=512}$	570.06	60.65	14.78

Table 7: Validation Perplexity (PPL) comparison. The column labeled Params (M) reports the total number of trainable parameters for each optimizer, where

b

is the block size and

\bm{L}_{\max}\!=\!256

Method	Sequence Length 512				Sequence Length 1024
	1 $\times$ 1 H100		1 $\times$ 8 H100		1 $\times$ 1 H100		1 $\times$ 8 H100
	8B	13B	8B	13B	8B	13B	8B	13B
8-bit Q-APOLLO	2.52	1.55	18.63	11.47	4.24	2.63	31.85	19.72
8-bit Q-GaLore	2.25	1.32	18.66	11.49	3.86	2.30	31.81	19.70
$\text{POET-XQ}_{b=256}$	2.77	1.82	21.96	14.22	4.12	2.72	32.41	21.67
$\text{POET-XQ}_{b=512}$	2.30	1.54	17.90	11.93	3.56	2.37	28.27	18.75

Table 8: Throughput (k tokens/s) with Llama-8B and Llama-13B.

Method	Llama-3B			Llama-8B			Llama-13B
Method	$\bm{L}_{\max}\!=\!512$	$\bm{L}_{\max}\!=\!1024$	$\bm{L}_{\max}\!=\!2048$	$\bm{L}_{\max}\!=\!512$	$\bm{L}_{\max}\!=\!1024$	$\bm{L}_{\max}\!=\!2048$	$\bm{L}_{\max}\!=\!512$	$\bm{L}_{\max}\!=\!1024$	$\bm{L}_{\max}\!=\!2048$
AdamW	28.74	28.78	31.43	78.89	76.34	78.69	OOM	OOM	OOM
Muon	18.62	19.27	21.82	50.30	53.46	54.98	76.32	77.02	OOM
APOLLO	19.74	20.40	21.32	51.34	53.49	56.94	OOM	OOM	OOM
GaLore	19.06	20.17	21.16	44.52	45.62	54.71	67.15	67.86	73.37
$\text{LoRA}_{r=160}$	13.60	16.79	22.68	27.90	33.63	43.78	42.48	49.78	63.50

Table 9: Peak memory (GB) for different models and sequence lengths (

\bm{L}_{\max}

) on a single H100 (batch size 1, no gradient accumulation).

Figure 6: Throughput (k tokens/s) across different numbers of GPUs. The solid line denotes the actual throughput, and the dashed line denotes the ideal linear scaling throughput. The ideal throughput of

k

GPUs is defined as

T_{k,ideal}=T_{8,real}\times k/8

	Seq. Length 512 (Llama-8B)			Seq. Length 2048 (Llama-8B)			Seq. Length 512 (Llama-13B)			Seq. Length 2048 (Llama-13B)
	1x1 H100	8x8 H100	Ratio	1x1 H100	8x8 H100	Ratio	1x1 H100	8x8 H100	Ratio	1x1 H100	8x8 H100	Ratio
AdamW	7.60	186.16	24.5	OOM	437.21	N/A	OOM	105.98	N/A	OOM	273.02	N/A
$\text{LoRA}_{r=160}$	5.28	310.59	58.8	6.42	396.97	61.8	3.68	217.52	59.1	4.42	OOM	N/A
$\text{LoRA}_{r=320}$	4.03	228.80	56.8	4.82	298.41	61.9	2.81	166.78	59.3	3.26	OOM	N/A
$\text{POET-X}_{\text{mem},b=256}$	3.73	199.95	53.7	5.92	362.75	61.2	2.74	136.14	49.7	4.02	246.47	61.3
$\text{POET-X}_{\text{mem},b=512}$	3.16	166.92	52.9	5.26	315.18	59.9	2.22	112.10	50.5	3.58	228.07	63.7

Table 10: Throughput (k tokens/s) comparison between baselines and POET-X. Ratio: throughput ratio between 1x1 H100 and 8x8 H100.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Llama 3B	Llama 8B	Llama 13B
Hidden dimension	2560	4096	5120
Intermediate dimension	7168	14336	13824
Number of attention heads	32	32	40
Number of hidden layers	32	32	40

Appendix

Table 11: Model architectures for different Llama variants.

Model	Spec.	# GPU	lr (base)	lr (POET)	POET reset gap	training steps	batch size	grad acc.
Llama 3B	$b=256$	8	1e-3	$5\text{e-}4\cdot\gamma$	400	600,000	64	1
Llama 3B	$b=512$	8	1e-3	$5\text{e-}4\cdot\gamma$	400	600,000	64	1

Appendix

Table 12: Hyper-parameter setup of POET-X.

Model	Spec.	# GPU	lr (base)	lr (POET)	POET reset gap	training steps	batch size	grad acc.
Llama 3B	$b=256$	8	1e-3	$5\text{e-}4\cdot\gamma$	400	100,000	64	1
Llama 3B	$b=512$	8	1e-3	$5\text{e-}4\cdot\gamma$	400	100,000	64	1

Appendix

Table 13: Hyper-parameter setup of POET-XQ.

Llama-8B	Sequence Length 512			Sequence Length 1024			Sequence Length 2048
Llama-8B	1x1 H100	8x8 H100	Ratio	1x1 H100	8x8 H100	Ratio	1x1 H100	8x8 H100	Ratio
AdamW	7.60	186.16	24.5	9.34	277.53	29.7	OOM	437.21	N/A
LoRA - 160	5.28	310.59	58.8	6.20	374.87	60.5	6.42	396.97	61.8
LoRA - 320	4.03	228.80	56.8	4.58	273.35	59.7	4.82	298.41	61.9
$\text{POET-X}_{\text{mem},b=256}$	3.73	199.95	53.7	5.29	294.25	55.6	5.92	362.75	61.2
$\text{POET-X}_{\text{mem},b=512}$	3.16	166.92	52.9	4.42	261.32	59.2	5.26	315.18	59.9

Appendix

Table 14: Throughput (k tokens/s) comparison between baselines and POET-X. Ratio: throughput ratio between 1x1 H100 and 8x8 H100.

Figure 7: Throughput (k tokens/s) across different numbers of GPUs. The solid line denotes the actual throughput, and the dashed line denotes the ideal linear scaling throughput. The ideal throughput of

k

GPUs is defined as

T_{k,ideal}=T_{8,real}\times k/8

Architecture	Method	Val PPL ( $\downarrow$ )	Memory (GB) ( $\downarrow$ )
Qwen3 1B	AdamW	14.78	36.09
Qwen3 1B	POET-X	14.67	29.42
Qwen3 MoE 1B	AdamW	16.89	28.41
Qwen3 MoE 1B	POET-X	15.38	27.51

Appendix

Table 15: Validation perplexity and memory usage for Qwen3 1B dense and MoE models trained on 20B tokens. Best results are bolded.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 342 papers

Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching
Oct 3, 2026Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi +3Nyström Method Second-Order Optimization
Mila & DIRO, Université de Montréal · Critical ML, University of Waterloo · Independent Researcher +2
Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation
Feb 27, 2026Zhengbo Wang, Jian Liang, Ran He +2Memory-Efficient Training Low-Rank Compression
University of Science and Technology of China · NLPR & MAIS, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences +1
Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers
May 9, 2026Aditya RanganathDeep Learning Optimization LLM Input Optimization
Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA

Oct 3, 2026J/K move · Enter open · S save

Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching

Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi, Shikai Fang+2

Mila & DIRO, Université de Montréal · Critical ML, University of Waterloo · Independent Researcher+2

Nyström Method Second-Order Optimization Memory-Efficient Training

Open paper AskPDF arXiv

Training large language models (LLMs) entails a fundamental trade-off: memory-efficient optimizers such as Adam discard cross-parameter curvature, whereas full-curvature methods such as SOAP can accelerate convergence at prohibitive memory costs. We introduce Clean, a memory-efficient and full-curvature optimizer designed to resolve this bottleneck. Clean leverages the randomized Nystrom method to accurately approximate the left and right preconditioners in SOAP, and to reduce the optimizer's memory complexity from quadratic to linear in terms of model dimensions. We subsequently reintegrate the off-subspace components to capture curvature information beyond the low-rank approximation, preserving rich curvature at minimal memory cost. We further propose Q-Clean, a low-precision variant that aggressively compresses optimizer states. Q-Clean reduces optimizer memory consumption by over 50% compared to Muon when pre-training a LLaMA-1.3B architecture, all while maintaining strong and competitive predictive performance. Notably, Clean operates with a smaller optimizer-state footprint than standard AdamW while reaching AdamW's final performance 26% faster in wall-clock time. Furthermore, our methods uniquely enable the pre-training of a 13B-parameter model on a single 80GB GPU, providing a scalable, efficient, and accessible approach to large-scale model optimization.

Abstract

LLM Training

Explore similar work

Clean: Second-order LLM Training at Linear Memory Cost via Nystr{ö}m Sketching

Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers