Title: Reparameterized LLM Training via Orthogonal Equivalence Transformation
Venue: NeurIPS 2025
Accepted: yes
Keywords: large language models

pretraining

orthogonality
Abstract: While large language models (LLMs) are driving the rapid advancement of artificial intelligence, effectively and reliably training these large models remains one of the field's most significant challenges. To address this challenge, we propose POET, a novel reParameterized training algorithm that uses Orthogonal Equivalence Transformation to optimize neurons. Specifically, POET reparameterizes each neuron with two learnable orthogonal matrices and a fixed random weight matrix. Because of its provable preservation of spectral properties of weight matrices, POET can stably optimize the objective function with improved generalization. We further develop efficient approximations that make POET flexible and scalable for training large-scale neural networks. Extensive experiments validate the effectiveness and scalability of POET in training LLMs.

Reparameterized LLM Training via Orthogonal Equivalence Transformation

Abstract

While Large language models (LLMs) are driving the rapid advancement of artificial intelligence, effectively and reliably training these large models remains one of the field's most significant challenges. To address this challenge, we propose POET, a novel reParameterized training algorithm that uses Orthogonal Equivalence Transformation to optimize neurons. Specifically, POET reparameterizes each neuron with two learnable orthogonal matrices and a fixed random weight matrix. Because of its provable preservation of spectral properties of weight matrices, POET can stably optimize the objective function with improved generalization. We further develop efficient approximations that make POET flexible and scalable for training large-scale neural networks. Extensive experiments validate the effectiveness and scalability of POET in training LLMs.

1 Introduction

Recent years have witnessed the increasing popularity of large language models (LLMs) in various applications, such as mathematical reasoning [13] and program synthesis [3] and decision-making [77]. Current LLMs are typically pre-trained using enormous computational resources on massive datasets containing trillions of tokens, with each training run that can take months to complete. Given such a huge training cost, how to effectively and reliably train them poses significant challenges.

The de facto way for training LLMs is to directly optimize weight matrices with the Adam optimizer [37, 55]. While conceptually simple, this direct optimization can be computationally intensive (due to the poor scaling with model size) and requires careful hyperparameter tuning to ensure stable convergence. More importantly, its generalization can remain suboptimal even if the training loss is perfectly minimized [36]. To stabilize training and enhance generalization, various weight regularization methods [4, 10, 12, 47, 49, 79] and weight normalization techniques [28, 38, 39, 50, 52, 54] have been proposed. Most of these methods boil down to improving spectral properties of weight matrices ( i.e. , singular values) either explicitly or implicitly. Intuitively, the spectral norm of a weight matrix ( i.e. , the largest singular value) provides an upper bound on how much a matrix can amplify the input vectors, which connects to the generalization properties. In general, smaller spectral norms ( i.e. , better smoothness) are considered to be associated with stronger generalization, which inspires explicit spectrum control [33, 59, 67, 79]. Theoretical results [6] also suggest that weight matrices with bounded spectrum can provably guarantee generalization. Given the importance of the spectral properties of weight matrices, what prevents us from controlling them during LLM training?

Inefficacy of spectrum control : Existing spectrum control methods constrain only the largest singular value, failing to effectively regularizing the full singular value spectrum. Moreover, there is also no guarantee for spectral norm regularization to effectively control the largest singular value.

Computational overhead : Both spectral norm regularization [79] and spectral normalization [59] require computing the largest singular value of weight matrices. Even with power iteration, this still adds a significant overhead to the training process, especially when training large neural networks. Additionally, spectral regularization does not scale efficiently with increasing model size.

Project page:

spherelab.ai/poet

To achieve effective weight spectrum control without the limitations above, we propose POET, a reParameterized training algorithm that uses Orthogonal Equivalence Transformation to indirectly learn weight matrices. Specifically, POET reparameterizes a weight matrix W ∈ R m × n with RW 0 P where W 0 ∈ R m × n is a randomly initialized weight matrix, R ∈ R m × m and P ∈ R n × n are two orthogonal matrices. Instead of optimizing weight matrices directly, POET keeps the ran-

Singular value (log scale)

Singular value index (sorted in a descending order)

Singular value (log scale)

Singular value index (sorted in a descending order)

(a) Standard training

(b) POET

500

400

300

200

100

0

10

0

10

1

10

-1

6

10

0

10

0

2

Step 0

Step 10000

Step 20000

Step 30000

Step 40000

Step 50000

Step 60000

Step 70000

Step 80000

Step 90000

Step 100000

Step 0

Step 10000

Step 20000

Step 30000

Step 40000

Step 50000

Step 60000

Step 70000

Step 80000

Step 90000

Step 100000

500

400

300

200

100

0

domly initialized weight matrix W 0 unchanged during training and learns two orthogonal matrices R , P to transform W 0 . This reparameterization preserves the singular values of weights while allowing flexible optimization of the singular vectors. POET effectively addresses the above limitations:

Strong spectrum control : Because orthogonal transformations do not change the singular values of weight matrices, POET keeps the weight spectrum the same as the randomly initialized weight matrices (empirically validated by Figure 1 even with approximations). Through the initialization scheme, POET thus directly controls the singular value distribution of its weight matrices. As a result, and in contrast to standard LLM training, POET matrices avoid undesirable large singular values after training (Figure 1 and Appendix I). To further facilitate the POET algorithm, we introduce two new initialization schemes: normalized Gaussian initialization and uniform spectrum initialization, which can ensure the resulting weight matrices have bounded singular values.

Efficient approximation : While a naive implementation of POET can be computationally expensive, its inherent flexibility opens up opportunities for efficient and scalable training. To address the key challenge of optimizing large orthogonal matrices, we introduce two levels of approximations:

-Stochastic primitive optimization : The first-level approximation aims to reduce the number of learnable parameters when optimizing a large orthogonal matrix. To this end, we propose the stochastic primitive optimization (SPO) algorithm. Given a large orthogonal matrix R ∈ R m × m , SPO factorizes it into a product of primitive orthogonal matrices, each involving significantly fewer trainable parameters. These primitives are constructed by parameterizing randomly sampled submatrices of the full matrix. This factorization is implemented as a memory-efficient iterative algorithm that sequentially updates one primitive orthogonal matrix at a time. To improve the expressiveness of the sequential factorization, we adopt a merge-then-reinitialize trick, where we merge each learned primitive orthogonal matrix into the weight matrix, and then reinitialize the primitive orthogonal matrix to be identity after every fixed number of iterations.

-Approximate orthogonality via Cayley-Neumann parameterization : The second-level approximation addresses how to maintain orthogonality without introducing significant computational overhead. To achieve this, we develop the Cayley-Neumann parameterization (CNP) which approximates the Cayley orthogonal parameterization [48, 65] with Neumann series. Our merge-then-reinitialize trick can effectively prevent the accumulation of approximation errors.

POET can be viewed as a natural generalization of orthogonal training [48, 51, 65], wherein the model training is done by learning a layer-shared orthogonal transformation for neurons. Orthogonal training preserves the hyperspherical energy [47, 49] within each layer-a quantity that characterizes pairwise neuron relationships on the unit hypersphere. While preserving hyperspherical energy proves effective for many finetuning tasks [51], it limits the flexibility of pretraining. Motivated by this, POET generalizes energy preservation to spectrum preservation and subsumes orthogonal training as its special case. The better flexibility of POET comes from its inductive structures for preserving weight spectrum, rather than more learnable parameters. We empirically validate that POET achieves better pretraining performance than orthogonal training given the same budget of parameters.

To better understand how POET functions, we employ vector probing to analyze the learning dynamics of the orthogonal matrices. Vector probing evaluates an orthogonal matrix R using a fixed, randomly generated unit vector v by computing v ⊤ Rv which corresponds to the cosine similarity between Rv and v . By inspecting the cosine similarities of seven orthogonal matrices throughout training, we

observe that the learning process can be divided into three distinct phases (Figure 2): (1) conical shell searching : The cosine starts at 1 ( i.e. , R is the identity) and gradually converges to a stable range of [0 . 6 , 0 . 65] , which we observe consistently across all learnable orthogonal matri-

0

50K

100K

150K

200K

250K

300K

Iteration number

0.60

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Cosine Similarity between

Rv

and

v

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

Phase I:

Conical shell

searching

Stable learning

on the conical shell

Phase II:

0

100K

200K

300K

3

5

7

9

Training loss

Iteration number

28

32

36

40

0

100K

200K

300K

Iteration number

Eval perplexity

Final

adjusting

Phase III:

Phase I:

Conical shell searching

Phase II:

Stable learning

Phase III:

Adjusting

Init

Final

ces. This suggests that R transforms v into a thin conical shell around its original direction. (2) stable learning on the conical shell : The cosine remains within this range while the model begins to learn stably. Despite the cosine plateauing, validation perplexity continues to improve almost linearly. (3) final adjusting : Learning slows and eventually halts as the learning rate approaches zero. We provide an in-depth discussion and empirical results in Appendix A,G. Our major contributions are:

We introduce POET, a novel training framework that provably preserves spectral properties of weight matrices through orthogonal equivalence transformation.

To enhance POET's scalability, we develop two simple yet effective approximations: stochastic principal submatrix optimization for large orthogonal matrices and the Cayley-Neumann parameterization for efficient representation of orthogonal matrices.

We empirically validate POET's training stability and generalization across multiple model scales.

2 From Energy-preserving Training to Spectrum-preserving Training

Orthogonal training [48, 51, 65] is a framework to train neural networks by learning a layer-shared orthogonal transformation for neurons in each layer. For a weight matrix W = { w 1 , · · · , w n } ∈ R m × n where w i ∈ R m is the i -th neuron, the layer's forward pass is y = W ⊤ x with input x ∈ R m and output y ∈ R n . Unlike standard training, which directly optimizes W , orthogonal training keeps W fixed at its random initialization W 0 = w 0 1 , . . . , w 0 n and instead learns an orthogonal matrix R ∈ R m × m to jointly transform all neurons in the layer. The forward pass becomes y = ( RW 0 ) ⊤ x . The effective weight matrix is W R = { w R 1 , · · · , w R n } where w R i = Rw i . A key property is its preservation of hyperspherical energy . With ˆ w i = w i / ∥ w i ∥ , orthogonal training ensures

̸

̸

where hyperspherical energy HE ( · ) characterizes the uniformity of neurons on a unit hypersphere. Prior work [47-49, 76] has shown that energy-preserving training can effectively improve generalization. Orthogonal finetuning (OFT) [51, 65] also demonstrates that finetuning foundation models while preserving hyperspherical energy achieves a favorable trade-off between efficient adaptation to downstream tasks and retention of pretraining knowledge. While the hyperspherical energy preservation is effective for finetuning, it can be too restrictive for pretraining. To allow greater flexibility in the pretraining phase, we relax the constraint from preserving hyperspherical energy to preserving the singular-value spectrum instead. By inherently maintaining the spectrum, energy-preserving training is a special case of spectrum-preserving training. As a generalization, spectrum-preserving training learns a transformation T : R m × n → R m × n that preserves the spectrum:

