EMA state: define which tensors move and which clock advances

This teaching rule updates an averaged parameter after each successful optimizer step using the new parameter value. Its decay d multiplies the old average and one minus d multiplies the current parameter. Initialize the average explicitly, commonly from the starting parameters, and keep averaging outside gradient tracking. The optimizer continues to update its live parameters and moments; the EMA does not replace those moments. Variable-decay warmup is another algorithm and needs its own recurrence and saved counter.

Diffusers 0.35.1's EMAModel tracks an explicit iterable of parameters and has update-count-dependent decay logic; its default behavior is not simply this fixed-decay hand rule from the first call. PyTorch 2.14's AveragedModel supports configurable averaging and buffer behavior with separate semantics. Inspect the actual helper rather than infer its full state from an option named ema. Document which artifact supplies sampling weights and whether evaluation temporarily swaps parameters or uses a separate model copy.

Gradient accumulation creates several microbatches per optimizer update. Averaging after every microbatch, after every attempted update, or only after a successful update produces different clocks, especially when mixed-precision overflow skips a step. Pick and test one rule. A decay expressed per update also changes its example-scale memory when global batch size changes. Record update count, effective batch, and any warmup schedule before comparing smoothing between runs. A half-life in optimizer updates is not automatically a half-life in training examples.

Model buffers such as running statistics or counters are not the same as trainable parameters. Decide whether each is copied, averaged when meaningful, recomputed, or excluded under an explicitly supported policy. Integer counters should not accidentally receive floating interpolation. Many diffusion architectures avoid batch-running statistics, but that does not justify ignoring every registered buffer. Also verify device, dtype, distributed/sharded state handling, and condition-encoder versions. An averaged parameter list alone may not reconstruct a complete usable model.

Save live model parameters, optimizer and scaler state, EMA shadow tensors, counters, decay configuration, relevant buffers, and random/data-loader state under the desired resume contract. Check that resuming reproduces the next EMA update as well as the next optimizer update. If sampling swaps in averaged weights, restore the live training weights before the next optimizer step. Compare live and EMA checkpoints with the same data, prompts, sampler, and seed portfolio; an apparent improvement can otherwise come from a changed inference protocol.

The Markov lab exactly matches one restricted constant-input averaging family. Reset, set p=0.10 and q=0, predict state-one mass 0.19 after two Steps, then Step twice. Separately Reset, set p=0.20 and q=0, predict mass 0.36 after two Steps, then Step twice. Starting from zero, this equals averaging a constant value one with decays 0.9 and 0.8 respectively. The instrument has no changing model parameters, buffers, optimizer, or checkpoint state, so the match does not extend to a full EMA training run.