NB2 §4 — Optimizers as steepest descent under a norm

1. KL. For \theta\ne0, P_0 puts mass 1 on \{0\}, where P_\theta puts none, so P_0 is not absolutely continuous with respect to P_\theta and \mathrm{KL}(P_0\|P_\theta)=+\infty. At \theta=0 it is 0.

JS. For \theta\ne0, M puts mass \tfrac12 on each atom, so dP_0/dM=2 on \{0\} and \mathrm{KL}(P_0\|M)=\log2. The same holds for P_\theta, so \mathrm{JS}=\log2, the largest value JS can take. At \theta=0 it is 0.

W_1. The only coupling of \delta_0 and \delta_\theta is \delta_{(0,\theta)}, which moves unit mass a distance |\theta|. So W_1=|\theta|.

2. JS is constant on \theta\ne0, so its derivative there is zero and gradient descent never moves. KL is +\infty for every \theta\ne0, so there is no finite objective and no gradient. \frac{d}{d\theta}W_1=\mathrm{sign}(\theta) for \theta\ne0: every step moves \theta a distance \eta toward 0, ending in an oscillation of size at most \eta around 0, which a step size \eta_t\to0 with \sum_t\eta_t=\infty removes. Only W_1 is usable.

1. For \|\Delta\|\le\eta, write \Delta=-\eta d with \|d\|\le1; the ball is symmetric, so d ranges over the whole unit ball. Then \langle g,\Delta\rangle=-\eta\langle g,d\rangle\ge-\eta\max_{\|d\|\le1}\langle g,d\rangle=-\eta\|g\|_*, with equality at d=d^\star(g). (If g=0, every point of the ball is a minimizer.)

2. Write \Delta=c\,u with c\ge0, \|u\|=1. The objective is c\langle g,u\rangle+c^2/(2\eta). At fixed c the first term is smallest when u=-d^\star(g); for g\ne0 we can take \|d^\star\|=1, since a nonzero linear function on a ball is maximized on the boundary. This gives -c\|g\|_*+c^2/(2\eta), minimized at c=\eta\|g\|_* with value -\tfrac\eta2\|g\|_*^2. So \Delta_{\rm pen}=-\eta\|g\|_*d^\star(g): the trust-region direction, with length set by the dual norm of the gradient.

1. \langle g,d\rangle=\sum_ig_id_i\le\sum_i|g_i|\,|d_i|\le\|d\|_\infty\|g\|_1, with equality when every d_i=\mathrm{sign}(g_i). Coordinates with g_i=0 can take any value in [-1,1] (remark 2). The trust-region step -\eta\,\mathrm{sign}(g) moves every coordinate by exactly \eta, however small its gradient.

2. \langle g,d\rangle\le\sum_i|g_i||d_i|\le\max_j|g_j|\sum_i|d_i|\le\|g\|_\infty for \|d\|_1\le1, with equality when all the mass of d sits on a largest coordinate with the right sign: greedy coordinate descent. Ties give a whole face of maximizers.

3. \langle g,d\rangle=\langle M^{-1/2}g,M^{1/2}d\rangle\le\|M^{-1/2}g\|_2\,\|M^{1/2}d\|_2=\sqrt{g^\top M^{-1}g}\;\|d\|_M, with equality when M^{1/2}d\propto M^{-1/2}g, i.e. d\propto M^{-1}g. Normalizing, d^\star=M^{-1}g/\sqrt{g^\top M^{-1}g}, and the penalty step is -\eta\|g\|_*d^\star=-\eta M^{-1}g: preconditioned gradient descent. With M=F, the Fisher information, it is Amari’s natural gradient. Since \mathrm{KL}(p_\theta\|p_{\theta+\delta})=\tfrac12\delta^\top F\delta+O(\|\delta\|^3), \|\Delta\|_F^2/2 is, to second order, the KL divergence between the model before and after the step: natural gradient measures a step by how much the predictive distribution changes.

4. M=H, \eta=1. On the quadratic, g=H(\theta-\theta^*), so the step lands exactly on \theta^*. Newton is steepest descent in the norm of the loss’s own curvature. It is not usually described that way because away from a minimum H need not be positive definite, so \sqrt{\Delta^\top H\Delta} need not be a norm, and because this “norm” changes with the loss. For exponential-family likelihoods with the canonical link, the Fisher equals the generalized Gauss–Newton matrix, a positive semidefinite approximation to H (Martens); this is one reason natural gradient sits between gradient descent and Newton.

1. \langle G,D\rangle=\mathrm{tr}\big(\sum_i\sigma_iv_iu_i^\top D\big)=\sum_i\sigma_i\,u_i^\top Dv_i. Since u_i,v_i are unit vectors, |u_i^\top Dv_i|\le\|D\|_2, so \langle G,D\rangle\le\|D\|_2\|G\|_{\rm nuc}. For D=UV^\top, u_i^\top UV^\top v_i=1, which gives equality, and UV^\top has r singular values equal to 1, so \|UV^\top\|_2=1. Hence the dual of the spectral norm is the nuclear norm, d^\star(G)=UV^\top, the trust-region step is -\eta\,UV^\top and the penalty step -\eta\|G\|_{\rm nuc}UV^\top. Every nonzero singular value of the step is set to the same value, whatever its size in G.

