Contrastive Learning with Stronger Augmentations

In plain words

Contrastive self-supervised learning teaches a network by showing it two altered copies of the same photo and asking it to recognise them as one picture among many others. Those copies have to stay recognisable, so people have always used gentle edits such as crops and colour shifts. Strong edits (heavy distortions, shapes pushed out of place) would teach more, but the distorted copy no longer looks like the original, so the network can't match the two.

This paper does not force the distorted copy to match its clean twin. It asks a softer question: when the clean twin is compared against a large bank of other images, which ones does it look most like? The distorted copy is then trained to produce the same pattern of similarities. It can differ in appearance from the original and still agree on which other images are close to it.

The idea can be bolted onto existing contrastive methods as an extra loss term. The authors use it with a MoCo-style setup and report that it helps, mostly on ImageNet classification. It costs extra compute: the strong views go through the network too.

TL;DR

CLSA keeps a standard contrastive loss on gentle views and adds a divergence loss. The divergence loss makes a strongly augmented view's distribution over a bank of stored keys match the distribution of its weakly augmented twin, which acts as a teacher with gradients stopped. With ResNet-50 and 200 epochs, the single-crop variant reaches 69.4% ImageNet linear top-1 against MoCo v2 at 67.5%, all author-reported. Multi-crop CLSA* reaches 76.2% at 800 epochs, against 76.5% for supervised training.

Problem

Stronger augmentations expose patterns that gentle ones miss, and that can help a supervised or semi-supervised model §1. In contrastive learning, the loss treats the positive key as the single correct answer §3.1. A strongly distorted view is then pushed towards a key that no longer looks like it, and the authors report that naive use of strong augmentations does not help: a strong query alone gives 67.7 against 67.5, and strong query plus strong key gives 67.0 §1 tab. 5. They also show that, after MoCo v2 training, a strongly augmented view and its weakly augmented twin give visibly different similarity distributions fig. 3 §3.1.

Method

The network is a query encoder and a momentum key encoder, both ResNet-50 with a projection head §3.3. Each training image yields one key view, one or more weak query views, and one or more strong query views §3.2 §4.2. A memory queue keeps keys from earlier batches as negatives §3.1.

1. Contrastive term. A weak query is trained with the usual contrastive loss over its own key (positive) and the K queued keys (negatives) §3.1 eq. 1 eq. 2.
2. Relative-similarity distribution. For any query, the paper turns its similarities to the positive and to the K negatives into a softmax distribution at temperature τ eq. 3 §3.2.
3. Distributional divergence. The strong query gets its own distribution over the same candidates eq. 5. The loss L_D is a cross-entropy in which the weak query's distribution is the target, with stop-gradient on the target eq. 6 §3.2.
4. Combined loss. L = L_C + β · L_D, with β = 1.0 eq. 7 §3.2 §3.3.
5. Key update. The key encoder is an exponential moving average of the query encoder with momentum 0.999, and the new keys enter the queue, which drops the oldest entries §3.3.

The strong view S is built from 14 operations: shear and translate on each axis, rotate, auto-contrast, invert, equalise, solarise, posterise, contrast, colour, brightness and sharpness §3.3 tab. 1. Each of five rounds draws one operation at random and applies it with probability 0.5 §3.3. The weak view T is a standard crop to 224 px with colour jitter, grayscale, blur and flip §3.3.

Single-crop CLSA uses one weak query at 224 px, a 224 px key, and one strong view at 96 px §4.2. CLSA* (multi-crop) uses five strong views at 224, 192, 160, 128 and 96 px, and the DDM loss is the sum over them §4.2.

Pseudocode, one batch (the paper has no algorithm box; this follows eq. 1, eq. 3, eq. 5, eq. 6 and eq. 7):

