Source
Chen et al. (2020), LIT-591 — ARXIV-2002.05709, §4.2.
The claim
Put a small nonlinear MLP between the encoder and the contrastive loss, train
through it, and then throw it away. The representation you ship is the
encoder output h, not the projected z = g(h) the loss was actually
computed on.
Both halves are measured on the ImageNet linear probe:
- The head helps. Nonlinear beats linear by about 3 points, and beats no head at all by more than 10.
- The head’s own output is worse than its input, by more than 10 points — and this holds with the nonlinear head in place, so it is not an argument against having one.
Both measurements, and the recommendation to ship h rather than z, come
from SimCLR (LIT-591), whose framework introduced the nonlinear projection
head between encoder and NT-Xent loss and then ablated it in §4.2.
Why: the head is where the invariance gets absorbed
The loss demands invariance to the augmentations. Anything the two views differ by must be discarded somewhere, and a head gives the network a place to discard it that is not the representation.
The paper tests this rather than asserting it: train an MLP to recover the applied transformation from each layer.
| recover from | h | g(h) | chance |
|---|---|---|---|
| rotation | 67.6 | 25.6 | 25 |
| original vs corrupted | 99.5 | 59.6 | 50 |
| original vs Sobel | 96.6 | 56.3 | 50 |
| colour vs greyscale | 99.3 | 97.4 | 80 |
g(h) has genuinely lost the information; h still has it. The head is a
sacrificial layer, and its job is to be the thing that gets damaged.
Why it generalises past contrastive learning
The shape of the argument is: when a loss demands a property you do not want in your representation, give the model somewhere else to put it. That applies wherever the training objective is a proxy — an auxiliary head, a task-specific decoder, an alignment layer. Whenever a pipeline ships the tensor the loss was computed on, this is the question to ask.
Conditions
- Nonlinear, and small. Linear is measurably worse; output dimension barely matters once a head is present.
- It reduces the cost of a strong augmentation policy, it does not remove it (SOTA-361). Information destroyed before the encoder cannot be recovered after it.
- Measured under linear evaluation on ImageNet. Whether
hremains the better choice under full fine-tuning is not what this ablation asked. - Which layer, exactly, is an empirical question in a deeper stack. The paper compares two adjacent points; a pipeline with several projection-like stages has more than two candidates and should measure rather than assume.