← Research

Muonic: Input Conditioning Muon with Improved Diagonal Approximations

Pranavchandra Vivekananda
Sumukh Bettadapura, Ajan Subramanian

Optimizers

X LinkedIn Facebook WhatsApp

We have been working with a new kind of training objective over the past few weeks: epiplexity [1] and Zhang's closed form approximation of it. [2] It borrows from the idea of maximising learnable novelty from data. Instead of using loss for a wrong guess (as in cross-entropy) or trying to bring two probability distributions closer to each other (as in KL divergence), this training rewards for an objective that extracts structure from the samples it sees through a spectral readout.

Along the journey, we had a thought. What if we could use this idea in an optimiser? We started running toy experiments with spectral filters, changing how strongly different singular directions contributed to the update. Our first filter suppressed large singular values; a gentler version smoothly capped them. We naively realized that we were practically rederiving Muon. It takes a momentum matrix and approximately orthogonalises it, bringing its singular values towards a common scale. Different rules, but the broader idea of reshaping the update through its spectrum already had a very capable optimiser attached to it. [3]

Now, the spectrum only describes the loss locally to first order. It cannot tell us how quickly a promising direction becomes too steep as we move (basically, what a Hessian would provide). For a linear layer, that curvature depends partly on the inputs, because a weight change affects the output through those inputs.

The simplest way to do this is to take an optimiser and give it information on how input layer's features behave and use that information to bend the momentum to your will.

input statistics→A→MA−1→Muon(MA−1)\displaystyle \text{input statistics} \to A \to MA^{-1}\to \text{Muon}(MA^{-1})

But doing these calculations could potentially add more computation per step, which is where we had to use some smart linear algebra.

With the help of GPT, we came up with two ways do this and tested them on a toy transformer (4.84M-parameter Tiny Shakespeare transformer) over 5 seeds where it ended up beating vanilla Muon.

Methodology

1. What “input-aware” means

Consider one linear layer:

y=Wa,\displaystyle y=Wa,W∈Rm×d,\displaystyle W\in\mathbb R^{m\times d},a∈Rd.\displaystyle a\in\mathbb R^d.

Here, aa is the activation entering this particular layer. Each weight matrix gets statistics from its own inputs. Collect NN input vectors into X∈RN×dX\in\mathbb R^{N\times d}, with one input per row. Define

μ=1N∑i=1Nai,\displaystyle \mu=\frac1N\sum_{i=1}^N a_i,C=1NX⊤X=E[aa⊤].\displaystyle C=\frac1N X^\top X=\mathbb E[aa^\top].

CC is the uncentered second moment. Its diagonal contains feature energies, and its off-diagonal entries describe how features occur together:

Cjj=E[aj2],\displaystyle C_{jj}=\mathbb E[a_j^2],Cjk=E[ajak].\displaystyle C_{jk}=\mathbb E[a_ja_k].

Now why should you even care, just for an update?

Changing the weights by ΔW\Delta W changes the layer output by ΔWa\Delta Wa, so

E∥ΔWa∥22=tr⁡(ΔWCΔW⊤).\displaystyle \boxed{ \mathbb E\|\Delta Wa\|_2^2 = \operatorname{tr}(\Delta W C\Delta W^\top). }

We find that size doesn't matter when it comes to weight changes as they can have very different effects on the layer’s outputs, depending on the inputs. An input-aware optimizer uses CC, or an approximation to it, to account for this difference.

2. The weight update

The shared momentum calculation is

Ht=0.95Ht−1+0.05Gt,\displaystyle H_t=0.95H_{t-1}+0.05G_t,Mt=0.05Gt+0.95Ht,\displaystyle M_t=0.05G_t+0.95H_t,

where Gt=∇WLtG_t=\nabla_W L_t.

Ordinary Muon transforms MtM_t. Our variants first compute

Pt=MtAt−1.\displaystyle P_t=M_tA_t^{-1}.

and then we apply Muon’s spectral transformation to construct the update direction. [4] The final weight update rule becomes

Wt+1=(1−γt)Wt−ηtmax⁡(1,m/d) P5(Pt),\displaystyle W_{t+1} = (1-\gamma_t)W_t - \eta_t \sqrt{\max(1,m/d)} \,\mathcal P_5(P_t),

