Momentum Encoders and Stop-Gradients

A symmetric contrastive loss compares two branches that are trained by the same gradient. Both branches move toward each other, and in the limit where the objective is minimized over both, the cheapest solution is for both to output the same thing for every input. Self-distillation methods avoid that solution by breaking the symmetry: one branch receives gradients, the other does not, and the second branch’s parameters follow the first by an exponential moving average rather than by optimization.

The mechanism is small, a few lines of parameter arithmetic, but it changes the dynamics enough that several methods depend on it entirely. Understanding why requires separating what the target branch provides from what it costs.

The trivial solution and why symmetry permits it

Write the loss for a pair of views x1 and x2 with encoders f_θ and f_ξ. If θ and ξ are the same parameters and both receive gradients, the loss is minimized when f(x1) = f(x2) for every pair. That is satisfied trivially when f maps every input to a constant, because then every pair matches. The constant solution is a global optimum of the pairwise objective, and it is reachable by gradient descent: as the encoder loses discriminative structure, the positive pairs become easier to match, and the gradient on the negatives weakens.

Normalization and a uniformity term can exclude the constant solution, but they add terms and hyperparameters. The alternative is to make the two branches different objects. If ξ is not updated by the gradient, then the loss cannot be reduced by making ξ output constants; it can only be reduced by moving θ toward whatever ξ produces. The degeneracy is removed structurally rather than by a penalty.

The exponential moving average

The target branch’s parameters are updated after each step as

ξ ← m·ξ + (1 - m)·θ

with m close to one, typically in the range 0.99 to 0.999 and often ramped from a lower value toward a higher one over training. This is a low-pass filter on the parameter trajectory. It has no gradient, no optimizer state, and no memory beyond a second copy of the weights.

@torch.no_grad()
def update_target(online, target, m):
    for p_o, p_t in zip(online.parameters(), target.parameters()):
        p_t.mul_(m).add_(p_o.detach(), alpha=1 - m)
    # buffers (e.g. normalization statistics) are copied, not averaged
    for b_o, b_t in zip(online.buffers(), target.buffers()):
        b_t.copy_(b_o)

Two details matter in that loop. The .detach() prevents the update from participating in the graph, so no gradient flows into the target even indirectly. And buffers such as batch-normalization running statistics are copied rather than averaged, because averaging two running means is not the same as maintaining one.

The averaging coefficient has a time-constant interpretation. With a constant m, the target lags the online network by roughly 1/(1-m) steps: at m = 0.99 the effective window is about a hundred updates, and at m = 0.999 about a thousand. During that window the target changes slowly, so the regression problem the online branch solves is approximately stationary. A target that moves as fast as the online network provides a moving goal, which reintroduces the symmetric collapse through the back door.

The cost is the second copy of the model. For a large encoder this doubles parameter memory and doubles the forward-pass compute per step, since both branches process their views. The target forward pass runs under no_grad, so it stores no activations and needs no backward, but the forward arithmetic is real. In methods that maintain a queue of target representations, the target branch also produces the entries.

What the slow target actually provides

Three distinct benefits are often conflated.

The first is stability of the regression target. The online branch is trained to predict the target’s output for the same input under a different view. If the target is a slowly changing function, the prediction task has a consistent answer for many steps, and the online branch can make progress on it. This is the same argument as target networks in value-based reinforcement learning, where a rapidly moving bootstrap target destabilizes the regression.

The second is a better-behaved target for the hardest pairs. Because the target averages over recent parameter states, its output for a given input is smoother than the online output. Smoothing reduces the variance of the target without biasing it toward the online branch’s current errors, which is what an ensemble of recent models would do.

The third is implicit clustering. In methods that combine the target’s outputs with a centring or sharpening operation, the slow target accumulates structure over many steps, and the online branch is pushed to match a distribution that has had time to become non-uniform. The interaction between the two operations is delicate and is covered in Collapse and Redundancy Reduction in Self-Supervised Training.

Stop-gradient without momentum

Momentum is not the only way to break symmetry. A stop-gradient on one branch while both branches share parameters also works, and it is cheaper: one model, one forward pass per view, no second copy. The target for a pair is then the representation of one view computed under the current parameters but treated as a constant.

The difference from momentum is what the target is. With a stop-gradient and shared weights, the target is the current model’s output, so the regression target moves at full speed with the online branch. The collapse argument is weaker but not eliminated: the stop-gradient makes the constant solution a fixed point of the dynamics rather than a strict optimum, and whether training escapes it depends on the initialization and the loss terms. Methods that use this pattern typically add a variance or decorrelation term to prevent the fixed point from being attractive.

A stop-gradient also appears in the negative side of some losses, where it prevents the gradient from pushing negatives apart symmetrically. The effect is that the anchor moves away from the negative rather than both moving, which changes the effective learning rate per item and avoids oscillation when a negative is shared by many anchors in the batch.

Momentum schedules and the warm-up question

A constant high momentum at the start is a problem: the target initializes to a copy of the online network, and with m = 0.999 it takes a thousand steps before the target reflects any learning. During that window the target is effectively the initial random network, and the online branch is trained to reproduce random outputs, which wastes early training and can push the online branch toward a degenerate match. Raising the momentum over training, starting from a lower value, gives the target time to track early rapid change and then locks it down once the representation stabilizes.

The schedule interacts with the learning-rate schedule. If the learning rate decays, the online parameters move less per step, so a fixed momentum produces a target that lags by the same number of steps but a smaller parameter distance. Some recipes raise momentum as the learning rate falls to keep the effective lag in parameter space roughly constant.

Where the mechanism fails

The target branch drifts when the online branch changes regime. A change in augmentation strength, a restart of the learning-rate schedule, or a switch in the loss can move the online parameters faster than the target follows, and the regression target becomes a stale function that the online branch has already surpassed. The symptom is a temporary rise in loss followed by recovery, or a permanent loss of the structure the target had accumulated.

Buffer handling is a quiet source of error. Copying batch-normalization statistics from the online to the target at every step makes the target’s normalization track the online branch exactly, which is usually correct. Averaging them instead introduces a mismatch between the statistics used in the forward pass and the statistics used at evaluation. Either convention can be defended, but mixing them within one implementation produces a target whose outputs differ from the online outputs even when the parameters agree.

Evaluation must use one branch. The target and online branches differ slightly at all times, and which one is used for downstream evaluation changes the measured quality. Most methods evaluate the target, on the argument that its averaged parameters are less noisy; the difference is small but consistent enough to matter when comparing methods across papers.

Finally, the mechanism does not by itself define what the representation should contain. It only prevents the loss from being minimized by a constant. The content of the space is still determined by the augmentations, as Data Augmentation as an Invariance Specification describes, and by which pairs are treated as negatives, which is the subject of False Negatives and Batch Composition.