cs.LGAug 4, 2026

Joint Affine Spectral Shaping: Coupling Weight and Bias Updates Beyond Weight-Only Muon

Authors: Gongyue Zhang, Honghai Liu

Organizations: State Key Laboratory of Robotics and Systems, Harbin Institute of Technology, Shenzhen, Shenzhen, 518055, China.

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Matrix spectral optimizers reshape weight-update spectra but usually delegate vector-valued biases to a separate optimizer. We study whether this separation is neutral. We formulate each affine layer as a joint momentum matrix A=[MW,αmb]A=[M_W,αm_b] and apply a capped regularized-inverse spectral map to the complete matrix, producing both the weight and physical bias updates. A strict five-seed ablation on a four-layer BERT-mini trained from scratch on IMDb compares exact-SVD Muon, weight-only inverse shaping, affine-probe inverse shaping, and the proposed joint regularized inverse (JRI). Weight-only inverse shaping raises validation-loss-selected test accuracy from 84.903±0.242%84.903\pm0.242\% to 85.562±0.308%85.562\pm0.308\% and lowers selected test loss from 0.34790.3479 to 0.33450.3345. Allowing bias to alter the joint SVD while retaining an independent Adam bias update does not improve over weight-only inverse shaping. Using the transformed bias jointly raises selected test accuracy to 85.738±0.180%85.738\pm0.180\% and lowers test loss to 0.32910.3291, with all five seeds improving relative to the probe baseline. During the peak-performance window, JRI preserves the eligible weight-update norm while reducing the bias-update norm from 0.020950.02095 to 0.003010.00301, lowers boundary-function share from 86.58%86.58\% to 78.97%78.97\%, and changes the cosine between weight-induced boundary motion and explicit bias from +0.030+0.030 to −0.137-0.137. An independent 22-seed replication yields 85.743±0.203%85.743\pm0.203\% selected test accuracy. These results identify joint affine spectral allocation as a small but consistent extension to weight-only spectral optimization.

Explore similar work

CardsList
Jul 24, 2026cs.LG

Hidden Boundary Motion in Transformer Optimization: Function-Space Orthogonalization of Affine Weight and Bias Updates

Weights and biases are normally optimized as separate parameter tensors, yet they do not represent separate functions when the input to an affine layer has nonzero mean. For an affine map z=Wx+bz=Wx+b with input mean μμ, a weight update contains a sample-independent displacement ΔWμΔWμ that is functionally indistinguishable from a bias update. We call this hidden contribution \emph{boundary motion} and decompose each update into a centered, sample-varying \emph{shape} component and a shared \emph{boundary} component. On a four-layer Transformer trained from scratch on IMDb, the bias-like term gbμ⊤g_bμ^\top has a median norm equal to 0.664 of the raw weight-gradient norm across affine layers and training checkpoints. More strikingly, the median ratio \normΔWμ/\normΔb\norm{ΔWμ}/\norm{Δb} is 134.7, while \normΔWμ/\normΔb+ΔWμ\norm{ΔWμ}/\norm{Δb+ΔWμ} is 0.994. Thus, under AdamW, the observed boundary motion is almost entirely realized through the weight matrix rather than the explicit bias. We implement a diagnostic optimizer, Shape--Boundary Orthogonal AdamW (SBO-AdamW), that optimizes gW−gbμ⊤g_W-g_bμ^\top and gbg_b with independent Adam states and compensates the weight-induced boundary displacement. In a single-seed experiment, SBO-AdamW raises validation accuracy from 81.68% to 85.81% and validation-selected test accuracy from 78.73% to 82.73%, with the best validation checkpoint occurring at step 800 instead of step 3000. However, the moving-batch-center compensation produces severe bias-coordinate drift and strongly reduces boundary energy. The present evidence therefore supports hidden boundary motion as an important optimization mechanism, but it does not yet establish a final general-purpose optimizer. A stable centered-affine parameterization is identified as the required next step.
Zhang Gongyue, Sheng Yixuan, Liu donghan +3
AdamAffineTransformer ArchitecturesModel WeightsOrthogonalityBoundaries
AskPDF
Jun 7, 2026cs.LG

The Spectral Dynamics and Noise Geometry of Muon

Muon replaces a matrix gradient G=UΣV⊤G=UΣV^\top by its polar factor UV⊤UV^\top. This keeps the singular directions selected by the gradient, but makes the update spectrum flat. We study the optimization bias created by this operation. Under explicit alignment assumptions, we prove that the polar update is the one-step entropy-maximizing choice among bounded updates that use the gradient singular directions and do not adapt to the current weight spectrum. In an underdetermined regression model, we derive exact singular-value dynamics for continuous-time Muon and identify a measurement-dependent condition under which the normalized spectrum moves toward equal nonzero singular values. This geometry also rules out a common low-rank interpretation: at fixed Frobenius norm, Muon's distinguished state has a flat spectrum, whereas nuclear-norm minimization favors spectral concentration. Controlled matrix-sensing experiments separate the effect from simple gradient rescaling, show that norm-matched gradient descent does not reproduce Muon, and recover the predicted flattening trend across broad ablations. In small NanoGPT pretraining, Muon preserves stable rank, has a broad learning-rate plateau, and improves validation loss relative to AdamW; in a matched small-ViT control, the ranking reverses. The resulting picture is regime-dependent: Muon is not universally superior, but its flat-spectrum bias can help when many spectral directions need to remain active.
Pierfrancesco Beneventano, Mahmoud Abdelmoneum, Tomaso Poggio
MuonSpectrumPolarizationModel Pretraining
AskPDF
Jun 11, 2026cs.LG

Muonp^p: Muon with Fractional Spectral Powers

Muon is an increasingly widely used optimizer that replaces a gradient G=USV⊤G=USV^\top with its polar factor UV⊤UV^\top, thereby flattening the singular spectrum. However, full flattening discards singular-value information that may matter for adaptation. We introduce Muonp^p, a Muon-style optimizer that instead uses fractional spectral-power updates USpV⊤US^pV^\top for rational p∈(0,1)p\in(0,1), interpolating between Muon and gradient descent. To make it practical, we prove that fractional spectral powers cannot be computed by any fixed univariate polynomial iteration, and furthermore derive low-degree odd bivariate recurrences that approximate USpV⊤US^pV^\top using only matrix multiplications, preserving Muon's matrix-multiplication-only structure and compute complexity. We show that Muonp^p maximizes the linear improvement in loss under the Schatten qq-norm for q=1+1pq=1+\frac{1}{p}. Empirically, Muonp^p is especially effective for finetuning: on billion-scale models, Muonp^p improves validation perplexity and downstream task performance. We further analyze when Muonp^p is less suitable, through the lens of spectral geometry. Our results reveal important insights on when preserving the singular spectrum can bring significant gains, and introduce a principled way to achieve them.
Yihe Dong, Will Sawin
Muon OptimizerMuonGradient DescentMatrix Multiplication
AskPDF