where γt\gamma_t is the scheduled weight shrink and P5\mathcal P_5 denotes that approximate polar transformation.

For each layer, we construct an input geometry matrix from the activations aa entering that layer. Let

C=E[aa⊤],\displaystyle C=\mathbb E[aa^\top],s=tr⁡(C)d,\displaystyle s=\frac{\operatorname{tr}(C)}{d},

where dd is the input dimension. We approximate the normalized second moment C/sC/s, then add damping:

A=C^normalized+λI1+λ,λ=0.1.\displaystyle \boxed{ A=\frac{\widehat C_{\mathrm{normalized}}+\lambda I}{1+\lambda}, \qquad \lambda=0.1. }

Normalization removes the overall input-energy scale, while damping ensures that the inverse remains a good little boy and behaves. We explore two constructions.

Explore the mechanism / 01

Why should an update care about the inputs?

A tiny layer with two input features and one output. Follow a dot to see how changing the weights changes that sample’s output.

Horizontal position is feature 1; vertical position is feature 2. This batch spreads much farther horizontally. The layer sees more variation in feature 1.

Illustrative two-input, one-output layer; a weight-update direction stands in for momentum. No learning rate or weight decay is applied. Output position uses the same scale in both rows. Vertical jitter only separates the dots.

3. Input conditioning method A: Mean + diagonal

Mean + diagonal uses feature variances and the mean activation vector. The exact second moment decomposes as

C=Cov⁡(a)+μμ⊤,\displaystyle C=\operatorname{Cov}(a)+\mu\mu^\top,μ=E[a].\displaystyle \mu=\mathbb E[a].

We retain the diagonal of the centered covariance and the mean outer product. Writing vj=E[(aj−μj)2]v_j=\mathbb E[(a_j-\mu_j)^2],

Amean=[diag⁡(v)+μμ⊤]/s+λI1+λ.\displaystyle \boxed{ A_{\mathrm{mean}} = \frac{ \left[\operatorname{diag}(v)+\mu\mu^\top\right]/s +\lambda I }{1+\lambda}. }

The diagonal term captures each feature’s variability. The rank-one term μμ⊤\mu\mu^\top captures the shared structure implied by nonzero feature means. This approximation preserves individual feature second moments but omits centered cross-feature correlations.

Explore the mechanism / 02 · Move the data

Find what the mean misses.

Move the centre of this cloud. Then change its tilt without moving the centre. Which change can Mean + diagonal see?

Mean + diagonal retains each feature’s variance and the shared mean. Changing the centred correlation tilts the samples, but leaves this approximation unchanged.

Two-feature synthetic batch with fixed variances. The outline shows the axis-aligned centred spread the approximation retains; the arrow records its mean. Readouts are unnormalized second moments, before damping. No training results are simulated.
Loss curves / 01

Mean + diagonal

By training steps

MuonMean + diagonal
Mean + diagonal versus Muon: validation loss by training step Mean validation cross-entropy across five paired seeds; shaded bands show one standard deviation. Dashed line: Muon. Solid line: Mean + diagonal. Lower is better. Validation CE ↓ 1.4 1.6 1.8 2 2.2 200 400 600 800 1000 1200 Training steps

By training time

MuonMean + diagonal
Mean + diagonal versus Muon: validation loss by training time Mean validation cross-entropy across five paired seeds; shaded bands show one standard deviation. Dashed line: Muon. Solid line: Mean + diagonal. Lower is better. Validation CE ↓ 1.4 1.6 1.8 2 2.2 20 40 60 80 100 120 Training time (seconds)
Mean validation cross-entropy over five paired seeds; shaded bands show one standard deviation. Original full-size plot ↗ · Download plotted data ↓

4. Input conditioning method B: Nyström + diagonal

Nyström + diagonal captures a richer set of relationships between features. We choose a fixed random orthonormal matrix Ω∈Rd×r\Omega\in\mathbb R^{d\times r}, with r=64r=64, and measure the normalized second moment along those directions:

S=(C/s)Ω,\displaystyle S=(C/s)\Omega,K=Ω⊤S.\displaystyle K=\Omega^\top S.