where σ i ( W 0 ) denotes the i -th singular value of W 0 (sorted by descending order with σ 1 being the largest singular value). How we instantiate the transformation T results in different algorithms. Generally, T is a spectrum-preserving map, and can be either linear [42] or nonlinear [5]. If we only consider T to be a linear map, then Theorem 1 can fully characterize the form of T :

̸

Theorem 1 (Simplified results from [42]) . For a linear map T : R m × n → R m × n ( m = n ), if σ 1 ( T ( W )) = σ 1 ( W ) always holds for all W ∈ R m × n , then the linear map T must be of the following form: T ( W ) = RWP , for all W ∈ R m × n where R ∈ R m × m and P ∈ R n × n are some fixed elements in orthogonal groups O ( m ) and O ( n ) , respectively.

All parameterizations for the linear map T can be expressed as T ( W )= RWP , where R and P are orthogonal matrices. For instance, OFT is an energy-preserving method (a special case of spectrum-preserving training), where the map simplifies to T ( W )= RWI , with I as the identity.

POET preserves hyperspherical energy under isotropic Gaussian initialization . [48] shows that weight matrices that are initialized by zero-mean isotropic Gaussian distribution ( e.g. , [17]) are guaranteed to have small hyperspherical energy. Because zero-mean isotropic Gaussian is invariant under orthogonal transformation, RWP also has small energy (see Section 4 and Appendix B).

3 Reparameterized Training via Orthogonal Equivalence Transformation

This section introduces the POET framework, which reparameterizes each neuron as the product of a fixed random weight matrix and two learnable orthogonal matrices applied on both sides. POET serves as a specific implementation of spectrum-preserving training. Inspired by Theorem 1, it parameterizes the spectrum-preserving transformation T using a left orthogonal matrix that transforms the column space of the weight matrix and a right orthogonal matrix that transforms its row space.

3.1 General Framework

Following the general form of spectrum-preserving linear maps discussed in the last section, POET reparameterizes the neuron as RW 0 P , where W 0 ∈ R m × n is a randomly initialized weight matrix that remains fixed during training, and R ∈ R m × m , P ∈ R n × n are trainable orthogonal matrices. This reparameterization effectively applies an orthogonal equivalence transformation (OET) to random weight matrices. Specifically, OET is a double-sided transformation, defined as OET ( W ; R , P ) = RWP , where the input matrix W is multiplied on the left and on the right by orthogonal matrices R and P , respectively. The forward pass of POET can be thus written as

where R and P can be merged into a single weight matrix W RP = RW 0 P after training. Therefore, the inference speed of POET-trained neural networks is the same as conventionally trained ones.

Spectrum control . POET can be interpreted as learning weight matrices by simultaneously transforming their left singular vectors and right singular vectors while keeping the singular values unchanged. Given the singular value decomposition (SVD) W 0 = U Σ 0 V ⊤ , the reparameterized neuron weight matrix becomes W RP = RU Σ 0 V ⊤ P where both RU and V ⊤ P are orthogonal matrices. This effectively constitutes an SVD of W RP . It is also straightforward to verify that the spectral properties of W RP remain identical to those of the initial matrix W 0 .

Neuron initialization . Since POET preserves the spectral properties of the initial weight matrix W 0 , the choice of initialization plays a critical role. We consider two common schemes: (1) standard initialization , which samples from a zero-mean Gaussian with fixed variance (the default choice for LLaMA models); and (2) Xavier initialization [17], which uses a zero-mean Gaussian with variance scaled by the layer dimensions. To facilitate POET, we propose two new initialization schemes. The first method, uniform-spectrum initialization , applies SVD to a standard initialization and sets all singular values to 1 , balancing spectral properties throughout training. The second, normalized Gaussian initialization , normalizes neurons drawn from a zero-mean Gaussian with fixed variance. This

0

100

200

300

400

500

0.4

0.6

0.8

1

1.2

1.4

1.6

1.8

2

Singular value

Singular value index (sorted in a descending order)

Standard

Xavier

Uniform spectrum

is directly inspired by prior work showing that normalized neurons improve convergence [48, 50, 52]. To ensure that the POET-reparameterized network is statistically equivalent to a standard network at initialization, we always initialize both orthogonal matrices as identity matrices.

3.2 Efficient Approximation to Orthogonality

POET is conceptually simple, requiring only the optimization of two orthogonal matrices. However, these matrices are typically large, and naively optimizing them leads to significant computational challenges. We start by introducing the following efficient approximations.

3.2.1 Stochastic Primitive Optimization

The core idea of SPO is inspired by how QR factorization is performed using Givens rotations and Householder transformations. Both methods construct a large orthogonal matrix R by sequentially applying primitive orthogonal transformations ( e.g. , Givens rotations or Householder reflections), i.e. , R = ∏ c i =1 G i , where G i denotes the i -th primitive orthogonal matrix. While each G i is of the same size as R , it is parameterized by significantly fewer degrees of freedom. See Figure 4

(a) Givens rotation

(b) Householder reflection

(c) SPO (fully stochastic)

(d) SPO (block-stochastic)

for an illustration. Both Givens rotation and Householder reflection use relatively low-capacity parameterizations-for example, each Givens rotation G i involves only a single effective parameterwhich limits their efficiency in representing the full orthogonal matrix. SPO follows a similar idea of factorizing the original orthogonal matrix into multiple primitive orthogonal matrices. However, unlike Givens and Householder methods, SPO treats the number of effective parameters in each primitive matrix as a tunable hyperparameter and adopts a stochastic sparsity pattern.

Fully stochastic SPO . The basic idea of fully stochastic SPO is to randomly sample a small submatrix and enforce its orthogonality, allowing it to be easily extended to a full orthogonal matrix by embedding it within an identity matrix-a process similar to Givens or Householder transformations. To represent a large orthogonal matrix R ∈ R m × m , we start by defining c index sets S j = { s j 1 , · · · , s j b } ⊆ { 1 , · · · , m } ( j ∈ [1 , c ] ), where each set has cardinality | S j | = b , a hyperparameter controlling the number of effective parameters of a primitive orthogonal matrix. S j , ∀ j are randomly sampled from the full indices { 1 , · · · , m } . Let ˜ G j ∈ R b × b be a small orthogonal matrix, and D ( S j ) = { e ( s j 1 ) , · · · , e ( s j b ) } ∈ R m × b be a selection matrix, where e ( k ) is the standard basis vector with a 1 in the k -th position and 0 elsewhere. The factorization is given by

where D ( S i ) · ( A ) · D ( S i ) ⊤ is a projector that replaces the b × b sub-block with A . I m and I b are identity matrices of size m × m and b × b , respectively. To efficiently parameterize small orthogonal matrices ˜ G i , we can use the CNP introduced in the next section.

Block-stochastic SPO . While fully stochastic SPO is simple, it may fail to transform all neuron dimensions because the identity matrix leaves part of the space unchanged. See the blue blocks in Figure 4(c) as an example. To address this, we propose block-stochastic SPO, which first constructs a block-diagonal orthogonal matrix with small blocks for parameter efficiency, and then applies a random permutation to enhance expressiveness by randomizing the sparsity pattern. Block-stochastic SPO transforms all neuron dimensions simultaneously, as shown in Figure 4(d). Formally we have

where ˜ G j i ∈ R b × b is the j -th block of the block diagonal matrix, and Ψ i , ∀ i are all random permutation matrices. As long as each diagonal block ˜ G j i is an orthogonal matrix, both G i and R are also orthogonal matrices. We also use CNP to efficiently parameterize each orthogonal block ˜ G j i .

The merge-then-reinitialize trick . The factorizations in Equation (4) and (5) offer a simple approach to optimizing large orthogonal matrices by sequentially updating primitive orthogonal matrices. However, storing all previous primitives incurs high GPU memory overhead. To mitigate this, we propose the merge-then-reinitialize trick, where the learned primitive orthogonal matrix can be merged into the weight matrix after every certain number of iterations, and then reinitialized to the identity matrix. After reinitialization, stochastic sampling is repeated to select a new index set (in fully stochastic SPO) or generate a new permutation (in block-stochastic SPO). This trick allows only one primitive matrix to be stored at a time, substantially reducing GPU memory usage.

3.2.2 Cayley-Neumann Parameterization

The classic Cayley parameterization generates an orthogonal matrix R in the form of R = ( I + Q )( I -Q ) -1 where Q is a skew-symmetric matrix satisfying Q = -Q ⊤ . A minor caveat of this parameterization is that it only produces orthogonal matrices with determinant 1 ( i.e. , elements of the special orthogonal group), but empirical results in [48, 51, 65] indicate that this constraint does not hurt performance. However, the matrix inverse in the original Cayley parameterization introduces

numerical instability and computational overhead, limiting its scalability to large orthogonal matrices. To address this, we approximate the matrix inverse using a truncated Neumann series:

where a larger number of approximation terms k leads to a smaller approximation error. By avoiding matrix inversion, the training stability of POET is improved; however, this comes with a price-the approximation is valid only when the Neumann series converges in the operator norm. To initialize orthogonal matrices as identity, we set Q to a zero matrix in CNP, satisfying the convergence condition initially. As the training progresses, however, updates to Q may cause its operator norm to exceed 1 , violating this condition. Fortunately, our merge-then-reinitialize trick mitigates this issue by periodically resetting Q to a zero matrix, ensuring its operator norm remains small.

3.2.3 Overall Training Algorithm

Step 1: Initialization . We initialize the weight matrices using normalized Gaussian: W ← W 0 .

Step 2: Orthogonal matrix initialization . For fully stochastic SPO, we randomly sample an index set S , and parameterize ˜ G R ∈ R b × b and ˜ G P ∈ R b × b using CNP (Equation (6)). Both matrices are initialized as identity, so R and P also start as identity matrices. For block-stochastic SPO, we sample a random permutation matrix Ψ R , Ψ P , and parameterize { ˜ G 1 R , · · · , ˜ G ⌈ m b ⌉ R } and { ˜ G 1 P , · · · , ˜ G ⌈ n b ⌉ P } using CNP. Then we initialize them as the identity, so R and P again starts as identity matrices.

Step 3: Efficient orthogonal parameterization . For fully stochastic SPO, we have R = I m + D ( S )( ˜ G R -I b ) D ( S ) ⊤ and P = I n + D ( S )( ˜ G P -I b ) D ( S ) ⊤ . For block-stochastic SPO, we have R = Ψ ⊤ R Diag ( ˜ G 1 R , · · · , ˜ G ⌈ m b ⌉ R ) Ψ R and P = Ψ ⊤ P Diag ( ˜ G 1 P , · · · , ˜ G ⌈ n b ⌉ P ) Ψ P .

