Reparameterized LLM Training via Orthogonal Equivalence Transformation
Authors: Zeju Qiu
Simon Buchholz
Tim Z. Xiao
Maximilian Dax
Bernhard Schölkopf
Weiyang Liu
Venue: NeurIPS 2025 poster ()
Keywords: large language models
pretraining
orthogonality
Primary area: deep_learning

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
Zeju Qiu Simon Buchholz' Tim Z. Xiao Maximilian Dax! Bernhard Schölkopf! Weiyang Liu
Max Planck Institute for Intelligent Systems, Tübingen
Thee Chinese University of Hong Kong
Abstract
While Large language models (LLMs) are driving the rapid advancement of ar- tificial intelligence, effectively and reliably training these large models remains one of the field's most significant challenges. To address this challenge, we propose POET, novel reParameterized training algorithm that uses Orthogonal Equivalence Transformation to optimize neurons. Specifically, POET reparame- terizes each neuron with two learnable orthogonal matrices and fixed random weight matrix. Because of its provable preservation of spectral properties of weight matrices, POET can stably optimize the objective function with improved general- ization. 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.
Introduction
Recent years have witnessed the increasing popularity of large language models (LLMs) in various ap- plications, such mathematical reasoning [13] and program synthesis [3] and Docision-making [77]. Current LLMs are typically pre-trained using enormous computational resources on massive datasets containing trillions of tokens, with each training run that take months complete. Given such huge training cost, how to effectively and reliably train them poses significant challenges.
The de facto way for training I L.MSs is to directly optimize weight matrices with the Adam opti- mizer [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 regular- ization 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 a weight matrix (i.e., the largest singular value) provides upper bound on how much matrix amplify the input vectors, which connects to the generalization properties. In general, smaller spectral norms (i.e., better smoothness) 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?
considered be associated with stronger generalization, which inspires
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 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 weight matrices. Even with power iteration, this still adds significant overhead to the training process, especially when training large neural networks. Additionally, spectral regularization does not scale efficiently with increasing model size.
Prooecct lead Corresponding author
Project page: spherelab. a.c.b.ii/poo ai/po
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

To achieve effective weight spectrum control without the limitations above, we propose POET, reParameterized training algorithm that uses Orthogonal Equivalence Transformation to indi- rectly learn weight matrices. Specif- ically, POET reparameterizes weight matrix W E Ruxx with RWoP where Wo E R is randomly ini- tialized weight matrix, R E Rm*x and e R=*" are two orthogonal ma- trices. Instead of optimizing weight matrices directly, POET keeps the domly initialized weight matrix W, unchanged during training and learns two orthogonal matrices R, P transform Wn- This reparameterization preserves the singular values of weights while allow- ing flexible optimization of the singular vectors. POET effectively addresses the above limitations:
Ste
Stagaber vake indes Oatted as dearraing order)
Singsie vakae rndes (oeted decenczyy endr)
(4) Standad trxining
(b) POET
Figure I: Training dynamics singular values of the same weight matrix a.aMA model. Standard training on the left strictly follows the common practice for training LLMs (direct optimization with AdamW). POET on right the proposed upprosimation for large scale LLM training. The slight (almest negligible) singular value changes in POET are numerical and upproximation error.
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 and Appendix 1). 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 expen- sive, 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 E Rmmxm. SPO factorizes 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 implemented as a memory-efficient iterative algorithm that sequentially updates one primitive orthogonal matrix at time. To improve the expressiveness of the sequential factorization, we adopt a merge-then-reinitialiee 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-Neumand mation addresses how to maintain orthogonality without introducing significant computational overhead. To achieve this, we develop the Cayley- Neulman parameterization (CNP) which approximates the Cayley orthogonal parameterization [48, 65] with Neumann series. Our merge-thee--reinitiliiee trick can effectively prevent the accumulation of approximation errors.
parameterizaaton: The second-level approxi-
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 fiexibility 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 o by computing v Re which corresponds to the cosine similarity between Ro and v. By inspecting the cosine similarities of seven orthogonal matrices throughout training, we

---

observe that the learning pro- cess can be divided into three
T1ailing 1000
Phase f: Conkcal shell scandhing
Plane II:
Phuse III:
Stabie loamiise
Final
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 ob- serve consistently across all learnable orthogonal matri-
the comical shell
adjialing
Plane searciing
126
285, 280
Phase HC: Soiie burning Plane Mujualigg
Tocotion numbbr
Itention lumber
Figure 2: POET's three learning phases. Left. illustration; Middle: angle; Right: loss and validation.
ces. This suggests that R transforms into 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 in-depth discussion and empirical results in Appendix A,G. Our major contributions
We introduce POET, 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 parame- terization 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 framework to train neural networks by learning a layer-shared orthogonal transformation for neurons in each layer. For weight matrix W = {WI,-- w,, E Rm** where w; E " is the i-th neuron, the layer's 's forward pass is = W x with input x E I and output y E R". Unlike standard training, which directly optimizes W, orthogonal training keeps W fixed at its random initialization Wo w :.:...:: w and instead learns an orthogonal matrix Re E Rmxm to jointly transform all neurons in the layer. The forward pass becomes y = (RW)) x. The effective weight matrix is Wn (wff, preservation of hyperspherical energy. With WZ; = w: |wal]. orthogonal training ensures
w where wrf Rw, A key property is its
HE(Wo) = Ele =
Σ || Rii: Rů²l HE(W ()
(I)
it
where hyperspherical energy HE(.) characterizes the uniformity of neurons on a unit hypersphere. Prior work [47-49, 76] has shown that emergy-preserve training can effectively improve general- ization. Orthogonal finetuning (OFT) [51, 65] also demonstrates that finetuning foundation models while preserving hyperspherical energy achieves favorable trade-off between efficient adaptation to downstream tasks and retention of pretraining knowledge. While the hyperspherical energy preserva- tion 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 aagular-valuee spectrum instead. By inherently maintaining the spectrum, eneryy-preservie training training. As generalization, spectrum-preserviie training T Rmmxn - R m x that preserves the spectrum:
is a special case of learns a transformation
{ai (W (TWW)))) ... (m) (T(W))) = (ai (2o).. (.)iaan)))) (W))). (2)
where o, (W)) denotes the i-th singular value of Wo (sorted by descending order with o₁ being the largest singular value). How we instantiate the transformation T results in different algorithms. Generally, is pecctrum-preservii
map, and can be either linear [42] or nonlinear [5]. we only consider T to be linear map, then Theorem 1 can fully characterize the form of T:
Theorem 1 (Simplified results from [42]). For linear map T R + mm (m + ni), if (T))) = 01 (W) always holds for all W E RR,,, then the linear map T must be of the following form: T(W)= R.P.J. for all W Rm where R mm and P Boxn are some fixed elements in orthogonal groups (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 ppectruumppeseerviei training), where the map simplifies to T(W)= RW I, with I as the identity.
3

