Orthogonal Equivalence Transformation in LLMs

• Orthogonal Equivalence Transformation is a method that reparameterizes neural network weights by representing them as fixed random matrices bracketed by learnable orthogonal rotations, preserving singular value spectra.
• It enhances training stability and generalization by restricting learning to rotation adjustments, thereby maintaining desirable spectral properties and reducing gradient noise.
• POET leverages scalable techniques like SPO and Cayley-Neumann Parameterization to implement OET efficiently, demonstrating improved performance and lower computational overhead in LLM training.

Orthogonal Equivalence Transformation is a principled method for constraining the parameterization of neural network weight matrices—particularly in LLMs—to enhance stability, generalization, and efficiency during training. In the POET (reParameterization via Orthogonal Equivalence Transformation) framework, every neuron or layer’s weight matrix is reparameterized as a fixed random matrix bracketed by two learnable orthogonal matrices. This approach rigorously preserves singular value spectra, providing both strong theoretical guarantees and practical performance gains in large-scale training contexts.

1. Reparameterization with Orthogonal Equivalence Transformation

Orthogonal Equivalence Transformation (OET) refers to representing the trainable matrix $\bm{W}$ of a neural network layer as

$\bm{W} = \bm{R} \, \bm{W}_0 \, \bm{P},$

where:

• $\bm{W}_0$ is a fixed matrix, initialized randomly (e.g., via a normalized Gaussian distribution);
• $\bm{R}$ and $\bm{P}$ are trainable orthogonal matrices, satisfying $\bm{R}^\top \bm{R} = \bm{I}$ and $\bm{P}^\top \bm{P} = \bm{I}$.

In this formulation, only $\bm{R}$ and $\bm{P}$ are subject to optimization; $\bm{W}_0$ remains constant.

The transformation’s defining property is spectrum preservation: for any orthogonal $U$, $V$, if $\bm{W}_0 = U \Sigma_0 V^\top$ (SVD), then $\bm{W} = (\bm{R} U) \Sigma_0 (V^\top \bm{P})$ is another SVD, and the singular values $\Sigma_0$ do not change. This restricts the solution space to matrices with fixed singular values (isometric up to input and output rotations), while allowing optimization of the input/output directions (singular vectors).

2. Algorithmic Implementation in POET

POET (Parameterization via Orthogonal Equivalence Transformation) operationalizes OET in LLM training as follows:

Initialization:

• The base matrix $\bm{W}_0$ is randomly sampled and fixed.
• $\bm{R}$ and $\bm{P}$ are initialized to the identity.

Optimization:

• Direct gradient-based optimization of full orthogonal matrices is computationally prohibitive at scale. POET introduces two scalable parameterizations:
  • Stochastic Primitive Optimization (SPO): Decompose $\bm{R}$ and $\bm{P}$ into products of smaller orthogonal “primitives”, such as block-diagonal or low-dimensional (e.g., $b \times b$) orthogonal submatrices. These primitives are stochastically selected during optimization, maintaining global orthogonality while distributing parameter updates.

  • Cayley-Neumann Parameterization (CNP): Represent $\bm{R}$ (or $\bm{P}$) as the Cayley transform of a skew-symmetric matrix $\bm{Q}$, approximated by a truncated Neumann series:

    $$\bm{R} \approx (\bm{I} + \bm{Q})\left(\bm{I} - \bm{Q}\right)^{-1} \approx (\bm{I} + \bm{Q})\left(\bm{I} + \sum_{i=1}^k \bm{Q}^i\right),$$

  where $\bm{Q}^\top = -\bm{Q}$. This allows fast, memory-efficient updates of orthogonal parameters.

• Memory and Computation Reduction: After each training phase (e.g., after $J$ iterations), the current $\bm{R}$ and $\bm{P}$ matrices are merged into a single weight matrix to free GPU memory ("merge-then-reinitialize"), and fresh orthogonal blocks are started from identity.

Inference: Post-training, the composed weight $\bm{W}$ may be used directly, incurring no additional runtime overhead compared to conventionally trained models.

3. Theoretical Properties and Guarantees

Spectrum Preservation

The OET ensures: $$\sigma_i(\bm{W}) = \sigma_i(\bm{W}_0), \quad \forall i,$$ where $\sigma_i(\cdot)$ denotes the $i$th singular value. Consequently,

• Frobenius and spectral norms are preserved,
• The condition number is fixed to that of $\bm{W}_0$, avoiding ill-conditioning.

This spectral constraint links directly to model generalization: the class of transformations allowed by OET is provably no more complex than the initialization, making the learning process less prone to overfitting and more robust to overparametrization.