2. For semi-orthogonal O, \|O\|_F^2=\mathrm{tr}(OO^\top)=m and \|O\|_2=1. So \|O-G\|_F^2=m-2\langle O,G\rangle+\|G\|_F^2, and minimizing the distance means maximizing \langle O,G\rangle. By 1, \langle O,G\rangle\le\|G\|_{\rm nuc}, with equality at O=UV^\top, which is semi-orthogonal because U is m\times m orthogonal and V^\top V=I_m. It is the orthogonal factor of the polar decomposition G=(UV^\top)(V\Sigma V^\top).

3. Let U_\perp, V_\perp be orthonormal bases of the complements of the column spaces of U and V, and W any matrix with \|W\|_2\le1. Then D=UV^\top+U_\perp WV_\perp^\top has \|D\|_2=1 (the two blocks act on orthogonal subspaces) and \langle G,D\rangle=\|G\|_{\rm nuc} (because U^\top U_\perp=0). W=0 is the minimum-Frobenius-norm maximizer, and it is what the cubic Newton–Schulz iteration converges to (§4.4), since odd polynomial iterations leave zero singular values at zero.

4. With G=U\Sigma V^\top of full rank, (GG^\top)^{-1/4}=U\Sigma^{-1/2}U^\top and (G^\top G)^{-1/4}=V\Sigma^{-1/2}V^\top (on their ranges), so (GG^\top)^{-1/4}G(G^\top G)^{-1/4}=U\Sigma^{-1/2}\,\Sigma\,\Sigma^{-1/2}V^\top=UV^\top. Shampoo without accumulation is spectral-norm steepest descent, as Bernstein and Newhouse point out. Accumulation is where Shampoo and Muon part ways: Shampoo accumulates the gradient’s statistics, Muon accumulates the gradient itself (momentum) and orthogonalizes that.

H = np.diag([1.0, 25.0])
theta0 = np.array([0.93, 0.71])
loss = lambda th: 0.5 * th @ H @ th

def update(theta, method, eta):
    g = H @ theta
    if method == "gd":
        return theta - eta * g
    if method == "normalized":
        return theta - eta * g / np.linalg.norm(g)
    if method == "sign":
        return theta - eta * np.sign(g)
    if method == "precond":
        return theta - eta * g / np.diag(H)
    raise ValueError(method)

etas = {"gd": 0.07, "normalized": 0.05, "sign": 0.05, "precond": 0.7}
labels = {"gd": "GD (ℓ2 penalty)", "normalized": "normalized GD (ℓ2 trust)",
          "sign": "sign descent (ℓ∞ trust)", "precond": "precond. GD, M = diag(H)"}
fig, axs = plt.subplots(1, 3, figsize=(13, 3.6))
xx, yy = np.meshgrid(np.linspace(-1.2, 1.2, 200), np.linspace(-1.2, 1.2, 200))
axs[0].contour(xx, yy, 0.5 * (xx**2 + 25 * yy**2), levels=np.geomspace(1e-3, 15, 12), colors="0.8", linewidths=0.7)
for method, eta0 in etas.items():
    for ax, schedule in [(axs[1], "constant"), (axs[2], "decayed")]:
        th = theta0.copy(); losses, path = [loss(th)], [th]
        for t in range(200):
            eta = eta0 if schedule == "constant" else eta0 / np.sqrt(t + 1)
            th = update(th, method, eta); losses.append(loss(th)); path.append(th)
        ax.semilogy(losses, label=labels[method])
        if schedule == "constant":
            path = np.array(path)
            axs[0].plot(path[:40, 0], path[:40, 1], ".-", ms=3, lw=0.8, label=labels[method])
axs[0].set_title("first 40 steps, constant η"); axs[0].set_aspect("equal")
axs[1].set_title("constant η"); axs[2].set_title(r"$\eta_t=\eta_0/\sqrt{t+1}$")
for ax in axs[1:]:
    ax.set_xlabel("step"); ax.set_ylim(1e-12, 30)
axs[1].set_ylabel("loss"); axs[2].legend(fontsize=8)
plt.tight_layout(); plt.show()

At constant \eta, GD and preconditioned GD converge linearly: penalty-form steps shrink with the gradient. GD is slow because its stable step is capped by the largest curvature (\eta<2/25); preconditioning by \mathrm{diag}(H) equalizes the curvatures (exactly, since H is diagonal) and converges in a few steps. The trust-region methods take steps of fixed length \eta, so near the minimum they rattle at distance of order \eta and the loss stalls; sign descent’s approach, though, is not slowed by the ill-conditioning. With \eta_t\propto1/\sqrt{t+1} they keep improving, at a rate set by the schedule rather than the curvature, as in the practice of decaying Adam’s and Muon’s learning rate at the end of training.

1. The score transforms by the chain rule, \nabla_\xi\log p=J^\top\nabla_\theta\log p, so F_\xi=\mathbb E\big[J^\top\nabla_\theta\log p\,(\nabla_\theta\log p)^\top J\big]=J^\top F_\theta J. The same chain rule gives g_\xi=J^\top g_\theta.