Step 4: Inner training loop for updating orthogonal matrices . The equivalent weight matrix in the forward pass is RWP . Gradients are backpropagated through R and P to update ˜ G R , ˜ G P (fully stochastic) or ˜ G i R , ˜ G i P , ∀ i (block-stochastic). This inner loop runs for a fixed number of iterations.

Step 5: Merge-then-reinitialize . The learned orthogonal matrices R and P are merged into the weight matrix by W ← RWP . If not terminated, return to Step 2 for reinitialization.

4 Discussions and Intriguing Insights

Parameter and memory complexity . By introducing a hyperparameter b as the sampling budget, fully stochastic SPO decouples parameter complexity from the size of the weight matrices. With a small b , POET becomes highly parameter-efficient, though at the cost of slower convergence. This offers users a flexible trade-off between efficiency and speed. In contrast, block-stochastic SPO has

parameter complexity dependent on the matrix size ( i.e. , m + n ), making it more scalable than AdamW, which requires mn trainable parameters. In terms of memory complexity, both POET variants can be much more efficient than AdamW with a suitable sampling budget b . A comparison of parameter and memory complexity is given in Table 1.

Performance under a constant parameter budget . Since POET optimizes two orthogonal matrices R , P simultaneously, a natural question arises: which matrix should receive more parameter budget under a fixed total constraint? To investigate this, we conduct a controlled experiment where different ratios of trainable parameters are allocated to R and P under a fixed total budget. All other settings ( e.g. , architecture, data) remain unchanged, with full details provided in the Appendix. We use validation perplexity as the

28

30

32

34

36

Small budget

Medium budget

Large budget

Validation perplexity

R

: 0.5

P

: 0.5

R

: 1

P

: 0

R

: 0.9

P

: 0.1

R

: 0.75

P

: 0.25

R

: 0.25

P

: 0.75

R

: 0.1

P

: 0.9

R

: 0

P

: 1

evaluation metric. The total parameter budget matches that of fully stochastic POET with b = 1 h m for R and b = 1 h n for P , where h = 8 , 4 , and 3 correspond to small, medium, and large budgets, respectively. We explore seven allocation settings: R : P =1:0 ( i.e. , orthogonal training [48, 51, 65]), 0 . 9:0 . 1 , 0 . 75:0 . 25 , 0 . 5:0 . 5 ( i.e. , standard POET), 0 . 25:0 . 75 , 0 . 1:0 . 9 , and 0:1 . Results in Figure 5 show that POET with a balanced allocation between R and P yields the best performance.

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

Training Steps

0.93

0.94

0.95

0.96

0.97

SVD Entropy

(c) Average over all up projection layers

MUON

AdamW

POET

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

Training Steps

0.86

0.87

0.88

0.89

0.90

0.91

0.92

SVD Entropy

(b) Average over all Q projection layers

MUON

AdamW

POET

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

Training Steps

0.85

0.86

0.87

0.88

0.89

0.90

0.91

0.92

SVD Entropy

(a) Average over all K projection layers

MUON

AdamW

POET

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

Training Steps

0.88

0.89

0.90

0.91

0.92

SVD Entropy

(f) Average over all output projection layers

MUON

AdamW

POET

0.955

0.960

0.965

0.970

SVD Entropy

(e) Average over all down projection layers

MUON

AdamW

POET

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

Training Steps

0.91

0.92

0.93

0.94

0.95

0.96

0.97

SVD Entropy

(d) Gate projection layer

MUON

AdamW

POET

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

Training Steps

Guarantees of weight spectrum . For POET with standard and normalized Gaussian initializations, we have proved in Appendix C that the largest and smallest singular values of weights can be bounded. For normalized Gaussian, the bound is only dependent on the row-to-column ratio of the weight matrix. For standard Gaussian, the bound has an extra dependence on the neuron dimension.

Connection to generalization theory . Several generalization results [6, 63, 76] based on bounding the spectral norm of weight matrices. In particular, the spectrally-normalized margin analysis in [6] bounds the misclassification error in terms of a margin-based training loss and a complexity term. The complexity term is proportional to Q/ ( γn ) where γ and n are margin and sample size and Q bounds the spectral complexity. For an L -layer ReLU MLP and maximal width d , Q is bounded by

where ∥·∥ and ∥·∥ F denote spectral and Frobenius norm respectively. Those norms remain invariant when training the network with POET and at initialization they can be bounded with high probability using standard results from random matrix theory (Appendix C). The scale at initialization is typically chosen such that W ∈ R d × d satisfies ∥ W ∥ = O (1) and ∥ W ∥ = O ( √ d ) so that Q = O L ( d ) .

Approximation properties of SPO . We have seen in Theorem 1 that the factorization RWP with orthogonal matrices R and P is the most general spectrum preserving transformation of W . Here we express R and P as products of stochastic primitives, but as we state next, this does not reduce representation power when using sufficiently many primitives.

Lemma 1. If c ≥ αm ln( m )( m/b ) 2 for some α > 0 then with probability at least 1 -m -( α -2) over the randomness of the index sets S i we can express any orthogonal matrix R as a product of c primitives G i as in Eq. (4) . Moreover, the orthogonal matrix G i depends only on the sets S j and matrices G j selected in earlier steps.

The proof of this lemma can be found in Appendix D. The result extends to Block-stochastic SPO as this is strictly more expressive than fully stochastic SPO. The key idea of the proof is similar to the factorization of orthogonal matrices into a product of Givens rotations. Indeed, by multiplying R ⊤ with properly chosen primitive matrices G i we can create zeros below the diagonal for one column after another. Note that each G i has b ( b -1) / 2 parameters while R has m ( m -1) / 2 parameters, which implies that generally at least Ω(( m/b ) 2 ) primitives are necessary. In Appendix D we also provide a heuristic that with high probability for c = O (ln( m )( m/b ) 2 ) every orthogonal matrix can be written as a product of c orthogonal primitives G i .

Inductive bias . POET-reparameterized neurons result in neural networks that maintain identical architecture and parameter count during inference as conventionally trained networks. While standard training could technically learn equivalent parameters, they consistently fail to do so in practice. This indicates POET provides a unique inductive bias unavailable through standard training. POET also aligns with prior findings in [2, 18] that optimizing factorized matrices yields implicit inductive bias.

Dynamics of singular spectrum . Inspired by [46], we conduct a spectral analysis by comparing the singular spectrum among AdamW, Muon and POET. Following [1, 68], we compute SVD

entropy of the trained Llama-60M model at different iteration. Specifically, the SVD entropy, defined as H ( σ ) = 1 log n ∑ i σ 2 i ∑ j σ 2 j log σ 2 i ∑ j σ 2 j measures the diversity of singular values; higher entropy indicates a more uniform and diverse spectrum. [46] attributes the favorable performance of Muon over AdamW to the more diverse spectrum of weight matrices updates. As shown in Figure 6, POET consistently maintains high spectral diversity throughout training, owing to its orthogonal equivalence transformation. Therefore, POET can better explore diverse optimization directions.

POET minimizes hyperspherical energy . While POET generalizes energy-preserving training to spectrum-preserving training, it still ensures low hyperspherical energy when initialized with a zeromean isotropic Gaussian distribution. This is significant, as POET retains the generalization benefits of minimal hyperspherical energy [47-49]. This property comes from the invariance of zero-mean isotropic Gaussian to orthogonal transformation. [48] shows that for a weight matrix W initialized by zero-mean isotropic Gaussian, its neurons, after being normalized, are uniformly distributed on the unit hypersphere. This property provably leads to a small hyperspherical energy. Since zero-mean isotropic Gaussian is invariant to orthogonal transformation, OET ( W ; R , P ) = RWP does not change the distribution of W , i.e. , RWP = d W . We give a derivation in Appendix B. Therefore, the transformed neurons of the new weight matrix RWP are also uniformly distributed over the unit hypersphere after being normalized. This validates that POET provably preserves a small energy.

The property that POET preserves singular spectrum while retaining small hyperspherical energy well justifies why POET yields stable training and good generalization. Such a property only holds true when the weight initialization follows zero-mean isotropic Gaussian, which is perfectly satisfied by current weight initialization. It may also justify why zero-mean isotropic Gaussian initialization works better than the uniform spectrum initialization in our experiments. To further validate that POET maintains a small hyperspherical energy, we plot the hyperspherical energy of the Llama-60M model

MUON

AdamW

POET

0

10K

20K

30K

40K

50K

39.62

39.63

39.64

39.65

Training Steps

Hyperspherical Energy

at different iterations in Figure 7. We compare the sum of the hyperspherical energy of all the layers trained by AdamW, Muon or POET. The results again verify POET's energy-minimizing property. Different from previous orthogonal training [48], POET can not preserve exactly the same hyperspherical energy during training, but it can well minimize hyperspherical energy in an expectation sense. We also find that Muon can better minimize hyperspherical energy than AdamW.

5 Experiments and Results

We start by evaluating POET on large-scale LLaMA pretraining, followed by an extensive ablation study to justify our design choices. Detailed settings and additional results are given in Appendices.

5.1 LLMPretraining using LLaMA Transformers

We perform the pretraining experiments on the Llama transformers of varying sizes (60M, 130M, 350M, 1.3B) for POET. We use the C4 dataset [66], a cleaned web crawl corpus from Common Crawl, widely used for LLM pretraining [29, 56, 82]. For POET-BS, b is the block size of the block-diagonal orthogonal matrix. For POETFS, b in= bm for R and b out= bn

for P . We compare POET against GaLore [82], a low-rank pretraining method, and AdamW, the standard pretraining optimizer. We generally follow the settings in [82]. To better simulate the practical pretraining setting, we significantly increase the number of training tokens for all methods.

Table 2 shows that both POET-FS ( b =1/2) and POET-BS ( b =256) consistently outperform both GaLore and AdamW with significantly fewer parameters. For LLaMA-1B, POET-FS ( b =1/2) yields the best overall performance, achieving a validation perplexity of 13 . 70 , much better than AdamW ( 14 . 73 ) and GaLore ( 18 . 33 ). Block-stochastic POET with b =256 achieves the second-best performance

( 14 . 56 ), which still surpasses AdamW with only one-tenth of AdamW's trainable parameters. Similar patterns can be observed for models of smaller sizes. Moreover, we compare the training dynamics between AdamW and POET in Figure 8. The training dynamics of POET is quite different

0K

50K

100K

150K

200K

250K

300K

350K

400K

Steps

15

20

25

30

35

40

45

Perplexity

AdamW

POET(BS, 64)

POET(BS, 128)

POET(BS, 256)

POET(FS, 1/8)

POET(FS, 1/4)