---

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, RW P 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 specific implementation of
training. Inspired by Theorem 1, it parameterizes the prrerrriing transformation 7 using a left orthogonal matrix that transforms the column space of the weight matrix and right orthogonal matrix that transforms its row space.
3.1 General Framework
Following the general form of reparameterizes the neuron as P. where Wo E R' is a randomly initialized weight matrix that remains fixed during training, and R E Rm P E R are trainable orthogonal matrices. This reparameterization effectively applies an orthogonal equivalence transformation (OET) to random weight matrices. Specifically, OET a double-sided transformation, defined as OET(W; R, P) = RWP P, 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
linear maps discussed the last section, POET
Win (R x. s.t. {R' R = RR I, P P PP =I}.
where R and P can be merged into single weight matrix WRP RW, 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 simultancously transform- ing their left singular vectors and right singular vectors while keeping the singular values unchanged. Given the singular value decomposition (SVD) Wo UE, V the reparameterized neuron weight matrix becomes Wn = RU V T P where both RU and V are orthogonal matrices. This effectively constitutes an SVD of WRP. It is also straightforward to verify that the spectral properties of RP remain identical to those of the initial matrix Wo-
Neuron initialization. Since POET preserves the spectral proper- ties of the initial weight matrix Wo. 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,
1.5
Usitiom qucins
Xivia
0.4
1000 700
300
100
901
initialization, applies SVD to standard initialization and sets all singular values to 1, balancing spectral properties throughout train- ing. The second, normalized Gaussian initialization, normalizes neurons drawn from a zero-mean Gaussian with fixed variance. This is directly inspired by prior work showing that normalized neurons improve convergence [48, 50, 52]. To ensure that the repraaepareeeeeriiie network is statistically equivalent to a standard network at initialization, we always initialize both orthogonal matrices as identity matrices.
Singaber valaa index (scned descending onter)
Figure 3: Singular values of a weight matrix size 512 1376, randomly generated by different inttiaizaaionn schemes.
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.., R I G, where G; denotes the i-th primitive orthogonal matrix. While each G; of the same size as R, it is parameterized by significantly fewer degrees of freedom. See Figure 4

---

Givens rotation
(b) Householder ldr reflection
SPO (fully stochastic)
SPO (block-stochsste)
Figure Examples the primidive orthogonal transformation matrix different ortogonalizations (two examples for each method), Note that, bluc blocks represent 1. light purple blocks denote O and deep parple blocks are the actual orthogonal parameterization to be learned.
for an illustration. Both Givens rotation and Householder reflection use relatively low-capacity parameterizations-for example, each Givens rotation G; involves only single effective parameter- which limits their efficiency in representing the full orthogonal matrix. SPO follows 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 tunable hyperparameter and adopts stochastic sparsity pattern.
Fully stochastic SPO. The basic idea of fully stochastic SPO is randomly sample a small submatrix and enforce its orthogonality, allowing it to be easily extended to full orthogonal matrix by embedding it within an identity matrix-a a process similar to Givens or Householder transformations. To represent a large orthogonal matrix R E m we start by defining index sets S = {si.... . S (1. (j E [1, where each set has cardinality |S = b, a hyperparameter controlling the number of effective parameters a primitive orthogonal matrix. Si vj are randomly sampled from the full indices (1,.... .m.. Let G2 E R be a small orthogonal E mm be selection matrix, where e(k) is the standard
matrix, and D(S)) = {e(s)),.... )........ e(s))) basis vector with 1 in the k-th position and 0 elsewhere. The factorization given by
R = II Im + D(S)) (G. - - ))
S.L. G. G.G. = To. Vi,
(4)
s=1
Gy: The i-th primitive orthogoonal matrix
where D(S"). (A) D(S)) T is projector that replaces the x b sub-block with A. I, and I, are identity matrices of size m x m and b x b, respectively. To efficiently parameterize small orthogonal matrices G.,, 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 example. To address this, we propose block-stochasttic SPO, which first constructs a block-diagona 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
R = II C.. -Diag((", C.... G. T v.). s.t. (G)) (C) To, Vi,j. 1=1
(5)
G. The i-th primitive orthogonal matrix
where G E b b is the j-th block of the block diagonal matrix, and D., Vi are all random permuta- tion matrices. As long each diagonal block G/ is an orthogonal matrix, both G; and R are also orthogonal matrices. We also use CNP to efficiently parameterize each orthogonal block G².
The thereenttie-e 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- then--einitiillie 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 new index set (in fully stochastic SPO) or generate new permutation (in block-stochastic SPO). This trick allows only one primitive matrix to stored at 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 QI - (D(I-Q)- where Q is a skew-symmetric matrix satisfying Q Q- T minor caveat of this parameterization is that it only produces orthogonal matrices with determinant I (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
5

---

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:
R = (T + Q)=
x++ 2
(1+ Σ Q) (6)
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 identity, we set Q to zero matrix in CNP, satisfying the convergence condition initially. As the training progresses, however, updates to Q may cause its operator norm to exceed violating this condition. Fortunately, our merge-then 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 <- Wo
Step 2: Orthogonal matrix initialization. For fully stochastic SPO, we randomly sample an index set S, and parameterize G R E bxb and Pp E b using CNP (Equation (6)). Both matrices initialized as identity, so R and P also start as identity matrices. For block-stochastic SPO, we sample random permutation matrix R. L p. and parameterize { R. G R } and (Gp..... G.
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 = Im (G - D(S) and P = I, + D(S) (Gp D(S) For block-stochastic SPO, we G R n and P V Diag Gpp............... G. Dr..
have R V Diag
Step 4: Inner training loop for updating orthogonal matrices. The equivalent weight matrix in the forward pass is RWWP P. Gradients are backpropagated through R and P update GR, Gp (fully stochastic) G² G/p, Vi (block-stochastic). This inner loop runs for a fixed number of iterations.
Step 5: Mergeethen--reinittaliie 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 intro- ducing a hyperparameter as the sampling budget, fully stochastic SPO decouples parameter com- plexity from the size of the weight matrices. With a small b, POET becomes highly parameter-efficient, though at the cost of slower convergence. This of- fers users flexible trade-off between efficiency and speed. In contrast, block-stochastic SPO has parameter complexity dependent 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 suitable sampling budget b. A comparison of parameter and memory complexity is given in Table 1.
Method
# trainable params
Memory cost
AdamW
man
3mn
GaLore [82]
mn
mn + mr + 2nr
POET (FS)
bib - 1)
wn + 36(6 1)
POET (BS)
(mm + +) 1) mn + (m n) 1)
Tahle I: Comparison to existing methods. Assume W E Rm
(m I n)) GaLore with rank and POET with block size 6. FS denotes fully stochastic SPO, and denates block strchastcc SPO.
Performance under 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 evaluation metric. The total parameter budget matches that of fully stochastic POET with m for R and b = n for P, where = 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.255 0.5 (i.e., standard POET), 0.25 75, 0.1 :0.9, and 0: 1. Results in Figure show that POET with balanced allocation between R and P yields the best performance.
pepteess
Searl hodge Larse tradet
12
12
P.0 P.8. P.0
P.05 P.035 F.60 Fl
Figure 5: Performance of POET under constant total parameter budget on R. P.

---

SE00H FOET
0.00
3.00
5000
0.0
0.00
036
0.55
386 38, J6,
185. 386
488 450
56 I8K
38K 48K
(4) Averag K projection anyes
Avvegee over all projection layes
(4) Avesuge proeetion layezs
035
0.0
05
1005
POST
4155
090
6080
600
96.
10K.
156.
156. 336 596
ESC 96
18. 10K. 1SE 336 750
100 550 50
58 108
152
330
beyor
Avengeoore projecson byom
Avange oxtpes projection leyers
Figure 6: Dynamics comparison of uveraage singular value eniropy (singular value divensiy) between direct training (AdamwW Muon) and POET.
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-coe 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-normalizes margin analysis [6] bounds the misclassification error in terms of margin-based training loss and a complexity term. The complexity term is proportional to (y(n)) where y 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
Q = (TiIaa)) (
(vdllW 12)
3/2
(7)
where ||-|] and || 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 E Rddd satisfies
|| = O(1) and || = O(vd) so that Q = (d)
Approximation properties of SPO. We have seen in Theorem 1 that the factorization RW P 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 2 am lm(m)(m/b) for some a 0 then with probability least 1 (a-2) over the randomness of the index sets S" we can express any orthogonal matrix R as product of e primitives G; as in Eq. (4). Moreover, the orthogonal matrix G; depends only on the sets S and matrices Gi selected in earlier steps.
The proof of this lemma can be found in Appendix D. The result extends to Block- stochastcc SPO as this is strictly more expressive than fully stochastic SPO. The key idea of the proof similar to the factorization of orthogonal matrices into product of Givens rotations. Indeed, by multiplying R with properly chosen primitive matrices G; we can create zeros below the diagonal for one column after another. Note that each G; has b(b 1bb11)/2 parameters while R has m(m 1)/2 parameters, which implies that generally at least (((m/b)² primitives are necessary. In Appendix D we also provide heuristic that with high probability for e (In(m)(m//))) every orthogonal matrix can be written as product of e orthogonal primitives G,.
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 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
7

---

entropy of the trained Llama- 60M model at different iteration. Specifically, the SVD entropy. defined log E measures the diversity of singular values; higher entropy indicates 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.
as (o) - Ing E. 0j
POET minimizes hyperspherical energy. While POET generalizes mnergy-preservinn training to spectrum- spectrum-preserving training. it still ensures low hyperspherical energy when initialized with a zero- mean isotropic Gaussian distribution. This 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 cero-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, OE:(WW R, P) RWP does not change the distribution of W, i.e., RWF P =d W. We give a derivation in Appendix B. Therefore, the transformed neurons of the new weight matrix RW P 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 at different iterations in Figure 7. We compare the sum of the hyperspherical energy of all the layers trained by Adam W, Muon or POET. The results again verify POET's energy-minimg minimiznng 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.
-264
Adaa
Figure 7: Hyperspherical energy compari- AdamW, Muon and POET.
Experiments and Results 5
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 LLM Pretraining using LLaMA Transformers
We perform the pretraining ex-
Model (# tokens)
60M (30B)
130M (40B)
350M (4UB)
1.3B (50B)
periments on the Llama trans-
AdamW
26.68 (25.30M)
20.82 (84.9RM)
16.78 (302,
14.73 (1.218)
formers of varying sizes (60M,
Galore
29.81 (25.30M)
22.35 (84.93M)
17.99 382)58M)
18.33 (1.21R)
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 POET- FS, =bm for R and bon for P. We compare POET against GaLore [82]. 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.
LoRA,
39.70 (4.85M)
32.07 (11,21M)
25.19 (30.28M)
20.55 (59, 38M)
POETOS heai
(2.39M)
24.52 (5.52M)
20.29 (14.90M)
18.28 (29.22M)
POETaS Nerzs
(4.81M)
21.86 (11.12M)
18.05 (30.0M)
16.24 (58.91M)
POETES Nezw
25.29 (9.663)
19.88 (22.33M)
16.27 (60.32M)
14.56 (118.26M)
POETTx.aa
(0.53M)
29.67 (1.78M)
(6.34M)
18.46 (25.39M)
POET hats
28.69 (2.13M)
23.55 (7.13M)
19:42 (25.44M)
17.60 (101.66M)
POET ToETix
25.37 (8.54M)
19.94 (28.56M)
15.95 (101.86M)
13.70 (406.88M)
Table 2: Comparison of POET with popular pretraining mothods using different sizes of LLaMA models. Validation perplexity and the number of trainuble purameters reported.
Table shows that both POET-FS (1=1/2) and POET-BS (I=256) consistently outperform both GaLore and AdamW with significantly fewer parameters. For LLaMA- POET-FS (1=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-bes performance
8

---

(14.56), which still surpasses AdamW with only one-tenth of AdamW's train- able parameters. Similar patterns can be observed for models of smaller sizes. Moreover, we compare the training dynamics between AdamW W and POET in Figure 8. The training dynamics of POET is quite different
Medi
PITTEES 290
PCEDIL
1N
60%
140
Figure Validation perplexity dynamics on I.aMA-150M LLaMA- 1,3B.
from AdamW. W. After an initial rapid drop in perplexity, POET improves more slowly than Adam W. 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 could offer valuable insights into large-scale model training.
To highlight POET's non-trivial performance improvement, we in- crease 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 W is trained with almost triple the number of tokens. Re-
Adenti
POET with 2500 stpe
Adxaa rupe
sults are given in Figure 9. In this experiment, the Adam W learning rate was carefully tuned for the full training run, and no training tokens were repeated. Thus, the improvement is non-trivial and can- not 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- 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 cerammetrr-fffcienn 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 better coverage of weight parameters. Specifically, the -lochsstcccsttic 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.
cos is defadi paneny lars
Staps
Figure 9: Validation perplexity dynamics of POET (FS. b=1/2) and Adam on Llama- 60M, POET outperforms the A trained with almost twice the number of seen tokens.
Pretraining Llama-3B. To demonstrate the scalability of POET, we apply POET pretrain Llama with 3B parameters while reusing the same hyperpa- rameters 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.
Method AdamW
Perplexity
19.61
PORT
16.90
Table Validation perplex- ity of training LLMMA 3B model on billion tokens.
5.2 LLM Finetuning on Downstream Tasks
To better evaluate models beyond the validation per- plexity, we show the results of finetuning the trained model on the GIUE bench- mark [75]. This benchmark provides comprehensive assessment of a model's
FT
Model
CoLA
MNLI
MRPC
QNLI
QQP
RTE
SST-2
STS-B
AdamW
0,361
0.658
0.696
0.818
0.829
0.534
0.914
0.880
Full FT
POET
0,523
0
0.824
0.885
0.902
0.661
0.920
0.873
Adam 0,388
0.774
0.689
0.842
0.867
0.531
0.915
0.762
OFT
POET
0.437
0.812
0.740
0.924
0.791
Adam 0.435
0.804
0.806
0.856
0.889
0.653
0.904
0.878
POET
PORT
0.505
0.821
0.826
0.892
0.902
0.682
0.931
0.887
Table 4:
Downstream performance on GLUE between AdamW- and POET-pretrained models.
language understanding capabilities through 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 evalua- tion 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 -retrrrrta

---

model consistently outperforms the Adam p-rrtrriieed baseline across all tasks and finetuning meth- ods. 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 initial- ization 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.
Scheme
Perplexity
Standard
26.22
Xavier
spectrum
27.29
Normalized
25.37
Table Performance of dif- ferent initializations.
Tm
Perplexity
prrge-then-reinitaalie frequency. The proposed merge- thin--ennitiiallze trick allows POET to train only 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 Tm. which determines how often the orthogonal matrix is merged and reset the identity. The index set in POET-FS and the permutation matrix in POET-BS are also resampled at each reinitialization. Therefore, is quite important understand how this hyperparam- eter Tm, affects performance. Following the previous initialization experiment, we use POET-FS with 1=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 Tm 400 in all experiments by default.
30.29
25
27.27
50
25.99
200
25.37
400
25.31
1600
25.58
Table 6: perplex- ity different Tm.
Scheme
Perplexity
Neumann series approximation. CNP approximates the matrix inverse using 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 training divergence, highlighting the critical role of maintaining orthogonality. To balance overhead and performance, we find that 3 Neumann terms (k- 3) reach good trf
I: =
Not converged
& =
22.56
21.54
20.22
20.19
Table 7: Number of terms
in Neumann series.
Additionally, we evaluate the accu- racy of the Neumann approximation to understand how the number of Neumann terms affects the orthog- onality. The orthogonal approxi- mation error is defined by 6anh
122
Sknp
284
120
1003
0.53
2.0
3.33
1000
(a) Onthogonaa moztN
(33 Orthoggonal matrix P
IRR" - /| select a weight matrix and compute the approximation error of corresponding orthogonal matrices R and P. For clarity, we visualize the error 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-reinitaaliee trick can periodically reset the error. More results are given in Appendix L
rr. We randomly Figure 10: Approximation error of orthogonal matrices R and P of weight matrix.
6 Related Work, Concluding Remarks and Acknowledgement
Related work. Inspired by low-rank adaptation methods such as LoRA [26], a number of recent approaches [11, 20, 27, 30-32, 43-45, 53, 58, 71, 81, 82] have explored low-rank structures to enable efficient pretraining of large language models (LLMs). parallel, sparsity has also been extensively studied as a means improve training efficiency in neural networks [9, 14, 25, 72, 78). Compared to approaches that exploit low-rank structures, relatively few works have explored sparsity for pretraining. Our work broadly aligns with the sparse training paradigm, POET leverages
10

---

sparsely optimized orthogonal matrices to enhance training efficiency. A parallel line of research [34, 46, 60, 69. 80] focuses on developing efficient optimizers for large-scale neural networks. While our work also targets efficient training of large models, it is orthogonal to these efforts, as POET can be integrated with any optimizer. The way POET uses orthogonal matrices to transform neurons may also relate to preconditioned optimizers such as Muon [34], Shampoo [19] and SOAP [73], as well as to the broader field of manifold optimization (e.g.. [7]). POET-trained weight matrices remain statistically indistinguishable from randomly initialized ones due to the isotropy of zero- mean independent Gaussian distributions. This yields interesting connections to random neural networks [21, 40, 40, 62, 74], random geometry [22], and random matrix theory [16].
Concluding remarks. This paper introduces POET, a reparameterized training algorithm for large language models. POET models each neuron as the product of two orthogonal matrices and fixed random weight matrix. By efficiently learning large orthogonal transformations, POET achieves superior generalization while being much more parameter-effiiien than existing LLM pretraining methods. Experiments show that PÊET broadly applicable to both pretraining and finetuning tasks.
Acknowledgement. The authors would like to sincerely thank Lixin Liu, Han Shi, Gege Gao, Zhen Liu and many colleagues at Max Planck Institute for Intelligent Systems for many helpful suggestions. Additionally, the authors also sincerely thank all the anonymous NeurlPS reviewers for their constructive suggestions that greatly improved the quality of our work.
References
[1] 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
[2] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In NeurlPS, 2019. 7
[3] Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quac Le, et al. Program synthesis with large language models. arXiv preprint XXiv:2108.07733 2021. 1
[4] Nitin Bansal, Xiaohan Chen, and Zhangyang Wang. Can gain more from orthogonality regularizations in training deep networks? In NeurlPS, 2018.
[5] Line Baribeau and Thomas Ransford. Non-linear spectrum-preserving preserving maps. Bulletin of the London Mathematical Society, 32(1):8- 14, 2000. 3
[6] Peter I. Bartlett, Dylan J Foster, and Matus J Telgarsky. Spectrally-normalizeed margin bounds for neural networks. In NeurlPS, 2017. 1.7
[7] Silvere Bonnabel. Stochastic gradient descent on ricmannian manifolds. IEEE Transactions Automatic Control, 58(9):2217-222333 2013. 11
[8] Djalil Chafaı, Djalil Chafa, Olivier Guédon, Guillaume Lecue, and Alain Pajor. Singular values of random matrices. Lecture Notes, 13, 2009. 20
[9] Beidi Chen, Tri Dao, Fric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In NeurlPS, 2021. 10
[10] Tianlong Chen, Zhenyu Zhang, Yu Cheng, Ahmed Awadallah, and Zhangyang Wang. The principle of diversity: Training stronger vision transformers calls for reducing levels of redundancy, In CVPR, 2022.
[11] Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue, Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of Ilms under low-rank constraint? arXiv preprint arXiv:2410.01623, 2024.
10
[12] Moustapha Cisse, Piotr Bojanowski, Edouard Grave, Yann Dauphin, and Nicolas Usunier. Parseval networks:
Improving robustness to adversarial examples. In ICML, 2017.
[13] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jenry Tworek. Jacob Hilton, Reichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint 10.14468 2021.

---

[14] 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
[15] Tri Dao, Alhert Gu, Matthew Fichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. ICML, 2019. 10
[16] Alan Edelman and N Raj Rao. Random matrix theory. Acta numerica, 14:233-297, 2003... 11
[17] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedffowwrr neural networks. In A/STATS, 2010. 4, 10, 19
[18] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization matrix factorization. NeurlPS, 2017. 7
[19] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In ICML, 2018.
[20] Andi Han, Jiaxiang Li, Wei Huang, Mingyi Hong, Akiko Takeda, Pratik Jawanpuria, and Bamdev Mishra. Sltraim: a sparse plus low-rank approach for parameter and memory efficient pretraining. arXiv preprint arXiv:2406.00 02214, 14, 2024. 10
[21] Boris Hanin. Random neural networks in the infinite width limit as gaussian processes. The Annals of Applied Probability, 33(6A):779--481 2023. 11
[22] Boris Hanin and David Rolnick. Complexity of linear regions deep networks. In ICMI,, 2019. 11
[23] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In ICCV, 2015. 19
[24] Karl Moritz Hermann, Tomás Kočiský. Edward Grefenstette, Lasse Espeholt. Will Kay, Mustafa Suleyman. and Phil Blunsom. Teaching machines to read and comprehend. 2015. 40
[25] 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
[26] Edward 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 /CLR, 2022. 9, 10
[27] Jia-Hong Huang. Yixian Shen, Hongyi Zhu, Stevan Rudinac, and Evangelos Kanoulas. Gradient weight- normalized low-rank projection for efficient Ulm training. In AAAI, 2025. 10
[28] 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.
[29] Tianjin Huang. Ziquan Zhu, Gaojie Jin, Lu Liu. Zhangyang Wang, and Shiwei Liu. Spam: Spike-aware adam with momentum reset for stable IIm training. 2025. 8
[30] Weihao Huang, Zhenyu Zhang, Yushun Zhang, Zhi-Quan Luo, Ruoyu Sun, and Zhangyang Wang. Galore- mini: Low rank gradient leaming with fewer learning rates. In NeurlPS Workshop Modern Machine Learning: Principles and Scalabiliiy, 2024.
in
[31] Minyoung Huh, Brian Cheung, Jeremy Bersstein, Phillip Isola, and Pulkit Agrawal. Training neural networks from scratch with parallel low-rank adapters. arXiv preprint arXiv:2402. 16828, 2024.
[32] 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 2xXiv:2407.12330 2024. 10
[33] Haoming Jiang, Zhchui Chen, Minshuo Chen, Feng Liu, Dingding Wang, and Tuo Zhao. On computation and generalization of generative adversarial networks under spectrum control. In ICLR, 2019.
[34] 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
[35] 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
12

---

[36] Nitish Shirish Keskar and Richard Socher. Improving generalization performance by switching from adam to sgd. arXiv preprint arXiv:/712.07628, 2017. 1
[37] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[38] Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. nolf-norm neural networks. In NeurIPS, 2017. 1
[39] Hojoon Lee, Youngdo Lee, Takuma Seno, Donghu Kim, Peter Stone, and Jacgul Choo. Hyperspherical normalization for scalable deep reinforcement learning. arXiv preprint arXiv:2502. 152000 2025.
[40] Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohk-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1111.0165,, 2017. 11
[41] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, AbdeIrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising -e-euucec--- pre-training for natural language generation, translation, and comprchension. arXiv preprint arXiv: 13461, 2019. 40
[42] Chi-Kwong Li and Nam-Kiu Tsing. Linear operators preserving unitarily invariant norms of matrices.
Linear and Multilinear Algebra, (0(1-2):119-1333 1990. 3
[43] Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Relora: High-rank training through low-rank updates. arXiv preprint arXiv:2307.05695 2023. 10
[44] Kaizhao Liang, Bo Liu, Lizhang Chen, and Qiang Liu. Memory- Ilm training with online subspace descent, arXiv preprint (rXiv:2408.1285 2024.
[45] Xutao Liao, Shaohui Li, Yuhui Xu, Zhi Li, Yu Liu, and You He. Galore+: Boosting low-rank adaptation for lims with head projection. arXiv preprint 2i.::2212.1992000 2024. 10
[46] Jingyuan Liu, Jianlin Su, Xingeheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for IIm training. arXiv preprint arXiv:2502 16982, 2025. 7, 8, 11
[47] Weiyang Liu, Rongmei Lin, Zhen Liu, Lixin Liu, Zhiding Yu, Bo Dai, and Song. Learning towards minimum hyperspherical energy. In NeurlPS, 2018. 1,
[48] Weiyang Liu, Rongmei Lin, Zhen Liu, James M Rchg, Liam Paull, Li Xiong, Le Song, and Adrian Weller. Orthogonal war-parameterized training. In CVPR, 2021. 2, 3, 4, 5, 6, 8, 19
[49] Weiyang Liu, Rongmei Lin, Zhen Liu, Li Xiong, Bernhard Schölkopf, and Adrian Weller. Learning with hyperspherical uniformity. In AISTATS, 2021. 3, 20
[50] Weiyang Liu, Zhen Liu, Zhiding Yu, Bo Dai, Rongmei Lin, Yisen Wang, James M Rehg, and Le Song. Decoupled networks. CVPR, 2018. 1,4
[51] Weiyang Liu, Zeju Qiu, Yao Feng. Yuliang Xiu, Yuxuan Xue, Loughui Yu, Haiwen Feng, Zhen Liu, Juyeon Hco, Songyou Peng. Yandong Wen, Michael Black, Adrian Weller, and Berhhard Schölkopf. Parameter-efficient orthogonal finetuning via butterfly factorization. In ICLR, 2024. 2, 3, 5, 6, 9
[52] Weiyang Liu, Yan-Ming Zhang, Xingguo Li, Zhiding Yu, Bo Dai, Tuo Zhao, and Le Song. Deep hyperspherical learning. In NIPS, 2017. 1,4
[53] Ziyue Liu, Ruijie Zhang, Zhengyang Wang, Zi Yang, Paul Hovland, Bogdan Nicolae, Franck Cappello, and Zheng Zhang. Cola: Compute-efficient pre-training of Ilms via low-rank activation. arXiv preprint arXiv:2502. 10940, 2025. 10
[54] Ilya Loshchilov, Cheng-PPigg Hsich, Simeng Sun, and Boris Ginsburg. ngpt: Normalized transformer with representation learning on the hypersphere. arXiv preprint 13Xiv:2410.01131, 2024.
[55] Tlya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2019. 1,25
[56] Chao Ma, Wenbo Gong, Meyer Scetbon, and Edward Meeds. Swan: Sgd with normalization and whitening enables stateless Ilm training. 2025. 8
[57] Albert W Marshall, Ingram Olkin, and Barry C Arnold. Inequalities: theory of majorization and its applications, volume 143. Springer. 1979. 21
[58] Roy Miles, Pradyumna Reddy, Ismail Elezi, and Jiankang Deng. Velora: Memory efficient training using rank-1 I sub-token projections. arXiv preprint 27114 2024, 10
13

---

[59] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. arXiv preprint XXiv:1802.05957 2018.
[60] Zhanfeng Mo, Long-Kai Kai Huang, and Sinno Jialin Pan. Parameter and memory efficient pretraining via low-rank riemannian optimization. In ICLR, 2025. 11
[61] 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
[62] Radford M. Neal, Priors for infinite networks. Bayesian learning for neural networks, pages 29-53, 1996.
11
[63] Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-hayesian approach to spectrally- normalized margin bounds for neural networks. In ICLR, 2018. 7
[64] Sean O'Rourke, Van Vu, and Ke Wang. Bigenvectors of random matrices: survey. Journal of Combinato- rial Theory, Series A, 144:361-442, 2016. 20
[65] 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 finetuming. In NeurIPS, 2023. 2, 3,
5, 6, 9
[66] 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 unified text-to-text transformer. Journal of machine learning research, 21 21(140):1-67, 2020. 8, 25
[67] Mihacla Rosca, Theophane Weber, Arthur Gretton, and Shakir Mohamed. A case for new neural network smoothness constraints. arXiv preprint arXiv:2012.07969, 2020.
[68] Olivier Roy and Martin Vetterli. The effective rank: A measure effective dimensionality. In European signal processing conference, 2007. 7
[69] 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 2505:2200,, 2025. 11
[70] Jack W Silverstein et al. The smallest cigenvalue of a large dimensional wishart matrix. The Annals of Probability, 13(4):1364 134-1368, 1985. 20
[71] DiJia Su, Andrew Gu, Jane Xu, Yuandong Tian, and Jiawei Zhao, Galore 2: Large-scale Ilm pre- pre-training by gradient low-rank projection. arXiv preprint (rXiv:2504220437 2025. 10
172] Vithursan Thangarasa, Abhay Gupta, William Marshall, Tianda Li, Kevin Leong, Dennis DeCoste, Scan Lie, and Shreyas Saxena. Spdf: Sparse pre-training and dense fine-tuning for large language models. In UAI, 2023. 10
[73] 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. 1121,,12222112111111 2024,
[74] Gilles Wainrib and Jonathan Touboul. Topological and dynamical complexity of random neural networks. Physical review letters, 110(11):18101, 2013.
[75] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samucl R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 1804.0760,, 2018. 9
[76] Bo Xie, Yingyu Liang. and Le Song. Diverse neural network learns true target functions. In AISTATS, 2017. 3,7
[77] Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Picter Abbeel, and Dale Schuurmans, Foundation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.04 001.222933,, 2023.
[78] Xinyu Yang. Jixuan Leng. Geyang Guo, Jiawei Zhao, Ryumei Nakada, Linjun Zhang, Huaxiu Yao, and Beidi Chen. S2ft: Efficient, scalable and generalizable Im finc-tuning by structured sparsity. arXiv preprint arXiv:2412.0622 2024, 10
[79] Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generalizability of deep learning. arXiv preprint 1705:1705..0944 2017.
14

---

[80] 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 gain more. arXiv preprint arXiv:2406. 16793, 2024.
[81] 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.082%6, 2024. 10
[82 Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Nemory-efficienn Il training by gradient low-rank projection. ICML, 2024. 6, 8, 10, 25
15

---

Appendix
Table of Contents
A Delving into POET's Three Training Phases
17
A.1 More Details on Vector Probing
A.2 Geometric Interpretation of the Trace of Orthogonal Matrices
A.3 Empirical Observations
B Minimum Hyperspherical Energy in POET
19
Guarantees of Weight Spectrum under POET
20
D Proofs of Lemma
23
E Experimental Details
25
F Implementation and CUDA Acceleration
27
G Results of Vector Probing for R and P
28
H Weight Update Evenness of Different POET Variants
30
Training Dynamics of Singular Values
31
J Orthogonality Approximation Quality using Neumann Series
37
K Full Results of Training Dynamics
39
L More Results POET as a Finetuning Method
40
16

---

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 very straightforward. We generate constant vector that is randomly initialized. Then we let it to be transformed by the learned orthogonal matrices R and Finally, we compute the cosine of their angle: Ro and Pv. this process, the probing vector v is always fixed. The full results are given in Appendix G.
Beyond partieular constant probing vector, also consider set of randomly sampled probing vectors that follow our proposed normalized Gaussian initialization. Specifically, we consider the following expectation:
En ( Re),
where v is vector initialized by normalized Gaussian distribution (thus uniformly distributed on unit hypersphere $m--1). Because [vve' m. then we have that
gm-- (o Ro} 1 Tr(R).
m
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 g(x)) v. RV (we simplify the expectation over the unit hypersphere to E):
Var(1(+)) = B{(v" RR)?? (E Ro})
(10)
First we compute E{(u" R
E((v) Rv)²] = Tr(R)² +
RT+R
2
m (m + 2)
(11)
Tr(R)* + Tr( Tr(R²) (22 + m
m(m 2)
Then compute (E(vT Rv))²:
(E Ro))² TT(R22
(12)
m
Finally, we combine pieces and have the final variance:
Var(v? Rv) m7 + Tr(R²) + 2Tr
(13)
m(m + 2)
which shrinks at the order of O(1/m). Therefore, when the dimension of orthogonal matrices is large, even if we use 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 trans- formation 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 orthogonal matrix R IRmx can be written
Tr(R) = Σ c Re;
(14)
where e,, Vi are unit basis vectors. This expression reveals that the trace measures the sum of inner products between each original direction e; and its transformed version Re,, Since e Re; can be interpreted as the cosine of the angle between e; and Re,, the trace thus reflects how much the orthogonal transformation aligns with deviates from the original coordinate directions.
17

---

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, 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 Tr(R) with small variance.
m
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 conical-shell searching rotates each orthogonal matrix R and P smoothly away from identity while preserving their singular values, the cosine similarity between transformed and initial weight vectors falls from 1 to z 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, whether you train with fully-stochastic 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 cone, and validation perplexity drops almost linearly because spectra remain frozen and gradients act only on meaningful directions. In this phase, the trace the orthogonal matrices stay almost as a constant.
Specifically, we hypothesize that the orthogonal transforms have reached "good" cone; thereafter they mostly shear vectors inside that shell, leaving the angle 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 learning meaningful directions. Because the geometry is stabilized in this phase, the learning of patterns happen in 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 singular vectors, 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.
18

---

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-mi me-m isotropic Gaussian distribution. Both Xavier [17] and Kaiming [23] initializations satisfy such weight initialization condition. Orthogonal equivalence transformation given by OET(W; R P) = RW P. where the input matrix W is multiplied on the left and on the right by orthogonal matrices R and P. respectively.
When W initialized zero-mean isotropic Gaussian distribution, [48] has shown that these random neurons, if normalized, are uniformly distributed on the unit hypersphere. This leads to 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 small hyperspherical energy.
We consider weight matrix W E R where each entry is ii.d. sampled from a zero-mean Gaussian distribution with variance the same variance o², i.e., Wij ~ *N o²). After applying orthogonal equivalence transformation, we have Wnew = RW P where R and P are two orthogonal matrices. Then compute the distribution of Wanww Because linear maps preserve Gaussianity, each entry of W is a finite linear combination of W and hence, Wnew follows joint Gaussian which can be fully characterized by mean and covariance.
The mean of W neew is given by
W R E W
P= R P
P E..
(15)
For its covariance, we consider two generic entries Whow and Whey:
Weew
RaK W Pij,
(16)
Winww
Σ
Ry Winew Pag.
0,00
Then we compute the covariance between the two entries:
Cov(W
Wey)
Σ
Ra Pig Cov(WWW Wae)
1j
(17)
o² 2 (ar(PT P)j
σ dwp
which implies the following resutls:
The covariance matrix a diagonal matrix, so different entries of Whey are uncorrelated. Because W new is a joint Gaussian and different entries are uncorrelated, each entry Wnew is independent.
Each entry of Weew has identical variance o².
To sum up, each entry of Wreow i.i.d. ((,, which identical the distribution of each entry of W. Because we have Wnnw =d W. we can conclude that, similar to W, W Wiew also has provably small hyperspherical energy among neurons.
Despite being extremely simple, we find that this is in fact significant result. Under zero-mean isotropic Gaussian initialization, be achieved simultancously. It also partially explains why the proposed normalized Gaussian initialization achieves the best performance.
training and energy-preserving training can
19

---

C Guarantees of Weight Spectrum under POET
For standard Gaussian initialization where each element the weight matrix W is sampled with normal distribution, we have the following standard results [8, 70]:
max W)
U.s.
+ VÁ
110000
(18)
miin
WW)
VA
vd
n--000
which gives spectrum guarantees for weight matrices generated by the standard Gaussian initialization.
In the following, give the spectrum guarantees for the normalized Gaussian initialization. We start by stating the following theorem from [49]:
Theorem Let ...... o, ERd be i.i.d. random vectors where each element follows the Gaussian distribution with mean and variance Then 0 =1= . Un S. are uniformly distributed
on the unit hypersphere gd-1, If the ralio ă converges constant XE (0, 1), asymplotically have for W (Vi,... Vn RRdx:
lim mmax (W) S (Vd VAd) (max 120000
||00
(19)
1
lim miin (W) 2 (vd Vdd)
(min
22+00
0
where m and miin denote the largest and the smallest singular value of matrix, respectively.
Proof. We first introduce the following lemma as the characterization of a unit vector that is uniformly distributed on the unit hypersphere $
Lemma 2 ([64]). Let be random vector that is uniformly distributed on the unit hypersphere gd-1 Then v has the same distribution as the following:
U1
12
(20)
1-
Vi=
W
where U1, U2.". Ud are i.i.d standard normal random variables.
Proof. The lemma follows naturally from the fact that the Gaussian vector {u: is rotationally invariant.
Then we consider random matrix W . ün) where é; follows the distribution of (ur.... ma). Therefore, it is also equivalent to a random matrix with each element distributed normally. For such a matrix W, we have from [70] that
lim max (W)= Vd Xd n-o0
(21)
lim Jmin (W) vā Xdd
where J (.) and (miu(-) denote the largest and the smallest singular value, respectively.
Then we write the matrix W as follows:
W W Q
0
0
| I|a
0
(22)
W
02
0
20

---

which leads to
lim max W)
lim maxx (W... Q)
7-000
(23)
lim min (W)
lim min (WQQ)
7-200
We fist assume that for symmetric matrix E RRRD (1 2 ... An (A). Then we introduce the following inequalities for eigenvalues:
Lemma 3 (157). Let G, H Ruxa be positive semi-definite symmetric, and let S is <
k S n. Then we have that
II [=1
GH)
S II t=1
Nr.
(H)
(24)
and
II
X GH) 2 II
Aa (+)))
(25)
t=1
t=1
where X3 denotes the i-th largest eigenvalue.
We first let 1 S <
< i S n. Because W E Bdxx and Q E IRuxn, we have the following:
Il oh (WQ) II A. (WQQTWT
(26)
h A ( WQQT)
by applying Lemma 3 to the above equation, we have that
A ( WQQT) Z II Xr (W (QQT=--(QQ)))
(27)
oi (W ()nn-et1(Q)
II Nr w WQQT)
< X (W (TTAAAQQQT)
(28)
k
(W)
Therefore, have that
o WQ) 2 II oi W onn=++(1))
(29)
t=1
II oi )Q)=0 < II or (Q)
(30)
Suppose we have A and i = then Eq. (29) gives
on (WQ) > ox (Q)
(31)
Then suppose we have & and i = 1, then Eq. (30) gives
( < a ( (
(32)
21

---

Combining the above results with Eq. (21) and Eq. (23), we have that
lim max W
lim max W Q) I lim (amax W) Fmax (Q))
= (Vd + vdd) max
10:12
(33)
lim min (W) lim min (W Q) N lim (amin (W) min (Q))
2-+00
22030
(vd v Ad) min
which concludes the proof.
Combing with the fact that
||uile
lim max
lim min no00
(34)
vd
d
we essentially have that
lim max WV) +1 + Vλ,
(35)
lim min - VA. 200
which can be written to the following results:
max W
1 1 VX
(36)
VÀ
min
12-500
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.
22

---

D Proofs of Lemma
Proof of Lemma 1. We consider an orthogonal matrix R and orthogonal primitives Gi corresponding to uniformly random subsets S C [m] of size b as explained in the main text (see equation (4)). The main claim we need to prove that given any vector can find an orthogonal primitive matrix G corresponding to the set S such that
R and set S C [m] with is E [m]] we
(Gv): for i E S with / k
(Gu) 2 0
(37)
(Gv): ia for lorls $ S.
Moreover, for all w E I with w, 0 for 2 k the relation
Gw w
(38)
holds. We can assume that the matrix D(S) {e(s))....., e(())) contains the entries in ascending order. Then we write
D(S)
(39)
32
where ह1 Rh corresponds the entries s; with 8; and ह2 E R to the remaining entries, in particular $by+1 k because k E S. It is well known that for every vector v there is a rotation Q aligning with the first standard basis vector, i.e., such that Qu Xe(1) for some 2 Consider such a matrix Q for the vector 02 and then define the orthogonal matrix
G c=(1
1b1 xb2
Oh xb2
(40)
Careful inspection of (4) implies that the last part of (37) actually true for any & as the second term has rows with all entries equal to zero for all I & S. For the first part we find
D(S)ED(S) u D(S)
T1 82
D(S)
(e(1))
e(ss)(0i))+ +e(()) beb
Σ
(41)
Here we used Sb, k in the last step. Since in addition
((1m - D(S) 1, D(S) 0
(42)
for all conclude that indeed (Gv); for E S and k, (Gv); 2 0. The remaining statement (38) follows from the observation that when decomposing in (39) we find
T Obz
(D(S))
(43)
(because w, ( for i k) and therefore
-G- (D)(()) O2
(44)
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 reverse the indices, i.e., consider products G. Assume that we have chosen G; for i S C and some C% E N such that the product
G
Pk Ga
R
(45)
satisfies Pk = 0 for all & k and l' > K' and & M 2 20 for k' < k. Let Cb+1 2 Ch +
(m/ In(m). Then, can bound for any 1 >> k the probability that there is C < j < Ck+1 such that {k. 1) C S using that S follows a uniform i.i.d. distribution by
P(Ac cc <j S C411: k,le 8²) <
(1
m²
b²
exp
(-mp
m
b²
m²
62
In In(m)
m aa
(46)
23

---

The union bound implies that with probability at least 1 m -+ there is for all >> ka < < Ck-1 1 such that { C St. If this holds we set G3 for C₁ < j S Ck. as constructed above if k: S
and Gj Im otherwise. This then ensures that
p GGe41 ..... G1 R
(47)
satisfies P = 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 pm is an upper triangular matrix with non-negative
diagonal entries with probability at least 1 - mmn -++ = 1 - m ma-2),, But we also know that P orthogonal and therefore satisfies pm 1m and we thus find
........ ... =.
(48)
Next we give heuristic that actually (In(m) m²² /b2) terms are sufficient to express every orthogonal map as a product of stochsastic primitives. For fixed we consider the map
T : 0()) - 0mm) ( Go) [[G. j=1
(49)
i 2 (mmmm)mmm)mmmmmmmmm /b2 we have that with probability least 1 m (a-2) for all k, [m]] there is j such that k, S1. Assume that this is the case. Recall that the tangent space of ((k) at the identity is the space of skew-symmetric matrices. Consider tangent vector (X,,..., with
X; E Skew(k). Then
D$ (X) X.)
D(S)) X, D(S')T ==1
(50)
This is surjective map on Skew(m) under the condition that for all k. / [m]] there is S e such that k., E S". We can therefore conclude that the image of & contains a neighbourhood of the identity. Moreover, since & is polynomial map. D& is surjective everywhere except for variety of codimension one. While this is not sufficient to conclude that the image of & is ((d) or dense in ((d) it provides some indication that this is the case.
24

---

E
Experimental Details
Parameter
Llama 60M
Llama 130M
Llama 350M
Llama 1.3B
Hidden dimension
512
768
1024
2048
Intermediate dimension
1280
2048
2816
5376
Number attention heads
12
16
32
Number of hidden layers
8
12
24
24
Table 8: Model architectures for different Llama variants.
Model
Spec.
GPU
Ir (base)
Ir (POET)
training steps
batch size
grad acc.
1/2
Ie-2
le-3
300,000
256
Llama 60M
b 1/4
1
le-2
2e-3
300,000
256
2
b 1/8
1
le-2
4e-3
300,000
256
2
b = 1/2
1
5e-3
le-3
400,000
128
2
Llama 130M
b 1/4
1
5e-3
2c-3
400,000
128
2
b 1/8
5e-3
4o-3
400,000
128
6 1/2
4
5e-3
le-3
400,000
128
Llama 350M
b 1/4
4
Se-3
2e-3
400,000
128
b 1/8
4
5e-3
4e-3
400,000
128
h 1/2
8
1e-3
1e-3
500,000
64
Llama 1.3B
b 1/4
8
1c-3
20-3
500,000
64
b 1/8
8
Ic-3
4e-3
500,000
64
Table Hyper-parameter setup of POET-FS.
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]] (Apache 2.0 license), which we also reproduce the AdamW and Gal baselines. We will release 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 ra- tio of 0.01. We use the number of warmup steps of O, weight decay of 0.01 and gradient clipping of 0.1. For the Adam baseline, we report results for the optimal learning rate from 11x10-2, 5x 10-3, 10000,
1 10-A 50105 1x10-=,5x100- 1 10105). After each merge-then- reinitalize step, we additionally increase the gradient clipping for 10 training steps to improve training stability.
Model architecture. Our work utilized the Hugging Face Transformers' code base to construct the Ilama 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-Forwand Network (FFN) has been slightly modified for the POET-BS, compared 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 large, meticulously cleaned version of Common Crawl's web crawl corpus. It was
25

---

Model
Spec.
GPU
Ir (base)
Ir (POET)
training steps
batch size
grad acc.
256
Ie-2
le-3
300,000
256
Llama 60M
b 128
1
Tc-2
2c-3
300,000
256
2
64
1
le-2
4e-3
300,000
256
2
256
1
5e-3
le-3
400,000
256
2
LJama 130M
6 128
Se-3
2e-3
400,000
256
2
64
5e-3
4c-3
400,000
256
2
b 256
4
5e-3
le-3
400,000
128
Llama 350M
b 128
4
Se-3
2e-3
400,000
128
1
Se-3
4e-3
400,000
6 256
8
1e-3
le-3
500,000
64
Llama 1.3B
b 128
8
Ic-3
20-3
500,000
64
64
8
le-3
4e-3
500,000
64
Table 10: Hyper- pareme-ar setup POET-BS.
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 4 or 8 GPUs.
26

---

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-symmetry matrix Q E R where the diagonal entries are zero (Q 0) and off-diagonal elements satisfy Qis Q2i. This structure allows us to store only the upper triangular part of Q vector, reducing the number of trainable parameters from n22 to n(n 1)/2. During the forward pass, Q is reconstructed on-the-fly using 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 1.3B LLaMA model single H100 8-GPU node yields 3.8x speedup over the baseline (i.e. native implementation). Table summarizes the contribution of each component the overall training time.
Design
Speed-Up
Neumann approximation
1.5x
Skew-symmerric CUDA kernel
1.3x
Total
3.8x
Table 11: Method design and clock time speed-up.
27

---

G Results of Vector Probing for R and P
In this ablation study, we perform vector probing on the orthogonal matrices R Rm R for all linear layers for all blocks of 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 eliminate the need for resampling and reinitialization of the orthogonal matrices. Interestingly. we observe this three- three-phased learning dynamics across different types of linear layers and different-depth transformer blocks.
modelllayers.com calf
model.layers.co attn.k
model. self attn.y proj
Simiarty
Cosino mmiurty
505
1006
1506
2004
2536
3936
505
100K
2004
2534
3336
555
1006
1506
2004
2536
333
model.layers.mmpp www.srs.mmp............ pro) (P)
model.ayeersmlppug pro) (P)
www.eelllyrssss.ddooooo pro) (P)
Ousine
Step
Step
Step
model.layers.self attn.o pro)
Uyero
1yer
2.0
bywe
laye
2.6 1932
1100
2504
Figure 11: Cosine similarity for vector probing of P across the self-attention components (query, key, value, and output projections) and feed-forward network components (up-, down-, and gate- pateeprrocctionn) in all transformer blocks of POET-trained Llama 60M model.
model.layers.soe self welf_fattn.__pr__
moodel.lyyers.seff. pproj (R)
model.layers.c self attn.v_proj
Siilarit
Oraine Smilarit
0.4
Stop
Stop
htddelllyyers.ml..geg proj (R)
model layers, www.s.llp... pro) (R)
model. model. www.pyyerr..mm.oow..... pro) (R)
Sentartyy
Coline
156K Stap
2334
2004
5Ec
156K Stap
2334
Step
model leddel.layers.cm self proj
gmilaryy....
2.
1.
G
15500
Figure 12: Cosine similarity vector probing of across the self-attention components (query, key, value, and output projections) and feed-forward network components (up-, down-, and gate-projections) from all transformer blocks of POET-trained Llama 60M model.
28

---

model.byyers.off attn..9 pro) (P)
model.layers.co welf attn. proj
model.layers self atto. pro)
200
300
5001
150
150
250
Step
Step
Step
model. wtw.eellayysssmp.pp.... pro
model.Jayers.mlpp pro) (P)
model.layers.mlp down pro) (P)
1400
1400
1200
11.00
1100
400
000
500
30X
200K
2006
253K
10CK
150K
3236
100K
253K
2306
Spep
S0ep
Step
model.Jayers.sell pro)
bowe
602
lext
103K
Figure Trace of P across the self-attention components (query, key, value, and output projections) and feod- forward network components (up-, down- and gate- (ite-projections) from all transformer blocks of a POET-trained Llama 60M model.
model.layers.self attn.q, pro) (R)
model.layers.seif attn.k, proj (R)
model.layers.sef attn.v
500
300
200
loar
koar
lear
byar
425
250
20
200
Stop
Stop
Stop
model.layers..o www.gatteea.mmmmmmmmmmmmmmmmmmmm proj (R)
model model.layers. www.ppppppppppppmm _roj
model.Jayers.mll downdownnn proj (R)
1000
100
1200
1100
250
250
100K
2586 Stop
250K
100K
159C
2006
333
50C
300K
200K
2536
3056
Stap
model.layers.seff pro)
1ywe
402
Figure 14: Trace of R across the self-attention components (query, key, value, and output projections) and feed- forward network components (up- down-, and gate-projections) from all transformer blocks a POET-rainee Llama 60M model.
29

---

H
Weight Update Evenness of Different POET Variants
To understand the higher parameter efficiency POET-BS compared POET-FS, we employ 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, 64x64 matrix was randomly initialized and trained for 100 steps under various POET-BS and POET-FS configurations. The merge-then-reinittaiiz trick performed at each iteration, and the same set of weight elements was effectively updated between two successive thane-tthe-rrinttialie 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 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, can introduce unexpected training difficulties in earlier stages.
200
200
175
150
100
100
75
(a) POET-BS (b
(b) POET-BS 16)
(c) POET-BS (b 32)
Figure 15: Visualization of the weight update mechanism of POET-BS after 100 steps of update and T.,. 1.
40
60
100
90
40
10
70
5
10
20
(a) POET-FS (b = 1/8)
(b) POET-FS (b 1/4)
(c) POET-FS (b 1/2)
Figure 16: Visualization of the weight update mechanism of POET-FS after 100 steps of update and Ti
30

---

I Training Dynamics of Singular Values
We conduct ablation study compare the training dynamics of singular values of weight matrices between Adam W 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. 60M LLaMA model was trained for 50,000 iterations with an effective batch size of 512, using both Adam W and POET-FS = 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-then- reinitialize step was applied before each evaluation. Training is finished at 50,000 steps, as spectral norm of the AdamW-trained model plateaued at this point.
model.layers layers attn.q_pre proj
model.layers. attn..k [itnk_proj pro]
layers proj
10
10
100
10
10
10
224
10"4
234
10
10
455
524
mader. 0.self attn.o pro
model.layers.0. m. pro
model.layers. www.errs...m....o. m.mlp.down pro
156
10-2
300
10"2
10a
10"
450
10
350
355
10
50K
500
1.0
300
400
200
107
model.layers w.t.s...ffffff....e atto.._proj
model.layers. w.sef...ftfftfff www.eee_etttttt_rrrre
model.layers.i seseerss...es.ftttt attn.v_proj
10
10-2
325
10
22c
200
10"h
235
10-
10
424
10
425
325
42A
30
420
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model. lawerll..yrss......
model.Jnyers...mbl
model.tayers...mpp..e
30-4
115
10"4
10a
30
100
200
500
100
200
300
400
500
100
200
300
400
500
Figure 17: Training dynamics of the singular values of weight matrices within Blocks (-1 (the i-th row represents Block i) of 60M llama Transformer trained with AdamW.
31

---

10
model.layers. 2ellf attn.q pro)
medel.ayers.2.ell attn.k pro)
model.byyers. 2elff attn.y pro)
10
10-1
325
10
235
235
10
235
235
10
655
625
10"4
100
700
100
400
500
100
700
300
400
500
100
700
100
400
500
modell layers 2. salf attn.o proj
modal www.:.ppppppppppp...
model. Landel.ayyar.2222 m.p. dow. pro)
100
10K
1.00
10
10"0
030
10
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model.layers.33 attn.q pro]
lodel. 3.seef attn.k pro)
model.layers attn.y pro)
100
114
10-4
110
10
10
120
10
10-4
614
10-4
32X
414
10
45C
loodel.layerr.3..r self wwt.tttttpppmm proj
model.Jayers. wtw.a//rr....mm...ppp......
wtw.:layyerr..ml.ddoww
1.50
430
500
100
model laddel.layer..4 layers.4.sefff attn.q pro]
modell.ayerrs.o 4.coff attn.k ppoo
model 4.roors...sslll attn.v pro)
100
10
224
10-2
254
30
624
3.0
30
10
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model.layers.4.seem attn.o pro
model.layers.4mmpll
moode.layers.4.mlp.d
4.00
10
R00
1.00
200
1.00
200
300
100
200
300
400
500
Figure 18: Training dynamics of the singular values of weight matrices within Blocks 2-4 (the i-th row represents Block i) of 60M Llama Transformer trained with AdamW.
32