2. \dot\theta=J\dot\xi=-J\,(J^\top F_\theta J)^{-1}J^\top g_\theta=-J\,J^{-1}F_\theta^{-1}J^{-\top}J^\top g_\theta=-F_\theta^{-1}g_\theta, whatever \phi was. For gradient flow, \dot\theta=J\dot\xi=-JJ^\top g_\theta: preconditioned by JJ^\top, which depends on the parametrization. Example: \theta=c\,\xi gives \dot\theta=-c^2g_\theta, so rescaling a coordinate by c changes its learning rate by c^2.

3. The step in \xi is \Delta\xi=-\eta F_\xi^{-1}g_\xi=-\eta J^{-1}F_\theta^{-1}g_\theta. Taylor expanding, \phi(\xi+\Delta\xi)=\phi(\xi)+J\Delta\xi+\tfrac12D^2\phi(\xi)[\Delta\xi,\Delta\xi]+\dots=\theta-\eta F_\theta^{-1}g_\theta+O(\eta^2). When \phi is affine, D^2\phi=0 and the two iterations agree exactly. The O(\eta^2) term can be large when J is nearly singular: a moderate step in \theta can be an enormous step in \xi. Code demo 4.2.D shows both effects.

1. \partial_{\xi_a}L=\sum_k(g_\theta)_k\,\partial_a\phi_k. Differentiating again, \partial_{\xi_a}\partial_{\xi_b}L=\sum_{k,l}(H_\theta)_{kl}\,\partial_a\phi_k\,\partial_b\phi_l+\sum_k(g_\theta)_k\,\partial_a\partial_b\phi_k, which is the claim, since J_{ka}=\partial_a\phi_k. At a critical point g_\theta=0 and only J^\top H_\theta J survives. Away from one the Hessian does not transform as a tensor: the second term depends on how curved the coordinate change is. The Fisher has no such term; Code demo 4.2.D shows the consequence for Newton.

2. With J=cI, H_\xi=c^2H_\theta, so \lambda_{\max} scales by c^2 and can be anything. Layer rescalings in ReLU networks are reparametrizations of this kind: they change \lambda_{\max} without changing the function. Meanwhile F_\xi^{-1}H_\xi=(J^\top F_\theta J)^{-1}J^\top H_\theta J=J^{-1}F_\theta^{-1}H_\theta J, which is similar to F_\theta^{-1}H_\theta and has the same eigenvalues. Raw sharpness is curvature measured against a ruler, the Euclidean norm of whatever coordinates you use; reparametrizing changes the ruler and the number, not the function.

1. \mathrm{KL}=\mathbb E_{x\sim p_0}\big[(x-\theta)^2-x^2\big]/(2s^2)=\mathbb E[-2\theta x+\theta^2]/(2s^2)=\theta^2/(2s^2). For W_p: the coupling Y=X+\theta costs exactly |\theta| in W_1 and W_2, and any coupling has \mathbb E|X-Y|\ge|\mathbb E X-\mathbb EY|=|\theta| (and (\mathbb E|X-Y|^2)^{1/2}\ge\mathbb E|X-Y|), so W_1=W_2=|\theta|. As s\to0 at fixed \theta, \mathrm{KL}\to\infty, consistent with 4.0.A; W does not depend on s.

2. On L_{\rm KL}, \theta\leftarrow(1-\eta/s^2)\theta, which converges iff 0<\eta<2s^2: the usable step size is tied to the blur. On L_W, \theta\leftarrow(1-\eta)\theta, which converges iff 0<\eta<2, whatever s is.

3. F=\mathbb E[(x-\theta)^2/s^4]=1/s^2. The natural-gradient step is \theta\leftarrow\theta-\eta\,s^2\cdot\theta/s^2=(1-\eta)\theta, identical to GD on L_W. But the Fisher is the second-order expansion of KL, so it measures how much the probabilities change, not how far mass moves: as s\to0 it diverges, and between point masses there is no Fisher metric at all, while the Wasserstein geometry is unchanged. The two agree here only because, for a translated Gaussian, both reduce to rescaling the translation.

What the plot shows. GD follows different trajectories in the two coordinate systems (JJ^\top rescales the second coordinate by e^{2\xi_2}). Natural gradient in \xi tracks natural gradient in \theta closely, differing only at O(\eta^2) per step, as 4.2.A.3 predicts. Newton in \xi fails: away from the optimum the extra term g_2e^{\xi_2} in H_\xi is negative here and makes H_\xi indefinite, so the Newton step in \xi_2 goes uphill and drives \theta_2=e^{\xi_2} toward 0, where the \xi-gradient vanishes without \theta being optimal. In \theta coordinates Newton and natural gradient coincide because F=H for this model; in \xi coordinates they come apart, because F_\xi=J^\top F_\theta J holds everywhere while H_\xi=J^\top H_\theta J holds only at critical points.

1. With \beta_1=0, m_t=g_t and 1-\beta_1^t=1, so \hat m_t=g_t; likewise \hat v_t=g_t^2. The update is -\eta\,g_t/\sqrt{g_t^2}=-\eta\,\mathrm{sign}(g_t) per coordinate (taking 0/0=0).