POET(FS, 1/2)

0K

100K

200K

300K

400K

500K

Steps

15

20

25

30

Perplexity

AdamW

POET(BS, 256)

POET(FS, 1/2)

from AdamW. After an initial rapid drop in perplexity, POET improves more slowly than AdamW. As seen in Phase II (Figure 2), this slower but stable progress can lead to better performance in later stages. We attribute this intriguing phenomenon to the unique reparameterization of POET and How we efficiently approximate orthogonality. The exact mechanism behind this phenomenon remains an open question, and understanding it could offer valuable insights into large-scale model training.

To highlight POET's non-trivial performance improvement, we increase the training steps ( i.e. , effectively tokens seen) for AdamW, and find that POET-FS ( b =1/2) still outperforms AdamW even even if AdamW is trained with almost triple the number of tokens. Results are given in Figure 9. In this experiment, the AdamW learning rate was carefully tuned for the full training run, and no training tokens were repeated. Thus, the improvement is non-trivial and cannot be attributed to merely increasing training steps. Interestingly, we also observe from Table 2 that POET's performance appears strongly correlated with the parameter budget and larger budgets consistently yield better results across model scales. This is particu-

0K

100K

200K

300K

400K

500K

600K

700K

Steps

25

30

35

40

45

Perplexity

AdamW

POET(FS, 1/2)

POET with ~250K steps

AdamW with >700K steps

the same perplexity

300K is the default training

steps for all methods

larly important for model scaling law [35]. Another notable observation is that POET significantly outperforms LoRA [26] given a similar parameter budget. For instance, with approximately 30M trainable parameters, POET attains a validation perplexity of 18 . 05 , significantly better than LoRA's 25 . 19 . We further observe that the block-stochastic variant is more parameter-efficient than the fully stochastic one. On the 130M model, it achieves a validation perplexity of 19 . 88 with nearly 6M fewer trainable parameters, compared to 19 . 94 for the fully stochastic variant. We hypothesize that this is due to better coverage of weight parameters. Specifically, the block-stochastic variant ensures all corresponding weights are updated at each step, unlike the more uneven updates in the fully stochastic variant. Experimental details and results on weight update coverage are provided in Appendix H.

Pretraining Llama-3B . To demonstrate the scalability of POET, we apply POET to pretrain Llama with 3B parameters while reusing the same hyperparameters from the 1.3B model. This model was trained with only 1 / 10 of the tokens used for the 1.3B model. Compared to Table 2, the slightly higher final perplexity is attributed to the reduced training data (5B tokens). Table 3 shows that POET maintains the performance advantage over AdamW at the 3B scale,

consistent with the conclusion drawn from Table 2 for smaller models. Our findings confirm that the benefits of POET are not limited to smaller models but can extend robustly to larger scales.

5.2 LLMFinetuning on Downstream Tasks

To better evaluate models beyond the validation perplexity, we show the results of finetuning the trained model on the GLUE benchmark [75]. This benchmark provides a comprehensive assessment of a model's

language understanding capabilities through a diverse set of downstream tasks. The performance on GLUE serves as an important indicator of the transferability of the model's learned representations. Performance was measured using accuracy for MNLI, MRPC, QNLI, QQP, RTE, and SST-2; the Matthews Correlation Coefficient for CoLA; and the Pearson Correlation Coefficient for STS-B. For each task, we finetune the models for 10 epochs using a consistent learning rate of 2e-5. The evaluation covers several finetuning strategies: full finetuning (Full FT) with AdamW, orthogonal finetuning (OFT) [51, 65], and finetuning with POET. The results in Table 4 show that the POET-pretrained

model consistently outperforms the AdamW-pretrained baseline across all tasks and finetuning methods. The results further validate POET's generalizability. In this setting, OFT uses significantly less parameters than both Full FT and POET, so OFT generally performs worse. However, the comparison between Full FT and POET is fair, since both methods update the entire model.

5.3 Ablation Studies and Empirical Analyses

Initialization schemes . We empirically compare different random initialization schemes for POET, including two commonly used ones (standard Gaussian, Xavier [17]) and two proposed ones (uniform spectrum, normalized Gaussian). Specifically, we use fully stochastic POET with b =1/2 to train Llama-60M on 30B tokens and report the validation perplexity in Table 5. Results show that the normalized initialization will lead to the best final performance, and we stick to it as a default choice. Interestingly, uniform

spectrum initialization performs poorly. This suggests a trade-off between preserving good weight spectral properties and achieving strong expressiveness. it may limit its expressiveness. Finding the optimal singular value structure for weights remains an important open problem. T Perplexity

Merge-then-reinitialize frequency . The proposed merge-then-reinitialize trick allows POET to train only a small fraction of the large orthogonal matrices R and P per iteration, significantly reducing GPU memory usage. However, this trick also introduces a reinitialization frequency hyperparameter T m , which determines how often the orthogonal matrix is merged and reset to the identity. The index set in POET-FS and the permutation matrix in POET-BS are also resampled at each reinitialization. Therefore, it is quite important to understand how this hyperparameter T m affects performance. Following the previous initialization experiment, we

use POET-FS with b =1/2 to train Llama-60M on 30B tokens. We vary the reinitialization frequency from 5 to 1600 and report the validation perplexity in Table 6. Results show that both 200 and 400 perform well. Therefore, we set T m = 400 in all experiments by default. Scheme Perplexity

Neumann series approximation . CNP approximates the matrix inverse using a Neumann series. As the number of Neumann terms directly influences the approximation quality, understanding its impact on model performance is essential. To this end, we evaluate how varying the number of Neumann terms affects performance, using POET-FS with b = 1/2 to train LLaMA-130M. Results in Table 7 show that increasing the number of Neumann terms generally improves validation perplexity. However, this also leads to slower training.

Moreover, dropping all Neumann terms ( k = 0 ) leads to training divergence, highlighting the critical role of maintaining orthogonality. To balance overhead and performance, we find that 3 Neumann terms ( k = 3 ) reach a good trade-off.

Additionally, we evaluate the accuracy of the Neumann approximation to understand how the number of Neumann terms affects the orthogonality. The orthogonal approximation error is defined by e orth = ∥ RR T -I ∥ F / ∥ I ∥ F . We randomly

0

200

400

600

800

1000

Step

k=2

k=4

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

(a) Orthogonal matrix

R

(a) Orthogonal matrix

P

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

k=2

k=4

select a weight matrix and compute the approximation error of corresponding orthogonal matrices R and P . For clarity, we visualize the error in the initial 1000 training steps in Figure 10. We observe that, with more Neumann terms, the orthogonal approximation error is indeed lower. Moreover, the merge-then-reinitialize trick can periodically reset the error. More results are given in Appendix L

References

Orly Alter, Patrick O Brown, and David Botstein. Singular value decomposition for genome-wide expression data processing and modeling. Proceedings of the National Academy of Sciences , 2000. 7

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In NeurIPS , 2019. 7

Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732 , 2021. 1

Nitin Bansal, Xiaohan Chen, and Zhangyang Wang. Can we gain more from orthogonality regularizations in training deep networks? In NeurIPS , 2018. 1

Line Baribeau and Thomas Ransford. Non-linear spectrum-preserving maps. Bulletin of the London Mathematical Society , 32(1):8-14, 2000. 3

Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. Spectrally-normalized margin bounds for neural networks. In NeurIPS , 2017. 1, 7

Silvere Bonnabel. Stochastic gradient descent on riemannian manifolds. IEEE Transactions on Automatic Control , 58(9):2217-2229, 2013. 11

Djalil Chafaı, Djalil Chafä, Olivier Guédon, Guillaume Lecue, and Alain Pajor. Singular values of random matrices. Lecture Notes , 13, 2009. 20

Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In NeurIPS , 2021. 10

Tianlong Chen, Zhenyu Zhang, Yu Cheng, Ahmed Awadallah, and Zhangyang Wang. The principle of diversity: Training stronger vision transformers calls for reducing all levels of redundancy. In CVPR , 2022. 1

Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue, Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of llms under low-rank constraint? arXiv preprint arXiv:2410.01623 , 2024. 10

Moustapha Cisse, Piotr Bojanowski, Edouard Grave, Yann Dauphin, and Nicolas Usunier. Parseval networks: Improving robustness to adversarial examples. In ICML , 2017. 1

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 , 2021. 1

Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In ICML , 2022. 10

Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In ICML , 2019. 10

Alan Edelman and N Raj Rao. Random matrix theory. Acta numerica , 14:233-297, 2005. 11

Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In AISTATS , 2010. 4, 10, 19

Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In NeurIPS , 2017. 7

Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In ICML , 2018. 11

Andi Han, Jiaxiang Li, Wei Huang, Mingyi Hong, Akiko Takeda, Pratik Jawanpuria, and Bamdev Mishra. Sltrain: a sparse plus low-rank approach for parameter and memory efficient pretraining. arXiv preprint arXiv:2406.02214 , 2024. 10

Boris Hanin. Random neural networks in the infinite width limit as gaussian processes. The Annals of Applied Probability , 33(6A):4798-4819, 2023. 11

Boris Hanin and David Rolnick. Complexity of linear regions in deep networks. In ICML , 2019. 11

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In ICCV , 2015. 19

Karl Moritz Hermann, Tomáš Koˇ ciský, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines to read and comprehend, 2015. 40

Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. JMLR , 2021. 10

Edward J. Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In ICLR , 2022. 9, 10

Jia-Hong Huang, Yixian Shen, Hongyi Zhu, Stevan Rudinac, and Evangelos Kanoulas. Gradient weightnormalized low-rank projection for efficient llm training. In AAAI , 2025. 10

Lei Huang, Xianglong Liu, Bo Lang, Adams Yu, Yongliang Wang, and Bo Li. Orthogonal weight normalization: Solution to optimization over multiple dependent stiefel manifolds in deep neural networks. In AAAI , 2018. 1

Tianjin Huang, Ziquan Zhu, Gaojie Jin, Lu Liu, Zhangyang Wang, and Shiwei Liu. Spam: Spike-aware adam with momentum reset for stable llm training, 2025. 8

Weihao Huang, Zhenyu Zhang, Yushun Zhang, Zhi-Quan Luo, Ruoyu Sun, and Zhangyang Wang. Galoremini: Low rank gradient learning with fewer learning rates. In NeurIPS Workshop on Fine-Tuning in Modern Machine Learning: Principles and Scalability , 2024. 10

Minyoung Huh, Brian Cheung, Jeremy Bernstein, Phillip Isola, and Pulkit Agrawal. Training neural networks from scratch with parallel low-rank adapters. arXiv preprint arXiv:2402.16828 , 2024.

Ajay Jaiswal, Lu Yin, Zhenyu Zhang, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. From galore to welore: How low-rank weights non-uniformly emerge from low-rank gradients. arXiv preprint arXiv:2407.11239 , 2024. 10

