A moving average of the weights (EMA)

The averaged copy moves a small fraction of the way each step

Call a weight after step t w_t, and its averaged copy e_t. Pick a number β (the Greek letter beta) a little below 1, such as 0.999. After every training step, update the copy:

e_t = β × e_t−1 + (1 − β) × w_t

The copy keeps a share β of its old value and takes a share 1 − β from the new weight. A model has millions of weights; each one gets its own copy, updated by the same rule.

Old weights fade by a factor β each step

Unroll the rule and the copy is a weighted average of all past weights. The newest weight counts 1 − β. The one before counts (1 − β) × β, the one before that (1 − β) × β², and so on. With β = 0.9 the shares are 0.1, 0.09, 0.081, … . The shares shrink by the same factor at each step back; that is what "exponential" means here (see exponentials).

Two numbers describe how far back the average reaches:

• It covers roughly the last 1 / (1 − β) steps: 10 steps for 0.9, 1,000 steps for 0.999.
• A weight's share halves every ln 2 / (−ln β) steps, about 0.69 / (1 − β) when β is close to 1. For 0.999 that is about 693 steps.

A larger β smooths more but lags behind

The figure shows one weight over 1,500 steps of a made-up training run. The dashed line is where the weight is heading: it climbs at first, then levels off. The grey line is the trained weight, jittering around the dashed line. The orange line is the averaged copy.

Predict first. Is the largest β always the best? Then move the slider to find the β that keeps the orange line closest to the dashed line over the whole run. The readout gives the average distance; make it as small as you can.

In this short run, the best β is near 0.95. Real training runs last hundreds of thousands of steps, and their weights drift slowly compared with their jitter. So values such as 0.999 or 0.9999 are common there.

In code, the copy is a second model

This PyTorch snippet trains a model with one weight and one bias (a linear layer) on minibatches of 4. After each step it updates the averaged copy. Then it compares how much each one jitters over the last 200 steps.

ema.py

import copy
import torch

torch.manual_seed(0)                          # set the seed of PyTorch's random numbers
x = torch.randn(1000, 1)                      # 1,000 random inputs, one per row
y = 3 * x + 0.5 * torch.randn(1000, 1)       # the best slope is close to 3

model = torch.nn.Linear(1, 1)                 # one weight and one bias
ema = copy.deepcopy(model)                    # the averaged copy; never trained
opt = torch.optim.SGD(model.parameters(), lr=0.1)  # gradient descent on the weights
beta, raw, avg = 0.95, [], []

for step in range(400):
    i = torch.randint(0, 1000, (4,))          # a minibatch of 4 examples
    loss = ((model(x[i]) - y[i]) ** 2).mean()
    opt.zero_grad(); loss.backward(); opt.step()  # zero the gradients, compute them, step
    with torch.no_grad():                     # the EMA update, for every weight
        for pe, p in zip(ema.parameters(), model.parameters()):
            pe.mul_(beta).add_(p, alpha=1 - beta)
    # Add each copy's weight to the end of its list, as a plain number from .item().
    raw.append(model.weight.item()); avg.append(ema.weight.item())

r, a = torch.tensor(raw[200:]), torch.tensor(avg[200:])   # after both settle
print(f"raw weight: mean {r.mean():.3f}, spread {r.std():.3f}")
print(f"EMA weight: mean {a.mean():.3f}, spread {a.std():.3f}")

raw weight: mean 3.005, spread 0.092
EMA weight: mean 3.002, spread 0.037

Both copies center on the same value, but the averaged copy jitters less than half as much. pe.mul_(beta).add_(p, alpha=1 - beta) is the update rule above, applied in place. Inside torch.no_grad(), PyTorch records no steps for automatic gradients, which lets the code change the averaged weights in place.