---

Sttel..lyyrr.S.selle attn.q pro]
model.ayerr. 5.self attn.k pro)
model. S.lff attn.y pro)
10F
10
10
10
325
10"1
325
10"4
325
255
235
10
10h4
10
355
60%
625
42c
10
10
100
700
100
100
500
100
700
300
400
500
100
700
100
400
500
modell layers. S. salf attn.o proj
modal. wwp.p..r......ppppppppp.. wwwpppppppppp.p...
model. londel.ayyer..So www..o.ww....wpp.mmmmmmmmmm pro)
10F
10K
1.00
10"4
107
10
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model.layers. attn.q pro]
layer. attn.k pro)
model.layers. attm.y pro)
10
10
10
10
114
10
11
10
10a4
120
226
10-4
414
414
614
10
32A
10
model layers. self attmn.......ooooo.. proj
model.layers. mmpppupppomppp.ppoojoooo..
model.layers. www.p.dwwnpppppppp.mmmm _oppdoow__rooo
10
1.50
10
10
10"4
500
436
100
100
model layers layers.7 7. sel attn.q pro]
moddll.aayrr..7 7.sof attn. ppro
model. 7.lelf attn.v pro]
10
100
10-1
224
10
254
10-2
10-1
604
622
62c
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
ladel.layers.7..oo pro
model.layers.7.mll
layel. mlwp.down
10
126
10-1
256
10-3
10a
10a
10-3
1.00
200
300
100
200
300
400
500
Figure 19: Training dynamics of the singular values of weight matrices within Blocks 5-7 (the i-th row represents Block i) of 60M Llama Transformer trained with AdamW.
33

