PaperlayerView exact revision on arXiv

Source-linked AI summary

Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation

Kexuan Shi, Hanxuan Li, Zeju Qiu, Yandong Wen, Simon Buchholz, Weiyang Liu

arXiv:2605.12492v1cs.LGstat.ML

TL;DR

LLM training needs optimizers that are both efficient and stable, while weight spectral norms may drift during Muon-based optimization. Pion instead uses left and right orthogonal transformations to preserve weight spectra, and performs competitively across pretraining, supervised finetuning, and reinforcement learning.

  • Problem

    Designing optimizers that are efficient and stable remains difficult because weight spectral norms may drift during training.

  • Method

    Pion updates weight matrices through left and right orthogonal transformations on the iso-spectral manifold, preserving their singular-value spectra without explicit reparameterization.

  • Results

    Pion is competitive with standard optimizers across pretraining, supervised finetuning, and reinforcement learning.

  • Takeaways & Limitations

    Pion offers a spectrum-preserving optimizer for stable training across diverse LLM optimization settings.

  • Takeaways & Limitations

    The convergence result assumes smoothness, a lower-bounded objective, bounded stochastic-gradient variance, and iterates remaining on the iso-spectral manifold.

Abstract

from arXiv · show

We introduce Pion, a spectrum-preserving optimizer for large language model (LLM) training based on orthogonal equivalence transformation. Unlike additive optimizers such as Adam and Muon, Pion updates each weight matrix through left and right orthogonal transformations, preserving its singular values throughout training. This yields an optimization mechanism that modulates the geometry of weight matrices while keeping their spectral norm fixed. We derive the Pion update rule, systematically examine its design choices, and analyze its convergence behavior along with several key properties. Empirical results show that Pion offers a stable and competitive alternative to standard optimizers for both LLM pretraining and finetuning.

1 Introduction

Pion addresses the growing need for efficient and stable optimizers for scaling LLMs by directly preserving weight spectra during training. It applies left and right orthogonal transformations without auxiliary reparameterization or explicit normalization.

  • Motivation: Scaling LLMs makes training increasingly difficult, intensifying the need for optimizers that are both efficient and stable.Training stability can be partially characterized by µP, which constrains weight and update spectral norms to maintain width-invariant activation scale.
  • Algorithmic spectrum control: Pion derives updates on the iso-spectral manifold, preserving the weight spectrum without explicit normalization.Its weight spectral norm is upper bounded, and its update spectral norm is also guaranteed to be upper bounded, making it compatible with µP.
  • Relation to POET: Unlike POET, Pion removes auxiliary orthogonal parameterization and turns spectrum preservation into a direct optimizer.POET learns two orthogonal factors around a randomly initialized base weight, whereas Pion updates the weights themselves.
  • Core mechanism: Pion updates each weight matrix through left and right orthogonal transformations, preserving its singular-value spectrum throughout training.This mechanism avoids the auxiliary parameterization associated with complications such as loss spikes and careful momentum design.

2 Pion: A Spectrum-Preserving Optimizer for LLM Training

Pion updates weight matrices through orthogonal transformations that preserve their spectra, interpreting changes as rotational motion rather than rescaling. Experiments motivate scale-consistent updates, geometry-aligned momentum, alternate transformations, and second-order exponential approximations as core design choices.

  • Update rule: Pion updates identity factors directly on the orthogonal group, avoiding explicit reparameterization while preserving the weight matrix spectrum.The update operates on the left and right orthogonal factors of each weight matrix.
  • Geometric interpretation: Orthogonal transformations preserve row and column ℓ2 norms, so Pion’s update norm measures angular deviation rather than magnitude rescaling.This distinguishes Pion’s rotational motion from vanilla gradient descent, which generally mixes magnitude and direction changes.
  • Scale consistency: RMS-controlled scale consistency stabilizes training at larger learning rates, whereas bilateral normalization adds no further stability or final-performance gains.The naive update performs well only at small learning rates and diverges when the learning rate is large.
  • Momentum: Lie-algebra momentum converges fastest, slightly outperforming transported ambient-space momentum, while Lie+Lie combinations outperform ambient-space and mixed variants.All first-order momentum formulations remain spectrum-preserving; ambient-space momentum is most efficient but has biased estimates from tangent-space mismatch.
  • Efficient updates: Alternate updates retain performance close to bilateral updates while reducing per-step computation by decoupling input-side and output-side transformations across iterations.For Lie+Lie, alternate update reaches a final loss of 3.3654, about 0.23% higher than bilateral update.
  • Exponential approximation: Pion adopts a second-order exponential approximation because it nearly preserves the spectrum, avoids repeated error compounding, and makes higher-order approximations unnecessary.The first-order approximation degrades convergence and singular-value preservation, while the Cayley transform provides only modest gains.