Haoming Jiang, Zhehui Chen, Minshuo Chen, Feng Liu, Dingding Wang, and Tuo Zhao. On computation and generalization of generative adversarial networks under spectrum control. In ICLR , 2019. 1

Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. 11

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020. 9

Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint arXiv:1712.07628 , 2017. 1

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR , 2015. 1

Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-normalizing neural networks. In NeurIPS , 2017. 1

Hojoon Lee, Youngdo Lee, Takuma Seno, Donghu Kim, Peter Stone, and Jaegul Choo. Hyperspherical normalization for scalable deep reinforcement learning. arXiv preprint arXiv:2502.15280 , 2025. 1

Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.00165 , 2017. 11

Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension. arXiv preprint arXiv:1910.13461 , 2019. 40

Chi-Kwong Li and Nam-Kiu Tsing. Linear operators preserving unitarily invariant norms of matrices. Linear and Multilinear Algebra , 26(1-2):119-132, 1990. 3

Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Relora: High-rank training through low-rank updates. arXiv preprint arXiv:2307.05695 , 2023. 10

Kaizhao Liang, Bo Liu, Lizhang Chen, and Qiang Liu. Memory-efficient llm training with online subspace descent. arXiv preprint arXiv:2408.12857 , 2024.

Xutao Liao, Shaohui Li, Yuhui Xu, Zhi Li, Yu Liu, and You He. Galore+: Boosting low-rank adaptation for llms with cross-head projection. arXiv preprint arXiv:2412.19820 , 2024. 10

Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for llm training. arXiv preprint arXiv:2502.16982 , 2025. 7, 8, 11

Weiyang Liu, Rongmei Lin, Zhen Liu, Lixin Liu, Zhiding Yu, Bo Dai, and Le Song. Learning towards minimum hyperspherical energy. In NeurIPS , 2018. 1, 2, 3, 8

Weiyang Liu, Rongmei Lin, Zhen Liu, James M Rehg, Liam Paull, Li Xiong, Le Song, and Adrian Weller. Orthogonal over-parameterized training. In CVPR , 2021. 2, 3, 4, 5, 6, 8, 19

Weiyang Liu, Rongmei Lin, Zhen Liu, Li Xiong, Bernhard Schölkopf, and Adrian Weller. Learning with hyperspherical uniformity. In AISTATS , 2021. 1, 2, 3, 8, 20

Weiyang Liu, Zhen Liu, Zhiding Yu, Bo Dai, Rongmei Lin, Yisen Wang, James M Rehg, and Le Song. Decoupled networks. In CVPR , 2018. 1, 4

Weiyang Liu, Zeju Qiu, Yao Feng, Yuliang Xiu, Yuxuan Xue, Longhui Yu, Haiwen Feng, Zhen Liu, Juyeon Heo, Songyou Peng, Yandong Wen, Michael J. Black, Adrian Weller, and Bernhard Schölkopf. Parameter-efficient orthogonal finetuning via butterfly factorization. In ICLR , 2024. 2, 3, 5, 6, 9

Weiyang Liu, Yan-Ming Zhang, Xingguo Li, Zhiding Yu, Bo Dai, Tuo Zhao, and Le Song. Deep hyperspherical learning. In NIPS , 2017. 1, 4

Ziyue Liu, Ruijie Zhang, Zhengyang Wang, Zi Yang, Paul Hovland, Bogdan Nicolae, Franck Cappello, and Zheng Zhang. Cola: Compute-efficient pre-training of llms via low-rank activation. arXiv preprint arXiv:2502.10940 , 2025. 10

Ilya Loshchilov, Cheng-Ping Hsieh, Simeng Sun, and Boris Ginsburg. ngpt: Normalized transformer with representation learning on the hypersphere. arXiv preprint arXiv:2410.01131 , 2024. 1

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR , 2019. 1, 25

Chao Ma, Wenbo Gong, Meyer Scetbon, and Edward Meeds. Swan: Sgd with normalization and whitening enables stateless llm training, 2025. 8

Albert W Marshall, Ingram Olkin, and Barry C Arnold. Inequalities: theory of majorization and its applications , volume 143. Springer, 1979. 21

Roy Miles, Pradyumna Reddy, Ismail Elezi, and Jiankang Deng. Velora: Memory efficient training using rank-1 sub-token projections. arXiv preprint arXiv:2405.17991 , 2024. 10

Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. arXiv preprint arXiv:1802.05957 , 2018. 1

Zhanfeng Mo, Long-Kai Huang, and Sinno Jialin Pan. Parameter and memory efficient pretraining via low-rank riemannian optimization. In ICLR , 2025. 11

Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization, 2018. 40

Radford M Neal. Priors for infinite networks. Bayesian learning for neural networks , pages 29-53, 1996. 11

Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spectrallynormalized margin bounds for neural networks. In ICLR , 2018. 7

Sean O'Rourke, Van Vu, and Ke Wang. Eigenvectors of random matrices: a survey. Journal of Combinatorial Theory, Series A , 144:361-442, 2016. 20

Zeju Qiu, Weiyang Liu, Haiwen Feng, Yuxuan Xue, Yao Feng, Zhen Liu, Dan Zhang, Adrian Weller, and Bernhard Schölkopf. Controlling text-to-image diffusion by orthogonal finetuning. In NeurIPS , 2023. 2, 3, 5, 6, 9

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research , 21(140):1-67, 2020. 8, 25

Mihaela Rosca, Theophane Weber, Arthur Gretton, and Shakir Mohamed. A case for new neural network smoothness constraints. arXiv preprint arXiv:2012.07969 , 2020. 1

Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In European signal processing conference , 2007. 7

Ishaan Shah, Anthony M Polloreno, Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh J Shah, et al. Practical efficiency of muon for pretraining. arXiv preprint arXiv:2505.02222 , 2025. 11

Jack W Silverstein et al. The smallest eigenvalue of a large dimensional wishart matrix. The Annals of Probability , 13(4):1364-1368, 1985. 20

DiJia Su, Andrew Gu, Jane Xu, Yuandong Tian, and Jiawei Zhao. Galore 2: Large-scale llm pre-training by gradient low-rank projection. arXiv preprint arXiv:2504.20437 , 2025. 10

Vithursan Thangarasa, Abhay Gupta, William Marshall, Tianda Li, Kevin Leong, Dennis DeCoste, Sean Lie, and Shreyas Saxena. Spdf: Sparse pre-training and dense fine-tuning for large language models. In UAI , 2023. 10

Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.11321 , 2024. 11

Gilles Wainrib and Jonathan Touboul. Topological and dynamical complexity of random neural networks. Physical review letters , 110(11):118101, 2013. 11

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 , 2018. 9

Bo Xie, Yingyu Liang, and Le Song. Diverse neural network learns true target functions. In AISTATS , 2017. 3, 7

Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foundation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.04129 , 2023. 1

Xinyu Yang, Jixuan Leng, Geyang Guo, Jiawei Zhao, Ryumei Nakada, Linjun Zhang, Huaxiu Yao, and Beidi Chen. S2ft: Efficient, scalable and generalizable llm fine-tuning by structured sparsity. arXiv preprint arXiv:2412.06289 , 2024. 10

Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generalizability of deep learning. arXiv preprint arXiv:1705.10941 , 2017. 1

Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.16793 , 2024. 11

Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. Q-galore: Quantized galore with int4 projection and layer-adaptive low-rank gradients. arXiv preprint arXiv:2407.08296 , 2024. 10

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient llm training by gradient low-rank projection. In ICML , 2024. 6, 8, 10, 25

Table of Contents

Appendix

A Delving into POET's Three Training Phases

A.1 More Details on Vector Probing

The three training phases of POET are summarized from the empirical observation of the vector probing results. The idea of vector probing is very straightforward. We generate a constant vector v that is randomly initialized. Then we let it to be transformed by the learned orthogonal matrices R and P . Finally, we compute the cosine of their angle: v ⊤ Rv and v ⊤ Pv . In this process, the probing vector v is always fixed. The full results are given in Appendix G.

Beyond a particular constant probing vector, we also consider a set of randomly sampled probing vectors that follow our proposed normalized Gaussian initialization. Specifically, we consider the following expectation:

where v is a vector initialized by normalized Gaussian distribution (thus uniformly distributed on a unit hypersphere S m -1 ). Because E { vv ⊤ } = 1 m , then we have that

where Tr ( · ) denotes the matrix trace. Its geometric interpretation is the cosine of the rotation angle between v and Rv .

Next, we look into the variance of q ( x ) = v ⊤ Rv (we simplify the expectation over the unit hypersphere to E ):

First we compute E { ( v ⊤ Rv ) 2 } :

Then we compute ( E { v ⊤ Rv } ) 2 :

Finally, we combine pieces and have the final variance:

which shrinks at the order of O (1 /m ) . Therefore, when the dimension of orthogonal matrices is large, even if we use a fixed random probing vector v , this rotation angle is quite consistent.

A.2 Geometric Interpretation of the Trace of Orthogonal Matrices

Let's delve deeper into the trace of orthogonal matrices. It generally represents how much a transformation preserves vectors in their original directions. Specifically, the trace indicates how much 'alignment' or similarity there is between the original vectors and their images after transformation.

The trace of an orthogonal matrix R ∈ R m × m can be written as

where e i , ∀ i are unit basis vectors. This expression reveals that the trace measures the sum of inner products between each original direction e i and its transformed version Re i . Since e ⊤ i Re i can be interpreted as the cosine of the angle between e i and Re i , the trace thus reflects how much the orthogonal transformation aligns with or deviates from the original coordinate directions.

We also plot the trace of both R and P during the POET training. The results are shown in Figure 13 and Figure 14. After dividing the trace by the orthogonal matrix dimension, we obtain that the result is generally in the range of [0 . 6 , 0 . 65] after training. This is similar to the results of vector probing. Therefore, we empirically verify the conclusion that the expectation of vector probing results is Tr ( R ) m with a small variance.

A.3 Empirical Observations

The training dynamics of POET presents three geometry-driven phases. We note that these phase changes are based on empirical observation, and further theoretical understanding of this process remains an open problem.

Phase I: conical-shell searching rotates each orthogonal matrix R and P smoothly away from the identity while preserving their singular values, so the cosine similarity between transformed and initial weight vectors falls from 1 to ≈ 0 . 6 ; this provides a spectrally well-conditioned 'cone' in which learning can proceed safely. this phase serves the role of 'spectral warm-up'. By plotting the cosine similarity of any one layer, we always see the same graceful slide towards 0 . 6 -0 . 65 , independent of model size, layer type, or whether you train with fully-stochastic or block-stochastic SPO. This phase carves out the thin 'shell' in which subsequent learning lives.