2. m_t is linear in the gradients, so it becomes c\,m_t; v_t is quadratic, so it becomes c^2v_t, and \sqrt{c^2v_t}=c\sqrt{v_t}. The ratio \hat m_t/\sqrt{\hat v_t} is unchanged. SGD’s update is multiplied by c: a loss multiplied by 10 needs SGD’s learning rate divided by 10, and no change for Adam (until \epsilon matters).

3. For SGD, \Delta\theta=-\eta\,\partial L/\partial\theta, so [\eta]=[\theta]^2/[L]. For Adam, \hat m/\sqrt{\hat v} is dimensionless, so [\eta]=[\theta]: a distance per step in parameter units, which can be set from the expected scale of the parameters (Kingma and Ba). So Adam’s \eta inherits the parametrization: if a weight’s natural scale changes with width, so must \eta (§4.5).

4. With \beta_2\to1, \hat v averages g^2 over a long window, so \hat v\approx\mathbb Eg^2=\mu^2+\sigma^2, nearly constant, and \hat m/\sqrt{\hat v}\approx\hat m/\sqrt{\mu^2+\sigma^2}. For a stationary sequence and many steps, \mathbb E\hat m=\mu and \mathrm{Var}\,m=(1-\beta_1)^2\sum_k\beta_1^{2k}\sigma^2=\frac{1-\beta_1}{1+\beta_1}\sigma^2. Dividing by \sqrt{\mu^2+\sigma^2} gives both claims. For \beta_1=0.9 the factor \sqrt{(1-\beta_1)/(1+\beta_1)}\approx0.23: on a pure-noise coordinate (\mu=0) the mean step is zero but individual steps are about 0.23\eta; on a noiseless one (\sigma=0) every step is \pm\eta, pure sign descent. Code demo 4.3.C checks both.

1. By Cauchy–Schwarz, \Big|\sum_k\beta_1^kg_{t-k}\Big|=\Big|\sum_k\frac{\beta_1^k}{\beta_2^{k/2}}\cdot\beta_2^{k/2}g_{t-k}\Big| \le\Big(\sum_k\frac{\beta_1^{2k}}{\beta_2^k}\Big)^{1/2}\Big(\sum_k\beta_2^kg_{t-k}^2\Big)^{1/2} =\frac{1}{\sqrt{1-\beta_1^2/\beta_2}}\sqrt{\frac{v_t}{1-\beta_2}}, where the geometric series converges because \beta_1^2<\beta_2. Multiplying by 1-\beta_1 gives the bound. Equality holds when \beta_2^{k/2}g_{t-k}\propto\beta_1^k\beta_2^{-k/2}, i.e. g_{t-k}\propto(\beta_1/\beta_2)^k: a gradient that has been growing geometrically, by a factor \beta_2/\beta_1\approx1.11 per step. For \beta_1=0.9, \beta_2=0.999 the bound is 0.1/\sqrt{0.001\times0.18919}\approx7.27. For \beta_1=\beta_2=\beta it is (1-\beta)/\sqrt{(1-\beta)(1-\beta)}=1, attained by a constant gradient. (For 0<\beta_2<1 there is no history-independent bound once \beta_1^2\ge\beta_2: if \beta_1^2>\beta_2, v forgets a past spike faster than m does and the ratio grows; at \beta_1^2=\beta_2 the weights in the Cauchy–Schwarz sum stop decaying, and ever longer aligned histories make the ratio grow like the square root of their length.)

2. A single spike gives m_t=(1-\beta_1)g and v_t=(1-\beta_2)g^2, so the ratio is (1-\beta_1)/\sqrt{1-\beta_2}\approx3.16. That history is not the worst one, so this is not an upper bound: a steadily growing gradient gives up to 7.27, because v forgets more slowly than m and lags behind it. Adam’s step is an \ell_\infty trust region in a loose sense, with radius up to about 7.3\eta for the default \beta’s; in practice typical steps are below \eta (4.3.A.4).

• Noisy coordinates. The mean signed update is about \eta|\mu|/\sigma\ll\eta (4.3.A.4), with fluctuations of about 0.23\eta around it, not steps of \pm\eta. “Sign of the mean, weighted by signal-to-noise, plus noise” describes it better.
• Oscillating gradients. At the edge of stability the gradient along the top curvature directions flips sign every step (NB2 §3). Then m averages toward zero while v does not, and the right description is EMA momentum with a slowly varying diagonal preconditioner P. With P\succ0 frozen, on a quadratic with H\succ0 and 0\le\beta_1<1, the iteration converges for every start iff \eta\lambda_{\max}(P^{-1}H)<2(1+\beta_1)/(1-\beta_1), which is 38 for \beta_1=0.9 (Cohen et al. 2022).
• Tiny gradients. When \sqrt{\hat v}\ll\epsilon, for example for rarely used embedding rows, the update is about -\eta\hat m/\epsilon (at \sqrt{\hat v}=\epsilon the denominator is already 2\epsilon): SGD with momentum and learning rate \eta/\epsilon, neither sign-like nor scale-invariant.

A fourth, milder one is a growing gradient, where v lags and steps exceed \eta (4.3.B).