# CLSA loss for one batch  [§3.2]
# Inputs:  key_view (N, C, H, W); weak_views, strong_views: lists of crops
# Output:  scalar loss L = L_C + beta * L_D  [eq. 7]
# queue (K, d): keys from earlier batches, FIFO  [§3.1]
# row_dot(a, b): per-row dot product, (N, d) x (N, d) -> (N, 1)
def clsa_loss(key_view, weak_views, strong_views, queue, tau, beta):
    # Key z from the momentum encoder, no gradient  [§3.1] [§3.3]
    keys = stop_gradient(normalize(key_encoder(key_view)))  # (N, d)
    loss_contrastive = 0
    loss_distill = 0
    for weak_view in weak_views:
        # Weak query z': positive is column 0, then K negatives  [eq. 1]
        weak = normalize(query_encoder(weak_view))  # (N, d)
        pos_weak = row_dot(weak, keys) / tau  # (N, 1)
        neg_weak = weak @ queue.T / tau  # (N, K)
        weak_logits = concat(pos_weak, neg_weak, axis=1)  # (N, 1 + K)
        # Contrastive term L_C: cross-entropy with the positive label  [eq. 1]
        loss_contrastive += mean(-log_softmax(weak_logits)[:, 0])
        # Teacher distribution from the weak query, no gradient  [eq. 3]
        teacher = stop_gradient(softmax(weak_logits))  # (N, 1 + K)
        for strong_view in strong_views:
            # Strong query z'': same candidate keys  [eq. 5]
            strong = normalize(query_encoder(strong_view))  # (N, d)
            pos_strong = row_dot(strong, keys) / tau  # (N, 1)
            neg_strong = strong @ queue.T / tau  # (N, K)
            strong_logits = concat(pos_strong, neg_strong, axis=1)
            # Divergence L_D: teacher as target for strong distribution  [eq. 6]
            loss_distill += mean(-sum(teacher * log_softmax(strong_logits), 1))
    # Combine the two terms  [eq. 7]
    # After the step: momentum-update key encoder, enqueue keys  [§3.3]
    return loss_contrastive + beta * loss_distill

Key equations, as the paper states them:

\mathcal L_{C} = \mathbb{E}_{i \in B}\left[-\log \frac{Q(i,i^{+})}{Q(i,i^{+}) + \sum_{k=1}^{K} Q(i,k)}\right] eq. 1

\mathcal L_{D} = \mathbb{E}_{i \in B}\left[-p(z_i|z_i')\log p(z_i|z_i'') - \sum_{k=1}^{K} p(z_k|z_i')\log p(z_k|z_i'')\right] eq. 6

\mathcal L = \mathcal L_{C} + \beta \, \mathcal L_{D} eq. 7

Hyperparameters

Results

Author-reported. Setup: ImageNet (ILSVRC) pretraining with ResNet-50, then a frozen-feature linear classifier, top-1 accuracy §4.2. Downstream: VOC07 linear classification, VOC07+12 and COCO detection with detectron2 §4.3. Baselines are the methods listed in each table; the hardware is 8 V100 GPUs §4.1.

ImageNet, 200 epochs tab. 2:

ImageNet, 800 epochs tab. 3:

Transfer tab. 4:

Ablation, 200 epochs tab. 5: MoCo v2 67.5; MoCo v2 with a strong query 67.7; MoCo v2 with strong query and strong key 67.0; CLSA with a contrastive loss in place of L_D 68.0; CLSA 69.4 §4.4.

Wall-clock tab. 6: MoCo v2 53 h, 200 epochs, 67.5; CLSA 52.5 h, 150 epochs, 68.3; CLSA 70 h, 200 epochs, 69.4 §4.4.

Augmentation strength s, k-NN, 100 epochs tab. 7: s = 0 gives 51.9; s = 3 gives 52.8; s = 5 (default) gives 53.0; s = 7 gives 52.6 §4.4.

Other methods tab. 8: MoCo v2 (reproduced) 67.5 → 69.4 with CLSA (+1.9); BYOL (reproduced) 66.5 → 67.3 with CLSA (+0.8), both 200 epochs §4.4.

Strongly augmented views, k-NN fig. 5 §4.4: the authors report +8.2 points over MoCo v2 at 200 epochs and +9.9 at 800 epochs on strong views.

Paper vs code

Code: official repository maple-research-lab/CLSA at code_commit. The LICENSE file is MIT and the copyright holder is the MAPLE lab [code: LICENSE]. The README calls the license "MIT for academic use" and shows a GNU badge README.md. We cite the code by path and line only; no code is quoted.