Phase II: stable learning on the conical shell occupies the bulk of training: the angles to the initial vectors stay locked in that narrow band, optimization now shears weights within the cone, and validation perplexity drops almost linearly because spectra remain frozen and gradients act only on meaningful directions. In this phase, the trace of the orthogonal matrices stay almost as a constant.

Specifically, we hypothesize that the orthogonal transforms have reached a 'good' cone; thereafter they mostly shear vectors inside that shell, leaving the angle to the original vector unchanged. The spectrum continues to be exactly that of the random initial matrix, so gradients can no longer distort singular values and instead devote capacity to learning meaningful directions. Because the geometry is stabilized in this phase, the learning of patterns happen in a stable subspace. This stable learning phase takes up 80% of the training time.

Phase III: final adjusting coincides with learning-rate decay; the orthogonal transforms barely move, making only tiny refinements to singular vectors, so additional steps yield diminishing returns. This phase is merely the LR cooldown; weights and spectra are already near their final configuration, so progress naturally slows.

B Minimum Hyperspherical Energy in POET

We start by showing that orthogonal equivalence transformation in POET can provably obtain small hyperspherical energy for its transformed weight matrices. This result holds when the weight matrices are independently initialized by zero-mean isotropic Gaussian distribution. Both Xavier [17] and Kaiming [23] initializations satisfy such a weight initialization condition. Orthogonal equivalence transformation is given by OET ( W ; R , P ) = RWP , where the input matrix W is multiplied on the left and on the right by orthogonal matrices R and P , respectively.

When W is initialized by zero-mean isotropic Gaussian distribution, [48] has shown that these random neurons, if normalized, are uniformly distributed on the unit hypersphere. This leads to a provably small hyperspherical energy for the randomly initialized weight matrix. In the following, we will show that after orthogonal equivalence transformation, the weight matrix still maintains a small hyperspherical energy.

We consider a weight matrix W ∈ R m × n where each entry is i.i.d. sampled from a zero-mean Gaussian distribution with variance the same variance σ 2 , i.e. , W ij ∼ N (0 , σ 2 ) . After applying orthogonal equivalence transformation, we have W new = RWP where R and P are two orthogonal matrices. Then we compute the distribution of W new . Because linear maps preserve Gaussianity, each entry of W new is a finite linear combination of W ij , and hence, W new follows a joint Gaussian which can be fully characterized by mean and covariance.

The mean of W new is given by

For its covariance, we consider two generic entries W new ij and W new i ′ j ′ :

Then we compute the covariance between the two entries:

which implies the following resutls:

The covariance matrix is a diagonal matrix, so different entries of W new are uncorrelated.

Because W new is a joint Gaussian and different entries are uncorrelated, each entry of W new is independent.

Each entry of W new has identical variance σ 2 .

To sum up, each entry of W new is i.i.d. N (0 , σ 2 ) , which is identical to the distribution of each entry of W . Because we have W new = d W , we can conclude that, similar to W , W new also has provably small hyperspherical energy among neurons.

Despite being extremely simple, we find that this is in fact a significant result. Under zero-mean isotropic Gaussian initialization, spectrum-preserving training and energy-preserving training can be achieved simultaneously. It also partially explains why the proposed normalized Gaussian initialization achieves the best performance.

C Guarantees of Weight Spectrum under POET

For standard Gaussian initialization where each element of the weight matrix W ∈ d × n is sampled with a normal distribution, we have the following standard results [8, 70]:

which gives spectrum guarantees for weight matrices generated by the standard Gaussian initialization.

In the following, we give the spectrum guarantees for the normalized Gaussian initialization. We start by stating the following theorem from [49]:

Theorem 2. Let ˜ v 1 , · · · , ˜ v n ∈ R d be i.i.d. random vectors where each element follows the Gaussian distribution with mean 0 and variance 1 . Then v 1 = ˜ v 1 ∥ ˜ v 1 ∥ 2 , · · · , v n = ˜ v n ∥ ˜ v n ∥ 2 are uniformly distributed on the unit hypersphere S d -1 . If the ratio n d converges to a constant λ ∈ (0 , 1) , asymptotically we have for W = { v 1 , · · · , v n }∈ R d × n :

where σ max ( · ) and σ min ( · ) denote the largest and the smallest singular value of a matrix, respectively.

Proof. Wefirst introduce the following lemma as the characterization of a unit vector that is uniformly distributed on the unit hypersphere S d -1 .

Lemma 2 ([64]) . Let v be a random vector that is uniformly distributed on the unit hypersphere S d -1 . Then v has the same distribution as the following:

where u 1 , u 2 , · · · , u d are i.i.d. standard normal random variables.

Proof. The lemma follows naturally from the fact that the Gaussian vector { u i } d i =1 is rotationally invariant.

Then we consider a random matrix ˜ W = { ˜ v 1 , · · · , ˜ v n } where ˜ v i follows the same distribution of { u 1 , · · · , u d } . Therefore, it is also equivalent to a random matrix with each element distributed normally. For such a matrix ˜ W , we have from [70] that

where σ max ( · ) and σ min ( · ) denote the largest and the smallest singular value, respectively.

Then we write the matrix W as follows:

which leads to and

We fist assume that for a symmetric matrix A ∈ R n × n λ 1 ( A ) ≥ · · · ≥ λ n ( A ) . Then we introduce the following inequalities for eigenvalues:

Lemma 3 ([57]) . Let G , H ∈ R n × n be positive semi-definite symmetric, and let 1 ≤ i 1 < · · · < i k ≤ n . Then we have that

where λ i denotes the i -th largest eigenvalue.

We first let 1 ≤ i 1 < · · · < i k ≤ n . Because ˜ W ∈ R d × n and Q ∈ R n × n , we have the following:

by applying Lemma 3 to the above equation, we have that

Therefore, we have that

Suppose we have k = 1 and i 1 = n , then Eq. (29) gives

Then suppose we have k = 1 and i 1 = 1 , then Eq. (30) gives

Combining the above results with Eq. (21) and Eq. (23), we have that

which concludes the proof.

Combing with the fact that

we essentially have that

which can be written to the following results:

which shows that under our proposed normalized Gaussian initialization, the maximal and minimal singular values are well bounded by a constant that is only dependent on the size of weight matrix. These results justify the effectiveness of our proposed normalized Gaussian initialization in POET.

D Proofs of Lemma 1

Proof of Lemma 1. We consider an orthogonal matrix R and orthogonal primitives G i corresponding to uniformly random subsets S j ⊂ [ m ] of size b as explained in the main text (see equation (4)). The main claim we need to prove is that given any vector v ∈ R m and a set S ⊂ [ m ] with k ∈ [ m ] we can find an orthogonal primitive matrix G corresponding to the set S such that

Moreover, for all w ∈ R m with w i = 0 for i ≥ k the relation

holds. We can assume that the matrix D ( S ) = { e ( s 1 ) , . . . , e ( s b ) } contains the entries s i in ascending order. Then we write

where ˜ v 1 ∈ R b 1 corresponds to the entries s i with s i < k and ˜ v 2 ∈ R b 2 to the remaining entries, in particular s b 1 +1 = k because k ∈ S . It is well known that for every vector v there is a rotation Q aligning v with the first standard basis vector, i.e., such that Qv = λ e (1) for some λ ≥ 0 . Consider such a matrix ˜ Q for the vector ˜ v 2 and then define the orthogonal matrix

Careful inspection of (4) implies that the last part of (37) is actually true for any ˜ G as the second term has rows with all entries equal to zero for all l / ∈ S . For the first part we find

Here we used s b 1 +1 = k in the last step. Since in addition

for all l ∈ S we conclude that indeed ( Gv ) l = 0 for l ∈ S and l > k , ( Gv ) k ≥ 0 . The remaining statement (38) follows from the observation that when decomposing as in (39) we find

(because w i = 0 for i ≥ k ) and therefore

by definition of ˜ G and we find Gw = w .

The rest of the proof is straightforward by induction combined with a simple coin collector problem. For the rest of the proof it is convenient to reverse the indices, i.e., to consider products G c · . . . · G 1 Assume that we have chosen G i for i ≤ c k and some c k ∈ N such that the product

satisfies P k l ′ ,k ′ = 0 for all k ′ < k and l ′ > k ′ and P k k ′ ,k ′ ≥ 0 for k ′ < k . Let c k +1 ≥ c k + α ( m/b ) 2 ln( m ) . Then, we can bound for any l > k the probability that there is no c k < j ≤ c k +1 such that { k, l } ⊂ S j using that S j follows a uniform i.i.d. distribution by

The union bound implies that with probability at least 1 -m -α +1 there is for all l > k a c k < j ≤ c k +1 such that { k, l } ⊂ S j . If this holds we set G j for c k < j ≤ c k 1 as constructed above if k ∈ S j and G j = 1 m otherwise. This then ensures that

satisfies P k +1 l ′ ,k ′ = 0 for k ′ ≤ k and l ′ > k ′ . For k ′ < k this follows from (38) and for k ′ = k from (37). We conclude by the union bound that P m is an upper triangular matrix with non-negative diagonal entries with probability at least 1 -mm -α +1 = 1 -m -( α -2) . But we also know that P m is orthogonal and therefore satisfies P m = 1 m and we thus find

Next we give a heuristic that actually O (ln( m ) m 2 /b 2 ) terms are sufficient to express every orthogonal map as a product of stochsastic primitives. For fixed c we consider the map

If c ≥ α ln( m ) m 2 /b 2 we have that with probability at least 1 -m -( α -2) for all k, l ∈ [ m ] there is j ≤ c such that k, l ∈ S j . Assume that this is the case. Recall that the tangent space of O ( k ) at the identity is the space of skew-symmetric matrices. Consider a tangent vector ( X 1 , . . . , X c ) with X i ∈ Skew( k ) . Then

This is a surjective map on Skew( m ) under the condition that for all k, l ∈ [ m ] there is j ≤ c such that k, l ∈ S j . We can therefore conclude that the image of Φ contains a neighbourhood of the identity. Moreover, since Φ is a polynomial map, D Φ is surjective everywhere except for a variety of codimension one. While this is not sufficient to conclude that the image of Φ is O ( d ) or dense in O ( d ) it provides some indication that this is the case.

E Experimental Details

This section outlines our experimental setup, including the codebase, datasets, and computational resources used.

Code framework. Our method is implemented on top of the codebase from [82] 1 (Apache 2.0 license), which we also use to reproduce the AdamW and GaLore baselines. We will release our code for reproducing all training results prior to publication.

Training details. We employed the AdamW optimizer [55] for all our training runs. The specific hyperparameters used for each experiment are detailed in the Table 9 and Table 10 referenced below. We use the consine learning rate scheduler with the minimum learning ratio of 0 . 01 . We use the number of warmup steps of 0 , weight decay of 0 . 01 and gradient clipping of 0 . 1 . For the AdamW baseline, we report results for the optimal learning rate from [1 × 10 -2 , 5 × 10 -3 , 1 × 10 -3 , 5 × 10 -4 , 1 × 10 -4 , 5 × 10 -5 , 1 × 10 -5 ] . After each merge-thenreinitalize step, we additionally increase the gradient clipping for 10 training steps to improve training stability.