1. \|\theta_{t+1}\|=\|(1-\eta_t\lambda)\theta_t+\eta_to_t\|\le(1-\eta_t\lambda)\|\theta_t\|+\eta_tb. Subtracting b/\lambda from both sides, \|\theta_{t+1}\|-b/\lambda\le(1-\eta_t\lambda)\|\theta_t\|+\eta_tb-b/\lambda=(1-\eta_t\lambda)(\|\theta_t\|-b/\lambda). If \|\theta_t\|\le b/\lambda, the right side is \le0, so the ball is invariant. From outside, iterating gives an excess of at most \prod_{s<t}(1-\eta_s\lambda) times the initial one. The product tends to 0 if \sum_s\eta_s=\infty (with summable steps and every \eta_s\lambda<1 the weights can stop short of the ball), and it decays geometrically if \eta_s\ge\eta_{\min}>0. Even then the ball is approached, not necessarily entered in finite time.

2. Sign descent has \|\mathrm{sign}(g)\|_\infty\le1, so with decoupled decay the weights approach the box \max_i|\theta_i|\le1/\lambda. Exactly orthogonalized updates have \|UV^\top\|_2=1, so the weights approach the spectral-norm ball \|W\|_2\le1/\lambda. Weight decay picks the radius of a constraint set in the same norm the optimizer uses for its steps. For Adam itself the step bound is b\approx7.3 rather than 1 (4.3.B), so the box is only loosely 1/\lambda; practical Muon, with its approximate quintic orthogonalization (singular values up to about 1.2, §4.4) and shape factor, has a correspondingly adjusted radius.

For idealized versions of these algorithms the constraint picture is a theorem: for Lion (Chen, Liu, Liang and Liu), for full-batch AdamW with \epsilon=0, whose limits under a non-increasing schedule with \sum_t\eta_t=\infty are KKT points of the loss subject to \|\theta\|_\infty\le1/\lambda (Xie and Li, ICML 2024; with \epsilon>0 AdamW can stall at an interior non-stationary point), and for Muon with exact orthogonalization, via part 1’s inequality (Chen, Li and Liu, Proposition 1).

1. XX^\top=U\Sigma V^\top V\Sigma U^\top=U\Sigma^2U^\top (using V^\top V=I), so for k\ge1, (XX^\top)^k=U\Sigma^{2k}U^\top and (XX^\top)^kX=U\Sigma^{2k}U^\top U\Sigma V^\top=U\Sigma^{2k+1}V^\top; k=0 is trivial. A linear combination of such terms is U\varphi(\Sigma)V^\top, with the same singular vectors, and iterating gives U\varphi^N(\Sigma)V^\top. (A negative \varphi^N(\sigma_i) would give singular value |\varphi^N(\sigma_i)| with the sign absorbed into u_i; for the cubic and the quintic on the ranges used here the values stay positive.) Zero singular values stay zero, which is why the limit is the reduced UV^\top (4.1.C.3).

2. The analysis needs every initial singular value in (0,1]. \|G\|_F=\big(\sum\sigma_i^2\big)^{1/2}\ge\sigma_{\max} guarantees this at the cost of one cheap reduction, while \|G\|_2 would need its own iterative computation. The cost: if G has r comparable singular values, \|G\|_F\approx\sqrt r\,\sigma_{\max}, so even the largest singular values start near 1/\sqrt r and the small ones much lower. For a 256\times512 Gaussian matrix, all singular values of G/\|G\|_F lie in about [0.019,0.105]. This is why the growth rate near zero, \varphi'(0), is what Muon optimizes.

0. p(1)=a+b=1 and p'(1)=a+3b=0 give b=-\tfrac12, a=\tfrac32: p(\sigma)=\tfrac32\sigma-\tfrac12\sigma^3, the classical iteration.

1. p(\sigma)=\sigma\iff\sigma(1-\sigma^2)=0, so the fixed points are 0,\pm1. For \sigma\in(0,1): p(\sigma)-\sigma=\tfrac12\sigma(1-\sigma^2)>0 and 1-p(\sigma)=\tfrac12(1-\sigma)^2(2+\sigma)>0, so the iterates increase and stay below 1; a bounded increasing sequence converges, and its limit is a fixed point in (0,1], which can only be 1. On (1,\sqrt3), p'(\sigma)=\tfrac32(1-\sigma^2)<0, so p decreases from p(1)=1 to p(\sqrt3)=0 and maps the interval into (0,1). Frobenius normalization keeps every \sigma_0\le1, so larger starting values never arise.

2. p(1+e)=\tfrac12\big(3+3e-(1+3e+3e^2+e^3)\big)=1-\tfrac32e^2-\tfrac12e^3. So e_{k+1}\approx-\tfrac32e_k^2: quadratic convergence, from p'(1)=0. (“Cubic” refers to the polynomial, not the order.)

3. The singular value grows by about 3/2 per step until it is order one, so it needs about \log(500)/\log(1.5)\approx15 steps to reach 0.5; iterating exactly gives 16. With five steps, only singular values that start above about 0.09 reach 0.6. The slow start, not the final convergence, is the bottleneck.

1. \varphi(1)=a+b+c=0.701 and \varphi'(0)=a=3.4445. Since \varphi(1)\ne1, the iteration does not converge to UV^\top.

