Research map / Large language models / Language, safety and interpretability
Reparameterized LLM Training via Orthogonal Equivalence Transformation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes POET, which reparameterizes each neuron weight matrix as R W0 P with a fixed random W0 and two learnable orthogonal matrices, using stochastic primitive optimization and Cayley-Neumann parameterization to scale LLM pretraining.
Problem
Training LLMs with direct AdamW optimization is costly and has suboptimal generalization and poorly controlled weight spectra; existing spectrum regularization is computationally expensive and only constrains the largest singular value.
Approach
Reparameterize each weight matrix as R W0 P with fixed random W0 and learnable orthogonal R and P, so singular values are preserved exactly; scale via stochastic primitive optimization (sampled sub-block orthogonal primitives merged periodically) and a truncated Neumann-series Cayley parameterization for orthogonality.
Key results
- POET-FS (b=1/2) reaches validation perplexity 13.70 on LLaMA-1.3B vs AdamW 14.73 and GaLore 18.33.
- POET-BS (b=256) reaches 14.56 on 1.3B with about one-tenth of AdamW's trainable parameters (118.26M vs 1.21B).
- On GLUE full finetuning, POET-pretrained model outperforms AdamW-pretrained model on all eight tasks (e.g., CoLA 0.523 vs 0.361).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace AdamW with POET fully stochastic, b=1/2 | LLaMA pretraining on C4 · C4 validation, LLaMA-60M to LLaMA-1.3B | perplexity: 26.68 -> 25.37; 20.82 -> 19.94; 16.78 -> 15.95; 14.73 -> 13.70 | Helped |
| replace AdamW with POET fully stochastic, b=1/2 | LLaMA-3B pretraining · LLaMA-3B, 5B tokens | perplexity: 19.61 -> 16.90 | Helped |
| use POET pretrained weights instead of AdamW pretrained weights | GLUE downstream finetuning · GLUE tasks across Full FT, OFT, and POET finetuning | task scores | Helped |
| allocate the parameter budget evenly between left and right orthogonal matrices | POET · Validation perplexity under a fixed parameter budget | perplexity | Helped |
| use normalized Gaussian initialization rather than standard Gaussian, Xavier, or uniform-spectrum initialization (ablation) | POET fully stochastic, b=1/2 · LLaMA-60M, C4 validation | perplexity: 26.22 -> 25.37; 25.79 -> 25.37; 27.29 -> 25.37 | Helped |
| use uniform-spectrum initialization instead of normalized Gaussian initialization (ablation) | POET fully stochastic, b=1/2 · LLaMA-60M, C4 validation | perplexity: 25.37 -> 27.29 | Hurt |
| set merge-then-reinitialize frequency to 5 rather than 400 (ablation) | POET fully stochastic, b=1/2 · LLaMA-60M, C4 validation | perplexity: 25.31 -> 30.29 | Hurt |
| set merge-then-reinitialize frequency to 25 rather than 400 (ablation) | POET fully stochastic, b=1/2 · LLaMA-60M, C4 validation | perplexity: 25.31 -> 27.27 | Hurt |
| set merge-then-reinitialize frequency to 50 rather than 400 (ablation) | POET fully stochastic, b=1/2 · LLaMA-60M, C4 validation | perplexity: 25.31 -> 25.99 | Hurt |
| set merge-then-reinitialize frequency to 200 rather than 400 (ablation) | POET fully stochastic, b=1/2 · LLaMA-60M, C4 validation | perplexity: 25.31 -> 25.37 | No clear effect |
| set merge-then-reinitialize frequency to 1600 rather than 400 (ablation) | POET fully stochastic, b=1/2 · LLaMA-60M, C4 validation | perplexity: 25.31 -> 25.58 | Hurt |
| retain Neumann-series terms rather than drop all terms (k=0) (ablation) | POET fully stochastic, b=1/2 · LLaMA-130M, C4 validation | perplexity: Not converged -> 22.56 | Helped |
Surprising findings
- Uniform-spectrum initialization performs poorly (perplexity 27.29) compared with normalized Gaussian (25.37), suggesting a tradeoff between spectral preservation and expressiveness.
- POET trains more slowly than AdamW after an initial drop but still reaches better final perplexity; POET-FS with b=1/2 beats AdamW trained on nearly triple the tokens.
Limitations
- Slower convergence than AdamW in early training; smaller sampling budgets converge more slowly.
- Spectrum preservation depends on initialization; uniform-spectrum initialization hurts expressiveness.
- Neumann series approximation requires periodic reinitialization and can diverge without enough terms (k=0 diverges).
Experimental setup
Datasets: C4 GLUE
Benchmarks: C4 validation perplexity GLUE
Metrics: validation perplexity accuracy Matthews correlation Pearson correlation hyperspherical energy singular value entropy
Baselines: AdamW GaLore LoRA OFT Muon
Builds on: orthogonal fine-tuning (OFT) GaLore LoRA Cayley parameterization
Closest papers by meaning
- ShiftAddLLM: Accelerating Pretrained LLMs via Post-Training Multiplication-Less Reparameterization (NeurIPS 2024 · similarity 0.87)
- Building on Efficient Foundations: Effective Training of LLMs with Structured Feedforward Layers (NeurIPS 2024 · similarity 0.86)
- Learn To be Efficient: Build Structured Sparsity in Large Language Models (NeurIPS 2024 · similarity 0.86)
- Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA (ICLR 2025 · similarity 0.86)
- Transformer-Squared: Self-adaptive LLMs (ICLR 2025 · similarity 0.86)
- Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.