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.

[GXYM]

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 ← μ B t − 1 + G t

(where B 0 = G 0 = 0 )

And the updated momentum expressed as EMA:

B t ′ ← μ B t − 1 ′ + ( 1 − μ ) G t

B T ′ may be expressed in terms of B T times a constant

B T = ∑ k = 0 T μ T − k G k B T ′ = ( 1 − μ ) ∑ k = 0 T μ T − k G k B T ′ = ( 1 − μ ) B T

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

Hello, have any specific experimental verification comparisons been conducted here?

[RobotSail]

@GXYM I only checked it empirically, did not explore further. The EMA was adopted by the NanoGPT speedrun due to achieving faster convergence, which is the primary motivation for choosing it over decaying sum here.

[GXYM]

@GXYM I only checked it empirically, did not explore further. The EMA was adopted by the NanoGPT speedrun due to achieving faster convergence, which is the primary motivation for choosing it over decaying sum here.

thansks