2. The positive fixed points solve c\sigma^4+b\sigma^2+(a-1)=0, a quadratic in \sigma^2, giving \sigma\approx0.868 and \sigma\approx1.264. At the first, \varphi'(0.868)=a+3b\sigma^2+5c\sigma^4\approx-1.58: |\varphi'|>1, so it repels, and the negative sign means iterates overshoot it from side to side. Numerically, generic starting points never settle: iterated further, they keep cycling through values between about 0.68 and 1.13.

3. After five steps, starting points in [10^{-4},1] end up in about [0.048,1.20]; the top of this range is the maximum of \varphi on [0,1], \approx1.202 at \sigma\approx0.55. Every starting value above about 1.3\times10^{-3} is at least 0.6 after five steps.

4. Near 0 the quintic multiplies by a=3.4445 per step, so by a^5\approx485 in five steps; the cubic multiplies by 1.5^5\approx7.6. In five steps the quintic brings singular values from about 10^{-3} up to order one, which the cubic needs about 16 steps to do. The price is accuracy: the output is U\tilde\Sigma V^\top with \tilde\Sigma in a band around 1 rather than UV^\top exactly. For an optimizer step, whose overall size is set by \eta anyway, Muon’s authors found this an acceptable trade.

def newton_schulz(G, coeffs, steps):
    a, b, c = coeffs
    X = G / np.linalg.norm(G)
    out = [np.linalg.svd(X, compute_uv=False)]
    for _ in range(steps):
        A = X @ X.T
        X = a * X + (b * A + c * A @ A) @ X
        out.append(np.linalg.svd(X, compute_uv=False))
    return out

CUBIC = (1.5, -0.5, 0.0)
MUON = (3.4445, -4.7750, 2.0315)

m, n = 128, 256
U0, _ = np.linalg.qr(rng.standard_normal((m, m))); V0, _ = np.linalg.qr(rng.standard_normal((n, m)))
G = U0 @ np.diag(np.logspace(0, -2, m)) @ V0.T + 1e-3 * rng.standard_normal((m, n))

cub = newton_schulz(G, CUBIC, 20); quin = newton_schulz(G, MUON, 5)
fig, axs = plt.subplots(1, 3, figsize=(12, 3.3))
s = np.linspace(0, 1.25, 400)
axs[0].plot(s, (3 * s - s**3) / 2, label="cubic p(σ)")
axs[0].plot(s, MUON[0] * s + MUON[1] * s**3 + MUON[2] * s**5, label="Muon quintic φ(σ)")
axs[0].plot(s, s, "k:", lw=1); axs[0].axhspan(0.7, 1.3, color="0.85", alpha=0.5, lw=0)
axs[0].set_xlabel("σ"); axs[0].legend(fontsize=8); axs[0].set_title("scalar maps (band: 0.7–1.3)")
bins = np.linspace(-4, 0.3, 60)
for ax, seq, ks, title in [(axs[1], cub, [0, 5, 10, 20], "cubic"), (axs[2], quin, [0, 1, 2, 3, 5], "Muon quintic")]:
    for k in ks:
        ax.hist(np.log10(seq[k]), bins=bins, histtype="step", lw=1.5, label=f"step {k}")
    ax.set_xlabel("log10 singular value"); ax.set_title(title); ax.legend(fontsize=8)
plt.tight_layout(); plt.show()
print(f"cubic after 5 steps:   σ in [{cub[5].min():.3f}, {cub[5].max():.3f}]")
print(f"cubic after 20 steps:  σ in [{cub[20].min():.3f}, {cub[20].max():.3f}]")
print(f"quintic after 5 steps: σ in [{quin[5].min():.3f}, {quin[5].max():.3f}]")

cubic after 5 steps:   σ in [0.025, 0.988]
cubic after 20 steps:  σ in [1.000, 1.000]
quintic after 5 steps: σ in [0.682, 1.202]

1. \|Wx\|_{\rm RMS}/\|x\|_{\rm RMS}=\big(\|Wx\|_2/\sqrt{n_{\rm out}}\big)\big/\big(\|x\|_2/\sqrt{n_{\rm in}}\big)=\sqrt{n_{\rm in}/n_{\rm out}}\;\|Wx\|_2/\|x\|_2; maximizing over x gives \sqrt{n_{\rm in}/n_{\rm out}}\,\|W\|_2. So \|W_l\|_{\rm RMS\to RMS}=\Theta(1) is the same as \|W_l\|_2=\Theta(\sqrt{n_l/n_{l-1}}), with n_{\rm in}=n_{l-1}, n_{\rm out}=n_l.

2. The norm is a constant multiple of the spectral norm, so it has the same steepest direction, UV^\top (4.1.C), and the unit ball \|D\|_{\rm RMS\to RMS}\le1 is \|D\|_2\le\sqrt{n_{\rm out}/n_{\rm in}}. The trust-region step is therefore -\eta\sqrt{n_{\rm out}/n_{\rm in}}\,UV^\top. Muon’s factor \sqrt{\max(1,n_{\rm out}/n_{\rm in})} agrees when the layer widens or is square, and differs by the constant \sqrt{n_{\rm in}/n_{\rm out}} when it narrows; at a fixed aspect ratio that does not change the width scaling.

