Skip to content

EMA

Exponential moving average of model weights, updated once per optimizer step with a warmup ramp. Validation and checkpointing run under the averaged weights — these are the weights you evaluate and ship.

ema

Exponential Moving Average of Model Weights

Flow-matching and diffusion training produce noisy gradient estimates: the loss regresses a velocity target that is only conditionally determined by the network input, so consecutive iterates bounce around the optimum rather than settling into it. Averaging weights over training suppresses that jitter and is standard practice for this model family — the averaged weights are what you evaluate and ship, while the raw weights keep training.

Usage:

ema = EMA(model, decay=0.999)

for batch in loader:
    loss = objective.compute_loss(batch)
    loss.backward()
    optimizer.step()
    ema.update()              # after every optimizer step

with ema.average_parameters():
    metrics = evaluate(model)  # model temporarily holds EMA weights

EMA

Maintains an exponential moving average of a model's parameters.

The shadow weights follow

\[ \theta^{\text{EMA}} \leftarrow d \, \theta^{\text{EMA}} + (1 - d) \, \theta \]

with decay \(d\). Early in training the shadow is dominated by its (arbitrary) initial value, so warmup ramps the effective decay up from 0 following \(\min(d, (1 + n) / (10 + n))\) at step \(n\). This is the schedule used by most diffusion implementations and it removes the need to tune a separate "start EMA at step N" threshold.

Buffers (e.g. BatchNorm running statistics) are copied rather than averaged, matching the reference implementations.

Parameters:

Name Type Description Default
model Module

Model whose parameters are tracked.

 required
decay float

Target decay rate. Typical values are 0.999 (short runs) to 0.9999 (long runs). Higher means smoother and slower to adapt.

 0.999
warmup bool

Ramp the effective decay early in training (default: True).

 True
device Optional[device]

Optional device for the shadow copy. Defaults to keeping each shadow tensor on the same device as its source parameter.

 None
Example

model = nn.Linear(2, 2) ema = EMA(model, decay=0.9, warmup=False) ema.update() with ema.average_parameters(): ... _ = model(torch.zeros(1, 2)) # uses averaged weights