Model architecture. Our work utilized the Hugging Face Transformers 2 code base to construct the Llama model for pretraining, which is under the Apache 2.0 license. The specific layer setups for the different scaled Llama models are summarized in Table 8. Note, the intermediate dimension of the Feed-Forward Network (FFN) has been slightly modified for the POET-BS, compared to the configs in [82], because the linear layer dimensions have to be divisible by the POET-BS block size b .

Dataset. We use the Colossal Clean Crawled Corpus (C4) dataset [66] for pretraining. The C4 data is a large-scale, meticulously cleaned version of Common Crawl's web crawl corpus. It was

originally introduced for training the Text-to-Text Transfer Transformer (T5) model and has since become a standard pre-training dataset for testing training algorithms for pre-training large language models. The dataset is released under the ODC-BY license.

Compute Resources. All the training tasks are performed on a NVIDIA HGX H100 8-GPU System node with 80GB memory each. Depending on the model scale, we train on 1, 4 or 8 GPUs.

F Implementation and CUDA Acceleration

To enable efficient POET training, we implement the Cayley-Neumann parameterization. To reduce memory usage, we leverage the structure of the skew-symmetric matrix Q ∈ R n × n , where the diagonal entries are zero ( Q ii = 0 ) and off-diagonal elements satisfy Q ij = -Q ji . This structure allows us to store only the upper triangular part of Q as a vector, reducing the number of trainable parameters from n 2 to n ( n -1) / 2 . During the forward pass, Q is reconstructed on-the-fly using a specialized CUDA kernel, significantly accelerating this process. In addition, the Neumann approximation removes the need for costly and numerically unstable matrix inversion, offering further computational gains. Overall, training a 1.3B LLaMA model on a single H100 8-GPU node yields a 3.8× speedup over the baseline ( i.e. , native implementation). Table 11 summarizes the contribution of each component to the overall training time.

G Results of Vector Probing for R and P

In this ablation study, we perform vector probing on the orthogonal matrices R ∈ R m × m , P ∈ R n × n for all linear layers for all blocks of a 60M Llama model trained with POET-FS. The cosine similarity results are reported in Figure 11 and Figure 12, and the trace results are reported in Figure 13 and Figure 14. Since we want to understand the learning dynamics of the orthogonal matrices, we employ b = 1 with POET learning rate of 5 × 10 -4 to eliminate the need for resampling and reinitialization of the orthogonal matrices. Interestingly, we observe this three-phased learning dynamics across different types of linear layers and different-depth transformer blocks.

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.q_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.k_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.v_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.mlp.gate_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.mlp.up_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.mlp.down_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.o_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.q_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.k_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.v_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.mlp.gate_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.mlp.up_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.6

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.mlp.down_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

0.7

0.8

0.9

1.0

Cosine Similarity

model.layers.self_attn.o_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.self_attn.q_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.self_attn.k_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

325

350

375

400

425

450

475

500

Trace

model.layers.self_attn.v_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

800

900

1000

1100

1200

1300

1400

Trace

model.layers.mlp.gate_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

800

900

1000

1100

1200

1300

1400

Trace

model.layers.mlp.up_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.mlp.down_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.self_attn.o_proj (P)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.self_attn.q_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.self_attn.k_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

325

350

375

400

425

450

475

500

Trace

model.layers.self_attn.v_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.mlp.gate_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.mlp.up_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

800

900

1000

1100

1200

1300

1400

Trace

model.layers.mlp.down_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

0

50K

100K

150K

200K

250K

300K

Step

350

400

450

500

Trace

model.layers.self_attn.o_proj (R)

layer 0

layer 1

layer 2

layer 3

layer 4

layer 5

layer 6

layer 7

H Weight Update Evenness of Different POET Variants

To understand the higher parameter efficiency of POET-BS compared to POET-FS, we employ a toy example to visualize their different weight update mechanisms by counting the total number of updates for each element of the weight matrix. The visualization results are given in Figure 15 and Figure 16. Specifically, in this experiment, a 64 × 64 matrix was randomly initialized and trained for 100 steps under various POET-BS and POET-FS configurations. The merge-then-reinitialize trick is performed at each iteration, and the same set of weight elements was effectively updated between two successive merge-then-reinitialize operations. For each weight element, we compute its total number of update in these 100 steps.

Given 100 training steps and updates from both R and P , each element of the weight matrix can be updated at most 200 times. This target is consistently achieved by POET-BS, and it is also agnostic to the block size. All POET-BS variants can enable the maximal number of updates for each weight element to be 200. In contrast, POET-FS results in significantly fewer updates per weight element, with updates also unevenly distributed. This unevenness arises from stochasticity, causing certain weights to be updated more frequently than others. While this is less problematic at large iteration counts, it can introduce unexpected training difficulties in earlier stages.

0

25

50

75

100

125

150

175

200

(a) POET-BS (

b

= 8

)

0

25

50

75

100

125

150

175

200

(b) POET-BS (

b

= 16

)

0

25

50

75

100

125

150

175

200

(c) POET-BS (

b

= 32

)

0

5

10

15

20

25

30

35

40

(a) POET-FS (

b

= 1

/

8

)

0

10

20

30

40

50

60

70

(b) POET-FS (

b

= 1

/

4

)

0

20

40

60

80

100

120

(c) POET-FS (

b

= 1

/

2

)

I Training Dynamics of Singular Values

We conduct an ablation study to compare the training dynamics of singular values of weight matrices between AdamW and POET . The results of AdamW are given in Figure 17, Figure 18 and Figure 19. The results of POET are given in Figure 20, Figure 21 and Figure 22. A 60M LLaMA model was trained for 50,000 iterations with an effective batch size of 512, using both AdamW and POET-FS ( b = 1 / 2 ). The model was evaluated every 5,000 steps, and the singular value dynamics are computed by performing singular value decomposition on the weight matrices. For POET, a merge-thenreinitialize step was applied before each evaluation. Training is finished at 50,000 steps, as the spectral norm of the AdamW-trained model plateaued at this point.

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.0.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.0.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.0.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.0.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.0.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.0.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.1.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.1.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.1.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.1.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.1.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.1.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.2.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

10 1

model.layers.2.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.2.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.2.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.2.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.2.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.3.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.3.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.3.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.3.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.3.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.3.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.4.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.4.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.4.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.4.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.4.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.4.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.5.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.5.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.5.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.5.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.5.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.5.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.6.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.6.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.6.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.6.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.6.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

model.layers.6.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.7.self_attn.q_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

10 1

model.layers.7.self_attn.k_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.7.self_attn.v_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.7.self_attn.o_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

10 1

model.layers.7.mlp.up_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 0

10 1

model.layers.7.mlp.down_proj

0

5K

10K

15K

20K

25K

30K

35K

40K

45K

50K

0

100

200

300

400

500

10 -2

10 -1

10 0

model.layers.0.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.0.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.0.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -2

10 -1

10 0

model.layers.0.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.0.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.0.mlp.down_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.1.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.1.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -2

10 -1

10 0

model.layers.1.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.1.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.1.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.1.mlp.down_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.2.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.2.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -2

10 -1

10 0

model.layers.2.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.2.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.2.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.2.mlp.down_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.3.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.3.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.3.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.3.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.3.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.3.mlp.down_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.4.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.4.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.4.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.4.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.4.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.4.mlp.down_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.5.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.5.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.5.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.5.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.5.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.5.mlp.down_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.6.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -2

10 -1

10 0

model.layers.6.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.6.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -5

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.6.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.6.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.6.mlp.down_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -4

10 -3

10 -2

10 -1

10 0

model.layers.7.self_attn.q_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -2

10 -1

10 0

model.layers.7.self_attn.k_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -2

10 -1

10 0

model.layers.7.self_attn.v_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 -3

10 -2

10 -1

10 0

model.layers.7.self_attn.o_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

6×10 -1

2×10 0

model.layers.7.mlp.up_proj

0

10K

20K

30K

40K

50K

0

100

200

300

400

500

10 0

4×10 -1

6×10 -1

model.layers.7.mlp.down_proj

0

10K

20K

30K

40K

50K

J Orthogonality Approximation Quality using Neumann Series

In this ablation study, we evaluate the approximation error of the orthogonal matrices R ∈ R m × m and P ∈ R n × n across all linear layers in Block 0 of a 130M LLaMA model trained with POET-FS ( b = 1 / 2 ) for 10,000 steps. Figure 23 and Figure 24 show the approximation error over the first 1,000 steps. Since the error difference between k = 3 and k = 4 was negligible, we used k = 3 for better computational efficiency. Empirically, while k = 1 or k = 2 suffices for smaller LLaMA models, larger k values are needed to avoid training divergence caused by exploding gradients due to approximation error.

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.q_proj (R)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.k_proj (R)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.v_proj (R)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.mlp.gate_proj (R)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.mlp.up_proj (R)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.mlp.down_proj (R)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.o_proj (R)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.q_proj (P)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.k_proj (P)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.v_proj (P)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

10 0

Approximation error

model.layers.0.mlp.gate_proj (P)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 5

10 - 4

10 - 3

10 - 2

10 - 1

10 0

Approximation error

model.layers.0.mlp.up_proj (P)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.mlp.down_proj (P)

k = 0

k = 1

k = 3

k = 4

k = 2

0

200

400

600

800

1000

Step

10 - 4

10 - 3

10 - 2

10 - 1

Approximation error

model.layers.0.self_attn.o_proj (P)

k = 0

k = 1

k = 3

k = 4

k = 2

Additionally, Figure 25 shows the orthogonality approximation error of Neumann series with different k over the first 10,000 training steps, illustrating how it decreases as training progresses. We observe a general downward trend in approximation error, indicating improved approximation over time. The results also suggest that using too few Neumann series terms can lead to training divergence in POET.

0

1000

2000

3000

4000

Step

0.0

0.1

Approximation error

model.layers.0.mlp.down_proj (P)

k = 0

k = 1

k = 2

k = 3

k = 4

K Full Results of Training Dynamics

We provide the full training dynamics of different POET variants under Llama 60M, Llama 130M, Llama 350M and Llama 1.3B in Figure 26. This figure is essentially an extended result of Figure 8. One can observe that the training dynamics of POET is quite different from AdamW, and more importantly, POET consistently yields better parameter-efficiency and generalization.

0K

50K

100K

150K

200K

250K

300K

Steps

30

40

50

60

Perplexity

AdamW

POET(BS, 64)

POET(BS, 128)

POET(BS, 256)

POET(FS, 1/8)