Moonlight uses a different convention: it scales the full-rank UV^\top of an A\times B matrix (RMS entry 1/\sqrt{\max(A,B)}) by 0.2\sqrt{\max(A,B)}, so every update has RMS entry about 0.2, like AdamW’s, and AdamW’s hyperparameters can be reused. Its spectral norm 0.2\sqrt{\max(A,B)} grows with width for a square layer: a choice about matching AdamW at a given size, not the width scaling of Condition 1.

1. \|h_l\|=\|W_lh_{l-1}\|\le\|W_l\|_2\|h_{l-1}\|=\|W_l\|_2\,\Theta(\sqrt{n_{l-1}}), which is O(\sqrt{n_l}) if \|W_l\|_2=O(\sqrt{n_l/n_{l-1}}); likewise for \Delta W_l\,h_{l-1}. Conversely, if the bound must hold for every input, including one aligned with the top right singular vector of W_l (where the inequality is an equality), the conditions are necessary. For the initialization, \sigma_l(\sqrt{n_l}+\sqrt{n_{l-1}})=\Theta(\sqrt{n_l/n_{l-1}}) gives \sigma_l=\Theta\big(\tfrac1{\sqrt{n_{l-1}}}\min(1,\sqrt{n_l/n_{l-1}})\big), Yang–Simon–Bernstein’s Parametrization 1; the \min matters only when the layer narrows. The full change \Delta h_l=\Delta W_lh_{l-1}+W_l\Delta h_{l-1}+\Delta W_l\Delta h_{l-1} also carries the previous layer’s change forward; the other two terms are controlled by the same bounds recursively, and the paper assumes the terms do not cancel.

2. \Delta W\,h=-\eta\,\delta\,(h^\top h)=-\eta\|h\|^2\delta, so \|\Delta W\,h\|=\eta\|\delta\|\|h\|^2. Also \|\Delta W\|_2=\eta\|\delta\|\|h\| for a rank-one matrix. Hence \|\Delta W\,h\|=\|\Delta W\|_2\|h\|: the inequality is an equality, because h is exactly the top right singular vector of the update it produced. So the direct contribution is as large as \|\Delta W\|_2 allows, and its scaling is set by \|\Delta W\|_2.

3. \|\Delta W\,h\|=s\,|v\cdot h|. For a random unit v\in\mathbb R^{n_{l-1}} independent of h, v\cdot h has mean zero and variance \|h\|^2/n_{l-1}, so \|\Delta W\,h\|\approx\|\Delta W\|_2\,\|h\|/\sqrt{n_{l-1}}: smaller than the bound by a factor \sqrt{n_{l-1}}. A low-rank update affects the features at full strength only because it is aligned with them, and gradient updates are aligned by construction. This is why the spectral norm, and not, say, the Frobenius norm, measures a gradient update’s effect.

4. Layer l’s first-order contribution is \langle\nabla_{W_l}L,\Delta W_l\rangle=-\eta_l\|\nabla_{W_l}L\|_F^2<0, and by assumption its magnitude is \Theta(1)=\Theta(\|\Delta W_l\|_2\|\nabla_{W_l}L\|_2). (The assumption is per layer: a total change of order one would not by itself force every layer’s share to be of order one.) With \|\Delta W_l\|_2=\Theta(\sqrt{n_l/n_{l-1}}) this forces \|\nabla_{W_l}L\|_2=\Theta(\sqrt{n_{l-1}/n_l}). An SGD step has \|\Delta W_l\|_2=\eta_l\|\nabla_{W_l}L\|_2, so Condition 1 needs \eta_l=\Theta\big(\sqrt{n_l/n_{l-1}}\big/\sqrt{n_{l-1}/n_l}\big)=\Theta(n_l/n_{l-1}). For a square hidden layer this is \Theta(1); for the input layer (fixed n_0) it grows like the width; for the output layer (fixed n_L) it shrinks like 1/\text{width}. These are muP’s SGD rates (TP-V, Table 3) and Yang–Simon–Bernstein’s eqs. (18)–(20).

1. \mathrm{sign}(a_ib_j)=\mathrm{sign}(a_i)\,\mathrm{sign}(b_j), so \mathrm{sign}(G)=s\,t^\top with s,t\in\{\pm1\}^n (almost surely no zeros). A rank-one matrix has spectral norm \|s\|\|t\|=\sqrt n\sqrt n=n. For an m\times n matrix it is \sqrt{mn}, equal to its Frobenius norm.

2. Upper bound: \|S\|_2\le\|S\|_F=\sqrt{n^2}=n for any \pm1 matrix. Lower bound: \|S\|_2\ge u^\top Sv for unit u,v. With u=a_1/\|a_1\|, v=b_1/\|b_1\|, u^\top Sv=\frac{\sum_{ij}a_{i1}b_{j1}\,\mathrm{sign}(G_{ij})}{\|a_1\|\,\|b_1\|}. The numerator is a sum of n^2 terms, each distributed as X\,\mathrm{sign}(X+Y), with mean c_r. The terms are dependent (entries in the same row share a_i), but only weakly, and a law of large numbers gives c_rn^2(1+o(1)); we do not make this rigorous here. The denominator is n(1+o(1)). So \|S\|_2\ge c_rn(1+o(1)). Positivity: conditional on X=x>0, \mathbb E[\mathrm{sign}(x+Y)]=P(Y>-x)-P(Y<-x)=P(|Y|<x) by the symmetry of Y, and similarly for x<0, so \mathbb E[X\,\mathrm{sign}(X+Y)\mid X]=|X|\,P(|Y|<|X|)\ge0, strictly positive with positive probability. For r=1, Y=0 and c_1=\mathbb E|a||b|=2/\pi. Numerically c_4\approx0.37 and c_{16}\approx0.20. Together, c_rn\lesssim\|S\|_2\le n.

