MoCo

Decoupling the dictionary from the batch

"Contrastive learning is a way of building a discrete dictionary on high-dimensional continuous inputs such as images. The dictionary is dynamic in the sense that the keys are randomly sampled, and that the key encoder evolves during training. Our hypothesis is that good features can be learned by a large dictionary that covers a rich set of negative samples, while the encoder for the dictionary keys is kept as consistent as possible despite its evolution."

He et al. (2020) propose to view contrastive learning as dictionary lookup. An encoded query is compared against a set of encoded keys . One key, , is the positive (it matches the query); the remaining keys are negatives. Training succeeds when the query is most similar to its positive key.

From the InfoNCE section, we know that more negatives raise the ceiling on how much mutual information the loss can capture. SimCLR addressed this by using very large batch sizes (up to 8192), but this requires significant GPU memory. MoCo takes a different approach: decouple the dictionary size from the batch size by maintaining a queue of encoded keys from recent mini-batches.

This raises a consistency problem. If the key encoder changes rapidly during training, the oldest keys in the queue were produced by a very different encoder than the newest ones. The dictionary would contain representations that are not comparable to each other. MoCo solves this with a momentum-updated encoder that evolves slowly, keeping the dictionary keys consistent even as they span many training steps.

Contrastive loss as dictionary lookup

Given a query and a dictionary of keys , where denotes the single positive key, MoCo uses the InfoNCE loss:

where is the temperature. The sum in the denominator runs over the positive and all negatives, so this is a -way softmax classifier that tries to identify .

The query and keys are produced by encoder networks. Given an input image , two random augmentations yield and . A query encoder produces , and a key encoder produces . The negative keys come from other images encoded by in recent mini-batches, stored in a queue.

For this dictionary lookup to produce good representations, the authors argue two properties are needed:

• Large. A bigger dictionary samples the underlying visual space more densely, providing harder negatives that force the encoder to learn finer distinctions.
• Consistent. The keys should be produced by the same or a similar encoder so that the query can meaningfully compare against them.

The MoCo mechanism

Dictionary as a queue

The dictionary is maintained as a first-in, first-out (FIFO) queue of size . At each training step, the current mini-batch of encoded keys is enqueued and the oldest mini-batch is dequeued.

This decouples the dictionary size from the batch size: the dictionary can be much larger than a typical mini-batch, allowing keys even with a batch size of just 256.

The oldest keys are removed because they were encoded by the most outdated version of the key encoder and are therefore the least consistent with the current keys.

Momentum update

Let denote the parameters of the query encoder and the parameters of the key encoder.

• Only is updated by back-propagation.

The key encoder parameters are updated as an exponential moving average of the query encoder:

where is the momentum coefficient. With (the value used in the paper), only 0.1% of the query encoder's weights are blended into the key encoder at each step. This makes evolve very slowly, so keys encoded many steps apart remain comparable.

The authors find that momentum is critical. Setting (a faster update) causes a significant accuracy drop, because the key encoder changes too rapidly and the queue keys become inconsistent. At the other extreme, means the key encoder never updates at all, which also fails because drifts away from the improving .

Putting it together

The full training step is compact. For each mini-batch:

1. Augment each image twice to obtain a query view and a key view.
2. Encode the query view with to get ; encode the key view with to get (with no gradient through ).
3. Compute the InfoNCE loss using , the positive , and the negatives from the queue.
4. Back-propagate through only and update .
5. Update the key encoder via momentum: .
6. Enqueue the current batch of keys; dequeue the oldest batch.

The gradient never flows through the key encoder or the queue. This is what makes the design efficient: the dictionary can hold tens of thousands of keys without any additional memory cost for storing their computation graphs.