POET(FS, 1/4)

POET(FS, 1/2)

(a) Llama 60M

0K

50K

100K

150K

200K

250K

300K

350K

400K

Steps

20

25

30

35

40

45

50

55

Perplexity

AdamW

POET(BS, 64)

POET(BS, 128)

POET(BS, 256)

POET(FS, 1/8)

POET(FS, 1/4)

POET(FS, 1/2)

(b) Llama 130M

0K

50K

100K

150K

200K

250K

300K

350K

400K

Steps

15

20

25

30

35

40

45

Perplexity

AdamW

POET(BS, 64)

POET(BS, 128)

POET(BS, 256)

POET(FS, 1/8)

POET(FS, 1/4)

POET(FS, 1/2)

(c) Llama 350M

0K

100K

200K

300K

400K

500K

Steps

15

20

25

30

35

40

Perplexity

AdamW

POET(BS, 64)

POET(BS, 128)

POET(BS, 256)

POET(FS, 1/8)

POET(FS, 1/4)

POET(FS, 1/2)

(d) Llama 1.3B

L More Results of POET as a Finetuning Method

To demonstrate the applicability of POET to general finetuning tasks, we apply it to finetune a BART-large model [41] on the NLP task of text summarization. Specifically, we evaluate POET on the XSum [61] and CNN/DailyMail [24] datasets, reporting ROUGE-1/2/L scores in Table 12. We note that both LoRA and OFT are designed solely for parameter-efficient finetuning and are not applicable to pretraining. Our goal here is to demonstrate that POET is also effective as a finetuning method. For consistency, we use the same configuration as in the pretraining setup, resulting in a higher parameter count. Experimental results show that POET not only supports finetuning effectively but also outperforms both full-model finetuning and parameter-efficient methods.

Table 12: Finetuning BART-large on XSum and CNN/DailyMail for text summarization. We report ROUGE1/2/L results (higher is better).

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction reflect our work's goal, scope, idea and contribution.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We have included analysis in the paper to discuss our method's limitations.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Yes, we provide a detailed derivation of the theory and reference all theoretical results our method is based on.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Yes, all methods are described in details to ensure the reproducibility of the results. We will also release the code and data.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Yes, the code and data will be released to ensure full reproducibility.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have described the training and testing in details both in the main paper and also in the Appendix section to ensure reproducibility and ease the understanding.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We have described the experiments in detail in the Appendix sections, ensuring reproducibility and soundness of our results.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have described the experiments in the main paper, more focusing on its setting and its results and have included the specific type of compute we used to develop our methods. We have included detailed descriptions in our Appendix section.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Yes, our research conforms with the ethics guidelines of NeurIPS.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our method mainly focuses on methods and its theoretical implications and other interesting findings. Our work does not facilitate down-stream applications that will have a negative impact on the society.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We mainly work on methods, we will release the code but will not release any data or trained model.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We will credit the owners of the code, data and models in our Appendix section, where we describe the implementation more thoroughly. We already included the citations for the assets we used in the main paper. Here are an overview of the most important licenses: Bart (License: apache-2.0), Llama 2 (LLAMA 2 COMMUNITY LICENSE AGREEMENT), Code (apache-2.0), C4 dataset (ODC-BY).

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Our project does not release any new assets (dataset, model weights etc.).

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our project does not involve any crowd sourcing or involve human experiments.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our project does not involve any crowd sourcing or involve human experiments. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: We did not use LLMs to develop our methods but developed it independently.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a training method for LLMs that focuses on preserving the spectral properties of weight matrices to enhance training stability and generalization. The key idea is to reparameterize a weight matrix $\\mathbf{W}$ as three matrices, $\\mathbf{RWP}$, where $\\mathbf{W}$ is a fixed matrix and $\\mathbf{R}$ and $\\mathbf{P}$ are learnable orthogonal matrices. This reparameterization ensures that the singular values of $\\mathbf{W}$ are preserved throughout training. To improve the efficiency, the authors introduce two techniques: structured matrix sparsity (Stochastic Primitive Optimization) and efficient orthogonalization (Cayley-Neumann Parametrization). The experimental results show that the proposed method outperforms baseline methods such as AdamW, LoRA, and GaLore. Furthermore, the authors demonstrate that the proposed method is also effective when applied to fine-tuning.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The proposed method preserves the spectral properties by the reparameterization and controls the spectral properties via matrix initialization. The authors investigate four types of initialization in their ablation studies.\n2. The authors introduce a parameter-efficient sparse orthogonal matrix.\n3. The authors present an efficient approximation method for orthogonalization, reporting a 1.5x speed-up (Appendix E).\n4. The experiments are well-designed, and the ablation studies provide valuable insights.\n5. The paper is clearly written.\n\nWeaknesses:\n1. The pretraining experiments report only perplexity, without any downstream task evaluation. Since pretrained models are typically used with fine-tuning or zero-/few-shot prompting, it might be beneficial to evaluate the pretrained models on downstream tasks as well. \n2. The BOFT paper [49] presents a parameter-efficient sparse orthogonal matrix using butterfly factorization. Could this butterfly method be applied to POET in place of SPO? If so, it could serve as a reasonable baseline to assess the effectiveness of SPO.\n\nQUESTIONS:\n1. I would like to confirm that the parameters updated by the optimizer are $\\mathbf{Q}$ in Eq. (6), not $\\mathbf{G}$ itself.\n2. Could you provide the training time of the proposed method?\n3. Line 417: $R$ and $I$ should be in bold, as they represent matrices.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents POET, a novel reparameterized training algorithm for LLMs that utilizes orthogonal equivalence transformation. By reparameterizing each neuron with two learnable orthogonal matrices and a fixed random weight matrix, POET preserves the spectral properties of weight matrices, stabilizing the optimization of the objective function and enhancing generalization. The authors also develop stochastic principal submatrix optimization and the Cayley-Neumann parameterization to improve POET's scalability. Experiments demonstrate POET's effectiveness across multiple model scales, including LLaMA models up to 1.3B parameters and BART models fine-tuned for summary generation.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThe paper is technically sound with a solid theoretical foundation. The claims regarding POET's preservation of spectral properties are well-supported by experiments across multiple model scales and architectures. The authors demonstrate careful and honest evaluation of their work's strengths and limitations. The paper is clearly written and well-organized, offering a good understanding of the proposed method. It provides sufficient detail for expert readers to reproduce the results. The results are impactful for the LLM training community. POET addresses the challenge of training large models more stably and generalizably, making it likely for others to use or build on these ideas. The work provides new insights by leveraging orthogonal equivalence transformation for LLM training. It introduces a novel combination of existing techniques with clear articulation of the reasoning behind this combination.\n\nWeaknesses:\nThe experiments only validate POET on models up to 1.3B parameters, leaving its effectiveness on significantly larger models unverified. The computational complexity analysis is limited, making it difficult to assess resource consumption in large-scale training scenarios. Some sections could benefit from additional clarity, particularly in explaining the Cayley-Neumann parameterization and its potential numerical stability issues. The significance is somewhat limited by the lack of validation on larger models. Addressing this would further enhance the paper's impact.\n\nQUESTIONS:\n1. What are the bottlenecks in applying the current method to larger models such as those with 7B or 70B parameters?\n2. What is the impact of the Cayley-Neumann iteration on training time and efficiency?\n3. Does the Cayley-Neumann parameterization method encounter numerical stability issues when handling large-scale orthogonal matrices?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to train deep learning models with a novel method POET based on orthogonal reparameterization. Specifically, it periodically reinitialize the weight matrix $W$ to $RWP$, where $R$ and $P$ are orthogonal matrices. To further improve efficiency, this paper considers two distinct strategies—Full stochastic SPO and block-stochastic SPO—to break down the full orthogonal matrix into smaller pieces, and make them share parameters to further improve parameter efficiency. To improve numarical stability, the authors use a truncated Neumann series to approximante the matrix inverse operation inherited in the Cayley-Neumann parameterization for constructing orthogonal matrices. By periodically merging $R$ and $P$ into $W$ and reinitializing $W$ to $RWP$, the operator norm can be effectively controlled to guarantee such approximations. Besides, the paper has conducted experimental studies on pre-training LLaMA models, exhibiting advantages on both efficiency and training performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\ni) The technical details, particularly parameterization details and training procedures, are novel and interesting.\n\nii) The potential of this method is promising. With fewer trainable parameters, the method can achieve comparable performance and superior training efficiency.\n\nWeaknesses:\n\ni) The trade-off between training efficiency and training performance is not super clear. Although POET are more parameter-efficient, most training curves in the paper have shown that Adam is still much faster (2~3 times, roughly) than POET in the initial training stages. Consequently, it is unclear if we should use Adam to train in the early stages for faster convergence.\n\nii) Technical details missing. Maybe I missed it but it remains unclear to me whether the optimizer needs restart whenever the weights are merged and reinitialized.\n\niii) The experimental setup is insufficient. First, the models used are not bigger than 1.3B. It is recommended to include 7B results at least for fine-tuning tests. Second, given the motivation to enhance generalization ability by controlling the spectral norm, I expect to see more results illustrating the generalization ability of POET-trained models, which typically requires accuracy scores on downstream tasks such as GLUE.\n\nQUESTIONS:\nWhat is the major conclusion drawn from the empirical observations of the three learning stages? How can it help with developing the POET method, or justifying the use of it?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-08"}

comments: SUMMARY:
This paper is  concerned with stable and efficient training of LLMs by means of imposing orthogonality on  weights matrices. The orthogonal weight training for DNNs is already used in the literature to  improve generalization. The  two initilization schemes for orthogonal weight training is novel. Numerical results indicate that, when AdamW is trained with nearly three times more tokenswithout repeating any and using a well-tuned learning rate—the improvement is significant and not just due to more training steps. Additionally, results show that POET's performance scales with parameter budget: larger models consistently perform better, which has key implications for scaling laws.

STRENGTHS AND WEAKNESSES:
The Streght of  this paper is that,  extensive  numerical simulations are provided  to demonstrate the  effectiveness of orthogonal training in LLMs. Weaknesses are:
- The inductive bias of POET mentioned on page  7 is not clearly explained. As a reader I can not determine what is the inductive bias the POET has. This paragraph needs ellaboration or mathematical proof.
- No mathematical proof is  provided on  how the  SPO and CNP initialization affect the generalization bound. Moreover, the  main generalization bound used in Th.1 is for MLP and not an LLM.

QUESTIONS:
The connection of  spectral norm of weights  to generalization bound that you used in (7) is for MLP, is the connection to attention-based  architecture  studied in the literature or a direct result of  (7), please  clarify.

LIMITATIONS:
The proofs in the supplementary meterial  only cover gaussian  initialization and SPO and Cayley-Neumann parametrization  are not proved to be affecting the generalization,
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-07-22