Exponential Moving Average of Weights

Keeping an exponential moving average (EMA) of model weights during training gives a smoother, often better model for evaluation. It's standard in diffusion models and in YOLO.

Implement a class EMA(params, decay=0.999, warmup=True), where params is a dict of NumPy arrays:

• The shadow weights start as copies of params.
• update(params): for each key, shadow = d * shadow + (1 - d) * param. With warmup, use d = min(decay, (1 + step) / (10 + step)), where step counts previous updates (starting at 0), so early averages aren't dominated by the random initialisation. Without warmup, d = decay.
• shadow: the dict of averaged arrays
• state_dict(): return copies of the shadow weights (so callers can't mutate the EMA)