Source code in flowpde/trainers/ema.py 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class EMA:
    """
    Maintains an exponential moving average of a model's parameters.

    The shadow weights follow

    $$
    \\theta^{\\text{EMA}} \\leftarrow d \\, \\theta^{\\text{EMA}}
                                    + (1 - d) \\, \\theta
    $$

    with decay $d$.  Early in training the shadow is dominated by its
    (arbitrary) initial value, so `warmup` ramps the effective decay up from
    0 following $\\min(d, (1 + n) / (10 + n))$ at step $n$.  This
    is the schedule used by most diffusion implementations and it removes the
    need to tune a separate "start EMA at step N" threshold.

    Buffers (e.g. BatchNorm running statistics) are copied rather than
    averaged, matching the reference implementations.

    Args:
        model: Model whose parameters are tracked.
        decay: Target decay rate.  Typical values are 0.999 (short runs) to
            0.9999 (long runs).  Higher means smoother and slower to adapt.
        warmup: Ramp the effective decay early in training (default: True).
        device: Optional device for the shadow copy.  Defaults to keeping
            each shadow tensor on the same device as its source parameter.

    Example:
        >>> model = nn.Linear(2, 2)
        >>> ema = EMA(model, decay=0.9, warmup=False)
        >>> ema.update()
        >>> with ema.average_parameters():
        ...     _ = model(torch.zeros(1, 2))   # uses averaged weights
    """

    def __init__(
        self,
        model: nn.Module,
        decay: float = 0.999,
        warmup: bool = True,
        device: Optional[torch.device] = None,
    ):
        if not 0.0 <= decay <= 1.0:
            raise ValueError(f"decay must be in [0, 1], got {decay}")

        self.model = model
        self.decay = decay
        self.warmup = warmup
        self.device = device
        self.num_updates = 0

        self.shadow: Dict[str, torch.Tensor] = {}
        for name, param in model.named_parameters():
            if param.requires_grad:
                shadow = param.detach().clone()
                self.shadow[name] = shadow.to(device) if device else shadow

        self._backup: Optional[Dict[str, torch.Tensor]] = None

    @property
    def current_decay(self) -> float:
        """Effective decay at the current step, accounting for warmup."""
        if not self.warmup:
            return self.decay
        return min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))

    @torch.no_grad()
    def update(self, model: Optional[nn.Module] = None) -> None:
        """
        Update the shadow weights.  Call once after every optimizer step.

        Args:
            model: Optionally update from a different model instance;
                defaults to the model this EMA was constructed with.
        """
        model = model if model is not None else self.model
        decay = self.current_decay
        self.num_updates += 1

        for name, param in model.named_parameters():
            if not param.requires_grad:
                continue
            if name not in self.shadow:
                # A parameter that became trainable after construction.
                self.shadow[name] = param.detach().clone()
                continue
            shadow = self.shadow[name]
            shadow.mul_(decay).add_(param.detach().to(shadow.device), alpha=1.0 - decay)

    @torch.no_grad()
    def copy_to(self, model: Optional[nn.Module] = None) -> None:
        """Write the averaged weights into `model`, permanently."""
        model = model if model is not None else self.model
        for name, param in model.named_parameters():
            if name in self.shadow:
                param.copy_(self.shadow[name].to(param.device))

    @torch.no_grad()
    def store(self, model: Optional[nn.Module] = None) -> None:
        """Stash the model's live weights so they can be restored later."""
        model = model if model is not None else self.model
        self._backup = {
            name: param.detach().clone()
            for name, param in model.named_parameters()
            if name in self.shadow
        }

    @torch.no_grad()
    def restore(self, model: Optional[nn.Module] = None) -> None:
        """Put back the weights saved by `store()`."""
        if self._backup is None:
            raise RuntimeError("restore() called without a matching store()")
        model = model if model is not None else self.model
        for name, param in model.named_parameters():
            if name in self._backup:
                param.copy_(self._backup[name])
        self._backup = None

    @contextmanager
    def average_parameters(
        self, model: Optional[nn.Module] = None
    ) -> Iterator[nn.Module]:
        """
        Temporarily swap the averaged weights into the model.

        The live training weights are restored on exit, including when the
        body raises, so a failed validation pass cannot corrupt training.
        """
        model = model if model is not None else self.model
        self.store(model)
        self.copy_to(model)
        try:
            yield model
        finally:
            self.restore(model)

    def to(self, device: torch.device) -> "EMA":
        """Move the shadow weights to `device`."""
        self.device = device
        for name, tensor in self.shadow.items():
            self.shadow[name] = tensor.to(device)
        return self

    def state_dict(self) -> Dict[str, Any]:
        """Serializable state for checkpointing."""
        return {
            "decay": self.decay,
            "warmup": self.warmup,
            "num_updates": self.num_updates,
            "shadow": {k: v.detach().cpu().clone() for k, v in self.shadow.items()},
        }

    def load_state_dict(self, state: Dict[str, Any]) -> "EMA":
        """Restore state saved by `state_dict()`."""
        self.decay = state["decay"]
        self.warmup = state.get("warmup", True)
        self.num_updates = state["num_updates"]
        self.shadow = {k: v.clone() for k, v in state["shadow"].items()}
        if self.device is not None:
            self.to(self.device)
        return self

    def __repr__(self) -> str:
        return (
            f"EMA(decay={self.decay}, warmup={self.warmup}, "
            f"num_updates={self.num_updates}, tracked={len(self.shadow)})"
        )
current_decay property

Effective decay at the current step, accounting for warmup.

update(model=None)

Update the shadow weights. Call once after every optimizer step.

Parameters:

Name Type Description Default
model Optional[Module]

Optionally update from a different model instance; defaults to the model this EMA was constructed with.

 None