---

lodel.layers. 0.seff attn.q pro)
model.layers. 0.self attn.k pro
model. 0.self attn.v pro)
10
104
10
28C
226
23c
53c
10
306
53c
700
500
100
model model.layers. layers. O O.self attm.otttp.... proj
modal.layers. O. wtm.ppppppppp...
model model.layers. layers. mlmppdown proj
102
13X
120
430
10a
10
204
124
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model.layers. 1.sef
model.layers. 1.self attn.k
model.tayers. 1.seff pro)
10-4
10
10
10a4
10
326
1.0
500
300
400
lnodel.lyeers...see _pron
model.layers. www.erss...ll.ur.... www.mmpppppppp..
malel.lyyers...lppppp... proj
1.0F
10"4
10a
VSC
30K
10a
40K
Ihddll layers, 2ele attn.g pro]
lodel. 2.solf attn.k proj
model model. layers, 2-seff attn.y pro]
104
10f
10"4
204
204
10
20c
204
30
524
524
52c
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model layers. attn.o pro)
model.layers..mmlp.uu proj
model.layers.2.mm mllp..dww pro)
1.04
10
10a
204
100
200
300
400
500
100
200
300
400
500
11
100
200
300
400
500
Figure 20: This plot illustrates the singular value training dynamics for individual weight matrices within Blocks 0-2 of a 60M Llama transformer model trained with POET. For each block, the dynamics are shown for the self-attention components (query, key, value, and output projections) and the feed-forward network components puojeroec,,ion, down- pwwjeprojcction, and gate troje-rrooctin....
34