The matrix SS is a sketch: a compact collection of measurements. The standard Nyström reconstruction [5] is

C^Nys=SK†S⊤,\displaystyle \widehat C_{\mathrm{Nys}}=SK^\dagger S^\top,

where K†K^\dagger is the pseudoinverse. The sketch can be computed directly from batch activations without constructing the full matrix CC.

Our implementation removes an estimated isotropic background from the retained spectrum, producing a corrected low-rank factor LL. We then restore the diagonal energy that this factor has not represented:

ej=max⁡(Cjjs−∑kLjk2,  0).\displaystyle e_j= \max\left( \frac{C_{jj}}s-\sum_kL_{jk}^2,\;0 \right).

The result is

ANys=LL⊤+diag⁡(e)+λI1+λ.\displaystyle \boxed{ A_{\mathrm{Nys}} = \frac{ LL^\top+\operatorname{diag}(e)+\lambda I }{1+\lambda}. }
Explore the mechanism / 03 · Measure & reconstruct

Measure a little.
Reconstruct the rest.

Each tile describes a relationship between two input features. Measure one column, and Nyström uses it to estimate the entries you haven’t measured. Add another column to see what changes.

A sampled column gives the relationship of one feature to every other feature. Nyström combines these measurements to reconstruct the full second-moment matrix. More independent measurements can recover more of its structure.

Three-feature synthetic example with coordinate-column measurements, chosen to make the sketch visible. The experiment uses a random rank-64 sketch. Diagonal tiles show feature energies; off-diagonal tiles show cross-feature second moments. Error is relative Frobenius error, with matching damping in the final stage—not validation loss.
Loss curves / 02

Nyström + diagonal

By training steps

MuonNyström + diagonal
Nyström + diagonal versus Muon: validation loss by training step Mean validation cross-entropy across five paired seeds; shaded bands show one standard deviation. Dashed line: Muon. Solid line: Nyström + diagonal. Lower is better. Validation CE ↓ 1.4 1.6 1.8 2 2.2 200 400 600 800 1000 1200 Training steps

By training time

MuonNyström + diagonal
Nyström + diagonal versus Muon: validation loss by training time Mean validation cross-entropy across five paired seeds; shaded bands show one standard deviation. Dashed line: Muon. Solid line: Nyström + diagonal. Lower is better. Validation CE ↓ 1.4 1.6 1.8 2 2.2 20 40 60 80 100 120 Training time (seconds)
Mean validation cross-entropy over five paired seeds; shaded bands show one standard deviation. Original full-size plot ↗ · Download plotted data ↓

Results

We tested a 4.84M-parameter transformer on Tiny Shakespeare for 1,200 steps, using five fresh paired seeds per study.

Method Muon validation CE ↓ Candidate CE ↓ Loss reduction Additional training time
Nyström + diagonal 1.48923 1.46495 1.63% 12.0%
Mean + diagonal 1.48889 1.47410 0.99% 5.4%

Both methods achieved lower final loss than Muon on all five paired seeds. Nyström also retained its advantage within Muon’s recorded compute budget.

These results show that our input-aware preconditioning can improve validation quality in this setting. We acknowledge that we build on the shoulders of giants and we're just excited to post about some empirical work that we did in an engaging way. While the results are modest, it's genuinely interesting to see how this pans out at scale and what more we can do in this direction.

References

[1] Marc Finzi et al. (2026). From Entropy to Epiplexity: Rethinking Information for Computationally Bounded Intelligence ↗. The epiplexity framework.

[2] Yanbo Zhang and Michael Levin (2026). Intelligence from Learnable Novelty ↗. A closed-form estimator of learnable novelty.

[3] Keller Jordan (2024). Muon: An optimizer for hidden layers in neural networks ↗. Muon and its five-step Newton–Schulz transformation.

[4] Zhehang Du and Weijie Su (2026). The Newton-Muon Optimizer ↗. Prior work on input-second-moment preconditioning before the matrix-sign update.

[5] Zachary Frangella, Joel A. Tropp and Madeleine Udell (2021). Randomized Nyström Preconditioning ↗. Background on randomized low-rank approximation and preconditioned linear solvers.

© 2026 Kubo Technologies, Inc.