Sign compression for Muon: SignMuon, MuonSign, and the Limits of Error Feedback

Overview§

This paper studies extreme compression of the Muon optimizer, which uses a matrix-wise orthogonalization (Newton-Schulz) preconditioner on the gradient. The authors propose SignMuon: take the elementwise sign of the final Muon update, yielding one bit per parameter. They contrast this with placing the sign before the linear minimization oracle (LMO) or on both sides, and with error-feedback corrections.

Background: Muon update§

Muon computes a preconditioned update via the LMO:

$$M_t = \text{LMO}_G(G_t) = \text{NewtonSchulz}(G_t)$$

where $G_t$ is the gradient reshaped to a matrix, and then applies momentum and weight decay. The raw update is $\Delta_t = \mu u_{t} + M_t$ (with momentum state). SignMuon compresses this final update:

$$\text{SignMuon}(\Delta_t) = \text{sign}(\Delta_t) \in \{-1,+1\}^{d}$$

Only the signs are communicated, reducing bandwidth to 1/32 of a 32-bit float.

Sign placement variants§

The authors define three placements:

• MuonUSign: sign before the LMO, i.e., $\text{sign}(G_t)$ is fed to the LMO, then momentum applied.
• MuonSign: sign on both sides, $\text{sign}(\text{LMO}(\text{sign}(G_t)))$.
• SignMuon: sign after the LMO (the proposed heuristic).

They construct an explicit small instance on which both MuonUSign and MuonSign ascend on a linear function, proving no sign placement around the oracle is universally descending.

Error feedback§

Error feedback is the standard remedy for biased compressors. Applied to Muon's output (i.e., compressing the update and adding the error back to the next update), it fails: the paper shows divergence for every smoothness constant, step size, and momentum—a strong impossibility result.

In contrast, applying error feedback to the gradient before the LMO works. The resulting methods EF21-MuonUSign and EF21-MuonSign attain the standard $\mathcal{O}(T^{-1/2})$ rate for the squared gradient norm on smooth nonconvex problems. EF21-MuonSign sends one bit per direction, while EF21-MuonUSign sends signs of gradients before the LMO.

Experiments: practice reverses theory§

Across centralized CIFAR-10, federated CIFAR-10, and the nanoGPT speedrun, the strongest compressed method is consistently SignMuon (sign-after-the-LMO)—the very variant proven divergent. The provably convergent EF21 variants trail in accuracy and wall-clock time. This highlights that at practical scales, the heuristic placement matters more than formal convergence guarantees.

Implementation sketch§

The following pseudo-PyTorch illustrates the core update loop:

import torch

def newton_schulz(M, steps=5):
    a, b, c = 1.0, 0.55, -0.05  # optimal constants for 2D
    X = M
    for _ in range(steps):
        X = a * X + b * X @ X.transpose(-2, -1) @ X + c * X @ X @ X.transpose(-2, -1) @ X
    return X

def sign_muon_step(grad, momentum, lr, mu=0.9):
    G = grad.reshape(grad.shape[0], -1)
    M = newton_schulz(G)
    update = mu * momentum + M
    comp = torch.sign(update)  # 1 bit per element
    # apply compressed update (e.g., allreduce of signs)
    return comp, update  # update stored as momentum state

The key takeaway: compressing the preconditioned update (after the LMO) is a delicate heuristic that occasionally diverges in theory but excels in practice, whereas error-feedback-corrected gradient compression is theoretically sound yet practically weaker.