3 Experiments and Results

Experiments evaluate Pion against AdamW and Muon across LLM pretraining, supervised finetuning, and reinforcement learning. Results emphasize Pion’s training stability, including challenging normalization-free and ultra-deep settings, while remaining competitive in downstream tasks.

  • Experimental scope: Pion is evaluated against AdamW and Muon across pretraining, supervised finetuning, and reinforcement learning scenarios.The experiments cover both pretraining and post-training, including supervised finetuning and reinforcement learning.
  • Pretraining stability: Training stability is monitored through validation loss, maximum attention logits, and SwiGLU activation, weight, and output norms.These quantities are used as indicators of stability during LLaMA-based pretraining.
  • Normalization-free pretraining: In normalization-free pretraining, Pion stress-tests whether the optimizer alone can regulate scale when AdamW and Muon soon fail.The setting removes all normalization layers from a 60M LLaMA-based model.
  • Ultra-deep pretraining: 0.0892 mean standard deviation makes Pion the most stable optimizer in ultra-deep training, versus 0.0931 for AdamW and 0.0927 for Muon.The models scale from 8 to 200 layers, and stability is measured by the mean standard deviation of the local loss trajectory.
  • Supervised finetuning: Pion achieves a highly competitive stability-plasticity tradeoff against AdamW and Muon in supervised finetuning.Experiments use full-parameter finetuning of Qwen2.5-1.5B and Llama-3.2-3B on MetaMathQA and Magicoder-Evol-Instruct-110K.
  • Reinforcement learning with verifiable reward: Pion is tested for RLVR with GRPO because its spectrum-preserving updates align with the spectral structure observed in RLVR training.Experiments use Qwen3-1.7B and DeepSeek-R1-Distill-Qwen-1.5B on DeepMath, with evaluation on five mathematical reasoning benchmarks.

4 Related Work and Concluding Remarks · Appendix

The paper situates Pion within efforts to stabilize LLM training through scale control and matrix-aware optimization. It concludes that Pion preserves weight spectra and remains competitive across multiple training settings.

  • 4 Related Work and Concluding Remarks: Stable LLM training combines optimization dynamics with scale control.Related approaches include AdamW with normalization layers or scale-aware parameterizations.
  • 4 Related Work and Concluding Remarks: AdamW is commonly paired with normalization layers for stable LLM training.The passage also identifies scale-aware parameterizations as an alternative scale-control strategy.
  • 4 Related Work and Concluding Remarks: Matrix-aware optimization improves training stability by exploiting weight-matrix structure.The related-work discussion places Pion among recent structure-aware optimization methods.
  • 4 Related Work and Concluding Remarks: Pion uses orthogonal equivalence transformations to provide spectrum-preserving optimization for stable training.Its design avoids relying on explicit reparameterization.
  • 4 Related Work and Concluding Remarks: Pion provably preserves the weight spectrum and maintains minimal hyperspherical energy throughout training.These properties are achieved without explicit reparameterization.
  • 4 Related Work and Concluding Remarks: Pion is competitive with standard optimizers across pretraining, supervised finetuning, and reinforcement learning.The conclusion attributes this assessment to the reported empirical results.

A Geometric Structure of Pion’s Update

Pion updates preserve a weight matrix’s singular values by orthogonally rotating its left and right singular subspaces. The update’s magnitude reflects angular movement, with Lie-algebra norms controlling average and maximum rotation angles.

  • Singular-value preservation: Orthogonal left and right transformations preserve W_t’s singular values while rotating its left and right singular subspaces.The transformed factorization remains a valid singular value decomposition up to orthogonal basis changes.
  • Update magnitude: The Frobenius norm of the displacement measures the aggregate angular movement induced by the two rotations.Orthogonal multiplication preserves vector norms and changes row and column directions without scaling them.
  • Planar-rotation structure: Each 2×2 block of exp(−ηΣ) implements an independent planar rotation with angle determined by −ηθ_j.The right-side update represents the right singular space in the basis induced by G_in^t before rotating its corresponding 2D invariant subspaces.
  • Rotation magnitudes: The Frobenius norm of the input-side Lie-algebra element, scaled by √d_in, characterizes average rotation magnitude, while its spectral norm controls the maximum angle.The same characterization applies to G_out^t for the output-side update.

