Lipschitz Finite-Newton–Schulz Muon

(1) Integrate at the Muon optimizer update for every 2-D weight matrix W, after momentum or gradient accumulation and before applying the parameter step. Let M_t = beta M_{t-1} + (1-beta) grad_t, where beta is the momentum coefficient. Replace exact SVD-based polar(M_t), or the existing fixed-depth orthogonalization, by finite Newton–Schulz. This changes only the direction calculation; Adam-style scalar updates for biases and LayerNorm parameters remain unchanged.

(2) Pseudocode:

for each matrix parameter W with gradient G:
    M = beta*M + (1-beta)*G
    alpha = max(eps, estimate_operator_norm(M))
    X = M / alpha
    if rows(M) >= cols(M):
        repeat j = 1..t:
            X = 0.5 * X @ (3*I - X.T @ X)
    else:
        repeat j = 1..t:
            X = 0.5 * (3*I - X @ X.T) @ X
    direction = alpha * X
    W = W - lr * shape_scale(W, direction) * direction

Use power iteration with 2–3 matrix-vector products to estimate alpha = ||M||_op, and clamp alpha to at least 1e-6. Start with t=2 or 3. For an adaptive schedule, choose t_t = min(t_max, max(t_min, ceil(c log(1+q_t)))) where q_t is the normalized training-step fraction or a running inverse gradient-noise estimate; a simpler first test uses t=2 for the first 70% of training and t=4 thereafter.

(3) The paper supplies the spectral-map construction, polar/nuclear duality, and the logarithmic-depth smoothing guarantee. Estimate ||M||_op empirically by power iteration, and measure the actual orthogonalization residual r_t = ||X_t^T X_t-I||_F / sqrt(min(m,n)). Also estimate update smoothness by recording ||direction_t-direction_{t-1}||_F / ||M_t-M_{t-1}||_F. No SVD is needed during training; occasional SVDs on sampled layers can diagnose singular-value behavior.

(4) First experiment: train a 100M–300M decoder-only Transformer on a fixed small language-model corpus using the same batch size, learning rate, momentum, and FLOP budget. Compare baseline Muon with fixed t=5, exact-polar Muon if feasible, finite t=2, finite t=3, and the adaptive 2-to-4 schedule. Log training loss versus optimizer FLOPs, gradient/update cosine similarity, orthogonalization residual, and loss spikes after rank or scale changes. Success is lower loss at equal optimizer compute, fewer unstable spikes, or equal perplexity with fewer Newton–Schulz matrix multiplications. A falsification is that shallow finite-depth variants fail to match t=5 perplexity or show no stability improvement over exact polar updates on ill-conditioned or low-rank initialized layers.