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

SAI replication review · Referee report

Summary

POET-X is a systems-level engineering follow-up to POET (Qiu et al., 2025a), the spectrum-preserving reparameterization that trains each weight matrix $W$ via an orthogonal equivalence transformation $W \leftarrow R W P$ . POET's stability came from keeping the spectrum of $W_0$ intact and parameterizing $R,P$ via block-diagonal Cayley–Neumann orthogonalization, but its original weight-centric implementation materialized transformed weights and full block-diagonal matrices, making it memory- and compute-heavy relative to AdamW. The conceptual move here is to keep POET's objective and analytical guarantees unchanged while re-plumbing every step of the forward and backward passes so that the algorithm's inherent parameter sparsity finally shows up as memory sparsity: an input-centric matrix-free reformulation of $R W P x$ , a compact upper-triangular skew-symmetric representation, batch-wise block-diagonal matmul that never builds the sparse matrix, permutation reduction that absorbs two of the four permutations into $W$ in advance, and a fused Triton kernel for the Cayley–Neumann series. Layered on top are two variants — POET-Xfast and a checkpointed POET-Xmem — and a quantized POET-XQ that becomes possible because POET-Xmem never keeps a high-precision weight resident. The empirical story is that POET-X reaches LoRA-level memory (Table 9, Figure 4), keeps single-layer overhead close to nn.Linear (Figure 3), and lets a single H100 pretrain 8B–13B Llama models where AdamW OOMs; direct execution of the released code confirms the memory hierarchy at 13B (POET-Xmem fits at 51.05 / 57.78 GB vs AdamW OOM) and the single-layer POET-X latencies (1.29 / 1.87 ms vs paper 1.38 / 1.89 ms). The main conceptual limitation is that the paper reports a systems win but a mixed algorithmic story: on Table 6 POET-X_b=256 is worse than AdamW on validation PPL and Muon Pareto-dominates POET-X_b=512 on both PPL and memory, so the frequently repeated "consistently better-than-AdamW" phrasing is stronger than the numbers support. Combined with the absence of seed variance, an end-to-end ablation of the stacked optimizations, confounded wall-clock comparisons (POET-X on DDP vs AdamW on FSDP), and an 8x-runtime headline whose original-POET denominator is missing from the release (measured ~3.1x), the claims outrun the evidence in places even though the underlying engineering is impressive.

Strengths

Real engineering contribution. The five-part re-plumbing (input-centric matmul chain, compact skew-symmetric storage, batch-parallel block-diagonal matmul, permutation merging, Triton-fused CNP with reused $Q, Q^2$ ) is well-motivated at each step, and each step is measured in isolation in Tables 1–5. The Section 3.5 activation analysis and the fast/mem split are exactly the right things to include.
Substantive memory gains at scale, independently verified. Table 9's peak-memory numbers — POET-Xmem_b=256 fitting 13B at $L_{\max}=2048$ in 47.21 GB versus AdamW and original POET OOM at the same setting — are a genuine capability jump for single-GPU pretraining, and the executed peak-memory sweep confirms the hierarchy (POET-Xmem is the lowest-memory method at every scale, POET fits only at 3B, AdamW OOMs at 13B). Single-layer POET-X forward+backward latency also lands within a few percent of the paper.
Meaningful quantized-training path. Making POET-XQ implementable specifically through POET-Xmem (recompute enables never storing dequantized $W$ ) is a clean design decision, and Tables 7–8 show POET-XQ competitive with or better than 8-bit Q-GaLore and Q-APOLLO on both PPL and throughput at 8B/13B.
Preserved conceptual footprint. POET-X does not alter POET's spectrum-preserving objective, so the analytical guarantees of the original method (spectrum retention, small hyperspherical energy) carry over unchanged — a cleaner story than most efficient-training papers that trade off algorithmic behavior for memory.

Weaknesses