B Convergence Analysis

The convergence analysis characterizes Pion’s spectrum-preserving trajectory as lying on an isospectral manifold and derives its tangent space and first-order stationarity conditions. It then uses a combined stationarity measure under smoothness, lower-boundedness, and bounded-trajectory assumptions.

  • Geometry of Spectrum-Preserving Updates: Spectrum-preserving updates keep the idealized trajectory on the isospectral manifold MW0 of matrices sharing W0’s singular values.This manifold is defined through the initialization W0.
  • Geometry of Spectrum-Preserving Updates: The tangent space at any W ∈ MW0 consists of directions GoutW + WGin generated by skew-symmetric matrices on the left and right.Specifically, Gout ∈ so(m) and Gin ∈ so(n), where so(k) contains matrices satisfying G⊤ = −G.
  • Stationarity on the Isospectral Manifold: A point is first-order critical on MW0 exactly when GW⊤ − WG⊤ = 0 and W⊤G − G⊤W = 0.These conditions arise by requiring the gradient to be orthogonal to every tangent direction.
  • Convergence Conditions: For bilateral updates, convergence is assessed with a combined stationarity measure whose convergence to zero implies first-order convergence on the isospectral manifold.The analysis states that proving the measure tends to zero directly establishes first-order convergence.
  • Convergence Conditions: The analysis assumes Euclidean L-smoothness, a globally lower-bounded objective, and boundedness along the trajectory.The boundedness assumption is motivated by compact isospectral dynamics and bounded truncated updates for sufficiently small step size.

B.1 Simplified Single-Side Analysis · B.2 Deterministic Bilateral Update · B.3 Stochastic Bilateral Update