3. For an n\times n hidden layer Condition 1 asks for \|\Delta W\|_2=\Theta(1). A sign-based update has \|\Delta W\|_2=\eta\,\Theta(n) when G is low-rank, so \eta=\Theta(1/n); it would have \eta\cdot2\sqrt n if the sign matrix were dense and random, giving \eta=\Theta(1/\sqrt n). For an n_l\times n_{l-1} layer the low-rank case gives \eta\sqrt{n_ln_{l-1}}=\Theta(\sqrt{n_l/n_{l-1}}), i.e. \eta=\Theta(1/n_{l-1})=\Theta(1/\text{fan-in}): muP’s Adam rule. The dense case has \|S\|_2\approx\sqrt{n_l}+\sqrt{n_{l-1}}, giving \eta=\Theta\big(\sqrt{n_l/n_{l-1}}/(\sqrt{n_l}+\sqrt{n_{l-1}})\big), which is \Theta(n^{-1/2}) at a fixed aspect ratio. A width-independent Adam learning rate (standard parametrization) makes hidden-layer updates grow like the width.

from scipy.special import gammaln

def top_singular_value(S, iters=40):
    v = rng.standard_normal(S.shape[1]).astype(S.dtype)
    for _ in range(iters):
        u = S @ v; u /= np.linalg.norm(u)
        v = S.T @ u; s = np.linalg.norm(v); v /= s
    return s

ns = np.array([64, 128, 256, 512, 1024, 2048])
limit = lambda r: np.exp(gammaln(r / 2) - gammaln((r + 1) / 2)) / np.sqrt(np.pi)   # E|w_1|
fig, ax = plt.subplots(figsize=(5.6, 3.8))
for r in (1, 4, 16, "dense"):
    norms = []
    for n in ns:
        G = rng.standard_normal((n, n)) if r == "dense" else rng.standard_normal((n, r)) @ rng.standard_normal((r, n))
        norms.append(top_singular_value(np.sign(G).astype(np.float32)))
    norms = np.array(norms)
    slope = np.polyfit(np.log(ns), np.log(norms), 1)[0]
    line, = ax.loglog(ns, norms, "o-", ms=4, label=f"rank {r}: slope {slope:.2f}" if r != "dense" else f"dense: slope {slope:.2f}")
    ref = limit(r) * ns if r != "dense" else 2 * np.sqrt(ns)
    ax.loglog(ns, ref, ":", color=line.get_color(), lw=1)
    print(f"{str(r):>5}: ||sign(G)||_2 / limit at n = 2048: {norms[-1] / ref[-1]:.3f}")
ax.set_xlabel("width n"); ax.set_ylabel(r"$\|\mathrm{sign}(G)\|_2$")
ax.set_title("dotted: $\\mathbb{E}|w_1|\\,n$ (low rank), $2\\sqrt{n}$ (dense)", fontsize=9)
ax.legend(fontsize=8); plt.tight_layout(); plt.show()

    1: ||sign(G)||_2 / limit at n = 2048: 1.000
    4: ||sign(G)||_2 / limit at n = 2048: 1.043
   16: ||sign(G)||_2 / limit at n = 2048: 1.120
dense: ||sign(G)||_2 / limit at n = 2048: 0.983

Let N_k be the norm of the first k layers composed, with mass k and sensitivity 1. Then N_1=\|W_1\|_2 and, by the rule with \mu_1=k-1, \mu_2=1, s_2=1, N_k=\max\Big(\tfrac{k}{k-1}N_{k-1},\ k\,\|W_k\|_2\Big). If N_{k-1}=(k-1)\max_{l<k}\|W_l\|_2, this gives N_k=k\max_{l\le k}\|W_l\|_2, so by induction N_L=L\max_l\|W_l\|_2. A max of layer norms is the case of §4.1’s remark 3, so the steepest-descent step of modular radius \eta moves every layer by its own steepest step with spectral size \eta/L.

The bound behind it is a worst case: inside the well-normed region, a perturbation of size \epsilon in each of L layers changes the output by up to about L\epsilon if the perturbations add coherently. Steps of \eta/L per layer keep that of order \eta, which suggests per-layer step sizes (or branch multipliers) shrinking like 1/L. For residual networks, posed as the scaling L^{-\alpha} of each branch, this is debated: Tensor Programs VI and Bordelon et al. use \alpha=\tfrac12, while CompleteP (Dey et al. 2025) reports better hyperparameter transfer with \alpha=1 in pre-LN transformers. Masses let some layers (embeddings, say) take a larger or smaller share than the equal split.