Skip to content

Model Averaging — EMA, SWA

Overview

Averaging models over training is a free ensemble: it smooths the loss landscape, de-noises the last-epoch chatter, and routinely adds 0.5-2% accuracy or BLEU/ROUGE for nothing at inference time.

  • EMA (Exponential Moving Average): ema = decay * ema + (1-decay) * current, updated every step. Best for final deployment weights.
  • SWA (Stochastic Weight Averaging): uniform mean of checkpoints over the last few epochs. Better for training-end quality and flat minima.
  • EMA decay example: 0.999-0.9999 for long jobs; 0.99 for short bursts.

💡 Rule: use the EMA weights at eval/inference; keep the live weights for training (they adapt faster). Never train the EMA copy itself.

EMA — the Standard Implementation

import copy, torch, torch.nn as nn

class EMA:
    def __init__(self, model, decay=0.999):
        self.decay = decay
        self.ema = copy.deepcopy(model)   # shadow copy
        for p in self.ema.parameters():
            p.requires_grad_(False)

    @torch.no_grad()
    def update(self, model):
        ema_p = dict(self.ema.named_parameters())
        for name, p in model.named_parameters():
            if name in ema_p:
                ema_p[name].mul_(self.decay).add_(p.detach(), alpha=1 - self.decay)

    def state_dict(self):
        return self.ema.state_dict()

ema = EMA(model, decay=0.999)
for step in range(100):
    loss = (model(torch.randn(8, 4))).abs().mean()
    loss.backward(); opt.step()
    ema.update(model)                    # shadow follows live model

# switch to EMA at eval:
ema_model = ema.ema.eval()
with torch.no_grad():
    print(ema_model(torch.randn(1, 4)).shape)

⚠️ deepcopy needs the model on CPU or the same device — either deepcopy before .cuda() or move after copy.

Buffers & EMA (Batchnorm gotcha)

Copying only parameters leaves BatchNorm running stats stale. Copy buffers too:

@torch.no_grad()
def ema_update(self, model):
    for name, p in model.named_parameters():
        if name in self.ema_names:
            self.ema_params[name].mul_(self.decay).add_(p.detach(), alpha=1 - self.decay)
    for name, b in model.named_buffers():
        if name in self.ema_buffers:
            self.ema_buffers[name].copy_(b)   # buffers follow live copy

💡 Many opt for: after training, re-run a few eval forwards on EMA model so BatchNorm recalibrates (model.train(False) + forward pass on real data).

SWA — Stochastic Weight Averaging

from torch.optim.swa_utils import AveragedModel, SWALR

swa_model = AveragedModel(model)          # averages params over time
swa_sched = SWALR(opt, swa_lr=0.05)       # low, constant LR for the SWA phase

for epoch in range(10):
    train_one_epoch()
    if epoch >= 7:                        # start averaging near the end
        swa_model.update_parameters(model)
        swa_sched.step()

torch.optim.swa_utils.update_bn(swa_model, val_loader)   # recalibrate BN stats

EMA vs SWA — When to Use What

EMA SWA
Update frequency every step per epoch/checkpoint
Effective ensemble infinite memory last-N checkpoints only
Final weights long-run smoothed flattened minima
Best for chat/LLM deployment, RL value nets vision/classical-generalization tasks
Decay/lr 0.99-0.9999 SWA phase LR ~ peak/10

✅ EMA is the modern default for LLM/RL deployment; SWA shines on classification deadlines with fewer training steps.

Pitfalls

  1. Use EMA at inference, not in training metrics — train metrics reflect fast weights.
  2. torch.compile + EMA: compile the EMA copy separately or keep EMA un-compiled (shadows don't benefit).
  3. Don't EMA-then-distill: EMA-of-EMA compounds; one EMA layer is enough.
  4. Checkpoint both: save model and ema.state_dict() — resume keeps both consistent.
ckpt = {"model": model.state_dict(), "ema": ema.state_dict(), "step": step}
torch.save(ckpt, "ckpt.pt")

Key Takeaways

  • EMA/SWA are free ensembles: no extra forward passes, small memory (one shadow copy).
  • EMA: per-step decay, best for deployment; SWA: end-phase averaging, best for final checkpoints.
  • Copy buffers (BN stats) in EMA; recalibrate BN after SWA.
  • Serve EMA weights at inference; keep the fast live weights during training.
  • Save both in checkpoints for clean resumability.