---

layer.. attn.q pro)
model.layers.3 attn.k pro
model.layers..... attn.v pro)
10F
10
10"4
10-1
10
10-4
10"
20K
226
23c
506
305
53c
100
700
100
500
model model.layers. self attn.o. proj
modal.layers. 3.m.p.mmppppppp walpp.ppprrr
model.layers. layers. 3. down proj
model
10
13X
120
10"4
430
10
430
10a
204
10
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model.tayers. law.rs...erlffllf
model. attn.kk
model.tayers.4..u pro)
326
336
10
300
400
model.layers...uok
model.layers.4 htt.yyerr.....p..ppppr. mwp.mpppppoooojj..ppppp
modelllayers..
proj
10a
10
VSC
10
22c
10a
42C
10
modell layers, 5.self attn.q pro]
layer, 5.self attn.k proj
moddel.layers. S.self attn.y pro]
10F
10
204
204
22C
204
$2C
524
52c
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model lodel.layers. layers.. S.self attn.o pro
moodel.lyyersS.mlp.uu proj
madel www.:lyyer.....doow........ pro)
10
10a
10-4
2E4
10 4
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
Figure 21: This plot illustrates the singular value training dynamics for individual weight matrices within Blocks 3-5 of a 60M Llama transformer model trained with POET. For each block, the dynamics are shown for the self-attention components (query, key, value, and output projections) and the feed-forward network components (up- puojeroect,ion, down-prooeen pww-eproecction, and gate- projeerrocction..
35

---

lodel.layers...uee attn.q pro)
model.layers.6.om 6.seff attn.k proj
model.layers.6 sesif attn.v pro
10
10
10a4
10
10"0
306
324
35c
10
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
mydeel.byerr....oo selef attn.o pro]
model.layers.co www.p.p........ppppppppppppppppp mippup
model wandrr..ayers.....o w.w.p.ddllopwe...... pro)
23c
10
536
53c
10
10a
204
10
100
200
300
400
500
100
200
300
400
500
100
200
300
400
500
model.layers.7.ee attn.q pro]
model. attn.k pro)
model.bayers7.cc attm.y pro)
10F
105
10
10"4
10-1
10-4
104
32C
32c
10-4
464
404
43c
124
model.layers, attn.o pro]
madeel.ayers. wttp:rs..mmp..pppp.r.. prpproj
modelllayers 7.mlp.dowwr down pro)
10*
10-4
32c
10-3
53c
Figure 22: This plot illustrates singular value training dynamics individual weight matrices within Blocks 6-7 of a 60M Llama transformer model trained with POET, For each block, the dynamics are shown for the self-attend attentienoo components (query, key, value, and output projections) and the herd-forward network components (up-pre projection, down- projecroo,ction. and gate- projection).
36

