POET: Reparameterized Orthogonal Training
- The paper demonstrates how reparameterization via orthogonal transformations replaces direct dense matrix optimization, ensuring exact spectrum preservation.
- The framework provides a theoretically sound alternative with controlled spectral norms, leading to more stable training and improved generalization.
- POET incorporates efficient approximations like SPO and CNP to scale orthogonal training for large language models while reducing trainable parameters.
Searching arXiv for POET, POET-X, and related orthogonal-equivalence optimizer work. Reparameterized Orthogonal Equivalence Training (POET) is a reparameterized training framework for LLMs that replaces direct optimization of dense weight matrices with optimization over orthogonal transformations applied to a fixed random matrix. In POET, each weight matrix is represented through an Orthogonal Equivalence Transformation (OET), so training proceeds by learning how to rotate a base matrix on the left and right while preserving its singular values. The method is presented as a spectrum-preserving alternative to direct weight optimization with AdamW, with the stated aims of more stable optimization, stronger control over weight spectrum, and improved generalization, especially in large-scale pretraining (Qiu et al., 9 Jun 2025).
1. Motivation and problem setting
POET is motivated by two practical problems in standard LLM training. First, direct weight optimization is described as unstable and hard to tune at scale: AdamW on huge matrices is expensive, and training often requires careful hyperparameter tuning to avoid poor convergence. Second, even when training loss is minimized, generalization may still be suboptimal, and the method connects this issue to spectral properties of the learned weights (Qiu et al., 9 Jun 2025).
The framework is positioned against spectral regularization methods that often only constrain the top singular value and can be expensive because they rely on spectral norm estimation or power iteration. POET instead aims to provide exact spectral preservation by construction. It is also framed as a generalization of earlier orthogonal-training methods that preserve hyperspherical energy: those methods were reported as effective for finetuning but too restrictive for pretraining, whereas POET preserves the full singular-value spectrum while still allowing useful representation learning (Qiu et al., 9 Jun 2025).
This suggests that POET should be understood less as a conventional optimizer replacement and more as a geometric restriction on the admissible training trajectory. The restriction is not arbitrary: it is intended to keep the spectrum controlled throughout optimization rather than correcting it after unconstrained updates.
2. Orthogonal equivalence transformation and spectrum preservation
The central reparameterization is the Orthogonal Equivalence Transformation. For a weight matrix , POET represents the trainable matrix as
where is a fixed randomly initialized matrix, and and are trainable orthogonal matrices satisfying
The forward pass is
In this parameterization, the left matrix acts on the column space and the right matrix on the row space. Because orthogonal matrices preserve inner products and singular values, POET changes singular vectors while keeping singular values fixed. If
then
which remains an SVD-like decomposition with the same singular values (Qiu et al., 9 Jun 2025).
The core spectral-preservation claim is stated as follows: if 0 is a spectrum-preserving map, then
1
The paper further cites a theorem showing that for linear maps 2 with 3, if the largest singular value is always preserved, then the map must be of the form
4
with 5 and 6 (Qiu et al., 9 Jun 2025).
POET is contrasted with hyperspherical-energy-preserving orthogonal training, where a layer is reparameterized only on the left as 7. In that setting, hyperspherical energy is preserved:
8
POET generalizes this construction by preserving the full singular spectrum rather than only that geometric energy quantity (Qiu et al., 9 Jun 2025).
3. Optimization geometry, initialization, and training procedure
In standard training, the optimizer updates every entry of 9 directly. In POET, the optimizer updates 0 and 1 instead, so gradients act on a reparameterized manifold of matrices that all share the same singular values as the initial matrix. The stated consequence is that singular values are prevented from drifting arbitrarily, which is linked to more stable learning and better generalization (Qiu et al., 9 Jun 2025).
The generalization discussion is tied to bounds that depend on spectral norms. For an 2-layer ReLU MLP with maximal width 3, the cited complexity term is
4
The argument is that if the norms remain bounded, the generalization bound improves; because POET preserves the spectrum of each weight matrix and starts from carefully chosen initializations, it keeps these norms controlled throughout training (Qiu et al., 9 Jun 2025).
Initialization is therefore treated as a first-class component. The method considers standard Gaussian and Xavier initialization, and introduces uniform-spectrum initialization and normalized Gaussian initialization. Uniform-spectrum initialization sets all singular values of an initial matrix to 1 by SVD. Normalized Gaussian initialization normalizes random Gaussian vectors to lie on the unit hypersphere, and the paper proves the asymptotic bounds
5
These results are used to ensure that the preserved spectrum is well behaved at the start of training (Qiu et al., 9 Jun 2025).
The training algorithm is described in five steps. First, initialize the weight matrix with normalized Gaussian and set the orthogonal matrices to identity. Second, initialize the Cayley-Neumann Parameterization (CNP) parameters for 6 and 7. Third, construct the approximate orthogonal matrices using either fully stochastic or block-stochastic Stochastic Primitive Optimization (SPO). Fourth, run an inner loop where gradients update the orthogonal parameters through the transformed weight 8. Fifth, merge then reinitialize:
9
and reset the primitive factors to identity. This merge-then-reinitialize step is stated to be crucial for memory efficiency and for preventing approximation errors from accumulating (Qiu et al., 9 Jun 2025).
4. Scalable approximations: SPO and Cayley-Neumann parameterization
The main practical obstacle is that learning full orthogonal matrices 0 and 1 naively is expensive. POET addresses this with two approximation layers: SPO and CNP (Qiu et al., 9 Jun 2025).
SPO factorizes an orthogonal matrix into a product of primitive orthogonal matrices,
2
In fully stochastic SPO, each primitive selects a random subset of coordinates 3 of size 4 and embeds a small orthogonal block into the full identity:
5
with
6
In the block-stochastic variant,
7
with
8
The block-stochastic version is described as updating all dimensions more evenly (Qiu et al., 9 Jun 2025).
CNP provides an efficient orthogonal parameterization through a skew-symmetric matrix 9. The exact Cayley transform is
0
and POET approximates the inverse using a truncated Neumann series:
1
The paper notes that the approximation is valid when 2, and that periodic reset via reinitialization helps keep that condition satisfied (Qiu et al., 9 Jun 2025).
For a matrix 3, the reported trainable-parameter and memory formulas are:
| Variant | Trainable parameters | Memory |
|---|---|---|
| AdamW | 4 | roughly 5 |
| POET-FS | 6 | 7 |
| POET-BS | 8 | 9 |
The implementation also uses a CUDA kernel that stores only the upper triangular part of the skew-symmetric matrix and reconstructs it on the fly. Together with the Neumann approximation, this yields a reported 3.8× speedup over the naive implementation for a 1.3B LLaMA model, with 1.5× from Neumann approximation and 1.3× from the skew-symmetric CUDA kernel (Qiu et al., 9 Jun 2025).
5. Empirical behavior, scaling results, and training dynamics
The primary experimental setting is LLaMA-style pretraining at 60M, 130M, 350M, and 1.3B parameters on C4, using validation perplexity and comparing against AdamW, GaLore, and LoRA in some comparisons (Qiu et al., 9 Jun 2025).
The main reported result is that POET consistently outperforms AdamW and GaLore while using far fewer trainable parameters. On the 1.3B model, POET-FS with 0 reaches validation perplexity 13.70, compared with 14.73 for AdamW and 18.33 for GaLore. The block-stochastic version with 1 achieves 14.56, still better than AdamW while using about one-tenth the trainable parameters. Under similar parameter budgets, POET also outperforms LoRA; for example, around 30M trainable parameters on the 350M model, POET achieves perplexity 18.05 versus LoRA’s 25.19 (Qiu et al., 9 Jun 2025).
The paper emphasizes that POET’s training dynamics differ from AdamW. It may improve more slowly after an initial drop, but its later-stage stability leads to better final performance. It is also reported that POET can outperform AdamW even when AdamW is trained for substantially more tokens, so the gain is not merely from longer training (Qiu et al., 9 Jun 2025).
A qualitative account is given through vector probing. For a fixed random unit vector 2, the method tracks quantities such as 3 and 4. Training is described as proceeding through three phases: conical shell searching, where the cosine drops from 1 to around 5–6; stable learning on the conical shell, where the cosine remains in that range while validation perplexity continues to improve; and final adjusting, when the learning rate decays and the model makes only small refinements (Qiu et al., 9 Jun 2025).
The ablations identify several sensitive components. Normalized Gaussian works best among the tested initializations, whereas uniform-spectrum initialization performs worse. For merge-then-reinitialize frequency, 7 or 8 works well, with 9 used by default. For the Neumann approximation, more terms improve performance; with 0, training does not converge, while 1 or 2 performs much better. The orthogonality error decreases with more Neumann terms and is periodically reset by reinitialization. Update-coverage analysis further indicates that block-stochastic updates are more even than fully stochastic ones, which likely explains their better parameter efficiency at comparable sizes (Qiu et al., 9 Jun 2025).
Beyond pretraining, POET-FS is reported to outperform both LoRA and OFT on BART-large finetuning for XSum and CNN/DailyMail summarization, and to slightly exceed full fine-tuning on ROUGE metrics. The method is nevertheless described as especially useful in the pretraining regime, where standard orthogonal finetuning can be too restrictive (Qiu et al., 9 Jun 2025).
6. Successors, related formulations, and limitations
Subsequent work clarifies both the scalability issues of the original implementation and the broader methodological family built around orthogonal equivalence transformations. POET-X is introduced as a scalable and memory-efficient variant that keeps the same orthogonal-equivalence principle but replaces explicit large transformed-weight computation with an input-centric implementation, block-stochastic orthogonal factors, custom permutation handling, and fused Triton kernels. The stated motivation is that original POET, despite its stability, incurred high memory consumption and computational overhead due to intensive matrix multiplications; it could become prohibitive beyond roughly 3B parameters and could even use more memory than AdamW because of activation storage (Qiu et al., 5 Mar 2026).
The reported effect of these changes is substantial. In a single-layer profile, full forward+backward time drops from 10.59 ms in original POET to 1.38 ms for 3 and 1.89 ms for 4, and the method is described as exhibiting PEFT-like memory usage at full-model scale. In a 3B Llama pretraining setup, 5 achieves validation PPL 12.05 versus AdamW’s 12.69, and the system is reported to scale to 8B and 13B models on a single H100 while AdamW runs OOM under the same settings (Qiu et al., 5 Mar 2026). This suggests that the primary obstacle to using POET at very large scale was not the orthogonal-equivalence principle itself, but the original realization of that principle on GPU systems.
A related development, Pion, turns the geometric idea behind POET into a direct optimizer. Rather than parameterizing weights as 6, Pion updates the current weight directly by orthogonal left and right transformations,
7
thereby preserving the full singular spectrum while rotating row and column subspaces. The paper explicitly states that Pion is not merely POET itself, nor a direct reparameterization, but a related direct optimization formulation inspired by POET (Shi et al., 12 May 2026).
The limitations of POET remain those stated in the original work. The method depends on the assumption that preserving the spectrum of the initial weights is beneficial; this is not proved universally. Efficiency depends on hyperparameters such as block size 8, the number of Neumann terms 9, and merge frequency 0, all of which require tuning. Very small 1 reduces parameters but slows convergence. The fully stochastic version can lead to uneven weight updates, while the block version is more structured but less flexible in some settings. CNP only approximates orthogonality, so aggressive approximation can destabilize training. Although orthogonal transforms are merged into the weights at inference so runtime cost is removed, training still has overhead compared with plain AdamW because orthogonal structure must be maintained and the model must be periodically reparameterized. The exact mechanism behind POET’s improved training dynamics is also stated to remain an open question (Qiu et al., 9 Jun 2025).
A common misconception is that POET is simply an orthogonal finetuning method or that its sparse and structured parameterization automatically guarantees practical memory savings. The first is contradicted by its use in pretraining and by its preservation of the full singular spectrum rather than only hyperspherical energy; the second is contradicted by the original implementation experience that motivated POET-X. Within the available evidence, POET is best understood as a spectrum-preserving reparameterized training method whose mathematical constraint is clear, whose empirical behavior is favorable in several LLM settings, and whose practical viability at billion-parameter scale depends on the systems-level refinements introduced after the original formulation.