Source code in flowpde/trainers/ema.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
@torch.no_grad()
def update(self, model: Optional[nn.Module] = None) -> None:
    """
    Update the shadow weights.  Call once after every optimizer step.

    Args:
        model: Optionally update from a different model instance;
            defaults to the model this EMA was constructed with.
    """
    model = model if model is not None else self.model
    decay = self.current_decay
    self.num_updates += 1

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if name not in self.shadow:
            # A parameter that became trainable after construction.
            self.shadow[name] = param.detach().clone()
            continue
        shadow = self.shadow[name]
        shadow.mul_(decay).add_(param.detach().to(shadow.device), alpha=1.0 - decay)
copy_to(model=None)

Write the averaged weights into model, permanently.

Source code in flowpde/trainers/ema.py 
127
128
129
130
131
132
133
@torch.no_grad()
def copy_to(self, model: Optional[nn.Module] = None) -> None:
    """Write the averaged weights into `model`, permanently."""
    model = model if model is not None else self.model
    for name, param in model.named_parameters():
        if name in self.shadow:
            param.copy_(self.shadow[name].to(param.device))
store(model=None)

Stash the model's live weights so they can be restored later.

Source code in flowpde/trainers/ema.py 
135
136
137
138
139
140
141
142
143
@torch.no_grad()
def store(self, model: Optional[nn.Module] = None) -> None:
    """Stash the model's live weights so they can be restored later."""
    model = model if model is not None else self.model
    self._backup = {
        name: param.detach().clone()
        for name, param in model.named_parameters()
        if name in self.shadow
    }
restore(model=None)

Put back the weights saved by store().

Source code in flowpde/trainers/ema.py 
145
146
147
148
149
150
151
152
153
154
@torch.no_grad()
def restore(self, model: Optional[nn.Module] = None) -> None:
    """Put back the weights saved by `store()`."""
    if self._backup is None:
        raise RuntimeError("restore() called without a matching store()")
    model = model if model is not None else self.model
    for name, param in model.named_parameters():
        if name in self._backup:
            param.copy_(self._backup[name])
    self._backup = None
average_parameters(model=None)

Temporarily swap the averaged weights into the model.

The live training weights are restored on exit, including when the body raises, so a failed validation pass cannot corrupt training.

Source code in flowpde/trainers/ema.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
@contextmanager
def average_parameters(
    self, model: Optional[nn.Module] = None
) -> Iterator[nn.Module]:
    """
    Temporarily swap the averaged weights into the model.

    The live training weights are restored on exit, including when the
    body raises, so a failed validation pass cannot corrupt training.
    """
    model = model if model is not None else self.model
    self.store(model)
    self.copy_to(model)
    try:
        yield model
    finally:
        self.restore(model)
to(device)

Move the shadow weights to device.

Source code in flowpde/trainers/ema.py 
174
175
176
177
178
179
def to(self, device: torch.device) -> "EMA":
    """Move the shadow weights to `device`."""
    self.device = device
    for name, tensor in self.shadow.items():
        self.shadow[name] = tensor.to(device)
    return self
state_dict()

Serializable state for checkpointing.

Source code in flowpde/trainers/ema.py 
181
182
183
184
185
186
187
188
def state_dict(self) -> Dict[str, Any]:
    """Serializable state for checkpointing."""
    return {
        "decay": self.decay,
        "warmup": self.warmup,
        "num_updates": self.num_updates,
        "shadow": {k: v.detach().cpu().clone() for k, v in self.shadow.items()},
    }
load_state_dict(state)

Restore state saved by state_dict().

Source code in flowpde/trainers/ema.py 
190
191
192
193
194
195
196
197
198
def load_state_dict(self, state: Dict[str, Any]) -> "EMA":
    """Restore state saved by `state_dict()`."""
    self.decay = state["decay"]
    self.warmup = state.get("warmup", True)
    self.num_updates = state["num_updates"]
    self.shadow = {k: v.clone() for k, v in state["shadow"].items()}
    if self.device is not None:
        self.to(self.device)
    return self