The analysis derives descent identities for Pion’s single-side rotations, then extends them to deterministic and stochastic bilateral updates. Under smoothness, boundedness, and sufficiently small learning rates, the updates converge toward first-order stationary solutions or neighborhoods on the isospectral manifold.

  • B.1 Simplified Single-Side Analysis: The truncated-exponential expansion isolates first-order descent terms and groups the remaining terms into higher-order contributions.The expansion is used to establish the single-side descent identities later applied to bilateral updates.
  • B.1 Simplified Single-Side Analysis: Single-side in-side and out-side rotations align with the Riemannian descent direction, while controlled second-order terms yield convergence for the alternating update.The bilateral update applies both descent directions in the same step.
  • B.2 Deterministic Bilateral Update: The deterministic bilateral expansion combines both single-side descent identities and bounds the remainder under the stated assumptions.The remainder contains terms of order η2 and higher, with boundedness used in the final control.
  • B.2 Deterministic Bilateral Update: Under Assumptions B.4–B.6 and sufficiently small learning rate η with cη > 0, deterministic bilateral Pion converges to a first-order stationary point on the isospectral manifold.The result is stated in Theorem B.7 and supported by the preceding descent argument.
  • B.2 Deterministic Bilateral Update: The deterministic stationarity measure converges at rate O( 1, as stated in the convergence result.The supplied passage truncates the rate expression after O( 1.
  • B.3 Stochastic Bilateral Update: The stochastic analysis models the mini-batch gradient as an unbiased estimator with bounded variance and transfers the deterministic bilateral descent argument through conditional expectation.The stochastic noise is analyzed under the unbiasedness and bounded-variance assumptions.
  • B.3 Stochastic Bilateral Update: Under Assumptions B.4–B.6, unbiased stochastic gradients, bounded variance, and η=C/ T with sufficiently small C >0, Theorem B.8 gives stochastic convergence.The supplied theorem passage truncates the displayed bound or rate following the setup.
  • B.3 Stochastic Bilateral Update: The stochastic update converges to a stochastic first-order stationary neighborhood on the isospectral manifold and follows the standard stochastic nonconvex rate.The rate statement is associated with sufficiently small C.

C Additional Discussion on Computation Complexity · D Experimental Details

The discussion analyzes Pion’s computational and memory costs, showing that alternate updates halve dominant update-side computation while practical costs remain moderate. Experimental details are provided for the configurations used throughout the manuscript.

  • C Additional Discussion on Computation Complexity: Alternate updates reduce the dominant update-side computation by roughly a factor of two when averaged over consecutive steps.Only one exponential map is applied at each step.
  • C Additional Discussion on Computation Complexity: Pion’s persistent optimizer states use separate first- and second-moment buffers on the input- and output-side Lie algebras.Their relative overhead versus Adam depends on matrix shape: it is a small constant for nearly square matrices and grows with aspect ratio for highly rectangular matrices.
  • C Additional Discussion on Computation Complexity: With buffer reuse or in-place construction, additional peak temporary memory is O(d2out+doutdin).Naively materializing both RMS-scaling products can require two extra dout×din buffers.
  • C Additional Discussion on Computation Complexity: The practical cost experiments run on 8 NVIDIA H100 GPUs with NVLink, using distributed data parallelism without gradient accumulation and reporting peak per-GPU memory and wall-clock step time.The measured results are reported under the same training configuration.
  • C Additional Discussion on Computation Complexity: Full Pion uses 59,839 MB per GPU, 16.0% more than AdamW and 26.6% more than Muon.The extra memory mainly comes from input- and output-side Lie algebra states and temporary matrix products.
  • C Additional Discussion on Computation Complexity: Pion takes 0.5679 seconds per step versus 0.3932 seconds for AdamW and 0.5505 seconds for Muon.Thus, Pion is about 44.4% slower than AdamW but only about 3.2% slower than Muon; removing second-moment buffers does not noticeably change step time.
  • D Experimental Details: The experimental-details section specifies the detailed setups and configurations for each experiment presented in the main manuscript.It serves as the configuration reference for the reported experiments.

D.1 Experiments for Design Principles … D.4 Reinforcement Learning with Verifiable Reward

The experiments evaluate Pion’s design principles, pretraining, supervised fine-tuning, and reinforcement learning under specified optimizer, model, dataset, and evaluation settings. Across these sections, the setups use Megatron-LM or task-specific frameworks, controlled training budgets, and benchmark-based evaluation.

  • D.1 Experiments for Design Principles: All design-principle experiments use Megatron-LM with bf16 mixed precision, batch size 512, cosine decay to 0.01 times the initial learning rate, and lr=1e-3 as the consistently best tested value.The shared configuration also uses a second-order matrix-exponential approximation and the T5-base tokenizer.
  • D.1 Experiments for Design Principles: Design-principle ablations rescale bilateral-normalization updates to match the original norm, apply RMS scaling with coefficient 0.2, and use β1 =0.9 and β2 =0.95.Direct momentum without RMS scaling is stable only under extremely small learning rates.
  • D.1 Experiments for Design Principles: The alternate-update experiments retain momentum accumulation on both sides because disabling it on the non-updated side degrades convergence, while Cayley-transform computations use float32 and torch.linalg.solve.The Cayley transform follows the standard formulation (I −1 2S)−1(I + 1 2S).
  • D.2 Pretraining: Pretraining experiments use Megatron-LM with bfloat16 precision, the T5-base tokenizer, cosine decay to 0.01 times the maximum learning rate, and a maximum learning rate of 5e-4 for all methods.Muon and Pion use an RMS scaling coefficient of 0.2.
  • D.3 Supervised Fine-tuning: Supervised fine-tuning uses Qwen2.5-1.5B and Llama-3.2-3B on 50K-sample subsets of MetaMathQA and Magicoder-Evol-Instruct-110K.Training runs for 3 epochs with global batch size 64, learning rate 1×10−5, cutoff length 4096 tokens, and FP32 optimizer updates in LLaMA-Factory on NVIDIA H200 GPUs.
  • D.3 Supervised Fine-tuning: Fine-tuning evaluation uses LM Evaluation Harness, with GSM8K evaluated 5-shot, other benchmarks zero-shot, deterministic greedy decoding, and HumanEval reported as pass@1.These protocols standardize generation and reporting across the supervised fine-tuning benchmarks.
  • D.4 Reinforcement Learning with Verifiable Reward: RLVR experiments train Qwen3-1.7B and DeepSeek-R1-Distill-Qwen-1.5B on DeepMath with maximum context lengths of 4096 and 8192, respectively.The VeRL pipeline uses vLLM rollouts and GRPO, with learning rate 1×10−6, 12 rollouts per prompt, and FP32 optimizer updates on NVIDIA H200 GPUs.
  • D.4 Reinforcement Learning with Verifiable Reward: RLVR evaluation covers AIME24, AIME25, AMC23, Minerva Math, and OlympiadBench, reporting averaged accuracy across multiple independent samples.Maximum generation lengths are 4,096 tokens for Qwen3-1.7B and 8,192 tokens for DeepSeek-R1-Distill-Qwen-1.5B.

E Limitations

Pion introduces computational and memory overheads from Lie-algebra optimization operations. These costs are largely amortized during standard LLM pretraining with large token batches, minimizing their practical wall-clock impact.

  • E Limitations: Pion incurs computational and memory overheads despite its stability and competitive performance.The passage identifies these overheads as limitations of the optimizer.
  • E Limitations: Computing Lie-algebra gradients and applying the truncated matrix exponential mapping require additional FLOPs.These operations contribute directly to Pion’s computational overhead.
  • E Limitations: Large token batches in standard LLM pretraining largely amortize Pion’s overhead, making its practical wall-clock impact minimal.The passage attributes this amortization specifically to standard pretraining regimes with large token batches.

F Compatibility with Maximal Update Parametrization

Pion is compatible with maximal update parametrization because its orthogonal transformations preserve the weight spectral norm, leaving only the generator spectral scales to control the update condition. The section proposes spectral norm scaling and explicit orthogonalization, then tests hyperparameter transfer under width scaling in a LLaMA-based architecture.

  • Spectral Conditions: Pion’s left and right orthogonal transformations preserve W’s singular values, making its spectral norm invariant after initialization satisfies the Forward Spectral Condition.Unlike standard optimization frameworks, Pion therefore maintains this condition automatically throughout training.
  • Spectral Conditions: Controlling the generators’ spectral scales makes the update inherit W’s Θ-order, satisfying the µP Update Spectral Condition.Enforcing ∥G_out,t∥2 = Θ(1) ensures ∥∆W_t∥2 has the same Θ-order as ∥W_t∥2.
  • Generator Schemes: Pion enforces generator control either by directly scaling spectral norms or by Newton-Schulz orthogonalization, which pushes non-zero singular values toward 1.The second scheme structurally guarantees a bounded Θ(1) norm and uniform update magnitude across active spectral directions.
  • Width-Scaling Experiments: The µP verification scales hidden size, intermediate size, and attention heads while keeping head dimension fixed, sweeping learning rates for the lowest validation loss.The experiments use a LLaMA-based architecture to test hyperparameter transfer under width scaling for both proposed schemes.

G Additional Results … G.4 Reinforcement Learning with Verifiable Reward

The additional results introduce a transported ambient-space momentum variant of Pion and evaluate its spectrum preservation, training stability, fine-tuning update strategies, and RLVR performance. Across these analyses, Pion preserves weight spectra and maintains stable monitored quantities, while update-strategy preferences differ between supervised fine-tuning and reinforcement learning.

  • G.1 Another variant of Pion: Pion is extended with a transported ambient-space momentum variant using learning rate, momentum coefficients, RMS and stability constants, and an alternating-update flag.The algorithm initializes first- and second-moment states and optimizes the weight matrix W_t.
  • G.1 Another variant of Pion: The variant supports an alternating update procedure whose odd-step behavior is explicitly specified in the algorithm.The supplied algorithm excerpt identifies conditional alternating updates and separate handling when t is odd.
  • G.2 Pretraining: Additional pretraining analyses monitor Frobenius norms of representative weights, input and output activations, and maximum attention logits as stability indicators.These quantities are used to characterize optimization stability and activation amplification during large-scale training.
  • G.2 Pretraining: Pion preserves weight-matrix spectra close to initialization, whereas AdamW and Muon substantially distort the original spectra during optimization.This behavior is consistent with Pion’s spectrum-preserving design and reflects different optimization dynamics.
  • G.2 Pretraining: Pion keeps nearly all monitored stability quantities flat, while AdamW grows attention logits and activations and Muon suppresses logits but increases activations and down-projection norms.The comparison covers attention logits, activation magnitudes, and down-projection norms across training.
  • G.3 Supervised Fine-tuning: In supervised fine-tuning, bilateral updates consistently outperform alternating updates, although alternating updates remain competitive with AdamW and Muon.The analysis includes detailed benchmark breakdowns for Qwen2.5-1.5B and Llama-3.2-3B base models.
  • G.4 Reinforcement Learning with Verifiable Reward: In RLVR, alternating updates outperform bilateral updates, while the bilateral approach remains robust with overall performance closely comparable to the alternating strategy.Table 7 reports average accuracy over K generated samples per problem using the avg@K metric.
Loading 2605.12492v1…