🎓 A1Learner
🔍

PyTorch EMA of Weights

An exponential moving average of weights is a running average that trails the training weights and favours recent steps. The average often scores better than the raw final weights at eval time.

One EMA update by hand

Each update keeps most of the running average. A decay of 0.99 mixes in one percent of the new value.

average = torch.tensor(0.0)
value = torch.tensor(2.0)
average = 0.99 * average + 0.01 * value
print(average)

The average mirrors the model

The EMA dictionary holds one entry per tensor in the model's state_dict. The shapes match entry for entry.

print(list(ema.keys()))
print(ema["0.weight"].shape, model.state_dict()["0.weight"].shape)

The average at eval time

Loading the EMA weights into the model swaps the jittery final weights for the smoothed ones. The validation loss is the same number the closing example prints.

model.load_state_dict({name: value.to(backup[name].dtype) for name, value in ema.items()})
print(f"EMA validation loss: {validation_loss():.4f}")
model.load_state_dict(backup)

Example

import torch
from torch import nn


def make_blobs(count, seed, flip_share=0.0):
    g = torch.Generator().manual_seed(seed)
    first = torch.randn(count // 2, 2, generator=g) * 0.7 + torch.tensor([-1.0, -1.0])
    second = torch.randn(count // 2, 2, generator=g) * 0.7 + torch.tensor([1.0, 1.0])
    x = torch.cat([first, second])
    y = torch.cat([torch.zeros(count // 2), torch.ones(count // 2)]).long()
    perm = torch.randperm(count, generator=g)
    x, y = x[perm], y[perm]
    if flip_share > 0.0:
        flipped = torch.rand(count, generator=g) < flip_share
        y = torch.where(flipped, 1 - y, y)
    return x, y


train_x, train_y = make_blobs(120, seed=3, flip_share=0.15)
val_x, val_y = make_blobs(400, seed=4)
loss_fn = nn.CrossEntropyLoss()

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
optimizer = torch.optim.SGD(model.parameters(), lr=0.3)
decay = 0.99
ema = {name: value.detach().clone().float()
       for name, value in model.state_dict().items()}

batch_gen = torch.Generator().manual_seed(8)
for step in range(300):
    idx = torch.randint(0, len(train_y), (16,), generator=batch_gen)
    optimizer.zero_grad()
    loss = loss_fn(model(train_x[idx]), train_y[idx])
    loss.backward()
    optimizer.step()
    with torch.no_grad():
        for name, value in model.state_dict().items():
            ema[name].mul_(decay).add_(value.detach().float(), alpha=1 - decay)


def validation_loss():
    model.eval()
    with torch.no_grad():
        return float(loss_fn(model(val_x), val_y))


raw_loss = validation_loss()
backup = {name: value.detach().clone()
          for name, value in model.state_dict().items()}
model.load_state_dict({name: value.to(backup[name].dtype)
                       for name, value in ema.items()})
ema_loss = validation_loss()
model.load_state_dict(backup)

print(f"Raw weights validation loss: {raw_loss:.4f}")
print(f"EMA weights validation loss: {ema_loss:.4f}")

Key points

  • EMA: exponential moving average, a running average where each new value counts for a fixed small share
  • Decay: the share the average keeps of itself each step, 0.99 in the closing example
  • Use one shadow tensor per parameter tensor, updated after every optimizer step
  • Noisy labels: a seeded 15 percent of the train labels are flipped, so late steps chase wrong targets
  • The raw weights are the model's weights at the final step, jitter and all
  • Eval loads the EMA weights into the model, scores them, then restores the raw weights

In this example, a small MLP trains on 120 blob rows with SGD at a lively rate of 0.3. A seeded 15 percent of the train labels are flipped, so every batch carries wrong targets. The final weights land wherever the last noisy steps pushed them. Beside the training, an EMA with decay 0.99 updates after each of the 300 steps. The validation set holds 400 clean rows. The raw weights score a validation loss of 0.2972. The EMA weights score 0.2184 in the same model shell. The average cancels the jitter that the noisy steps added. The win is measured on data neither set of weights trained on.

Note

EMA costs one extra copy of the weights and a few lines, and it never changes the training itself.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why do the EMA weights beat the raw weights in this run?