"Consistently better than AdamW" is not what Table 6 shows. POET-X_b=256 reaches 12.76 val PPL vs AdamW's 12.69. Only POET-X_b=512 (12.05) actually beats AdamW. The paper elevates "better-than-AdamW" to a bolded contribution and repeats a similar phrase in the abstract, but one of the two shipped configurations underperforms.
Muon Pareto-dominates POET-X on Table 6. Muon achieves 11.45 PPL at 70.94 GB; POET-X_b=512 reaches 12.05 at 68.52 GB. The 2.4 GB memory delta is described as "much lower GPU memory" but the PPL gap of 0.60 is not marginal on this scale. The framing that POET-X is "only" slightly below Muon underplays that the direct competitor beats POET-X on the primary metric with only modestly more memory.
Wall-clock comparison in Figure 5 is confounded. POET-X runs DDP while AdamW is forced to FSDP because of memory, then a wall-clock win is presented as an efficiency finding. This is really a memory comparison. Either report the wall-clock at a scale where AdamW+DDP fits, or restate the finding as "the parallelism regime each method is actually forced into."
No ablation of the stacked optimizations. Five distinct optimizations combine to reach the headline 3x/8x. Tables 1–5 measure components individually, but the paper never shows a drop-one ablation on end-to-end training memory and runtime — so it is unclear which pieces are load-bearing at the target scale.
No seed variance or statistical error. All pretraining numbers are point estimates. Given PPL gaps of 0.07 (POET-X_b=256 vs AdamW) and 0.64 (POET-X_b=512 vs AdamW) on ~13-PPL baselines, and 0.11 on Qwen3 1B, the ordering could plausibly flip across seeds and there is no way for the reader to tell.
8B/13B numbers are single-run at a single hyperparameter setting. Appendix Tables 12–13 only document 3B hyperparameters. Larger models get no learning-rate sweep and no discussion of how sensitive perplexity is to the POET learning-rate scaling $\gamma = 0.5$ or the "POET reset gap = 400" at scale.
Stability claim inherited but not remeasured for POET-X. The abstract asserts POET-X "maintains the generalization and stability benefits of POET", but POET-X changes the storage layout for $Q$ , the CNP kernel, the block-matmul path and the permutation ordering — all things that could plausibly perturb loss dynamics. No direct stability check (loss variance, gradient-norm trajectory, divergence-free rate) is provided for POET-X itself.
Different token budgets for POET-X (60B) and POET-XQ (10B). Tables 6 and 7 are adjacent and reported in the same PPL units, but Table 6 uses 60B tokens and Table 7 uses 10B tokens on the same 3B model. The Table 7 numbers look substantially worse largely because those models are undertrained; the marginal cost of quantization cannot be read off this comparison.
Cayley approximation is mathematically incorrect as written. " $R \approx 2Q$ for small $Q$ " drops the identity: the correct first-order expansion is $R \approx I + 2Q$ , and $2Q \to 0$ is not orthogonal. The justification for $\gamma = 0.5$ is fine on the correct expansion, but the stated formula would mislead a careful implementer.
Notation gaps in the input-centric derivation. In $z = G_P^\top W G_R^\top x$ , the shapes of $G_R, G_P$ and the implicit transposes that reconcile them with the earlier $R\in\mathbb{R}^{m\times m}, P\in\mathbb{R}^{n\times n}$ convention are not stated; the bijection domain is written $\pi : I \to I$ with a single $I$ despite two distinct index sets $I_m, I_n$ ; the permutation subscripts $\pi_p, \pi_q$ cross-wire against dimensions $m, n$ ; the two permutations "merged into $W$ " are never named; and the POET-Xmem paragraph never says which tensor is recomputed.
Super-linear scaling ratios in Table 10 are unexplained. Several rows report 1-to-64-GPU ratios of 63.7 and 64.3 — at or above the 64x ceiling implied by the metric. Likely a 1-GPU-baseline overhead effect, but the paper frames the analysis against "ideal linear scaling" and never flags it. Figure 6 also uses an $T_8$ -based ideal while Table 10 uses a $T_1$ -based ideal, in the same paragraph, without noting the difference.
Fixed-offset memory saving from batch-parallel matmul is oversold. Table 4 shows the batch-parallel saving is a constant ~88 MB (relative benefit falls from 31.4% at seq 2048 to 9.6% at seq 16384). Fine as a micro-optimization, but a paper about scaling should note that this component is not itself a scaling contribution.
Chinchilla-optimality for the MoE run. Applying the 20 tokens/param dense-model rule to the 1B Qwen3 MoE by budgeting 20B tokens conflates total and active parameters. The claim of "Chinchilla-optimal token budgets" is looser than stated.
"POET-XQ_b=256 requires the overall least memory" is contradicted by one Table 9 cell. At 3B / $L_{\max}$ =2048, POET-XQ_b=256 = 16.15 GB while POET-Xmem_b=256 = 15.47 GB (and Q-GaLoRE 15.78 is bolded as the minimum). Either restate as "lowest in most settings" or verify the 15.47 figure.
Design justification for block-stochastic vs fully-stochastic rests on a 64x64 toy. Figure 1's coverage argument is used to motivate building POET-X on block-stochastic POET, but its only support is a 100-step run on a tiny matrix — nothing quantifies whether the coverage gap survives at LLM scale or translates into a PPL gap at all.
Minor wording/notation issues. Figure 1's caption overloads the symbol $b$ (integer block size vs $b = 1/8$ as sampling fraction); Section 4.2 slips "POET" where "POET-X" is intended; and "fully-stocahstic" is a typo.