---

J Orthogonality Approximation Quality using Neumann Series
In this ablation study, we evaluate the approximation error of the orthogonal matrices R R and P E _ across all linear layers in Block of a 130M LLaMA model trained with POET-FS 1/2) for 10,000 steps. Figure 23 and Figure 24 show the approximation error over the first 1,000 steps. Since the eπror difference between Io = 3 and k = 4 was negligible, we used i 3 for better computational efficiency. Empirically, while models, larger & values are needed to avoid training divergence caused by exploding gradients due to approximation eπror.
suffices for smaller LLaMA
modell setff
modil off
mnde Doaff
Stpp
Stap
Skpp
rocdd Inved Durslo peo) (8)
swdellaesssSmmppppp msi (R)
madelliger....n...oooo peoi (B)
Sep
Skyp
Sip
medd livers Bueff am 2ni (8)
Figure 23: For the transformer block 0, we show approximation error of orthogonal matrix R for the self- attention components (query, key. value, and output projections) and the rerward network components (up-projection, wwwn-projection,, and gate-projection).
109
modid Douff
Dugff
5x0
Figure 24: For the transformer block 0, we show approximation error of orthogonal matrix the self- attention components (query, key, value, and output projections) and the feed forward network components (up-projection, down- pwo-eproecction, and gate- sate-projection).
37

---

Additionally, Figure 25 shows the orthogonality approximation enror of Neumann series with different k over the first 10,000 training steps, illustrating how decreases 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 training divergence POET.
model waw.ellleesss.......d..o d.wlp. proj (P)
Approtimation enor
0.0
2000
3000
4000
Step
Figure 25: The approximation error of orthogonal matrix P randomly selected down-projection layer after training 10000 steps.
38

---

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-efficeeng and generalization.
60
Adamm
AdamW
POETIBS, 64]
POETIBS, 64)
Perplexity
50
POETIBS, 128)
POETIBS, 256)
Perplexity
45-
POETIBS, 256)
POETIFS, 1/8)
40
POETIFS, 1/81
40
POETIFS,
POETIFS, 1/4)
35
POETIFS. 1/21
POETIFS, 1/41
30
25-
20-
0K
50K
100K
130K Steps
200K
250K
300K
0K
50K
200K
150K
200K 250K Steps
300K
350K
400K
(a) Llama 60M
(b) Llama 130M
45
Adamm
AdammW
POETIBS, 64)
POETIBS, 64)
POETIBS, 128)
POETIBS, 128)
Perplexity
POETIBS, 256)
Perplexity
25
POETIFS, 1/8)
30
POETIFS, 1/81
30
POETIFS, 141)
POETIFS, 1/41
POETIFS, 1/2)
POETIFS, 1/21
20-
20
15
15
0K
50K
100K
150K
250K Steps
200K
300K 350K 400K
0K
100K
200K
300K
400K
500K
Steps
(c) Llama 350M
(d) Llama 1.3B
Figure 26: Validation perplexity during the training the LLama- based transformer with 60M, 130M, 350M and 1.3B parameters.
39

---

L More Results of POET as a Finetuning Method
To demonstrate the applicability of POET to general finetuning tasks, we apply to finetune BART-large model [41] on the NLP task of text summarization. Specifically, we evaluate POET on the XSum [61] and CNNi [24] datasets, reporting ROUGE-1/2/L. scores Table 12. We note that both LoRA and OFT are designed solely for parameter-efficien 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 the same configuration as in the pretraining setup, resulting in higher parameter count. Experimental results show that POET not only supports finetuning effectively but also outperforms both full-model finetuning and parameter-efficient methods.
Method
# Params
XSum
CNN/Daily
LoRA (r=32)
17.30M
43.38 / 20.20 / 35.25
43.17 / 20.31 29.72
(b=64)
8.52M
44.12
20.96
36.01
44.08
21.02 30.68
Full FT
406.29M
45.14
22.27
37.25
44.16
21.28 40.90
POET (FS,b=1/2)
144.57M
45.23
22.41
37.28
44.27
21.29 41.02
Table 12: Finetuning BART-large XSum and CNN/DailyMail for text summarization. We report ROUGE- 1/2/L results (higher better).
40

---

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 contribu-
tion.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have included analysis in the paper 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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
41

---

Yes, we provide a detailed derivation of the theory and reference all theoretical results our method is based on.
Justification:
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 cross- referenced.
All assumptions should be clearly stated or referenced 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 short proof sketch 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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: Yes, all methods are described 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 dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary 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 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 large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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 possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
42

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
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 NeurlPS code and data submission guidelines (https://nips. ca/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for 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 NeurlPS code and data submission guidelines (https:
heppublcccuuuidssssse guwdes/ddeeeee CCodeSumm ssi ssisssoniooooo ceiiinnooll Policy) 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 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 much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 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 appreciate the results and make sense of them.
The full details can 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, confi- dence intervals, or statistical significance tests, 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 with given experimental conditions).
43

---