• Contrastive denominator. The paper uses one positive and K queue negatives eq. 1 eq. 3. The code builds the positive logits against every key in the global batch, so the N in-batch keys also serve as negatives, next to the K queue keys. Its comment says "Nx1", but the shape is N×N [code: model/CLSA.py#L319-L333]. The same holds for the strong branch [code: model/CLSA.py#L346-L354].
• Multi-crop pairs. The paper gives one weak and one strong view per equation eq. 6, and says the DDM loss is summed over the strong augmentations §4.2. The code runs every weak query against every strong query, which gives 25 DDM terms in the five-crop run, and sums five contrastive terms [code: model/CLSA.py#L315-L358] run_multi.sh [code: training/train.py#L53-L58].
• Operation pool. The paper lists Invert and not Identity §3.3 tab. 1. The code uses the FixMatch pool with Identity and no Invert [code: data_processing/RandAugment.py#L148-L164].
• Magnitudes. The paper gives ranges tab. 1. The code draws the magnitude index uniformly from 1 to 9, never 10, and scales it by each op's maximum RandAugment.py. Contrast and related ops therefore range from about 0.14 to 0.86, not 0.05 to 0.95, and rotation reaches 27 degrees, not 30.
• Strong view pipeline. The code applies a random resized crop, colour jitter, grayscale, blur and flip before the five ops [code: data_processing/Multi_FixTransform.py#L47-L58]. The paper describes only the ops §3.3.
• Weak view parameters. Colour jitter (0.4, 0.4, 0.4, 0.1) with probability 0.8, grayscale 0.2, blur with sigma from 0.1 to 2 at probability 0.5, and ImageNet mean and std normalisation are in the code [code: data_processing/Multi_FixTransform.py#L26-L36]. The paper gives none of these values.
• Learning rate. The paper scales lr with batch size §4.1. The code does not main_worker.py. At batch 256 the two agree.
• VOC probe. The paper says batch 256 and no scheduler §4.1. The code defaults to batch 16 and a cosine schedule main.py main.py. The arguments used for the paper's numbers are not stated.
• Symmetric loss. The code has a sym branch that the paper does not describe. The run scripts use sym 0 run_single.sh [code: cmd/run_multi.sh]. The README marks sym 1 as unverified.
• Shuffle BN. Implemented only for DistributedDataParallel; single-GPU training is not supported CLSA.py README.md.

No discrepancies found in the momentum update, the temperature 0.2, the β of 1, the queue size of 65536, the 2048-d hidden MLP, the stop-gradient on the weak teacher, the linear-probe settings on ImageNet, or the single-crop and multi-crop run commands.

Pitfalls

• Denominator choice. The paper's formulas and the code differ on whether in-batch keys are negatives. The paper does not say which form produced its numbers. Pick one and say which in your write-up.
• Index in eq. 3. The denominator of eq. 3 adds the positive term separately and then sums negatives from k = 0 to K. Read literally, the positive may be counted twice, or the sum has K + 1 negatives. The code uses K negatives, as eq. 1 does.
• Loss scale grows with the number of views. The DDM loss and the contrastive terms are summed over views in the code. Each extra strong view adds a term, so with β = 1 the total loss scales with the view count. The paper does not state whether a mean would be used instead.
• Cost. Each strong view is an extra forward and backward pass through the query encoder. The single-crop CLSA takes 70 h for 200 epochs against 53 h for MoCo v2 tab. 6. The paper's comparison at equal time (150 epochs of CLSA against 200 of MoCo v2) is only a partial fix §4.4. Multi-crop takes about 166 h §4.2.
• Single-crop gap. At 200 epochs, single-crop CLSA (69.4) sits below InfoMin Aug (70.1) in the same table tab. 2. The headline claim is for the multi-crop variant.
• Stop-gradient. The gradient reaches the strong branch only. The weak teacher distribution is detached eq. 6 §3.2 CLSA.py.
• Frozen probes. Results rely on linear probes on frozen features §4.2. Fine-tuning results are not reported in the paper.
• Hardware and code. The code needs multiple GPUs with DistributedDataParallel and CUDA-only tensors CLSA.py README.md. Expect to port it before running it on a single machine.
• Licence ambiguity. The LICENSE file is MIT, but the README says "MIT for academic use" and shows a GNU badge. Check the licence before any code reuse.

Lineage

• Builds on: MoCo v2 (arXiv:2003.04297) with the same momentum encoder, queue and projection head §3.1 §3.3 tab. 2; SimCLR (arXiv:2002.05709) for augmentation design §2 §3.1.
• Strong augmentation ancestry: RandAugment and FixMatch-style pools §2 §3.3; FixMatch (arXiv:2001.07685) as the pool the code uses [code: data_processing/RandAugment.py#L148-L164].
• Compared against: MoCo v2, SimCLR, BYOL (arXiv:2006.07733), SWAV (arXiv:2006.09882), InfoMin Aug, PIRL and others tab. 2 tab. 3 tab. 4.
• Combined with: MoCo v2 and BYOL, as generalisation tests tab. 8.