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

Gongyue Zhang; Honghai Liu

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.

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=[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\pm0.242\%$ to $85.562\pm0.308\%$ and lowers selected test loss from $0.3479$ to $0.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\pm0.180\%$ and lowers test loss to $0.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.02095$ to $0.00301$ , lowers boundary-function share from $86.58\%$ to $78.97\%$ , and changes the cosine between weight-induced boundary motion and explicit bias from $+0.030$ to $-0.137$ . An independent 22-seed replication yields $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

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+b

with input mean

μ

, a weight update contains a sample-independent displacement

Δ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

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}

is 134.7, while

\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

g_W-g_bμ^\top

and

g_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

Adam Affine Transformer Architectures Model Weights Orthogonality Boundaries

Ask PDF

Jun 7, 2026cs.LG

The Spectral Dynamics and Noise Geometry of Muon

Muon replaces a matrix gradient

G=UΣV^\top

by its polar factor

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

Muon Spectrum Polarization Model Pretraining

Ask PDF

Jun 11, 2026cs.LG

Muon $^p$ : Muon with Fractional Spectral Powers

Muon is an increasingly widely used optimizer that replaces a gradient

G=USV^\top

with its polar factor

UV^\top

, thereby flattening the singular spectrum. However, full flattening discards singular-value information that may matter for adaptation. We introduce Muon

^p

, a Muon-style optimizer that instead uses fractional spectral-power updates

US^pV^\top

for rational

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

US^pV^\top

using only matrix multiplications, preserving Muon's matrix-multiplication-only structure and compute complexity. We show that Muon

^p

maximizes the linear improvement in loss under the Schatten

q

-norm for

q=1+\frac{1}{p}

. Empirically, Muon

^p

is especially effective for finetuning: on billion-scale models, Muon

^p

improves validation perplexity and downstream task performance. We further analyze when Muon

^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 Optimizer Muon Gradient Descent Matrix Multiplication

Ask PDF

Abstract

Explore similar work

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

The Spectral Dynamics and Noise Geometry of Muon

Muonp^pp: Muon with Fractional Spectral Powers

Muon $^p$ : Muon with Fractional Spectral Powers