PC Layer Cuts Llama-1B Training Tokens in Half With Zero Inference Cost

Researchers propose the PC Layer, a plug-in weight reparameterization that reshapes singular-value spectra during LLM pre-training with no inference overhead, achieving up to 2x token efficiency on Llama-1B.

·
·
·
PC Layer Cuts Llama-1B Training Tokens in Half With Zero Inference CostPRO
  • Researchers propose the PC Layer, a plug-in weight reparameterization that reshapes singular-value spectra during LLM pre-training with no inference overhead.
  • The method applies a low-degree polynomial to selected weight matrices each training step, amplifying small singular values and saturating large ones, without expensive SVD.
  • On Llama-1B with AdamW, PC reaches the baseline validation loss with 50% fewer tokens (2× speedup) and improves zero-shot accuracy by +2.06 points across 9 benchmarks.
  • PC also improves training under the Muon optimizer (1.13× token efficiency), and the two methods compose orthogonally since they act on different objects (weights vs. updates).
  • Compute overhead is under 0.4% FLOPs; memory increases ~9–10% due to autograd caching, with no inference cost after training.
  • Code is open-sourced at github.com/Empath-aln/PC-layer with training configs for 271M and 1B Llama-2 models.

Training large language models is expensive, and the community has spent years squeezing more out of every token. A new paper introduces the Preconditioning (PC) Layer, a lightweight architectural module that improves how efficiently a transformer learns during pre-training, by keeping the weight matrices in a healthier numerical state throughout the entire training run.

The problem: ill-conditioned weights slow everything down

Every linear layer in a transformer has a weight matrix whose singular values (think: the "stretching factors" of the transformation in different directions) tend to spread out during training. When the ratio between the largest and smallest singular values grows too large, the condition number of the matrix blows up. This makes gradients noisy in some directions and negligible in others, degrading signal propagation through depth and slowing convergence.

Existing fixes are partial. Normalization layers like RMSNorm and QK-Norm control activations, not weights. Spectral normalization caps the maximum singular value but ignores the rest of the spectrum. What's missing is a method that actively reshapes the entire spectrum throughout training, without adding inference cost.

The idea: reshape the spectrum with a polynomial, not an SVD

The key insight is that you can reshape a matrix's singular values using a low-degree polynomial, without ever computing an expensive singular value decomposition. If W has singular values σ₁ ≥ ... ≥ σₘ, then the matrix p(WW⊤)W has singular values |g(σᵢ)|, where g(σ) = p(σ²)σ is a scalar function. By choosing p to amplify small singular values and saturate large ones, you narrow the spectral spread without flattening it completely.

Pro article

This story is for Pro members

You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.

Trending
  • No trending articles

Comments

avatar

Next Reads