Reproducibility & code

8x runtime headline does not reproduce; the original-POET denominator is missing from the release. Executing the shipped code with the paper's dim=4096 / seq=2048 / block=256 configuration yielded a single-layer forward+backward speedup of ~3.1x (POET=4.01 ms vs POET-Xfast=1.29 ms), well below the paper's 8x (10.59→1.38 ms). The measured POET-Xfast (1.29 ms) and POET-Xmem (1.87 ms) latencies match the paper's 1.38 / 1.89 ms within ~6%, and the plain nn.Linear baseline lands at 1.12 ms vs 1.11 ms. The gap is on the POET side: the repo ships no fused, weight-centric original-POET path, so bench_single_layer.py compares against an un-fused PyTorch chain rather than the reference POET that produces 10.59 ms.
Headline memory story reproduces. The peak-memory sweep (63 cells, batch 1, bf16, on an 80 GB GPU) reproduces the Table 9 hierarchy: AdamW OOMs at 13B, original POET OOMs at 8B and 13B, POET-Xmem is the lowest-memory method at every scale, and POET-Xmem_b=256/b=512 fit 13B / $L_{\max}$ =2048 at 51.05 / 57.78 GB (paper 47.21 / 59.02). Most cells fall within ~15%.
Custom CUDA permutation kernel is absent from the release. poet_torch/core/ops.py's PermutationFunction is x[..., perm] — exactly the PyTorch-native baseline Table 1 claims to beat by up to 20x. No hand-written CUDA or Triton permutation kernel exists in the repo, so Table 1's "Ours" column cannot come from this code path.
Cayley Triton kernel (Table 5) does not reproduce cleanly. The shipped fused torch.ops.poet.cayley vs the shipped _cayley_transform_pytorch gave 0.58–2.38x speedups at block=256 and 1.05–1.54x at block=512 on an A100, versus the paper's tight 2.93–2.99x / 1.87–1.98x band. Part of this is A100-vs-H100 hardware, but the exact hidden dim and measurement config the paper used are not shipped and cannot be pinned down.
Batch-parallel block-diagonal tables (3 and 4) are numerically correct but not reproducible in shape. The two paths agree bitwise (max abs diff 0.0), but the flat ~2.3x speedup band and the 31.4%→9.6% memory-reduction curve did not reproduce — measurements were launch/overhead-bound at hidden=4096 and only qualitatively decreasing at hidden=16384. The paper does not specify the hidden dim or allocator/measurement protocol.
No benchmark harness for Tables 1–5 or Figures 3–4. These are the paper's most concrete numerical claims, but no benchmark scripts are shipped; Table 2's 4-permute-vs-2-permute comparison in particular requires exposing a 4-permute path that appears baked into merge_then_reinitialize.
Absolute perplexity (Table 6) not verifiable from the release. The training entrypoint runs end-to-end, but the real allenai/c4 dataset and t5-base tokenizer are not bundled. On a small offline substitute, POET-X_b=512 reached PPL 518.81 vs AdamW 704.98 — the POET-X < AdamW direction reproduces and the b=512 trainable-parameter count matches exactly (570.06M), but the absolute PPL numbers cannot be checked.
No 8B/13B pretraining scripts. examples/scripts/ covers only 3B (POET-X and Q-POET). Tables 7–10, 14 and Figures 3–6 are 8B/13B runs, and Appendix Tables 12–13 document only 3B hyperparameters, so there is no learning rate, reset gap, or batch schedule for the larger runs.
Baselines are not integrated. torchrun_main.py only dispatches adamw / muon / poet / q_poet. LoRA, GaLore, APOLLO, 8-bit Q-GaLore and 8-bit Q-APOLLO all appear as rows in Tables 6, 7, 9 and 10 but none have plumbing in the shipped harness or a README pointer to matched external configs.
Multi-node throughput / memory-profiling scripts are missing. The 100-warmup / 20-measure throughput protocol and the peak-memory sweep across 3 model sizes × 3 sequence lengths (Tables 8, 9, 10, 14, Figures 6–7) have no dedicated logger or orchestration script.
Qwen3 architectures (Table 15) are not in the code. torchrun_main.py hardwires LlamaForCausalLM; there is no Qwen3 config, no MoE routing and no 1B non-Llama config. Table 15's architectural-generalization claim is not runnable from the release as shipped.
10B-token budget in Table 7 is inconsistent with the shipped Q-POET script. pretrain_qpoet_3b.sh and Appendix Table 13 correspond to ~1.6B tokens (100k steps × batch 64 × seq 256), not the 10B tokens quoted for Table 7.
torch.compile regime for baselines is not documented, and a compile-related robustness gap surfaced. POETLinear.forward is decorated @torch.compile(fullgraph=True), which hard-fails when sweeping many shapes in one long-lived process — the peak-memory sweep had to run eager. The paper does not state whether AdamW/LoRA/GaLore/APOLLO ran with torch.compile at the same nesting level as POET-X.
Partial coverage that does exist. The Triton-fused Cayley kernel (cayley_forward_kernel, cayley_backward_kernel_grad_Q) matches Eq. 8/9 with $k=3$ ; POET-Xfast and POET-Xmem are both wired through forward_core with a mem_efficient_mode toggle; the compact skew-symmetric storage (n_elements = b*(b-1)/2) is implemented; and chain_layer_checkpoint_mem_o2 implements the checkpointed variant. So the paper's algorithmic pipeline is broadly present in code — what is missing is the measurement machinery, the baselines, the original-POET reference, and the 8B/13B / Qwen3 configs.