Inductive Bias and Optimization Landscape

By fixing the spectrum, OET introduces an inductive bias favoring solutions that are efficiently representable through rotations of initial directions (singular vectors). This allows exploration of a large, but regularized, portion of weight space, potentially avoiding pathological configurations with excessively large or small singular values, which are associated with gradient instability.

This suggests that OET creates a more uniform and stable optimization landscape compared to unconstrained training.

4. Practical Benefits and Experimental Observations

Tests on LLaMA models across scales (60M, 130M, 350M, 1.3B parameters) demonstrate:

• Improved generalization: POET achieves lower validation perplexity than AdamW and GaLore, even when the number of trainable parameters is reduced (e.g., using SPO block fraction $b=1/2$, POET achieves better perplexity than AdamW and GaLore using $1/3$ the parameters).
• Training stability: Training curves with POET are smoother, with reduced gradient noise and better convergence.
• Spectral stability: Plots of singular values during training confirm strict preservation, in contrast with AdamW where the spectrum drifts.
• Parameter efficiency: Through low-dimensional primitives and effective block strategies, POET reduces computation and memory, facilitating scalable training in large models.

Experiments also suggest that too restrictive an initialization (e.g., all singular values fixed to unity) degrades expressivity and generalization, indicating the importance of a well-chosen spectrum for $\bm{W}_0$.

5. Implementation Considerations and Limitations

Scalability and Efficiency

• SPO and CNP allow retention of the benefits of OET without prohibitive costs in large-scale contexts. Efficient CUDA implementations yield up to 3.8x speed-up over naive approaches.
• Frequent merging of orthogonal blocks, careful block size selection, and Neumann truncation depth are tunable hyperparameters impacting performance and efficiency.

Expressive Power and Initialization Sensitivity

• The expressive power of OET is bounded by the spectral properties of $\bm{W}_0$. OET cannot increase or decrease the spectrum; it only explores the manifold generated by orthogonal rotations.
• Initialization must be carefully considered, as it determines the upper bound of expressivity; random Gaussian matrices with diverse spectra perform best.

Potential Challenges

• Excessive restriction (uniform spectrum) can harm performance.
• Lower block size in SPO can slow convergence or limit expressivity.
• CNP's approximation quality depends on the depth of the Neumann series; too few terms can break orthogonality with subtle implications for training stability.

6. Broader Applications and Implications

The OET framework extends naturally to various parameter-efficient and robust neural training settings:

• Large-scale LLM pretraining: POET provides a scalable alternative to standard first-order optimizers, maintaining generalization while reducing parameter overhead.
• Fine-tuning and transfer learning: Demonstrated improvements over LoRA and OFT for adaptation tasks, with better perplexity and faster training dynamics.
• Continual and federated learning: The reliably bounded parameter space is attractive for distributed and federated environments.
• Manifold-based optimization methods: OET directly connects to approaches on the orthogonal (Stiefel) manifold, opening avenues for more advanced geometry-aware optimizations.

It is plausible that further theoretical refinement may enable even more expressive and efficient forms of OET, especially if the structure of $\bm{W}_0$ is adaptively learned or coordinated across layers or tasks.

7. Summary Table: POET and OET Properties

Property	OET/POET	Significance
Parameterization	$\bm{W} = \bm{R} \bm{W}_0 \bm{P}$	Separates spectrum (fixed) and directions (learned)
Orthogonality	$\bm{R}^\top\bm{R} = \bm{P}^\top\bm{P} = \bm{I}$	Preserves isometry, enables rotation-only updates
Spectrum Preservation	$\sigma_i(\bm{W}) = \sigma_i(\bm{W}_0)$	Training is confined to a fixed-spectral manifold
Scalability	SPO, CNP for efficient representation	Enables application in LLMs, large modules
Generalization	Empirically improved	Due to regularization by fixed spectrum
Expressivity	Determined by $\bm{W}_0$	Initialization must be rich, not uniform
Inference Cost	None (single $\bm{W}$ used at test time)	No loss of runtime efficiency

• POET: "Reparameterized LLM Training via Orthogonal Equivalence Transformation" (2506.08001)
• Connections to orthogonal optimization: "Parameterization of orthogonal matrices" (see references in the POET paper)

Orthogonal Equivalence Transformation, as instantiated in POET, presents a theoretically justified and practically effective approach for robust, scalable, and efficient LLM training. By preserving spectral properties and restricting learning to rotations in weight space, it establishes new directions for structured neural optimization and foundation model development.