The method for calculating the error bars should be explained (closed form formula, call 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.
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 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 range (e.g. negative enror 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
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 used develop our methods. We have included detailed descriptions in our Appendix section.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 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 failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https http://neuripppe wwtpppccl///eeeeieeeeeeeeee//eee htw.bliiie/thiccccceeeeeees.
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal 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 negative impact the society.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
44

---

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 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 (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
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 (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We mainly work on methods, we will release the code but will not release any data 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 2200 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 used and, possible, include URL.
45

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.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, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 available online, the authors are encouraged to reach out to the asset's 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 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurlPS 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 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 crowdsoureing nor research with human subjects.
46

---

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
Does the paper describe the usage of LLMs if 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, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We did not use LLMS to develop our methods but developed 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: / //tt://nn.uuueeeeeeeee wwt.c//eeeeeeneee/222222222 2022/LLM) for what should or should not be described.
47


Reviews:
[{"review_id": "jQ6JscPpFV", "reviewer": "Reviewer_hgyC", "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.", "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,", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "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:\n- 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.\n- 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.", "was_revised": true, "final_justification": "The authors  provided a detail rebuttal and clarified that some of the missing proofs are already in he literature, which I hope they will  be  properly cited and  explained in the paper.I am satisfied with the replies and explains and will keep my score.", "rebuttal": "We sincerely thank **Reviewer hgyC** for the recognition of our contribution and the constructive comments on our work. We take every comment seriously and hope our response can address the reviewer’s concerns. If there are any remaining questions, we are happy to address them. We summarize all the concerns into the following questions:\n\n---\n\n**Q1: 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.**\n\nA1: Thanks for the question. To be clear, the inductive bias we discussed here refers to constraints on the parameterized function. The unique weight update mechanism of POET introduces a few interesting explicit and implicit inductive biases, as detailed below. To address the reviewer’s concerns, we will add more context to better explain the inductive bias discussion.\n\nFor explicit inductive biases:\n\n**Spectrum-preserving weight matrices**: previous work [1] has shown that AdamW suffers from the lack of spectral diversity of the weight update, and gradient matrices, which have a more diverse spectrum, leads to favorable performance. We performed the same spectral analysis, by calculating the SVD entropy of the weight matrices during the training of a Llama 60M model and reporting the results in the following table. We show that POET, by guaranteeing to preserve the weight spectrum from the initialization, consistently maintains high spectral diversity throughout training, and therefore, POET can better explore diverse optimization directions.\n\n\n**Table A. Comparison of SVD entropy between a AdamW-pretrained model and a POET-pretrained model.**\n\n| Layer     | Initialization | AdamW  | POET   |\n|-----------|----------------|--------|--------|\n| Q proj    | 0.92           | 0.865  |**0.92**|\n| K proj    | 0.92           | 0.861  |**0.92**|\n| Up proj   | 0.97           | 0.946  |**0.971**|\n| Down proj | 0.97           | 0.861  |**0.972**|\n| Gate proj | 0.97           | 0.93  |**0.971**|\n| Out proj  | 0.92           | 0.897  |**0.92**|\n\n**(Probabilistic) energy-preserving weight matrices**: In addition, we compare the hyperspherical energy [2] of the model trained with AdamW and with POET. Hyperspherical energy, as introduced in our paper, is a quantity that characterizes the uniformity of neurons on the unit hypersphere.\n\nThe results are summarized in the Table B (in the Rebuttal)  below. Compared to AdamW, POET provides a much stronger guarantee of minimizing hyperspherical energy throughout training. It is important to note that the absolute value of the energy does not directly characterize neuron uniformity; rather, it is the relative difference from the minimum energy that reflects uniformity. Moreover, the initialization energy is already very close to the minimum energy (owing to the properties of the zero-mean isotropic Gaussian initialization).\n\n**Table B. Comparison of hyperspherical energy between a AdamW-pretrained model and a POET-pretrained model.**\n\n|     | Initialization | AdamW  | POET   |\n|-----------|----------------|--------|--------|\n| Energy    | 39.616          | 39.644  |**39.615**|\n\nThe energy difference between AdamW and initialization is about 0.03—roughly 30 times larger than the difference between POET and initialization. This gap is substantial and signals a significant deviation from uniformity, even though the raw values may appear small due to the high dimensionality. Notably, POET achieves an energy level almost identical to the initialization, demonstrating its energy-preserving (and energy-minimizing) inductive bias.\n\nFor implicit inductive biases, the multiplicative weight update can be viewed as a long sequence of matrix factorization of the final weight matrix. Specifically, if rolled out across multiple iterations, the final trained weight matrix can be equivalently expressed by $R_n \\cdots R_1 W_0 P_1 \\cdots P_n$ where $R_i$ and $P_i$ are orthogonal matrices. It may have some implicit inductive biases, as suggested by [3,4,5]\n\n[1] Liu, et al. Muon is scalable for LLM training. arXiv preprint arXiv:2502.16982 (2025).\n\n[2] Liu, et al. Learning towards minimum hyperspherical energy. NeurIPS 2018\n\n[3] Gunasekar, et al. Implicit regularization in matrix factorization. NeurIPS 2017\n\n[4] Arora, et al. Implicit regularization in deep matrix factorization. NeurIPS 2019\n\n[5] Li, et al. Algorithmic Regularization in Over-parameterized Matrix Recovery and Neural Networks with Quadratic Activations. COLT 2018\n\n---\n\n**Q2: No mathematical proof is provided on how the SPO and CNP initialization affect the generalization bound.**\n\nA2: Thanks for the question. Recall that we consider the initialization $W_{RP}=RW_0 P$ where $R$ and $P$ are initialized as identity matrices and $W_0$ follows a standard initialization scheme. Therefore, in terms of initialization, there is no difference between POET and standard training. \n\nIn particular, the spectral norms of our methods at initialization are exactly the same as standard training, so our method enjoys the same generalization guarantees in Eq.(7) of the main paper.\nSPO then refers to our approximation that we do not optimize over $R$ directly but only over suitable sub-blocks. We note that SPO does not compromise the orthogonality. Therefore, even with SPO, the spectral norm of weight matrices will not change, validating the generalization bound.\n\nCNP is an approximation scheme that avoids expensive matrix inversion and approximates the orthogonal matrix $R$ by a power series. Since CNP does not result in perfectly orthogonal matrices, the singular values of $W_{RP}$ change slightly but this change can be bounded by controlling the remainder of the Neumann series and Figure 1(b) from the main paper shows that it is very small in practice (almost neglectable). Therefore the generalization bounds do not deteriorate during training.\n\nHowever, we note that the generalization bound serves as an intuitive understanding of why controlling spectral norm is beneficial to generalization. There are no theoretical guarantees that can give a rigorous and tight generalization bound for practical LLMs. Therefore, a more involved theoretical analysis is beyond the scope of this paper. We do agree with the reviewer that more theoretical understanding will be useful and should be an important future direction.\n\n---\n\n**Q3: 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.**\n\nA3: Thanks for the question. This is a good point. At present, we cite well-known generalization bounds for MLPs to motivate the idea that controlling spectral norms generally supports generalization. To strengthen this argument, we will add references to additional works demonstrating similar results for transformer-based models. In fact, [1] and its extension [2] derive generalization bounds for transformer models that also involve the spectral norms of the weight matrices.\n\n\n[1] Edelman, et al.. Inductive Biases and Variable Creation in Self-Attention Mechanisms. ICML 2022\n\n[2] Trauger, et al. Sequence Length Independent Norm-Based Generalization Bounds for Transformers. AISTATS 2024\n\n\n---\n\n**Q4: the proofs in the supplementary material only cover gaussian initialization and SPO and Cayley-Neumann parametrization are not proved to be affecting the generalization.**\n\nA4: Thanks for the question. The statements on the generalization bounds are based on the norms of the weight matrices.  SPO does not change the norm because it only restricts the optimization to a subset of all orthogonal matrices. The CNP approximation results in small changes in the spectrum which can be bounded by $O(|Q|^{k+1})$ where $|Q|$ remains small due to the merge and reinitialize trick. We will add more discussion on this to address the reviewer’s concerns. \n\nMore importantly, we want to emphasize that our contribution does **NOT** lie in establishing a new generalization bound (nor did we make such a claim in the paper). Instead, we use existing generalization bounds as intuitive motivation and justification for why controlling spectral norms can be beneficial. Providing a rigorous theoretical generalization analysis for LLM training is beyond the scope of this work. Our contribution is the proposed reparameterization of weight matrices and the corresponding training algorithm for LLMs."}, {"review_id": "1eh3ECktlC", "reviewer": "Reviewer_sLfu", "summary": "This 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.", "questions": "What 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?", "limitations": "Yes.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "I have no further concerns or questions, and I will maintain my score as is.", "rebuttal": "We sincerely thank **Reviewer sLfu** for the recognition of our contribution and the constructive comments on our work. We take every comment seriously and hope our response can address the reviewer’s concerns. If there are any remaining questions, we are happy to address them. We summarize all the concerns into the following questions:\n\n---\n\n**Q1: 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\nA1: Great question! The reviewer’s main concern is that, based on the perplexity curves, AdamW appears to converge faster initially, raising doubts about whether POET can sustain its strong performance with fewer training steps.\n\nWe note that this convergence pattern reflects differences in training dynamics, particularly how each method responds to the learning rate scheduler, rather than demonstrating superior efficiency for AdamW. Specifically, POET exhibits distinct training phases, as illustrated in Figure 2 from the main paper. As long as the same learning rate scheduler is applied to both AdamW and POET, POET can still outperform AdamW even with significantly fewer training steps.\n\nTo fully address the reviewer’s concerns, we conducted an ablation study on the Llama-60M model, reducing the training to 1/6 and 1/3 of the total steps reported in Table  A (in the Rebuttal), while keeping all other hyperparameters unchanged.\n\nThe results show that even with significantly fewer training steps, POET consistently achieves a lower final perplexity than AdamW. This demonstrates that POET’s effectiveness is not limited to long training runs and that AdamW does not converge faster than POET, even under shorter training schedules. That said, we indeed find that the performance gap between POET and AdamW tends to widen with longer training steps, which aligns with the findings reported in our main paper.\n\n**Table A. Results of much fewer training steps for the 60M model.**\n\n| Training Steps | AdamW | POET (FS, b=½) |\n| :------------- | :---: | :--------------: |\n| 50,000         | 28.72 |      **28.68** |\n| 100,000        | 27.38 |      **27.03** |\n| 300,000 (Paper)| 26.68 |      **25.37** |\n\nAdditionally, we would like to share some insights on combining AdamW and POET training. We refer the reviewer to Table B (in the Rebuttal), which presents results on mixing strategies, such as pretraining with AdamW and then continuing (finetuning) with POET on downstream tasks, or vice versa. Our observations show no fundamental limitations in mixing AdamW and POET: a model pretrained with POET can be further trained with AdamW, and vice versa, while still achieving strong downstream performance. However, we find that using POET throughout the entire training always yields better performance. More importantly, using this mixed training strategy will destroy many nice properties such as weight spectrum and energy preservation. This may also be the reason that the mixed training strategy performs worse than using POET alone.\n\n\n\n---\n\n**Q2: 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\nA2: Thanks for the question. Yes, after performing the merge-then-reinitialize trick, we need to reset the optimizer states since the parameters being updated have changed. In developing this method, we also experimented with merge-then-reinitialize without resetting the optimizer states, but consistently observed worse performance.\n\nWe also note that this type of momentum reset has been used in prior methods [1] and does not harm final performance. We will add further explanation to the merge-then-reinitialize trick section in the revised paper to clarify this point.\n\n[1] Huang, et al. \"SPAM: Spike-aware adam with momentum reset for stable LLM training.\" ICLR 2025.\n\n---\n\n**Q3: 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\nA3: Great suggestion! We agree with the reviewer that whether a base model is good or not depends on its downstream performance. To address the reviewer’s concerns, we have conducted all the required experiments to further demonstrate the effectiveness of our proposed method. \n\n**Performance on downstream tasks**: We start by showing the results on the requested downstream task GLUE [2].\n\nFor comparison, we took the checkpoints of the 1.3B-AdamW baseline and our 1.3B-POET (FS, b=1/2) from Table 2 of the main paper and evaluated them on the standard GLUE benchmark [1]. For each task, we finetuned the models for 10 epochs using a consistent learning rate of 2e−5. The evaluation covers several finetuning strategies: full finetuning (FT) with AdamW, Low-Rank Adaptation (LoRA), and Orthogonal Finetuning (OFT), as well as finetuning with POET.\n\nPerformance 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. The results, summarized in the Table B (in the Rebuttal) below, show that the POET-pretrained model consistently outperforms the AdamW-pretrained baseline across all tasks and finetuning methods. This analysis further validates the effectiveness of POET, and we will include all these results in the revised paper.\n\n**Table B: Downstream performance comparison on the GLUE benchmark between POET- and AdamW-pretrained models.**\n\n| Metric | FT         |         | LoRA       |       | OFT        |        | POET       |       |\n|--------|------------|------------|------------|------------|------------|------------|------------|------------|\n| Model  | 1.3B-AdamW | 1.3B-POET  | 1.3B-AdamW | 1.3B-POET  | 1.3B-AdamW | 1.3B-POET  | 1.3B-AdamW | 1.3B-POET  |\n| COLA   | 0.361      | **0.523**  | 0.423      | **0.460**  | 0.388      | **0.437**  | 0.435      | **0.505**  |\n| MNLI   | 0.658      | **0.818**  | 0.772      | **0.822**  | 0.774      | **0.812**  | 0.804      | **0.821**  |\n| MRPC   | 0.696      | **0.824**  | 0.689      | **0.760**  | 0.689      | **0.740**  | 0.806      | **0.826**  |\n| QNLI   | 0.818      | **0.885**  | 0.843      | **0.870**  | 0.842      | **0.855**  | 0.856      | **0.892**  |\n| QQP    | 0.829      | **0.902**  | 0.865      | **0.880**  | 0.867      | **0.877**  | 0.889      | **0.902**  |\n| RTE    | 0.534      | **0.661**  | 0.534      | **0.567**  | 0.531      | **0.538**  | 0.653      | **0.682**  |\n| SST2   | 0.914      | **0.920**  | 0.908      | **0.927**  | 0.915      | **0.924**  | 0.904      | **0.931**  |\n| STSB   | **0.880**  | 0.873      | 0.818      | **0.865**  | 0.762      | **0.791**  | 0.878      | **0.887**  |\n\n**Performance on a larger model**: We now present results on pretraining with POET and AdamW at a larger model scale. Since this research is conducted in an academic setting, our computing resources are limited. As a result, the largest model we could feasibly pretrain for a controlled comparison within the short rebuttal period was a 3B-parameter model.\n\nTo this end, we trained a 3B-parameter model on a single 8×H100 node using 5 billion tokens with both AdamW and POET (FS, b=1/2). We emphasize that, due to the time constraints of the rebuttal phase, we were unable to train the 3B model on the same number of tokens as in the main paper. Instead, we used only 1/10 of the training tokens used for the 1.3B model, which explains why its performance is slightly lower than that of the 1.3B model. It is important to note that, to ensure a direct comparison focused purely on scalability, we did not perform any hyperparameter tuning. We simply reused the exact hyperparameters applied to the 1.3B models reported in Table 2 of the main paper.\n\nThe results, summarized in Table C (in the Rebuttal), report the final validation perplexity for both methods. We observe that POET maintains a similar performance advantage over AdamW at the 3B scale, consistent with the results shown in Table 2 of the main paper for smaller models. This demonstrates that POET’s advantages are not limited to the 1.3B scale but instead exhibit robust scalability to larger models.\n\n**Table C. Validation perplexity of training a Llama 3B model on 5 billion tokens.**\n\n|                                  | AdamW | POET(FS,b=1/2) |\n|----------------------------|-------------|-----------------------|\n| validation perplexity | 19.61      | 16.90                 |\n\n\n[2] Wang, et al. \"GLUE: A multi-task benchmark and analysis platform for natural language understanding.\" arXiv preprint arXiv:1804.07461."}, {"review_id": "N5aSg7cdi3", "reviewer": "Reviewer_UGEa", "summary": "The 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.", "questions": "1. 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?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "The authors addressed my concerns and questions and I still kept my previous scores.", "rebuttal": "We sincerely thank **Reviewer UGEa** for the recognition of our contribution and the constructive comments on our work. We take every comment seriously and hope our response can address the reviewer’s concerns. If there are any remaining questions, we are happy to address them. We summarize all the concerns into the following questions:\n\n---\n\n**Q1: What are the bottlenecks in applying the current method to larger models such as those with 7B or 70B parameters? The 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\nA1: We appreciate the reviewer’s helpful remark and agree that additional scaling experiments would further strengthen the impact of our work. To address the reviewer’s concerns, we provide an experiment to demonstrate POET's scalability to 3B models (due to the short rebuttal period, we have to use less training tokens than our main paper). The experiment details are given later.\n\nFirst, we want to clarify that there are no fundamental limitations to applying POET to larger models. Because POET’s parameter/memory complexity depends on a hyperparameter $b$, which can usually be set to a value much smaller than $m,n$, we can control the computational budget required by POET. In our experiments, we have tested $b=1/8,1/4,1/2m$ for $R$ and $b=1/8,1/4,1/2n$ for $P$. However, POET with smaller memory budget (i.e., smaller b) will require longer iterations to reach the same performance, but it provides a useful way to trade-off performance and compute budget. In practice, we find that POET demonstrates better parameter-efficiency than standard AdamW training.\n\nWhile there are no fundamental limitations for scaling POET, training bigger models requires (1) additional engineering efforts; and (2) more GPU compute budget. \n\nRegarding (1), pretraining models beyond 7B parameters would necessitate optimizing the implementation for multi-node setups (e.g., inter-node communication overhead). However, because POET shares conceptual similarities with adapter-based finetuning methods, i.e., adding adapter weights to the base layers, a design known for its scalability. We anticipate straightforward scaling to models of 70B parameters and beyond given enough compute resources. Our current experiments support POET’s smooth scalability up to the 1.3B model. The experiments on 1.3B models reuse the same hyperparameters from that on much smaller models, as we do not have the resources to perform hyperparameter search for the larger  models. See the main results in Table 2 of the main paper.\n\nRegarding (2), our experiments are unfortunately constrained by the GPU budget of an academic lab, making it infeasible for us to scale POET to substantially larger models within the scope of this work. \n\nAlthough we do not have the necessary resources for the scaling experiment to 7B models, we address the reviewer’s concerns by scaling POET to the 3B model during this rebuttal phase. We find that training POET on 3B models still shows significant and consistent performance gains compared to standard training with AdamW. \n\nThis experiment was designed to provide a direct comparison focused purely on scalability; therefore, we did not perform any hyperparameter tuning. Instead, we reused the exact hyperparameters from the 1.3B models reported in Table 2 of the main paper. We trained a 3B-parameter model on a single 8×H100 node with 5 billion training tokens using AdamW and POET (FS, b=1/2).\n\nWe emphasize that, due to the time constraints of the rebuttal phase, we could not train the 3B model on the same number of tokens as in the main paper. Instead, we used only 1/10 of the training tokens used for the 1.3B model, which explains why the performance is slightly lower than that of the 1.3B model. The results, summarized in the following table, report the final validation perplexity for both methods. We observe that POET maintains a similar performance advantage over AdamW at the 3B scale, consistent with the findings shown in Table 2 from the main paper for smaller models. This confirms that POET’s benefits are not limited to the 1.3B scale but extend robustly to larger models.\n\n**Table A. Validation perplexity of training a Llama 3B model on 5 billion tokens.**\n\n|                                  | AdamW | POET(FS,b=1/2) |\n|----------------------------|-------------|-----------------------|\n| validation perplexity | 19.61      | 16.90                 |\n\nWe hope these results address the reviewer’s concerns and will include them in the revised paper. We fully agree that demonstrating POET’s scalability to even larger models would further strengthen the impact of our work, and pursuing this will be a major direction for future research.\n\n---\n\n**Q2: What is the impact of the Cayley-Neumann iteration on training time and efficiency?.**\n\nA2: Thank you for raising this question. We assume the question is about the number of Neumann series terms, we have included an ablation study in the paper (Table 5) to evaluate how varying the number of Neumann terms affects performance, using POET-FS with b=1/2 to train LLaMA-130M. We report the final validation perplexity and additionally the training time in the table below.\n\n**Table B: Ablation study: validation perplexity and training time with different numbers of Neumann series.**\n\n| Scheme   | Perplexity   | Time       |\n|----------|--------------|------------|\n|    k=1   | Not converged| 3.51 iters/s  |\n|    k=2   | 22.56        | 3.41 iters/s  |\n|    k=3   | 21.54        | 3.32 iters/s  |\n|    k=4   | 20.22        | 3.27 iters/s  |\n|    k=5   | 20.19        | 3.22 iters/s  |\n\n---\n\n**Q3: Does the Cayley-Neumann parameterization method encounter numerical stability issues when handling large-scale orthogonal matrices?**\n\nA3: Thank you for raising this question. We understand the reviewer’s valid concern that performing orthogonalization on large matrices might encounter numerical stability issues. However, we did not encounter any numerical stability issues when dealing with large orthogonal matrices. \n\nImproving numerical stability is in fact one of the motivations for us to propose the Cayley-Neumann parameterization. To this end, SPO decomposes the training of large orthogonal matrices to many much smaller ones. In the Cayley-Neumann parameterization, Neumann series approximation gets rid of the matrix inverse, further enhancing its numerical stability.\n\nMoreover, we refer the reviewer to Table 2 of the main paper, where we compare the validation perplexity between POET and AdamW. For readability, the relevant results are summarized in Table C below. We find that POET with block-stochastic SPO demonstrates competitive performance against AdamW across different model scales. The key advantage of this approach is that the block size for orthogonalization is fixed (e.g., 256×256 in our settings) and does not grow with the layer dimension. Larger models can be simply trained with more of these small, fixed-size orthogonal matrices. Therefore, numerical stability is not a problem here. Specifically for the Neumann series, as long as we set the number of terms $k\\geq 3$, we never encounter any numerical stability issues in practice.\n\n\n**Table C. Validation perplexity comparison of POET-BS and AdamW.**\n\n| Model (# tokens)       | 60M (30B)               | 130M (40B)              | 350M (40B)              | 1.3B (50B)             |\n|-----------------------|-------------------------|-------------------------|-------------------------|------------------------|\n| AdamW                 | 26.68 (25.30M)          | 20.82 (84.93M)          | 16.78 (302.38M)         | 14.73 (1.21B)          |\n| POET(BS,b=256) | **25.29** (9.66M)       | **19.88** (22.33M)      | **16.27** (60.32M)          | **14.56** (118.26M)        |"}, {"review_id": "7Mx1WRy0DQ", "reviewer": "Reviewer_jtaX", "summary": "This 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.", "questions": "1. 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.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "My concerns have been addressed. I will maintain my positive score.", "rebuttal": "We sincerely thank **Reviewer jtaX** for the recognition of our contribution and the constructive comments on our work. We take every comment seriously and hope our response can address the reviewer’s concerns. If there are any remaining questions, we are happy to address them. We summarize all the concerns into the following questions:\n\n---\n\n**Q1: 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.**\n\nA1: Great suggestion! To address the reviewer’s concerns, we have conducted the requested experiments to further demonstrate the effectiveness of POET-trained models.\n\nFor comparison, we used the checkpoints of the 1.3B-AdamW baseline and our 1.3B-POET (FS, b=1/2) models from Table 2 of the main paper and evaluated them on the standard GLUE benchmark [1], which includes a diverse set of NLP tasks such as sentence classification, sentiment analysis, and question answering. For each task, we fine-tuned the models for 10 epochs using a consistent learning rate of 2e−5. The evaluation covers several finetuning strategies: full finetuning (FT) with AdamW, Low-Rank Adaptation (LoRA), and Orthogonal Finetuning (OFT), as well as finetuning with POET.\n\nPerformance 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. The results, summarized in Table A (in the Rebuttal), show that the POET-pretrained model consistently outperforms the AdamW-pretrained baseline across all tasks and finetuning methods. Moreover, we find that POET-pretrained models, if also finetuned by POET, achieved the best overall downstream performance. These results further validate the effectiveness of our proposed method, and we will include these results, along with the table, in the revised version of our paper.\n\n**Table A. Downstream performance comparison on the GLUE benchmark between POET and AdamW-pretrained models.**\n\n| Metric | FT |  | LoRA|  |OFT |  | POET|  |\n|--|--|--|--|--|--|--|--|--|\n| Model  | 1.3B-AdamW | 1.3B-POET  | 1.3B-AdamW | 1.3B-POET  | 1.3B-AdamW | 1.3B-POET  | 1.3B-AdamW | 1.3B-POET  |\n| COLA   | 0.361      | **0.523**  | 0.423      | **0.460**  | 0.388      | **0.437**  | 0.435      | **0.505**  |\n| MNLI   | 0.658      | **0.818**  | 0.772      | **0.822**  | 0.774      | **0.812**  | 0.804      | **0.821**  |\n| MRPC   | 0.696      | **0.824**  | 0.689      | **0.760**  | 0.689      | **0.740**  | 0.806      | **0.826**  |\n| QNLI   | 0.818      | **0.885**  | 0.843      | **0.870**  | 0.842      | **0.855**  | 0.856      | **0.892**  |\n| QQP    | 0.829      | **0.902**  | 0.865      | **0.880**  | 0.867      | **0.877**  | 0.889      | **0.902**  |\n| RTE    | 0.534      | **0.661**  | 0.534      | **0.567**  | 0.531      | **0.538**  | 0.653      | **0.682**  |\n| SST2   | 0.914      | **0.920**  | 0.908      | **0.927**  | 0.915      | **0.924**  | 0.904      | **0.931**  |\n| STSB   | **0.880**  | 0.873      | 0.818      | **0.865**  | 0.762      | **0.791**  | 0.878      | **0.887**  |\n\n[1] Wang, Alex, et al. \"GLUE: A multi-task benchmark and analysis platform for natural language understanding.\" arXiv preprint arXiv:1804.07461 (2018).\n\n---\n\n**Q2: 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\nA2: Great suggestion! The butterfly factorization used in BOFT is a parameter-efficient orthogonal parameterization, and we agree with the reviewer that it can be a good way to further improve POET’s parameter-efficiency. We identified two primary ways to incorporate this factorization into POET training. The first is to directly replace our block-sparse SPO with a fixed butterfly factorization. While this approach offers better parameter efficiency, it involves additional matrix-matrix multiplications (butterfly factorization requires $\\log(n)$ butterfly factors). This additional matrix-matrix multiplication, while being acceptable for finetuning (due to much shorter training runs), is computationally prohibitive for much longer pre-training schedules. Given this constraint, we could not feasibly train a model with this variant in the rebuttal period.\n\nTherefore, we implemented and tested the second variant. This approach adapts our \"merge-then-reinitialize\" step by replacing the random block permutation in POET-BS with a permutation ordered according to the butterfly factorization pattern. For a direct comparison, we trained two 60M parameter models on 5e10 tokens: one with our standard POET-BS and one with POET using this butterfly permutation. It effectively unrolls the butterfly factorization to different iterations of training. We ensured a fair comparison by using an identical parameter budget and the same set of hyperparameters for both methods.\n\nThe results in Table B (in the Rebuttal) below indicate that standard POET-BS training achieves better performance. We hypothesize that this advantage stems from the increased stochasticity in standard POET-BS, which seems particularly beneficial during the exploratory phase of pretraining. Moreover, unrolling the butterfly factorization across different training iterations can be suboptimal, as the butterfly factors are not learned simultaneously. Each factor is updated in separate steps, potentially reducing overall expressiveness. Finally, while the BOFT framework excels in parameter efficiency, strong regularization, and its constrained orthogonal matrix structure, making it highly effective for finetuning tasks. In contrast, pretraining tasks, as explored in our paper, demand greater flexibility.\n\nWe will include this new baseline comparison and the related discussion in the revised paper.\n\n**Table B. The validation perplexity comparison between POET and POET with the butterfly parameterization from the BOFT paper.**\n\n|| POET-BS| POET-BS (Butterfly)|\n|--|--|--|\n| BS=256   | **25.75** | 26.73          |\n| BS=128   | **27.47** | 29.16          |\n| BS=64     | **31.82** | 32.93          |\n\n---\n\n**Q3: I would like to confirm that the parameters updated by the optimizer are Q \n in Eq. (6), not G itself.**\n\nA3: Yes, the reviewer’s understanding is correct. The parameters updated by the optimizers are the Q in Equation 6.\n\n---\n\n**Q4: Could you provide the training time of the proposed method?**\n\nA4: Great suggestion! We are happy to perform the required training time comparison to address the reviewer’s concerns. \n\nTo provide a fair and comprehensive comparison of training efficiency between POET and AdamW, we evaluate the training efficiency using two distinct metrics, following the experimental setup in Figure 9 of our paper. First, we report both the absolute clock time per iteration to measure computational overhead. Second, we compare the total time each method requires to reach an equivalent target validation perplexity, thereby assessing overall convergence speed. We note that, while POET has a higher cost per iteration due to the orthogonalization step, its overall time to a target performance may in fact be superior to standard AdamW training. By presenting both metrics, we offer a more complete and transparent assessment of our method's practical efficiency. \n\nFollowing the experimental setup in Figure 9 of the main paper, we report two metrics: iteration-wise training time (where higher iters/s indicates less time per iteration, so larger values are better) and total training time (where smaller values are better) required to reach approximately the same validation perplexity of 25.2 when training a Llama‑60M model. The results are summarized in the following tables:\n\n**Table C. Iteration-wise training time of POET.**\n\n|  | AdamW   | POET (FS, b=½) | POET (FS, b=1/4) |\n|--|--|--|-----|\n| Iteration-wise Training time   | 3.97 iters/s    | 3.27 iters/s     | 3.35 iters/s     |\n\n**Table D. Total training until reaching the same validation perplexity of 25.2.**\n\n|   | AdamW       | POET (FS, b=½) |\n|----|---|--|\n| Training time to same perplexity  | 50.4 h   | 23.8 h   |\n\nIndeed POET’s per-iteration training time is currently slower than AdamW’s. This overhead stems from several operations, including non-trivial matrix orthogonalization and sparse matrix multiplications, which involve advanced indexing and matrix permutations. While we have taken steps to reduce this overhead, such as implementing low-level optimizations with custom CUDA kernels and introducing the efficient Cayley–Neumann parameterization. But we do want to recognize that there remains room for further improvements in clock-time efficiency.\n\nWe would also like to take this opportunity to clarify the scope of our work. Our primary goal is to demonstrate that the proposed principled energy/spectrum-preserving training paradigm can produce models that are consistently better than standard AdamW training up to a meaningful model scale (1.3B in the main paper, 3B in the rebuttal, see Table A in the response to Reviewer UGEa). At this stage, our focus has not been on heavily optimizing POET for large-scale, multi-node training, partly due to the resource constraints of our academic setting.\nHowever, we are confident that POET’s training efficiency can be significantly improved with more dedicated optimization to the low-level computations. We believe that improving POET’s efficiency and scalability represents a promising direction for future work.\n\n---\n\n**Q5: Line 417: R and I should be in bold, as they represent matrices.**\n\nA5: Thanks for the remark! We have updated it in the revised paper."}]

Author rebuttal:

Meta-review: This work introduces POET, a reparameterized training algorithm for large language models that optimizes neurons through orthogonal equivalence transformations. Each neuron is reparameterized with two learnable orthogonal matrices and a fixed random weight matrix, preserving the spectral properties of weight matrices to enable stable optimization and improved generalization. Efficient approximations further make the method scalable to large networks, and experiments validate its effectiveness for LLM training.

During the rebuttal, the authors provided sufficient evidence and clarification, demonstrating that POET scales effectively to larger models (3B), achieves lower validation perplexity, delivers superior downstream performance across multiple finetuning methods, and converges faster than AdamW. These results convincingly address the reviewers’ concerns.

**Recommendation**: Accept.
Decision: Accept (poster)