Muon momentum update appears to deviate from the paper; seems to be using a running average

[RobotSail]

In the write-ups for Muon, as well as other third-party sources such as Moonshot AI's Muon is Scalable for LLM Training paper, the update rule for Momentum is described as $B_t \gets \mu B_{t-1} + G_t$, per the following definition:

But the current muon_update function updates the momentum using muon.lerp_:

def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True):
    momentum.lerp_(grad, 1 - beta)
    update = grad.lerp_(momentum, beta) if nesterov else momentum
    if update.ndim == 4: # for the case of conv filters
        update = update.view(len(update), -1)
    update = zeropower_via_newtonschulz5(update, steps=ns_steps)
    update *= max(1, grad.size(-2) / grad.size(-1))**0.5
    return update

This calculation instead appears to have the effect of calculating Momentum as a running average instead of a decaying sum, i.e.:

$$ M_t \gets \mu B_{t-1} + (1 - \mu) G_t $$

Looking at the original implementation from commit b8dda, it looks like this has changed from the original implementation which computed it correctly:

                    if 'momentum_buffer' not in state:
                        state['momentum_buffer'] = torch.zeros_like(g)
                    buf = state['momentum_buffer']
                    buf.mul_(momentum).add_(g)
                    if group['nesterov']:
                        g = g.add(buf, alpha=momentum)

[RobotSail]

Did some more digging and reached out to You Jiacheng, who made the original change in NanoGPT, who pointed out that decaying sum and EMA are equivalent under a constant:

The original momentum term is defined as a decaying sum:

$$ B_{t} \gets \mu B_{t-1} + G_t $$

(where $B_0= G_0 = 0$)

And the updated momentum expressed as EMA:

$$ B'_{t} \gets \mu B'_{t-1} + (1 - \mu)G_t $$

$B'_T$ may be expressed in terms of $B_T$ times a constant

$$ \begin{align*} B_T&= \sum_{k=0}^T \mu^{T-k} G_k & \\ B'_T&= (1 - \mu) \sum_{k=0}^T \mu^{T-k} G_k & \\ B'_T&= (1 - \mu) B_T & \\ \end{align*} $$

But from an implementation POV, it seems like EMA has the effect of stabilizing the momentum: as $\mu$ becomes very large and carries past gradients forward, $1-\mu$ counteracts this by downscaling them. Similarly, if $\mu$ becomes very small and dramatically shrinks the past momentum, $1 - \mu$ serves to amplify these values so the final $B_T$ is still reasonable in magnitude.

[parlance-zz]

RobotSail is correct. Additionally, using lerp is faster and retains higher numerical precision.