Recommended Changes

Essential

Qualify the "consistently better-than-AdamW" claim. Restrict the bolded contribution and the abstract phrasing to POET-X_b=512, or otherwise state explicitly that b=256 trails AdamW on Table 6. (Weakness: "Consistently better than AdamW" is not what Table 6 shows.)
Re-frame the Muon comparison. State the numeric PPL gap (11.45 vs 12.05) and the small memory delta (2.4 GB) explicitly, and discuss whether POET-X's memory advantage grows enough at 8B/13B to flip the tradeoff — otherwise Muon Pareto-dominates on the Table 6 primary metrics. (Weakness: Muon Pareto-dominates POET-X on Table 6.)
Report seed variance for Table 6. Re-run AdamW and POET-X_b=512 for 2–3 seeds and report standard error so the reader can judge whether the 0.64-PPL ordering is stable. Extend to Table 15 if possible, since the Qwen3 1B gap is 0.11 PPL. (Weakness: No seed variance.)
Add a drop-one ablation of the stacked optimizations. Report end-to-end memory and runtime with each of {input-centric, compact skew, batch-parallel, permutation merge, Triton CNP, checkpointing} disabled in turn on Llama-3B, to substantiate that all five are load-bearing for the 3x/8x headline. (Weakness: No ablation of the stacked optimizations.)
Reconcile the wall-clock story in Figure 5. Add a scale where AdamW+DDP fits (e.g. 3B on 32 H100s) and re-run the comparison, or explicitly relabel the finding as a memory-driven parallelism outcome, not an algorithmic wall-clock win. (Weakness: Wall-clock comparison in Figure 5 is confounded.)
Fix the Cayley formula. Change " $R \approx 2Q$ " to " $R - I \approx 2Q$ " (or " $R \approx I + 2Q$ ") and re-derive the justification for $\gamma = 0.5$ from the correct expansion. (Weakness: Cayley approximation is mathematically incorrect as written.)
State shapes and transposes in the simplified forward pass, and clean up the permutation notation. Write out $\dim(G_R), \dim(G_P), \dim(x)$ in $z = G_P^\top W G_R^\top x$ ; use $\pi_p : I_m \to I_m$ / $\pi_q : I_n \to I_n$ rather than $\pi : I \to I$ ; align permutation subscripts with dimensions (or reuse the earlier $\Phi$ notation); name the two permutations merged into $W$ and the tensor $b = W G_R^\top x$ recomputed by POET-Xmem. (Weakness: Notation gaps in the input-centric derivation.)
Add a direct stability measurement for POET-X. Provide loss and gradient-norm trajectories at a matched scale for POET-X vs POET (and vs AdamW), so the abstract's inherited-stability claim is supported by POET-X's own data rather than by inheritance from POET. (Weakness: Stability claim inherited but not remeasured.)
Ship the original POET reference implementation (or the raw measurement scripts). The 3x/8x ratio, Figures 3–4, and Table 1's PyTorch baseline all need a comparable POET implementation in-repo (or cached measurement outputs). On execution, the measured runtime speedup is ~3.1x rather than 8x precisely because the shipped denominator is an un-fused PyTorch chain, not the fused POET that produces 10.59 ms. (Reproducibility: 8x runtime does not reproduce; Original POET absent from release.)
Ship the custom permutation kernel, or correct Table 1. Either release the CUDA/Triton kernel that produces the 15–20x speedup or clarify that the shipped x[..., perm] implementation is the one being benchmarked and update Table 1 accordingly. (Reproducibility: Custom CUDA permutation kernel is absent.)
Release the 8B / 13B pretraining scripts and hyperparameters. Add invocations, learning-rate schedules, POET reset gaps, and batch/accumulation settings for the 8B and 13B runs behind Tables 7–10 and Figures 3–6; extend Appendix Tables 12–13 accordingly. (Reproducibility: No 8B/13B pretraining scripts.)
Ship the benchmark harnesses used for Tables 1–5, 8–10 and Figures 3–7 with the exact hidden dim / measurement config. Include the 100-warmup / 20-step throughput protocol, the peak-memory sweep script, and the CNP / permutation / batch-parallel micro-benchmarks. (Reproducibility: Cayley kernel doesn't reproduce, Batch-parallel tables not reproducible in shape, No benchmark harness, Multi-node scripts missing.)

Suggested

Match the token budget between Table 6 and Table 7. Add either a 10B-token POET-X row or a 60B-token POET-XQ row so the marginal effect of quantization can be isolated from undertraining. (Weakness: Different token budgets for POET-X and POET-XQ.)
Reconcile the shipped Q-POET script with Table 7's 10B-token count. Update pretrain_qpoet_3b.sh (or Appendix Table 13) so the step count matches 10B tokens. (Reproducibility: 10B-token budget inconsistent with shipped Q-POET script; POET-XQ path not exercisable.)
Integrate (or externally document) the baseline optimizers. Add LoRA, GaLore, APOLLO, 8-bit Q-GaLore and 8-bit Q-APOLLO plumbing to torchrun_main.py or link to matched external configs with pinned commits. Document torch.compile regime per baseline. (Reproducibility: Baselines not integrated; torch.compile regime not documented.)
Ship Qwen3 dense + MoE integration for Table 15. Add the model wrapper, MoE routing and configs used for the Qwen3 experiments, or restate Table 15 as preliminary results whose artifacts will follow. (Reproducibility: Qwen3 not in the code.)
Document the C4 dataset and tokenizer setup. Pin a small reproducible C4 subset and tokenizer configuration that yields the paper's PPL numbers (or provide training logs / checkpoints), so Table 6's absolute perplexities can be audited. (Reproducibility: Absolute perplexity not verifiable from release.)
Explain the super-linear ratios in Table 10 and reconcile the two "ideal scaling" baselines. Add a footnote attributing 63.7 / 64.3 to 1-GPU-baseline overhead, and state explicitly that Figure 6 uses a $T_8$ -based ideal while Table 10 uses a $T_1$ -based ideal. (Weakness: Super-linear scaling ratios; two ideal baselines coexist.)
Clarify the fixed-offset nature of the batch-parallel memory saving. Note that Table 4's savings do not scale with sequence length, and quantify the fraction of the 3x memory reduction this component contributes. (Weakness: Fixed-offset memory saving oversold.)
Soften the "Chinchilla-optimal" MoE framing and disambiguate "1B". Distinguish total vs active parameters for the 1B Qwen3 MoE. (Weakness: Chinchilla-optimality for the MoE run.)
Add an LLM-scale check for block-stochastic vs fully-stochastic. Show coverage/PPL curves at a real layer size so Figure 1's toy no longer bears the whole justification for the block-stochastic design. (Weakness: Design justification rests on a 64x64 toy.)
Fix small wording issues. Distinguish $b$ (block size) from the fully-stochastic sampling fraction in the Figure 1 caption; change "POET" to "POET-X" in Section 4.2's baseline-comparison sentence; correct "fully-stocahstic" to "fully-stochastic"; restate "POET-XQ_b=256 requires the overall least memory" to reflect the Llama-3B / $L_{\max}$ =2048 exception. (Weaknesses: Minor wording/notation issues; POET-XQ overall least memory contradicted.)