An Empirical Study of Training Self-Supervised Vision Transformers

In plain words

Vision Transformers (ViTs) are now a common image model, but the training recipes for ViTs without labels are less settled than those for convolutional networks. This paper takes a contrastive learning method called MoCo v3 and applies it to ViTs. Each training picture is shown twice, as two random crops with colour and blur changes, and the model learns to give the two versions matching descriptions while keeping different pictures apart. A slowly changing copy of the model supplies the targets.

The paper is a study rather than a new method. Its main finding is that training ViTs this way is unstable. A run can look fine while quietly losing some accuracy, so the numbers can hide a partial failure. A simple change, keeping the first patch layer fixed at its random starting values, makes training steadier and more accurate in the settings the authors tried. The authors end with evidence in favour of ViTs for this kind of training, along with open questions.

TL;DR

• Empirical study, no new architecture: a ViT backbone trained with MoCo v3, a contrastive Siamese self-supervised method, on ImageNet-1k §1 §3 §5.
• Training instability hides partial failures. Freezing the random patch projection layer improves stability and accuracy in the tested cases §4.2.
• Author-reported: ViT-B/16 gets 76.7% linear-probing top-1 on ImageNet, and ViT-BN-L/7 gets 81.0% tab. 1 §6.3. The BN and /7 variants are not in the released code (see Paper vs code).
• Use the Hyperparameters table for the settings of the released ViT commands, which differ from some paper text (see Paper vs code).

Problem

Standard ViT training recipes do not carry over automatically to self-supervised contrastive training. The authors report that training with a ViT backbone is often unstable. Instability may not diverge. Instead it can cause mild accuracy loss, such as 1 to 3% §1. Sudden spikes in gradients produce dips in training curves, and a dense kNN monitor is needed to see them §4.1 app. A. Results are reported under ImageNet-1k self-supervised pretraining with linear probing evaluation §5 §6.

Method

MoCo v3 has two encoders with the same backbone and head structure. The query encoder f_q has a projection head and an extra prediction head. The key encoder f_k, the momentum encoder, has only the projection head. f_k is a moving average of f_q §3. Two random views of each image are encoded. Each query is matched against the keys of all other images in the batch, and the target is the key of the same image. The loss is symmetrised over the two views §3 alg. 1. The memory queue of MoCo v1/v2 is dropped, because the authors found its gain diminishing with a large batch §3.

The InfoNCE loss for one query is given by eq. 1:

\mathcal{L}_{q} = -\log \frac{\exp(q \cdot k^{+} / \tau)}{\exp(q \cdot k^{+} / \tau) + \sum_{k^{-}} \exp(q \cdot k^{-} / \tau)} eq. 1

Here k+ is the key of the same image and the k- are the keys of the other images. q and k are L2-normalised §3. Tau is the temperature §5.

The training step, in our own pseudocode (the paper's algorithm is alg. 1; the released code is cited for the same steps):

# One MoCo v3 training iteration  [alg. 1] [§3]
# Input:  images (N, 3, 224, 224); momentum m; temperature tau
# Output: scalar loss; f_q updated by the optimiser, f_k by momentum

def training_step(images, m, tau):
    # Two random augmented views of each image  [§5] [app. A]
    view1 = augment(images)  # (N, 3, 224, 224)
    view2 = augment(images)  # (N, 3, 224, 224)

    # Queries: backbone, projection head, then prediction head  [§3] [§5]
    queries1 = predictor(f_q(view1))  # (N, 256)
    queries2 = predictor(f_q(view2))  # (N, 256)

    # Keys: momentum encoder without the prediction head, no gradient  [§3]
    keys1 = f_k(view1)  # (N, 256)
    keys2 = f_k(view2)  # (N, 256)

    # Symmetrised contrastive loss  [eq. 1] [alg. 1]
    loss = contrastive(queries1, keys2, tau) + contrastive(queries2, keys1, tau)

    # Gradient step on f_q, then move f_k towards f_q  [alg. 1] [§3]
    optimiser_step(f_q, loss)
    f_k = m * f_k + (1 - m) * f_q
    return loss

The contrastive term over the batch, with the 2 tau scale used in the paper and in the code §5 alg. 1:

# InfoNCE over the in-batch keys, scaled by 2 * tau  [eq. 1] [§5]
# Input:  queries (N, 256), keys (N, 256), temperature tau
# Output: scalar loss

def contrastive(queries, keys, tau):
    # L2-normalise both sets of vectors  [eq. 1]
    queries = normalise(queries)  # (N, 256)
    keys = normalise(keys)  # (N, 256)

    # Score each query against every key in the batch  [alg. 1]
    logits = queries @ keys.T / tau  # (N, N)

    # The positive key of query i is key i, so the labels are the diagonal
    labels = range(N)  # [alg. 1]
    return 2 * tau * cross_entropy(logits, labels)  # [§5] [alg. 1]

The paper does not give the schedules in its text, so the released ViT code is the reference for them main_moco.py CONFIG.md. The learning rate is warmed up linearly per iteration, then decays with a cosine:

# Learning rate per iteration  [§5] [code: main_moco.py#L421-L426]
# Input:  epoch (fractional), warmup and total epochs, base lr, batch size
# Output: learning rate

def learning_rate(epoch, warmup, total, base_lr, batch):
    # Linear scaling: lr x batch / 256  [§4.1] [code: main_moco.py#L190]
    peak = base_lr * batch / 256
    if epoch < warmup:
        return peak * epoch / warmup
    progress = (epoch - warmup) / (total - warmup)
    return peak * 0.5 * (1 + cos(pi * progress))

The released ViT commands also raise the momentum coefficient towards 1 over training. The paper's text does not describe this for ViT (see Paper vs code):

# Momentum rising from m0 to 1 on a cosine  [code: main_moco.py#L431-L434]
# Input:  epoch (fractional), total epochs, starting coefficient m0
# Output: momentum m

def momentum_coefficient(epoch, total, m0):
    cosine = 0.5 * (1 + cos(pi * epoch / total))
    return 1 - (1 - m0) * cosine

Key design choices studied by the paper:

• Keys come from the same batch, so the negatives are the other images in the batch §3. The released code gathers keys from all devices, so the negatives span the global batch builder.py builder.py.
• The prediction head is on the query side only, as in BYOL §3 §6.2. Removing it costs about one point §6.2.
• The momentum coefficient m is 0.99 in the default setting; m = 0 behaves like SimCLR with a prediction head and a stop-gradient on the keys §6.2.
• The patch projection is made a fixed random layer by a stop-gradient right after it §4.2. This is the only change that targets stability directly.
• The ViT backbone follows the original ViT design: 2-D sin-cos position embeddings, a learned class token, and a final LayerNorm whose output is the feature §5. Learned position embeddings are an ablation option §6.2.

Hyperparameters

Results

Author-reported. Setup: ImageNet-1k pretraining (1.28M images) with MoCo v3 and ViT-B/16 unless stated; linear probing reports single-crop top-1 on the validation set §5 §6. Numbers are as printed in the paper.

Table 1 tab. 1 compares self-supervised transformers under linear probing and end-to-end fine-tuning.

Source: tab. 1; the masked patch prediction row is from prior work and is pre-trained on JFT-300M tab. 1.

Frameworks with ViT-S/16 and ViT-B/16 tab. 4, linear probing:

The caption says the R-50 results of the other frameworks come from the improved implementation in prior work tab. 4. Every entry uses two 224 crops, so multi-crop is not used tab. 4.

Ablations with ViT-B/16, linear probing §6.2 unless noted. The paper's tables are inline and not numbered.

Scaling and big ResNets fig. 8, linear probing (author-reported). ViT baseline: ViT-S 73.4, ViT-B 76.7, ViT-L 77.6, ViT-H 78.1. ViT-BN: 74.1, 77.5, 78.6, 79.1. ViT-BN/7: 77.0, 79.5, 81.0, not run. The text reports the previous best linear probing as 79.8 (SimCLR v2) and 79.6 (BYOL) §6.3.

Supervised comparison: at ViT-L, self-supervised linear probing (77.6) beats the supervised counterpart in prior work (76.53) when trained on ImageNet-1k §6.3.

Transfer learning tab. 6, end-to-end fine-tuning, ViT-B / ViT-L / ViT-H:

End-to-end fine-tuning on ImageNet-1k tab. 5, top-1 %:

ViT-B reaches 83.2 with 150 epochs of fine-tuning, and ViT-L reaches 84.1 with 100 epochs and drop path 0.5 §6.3. The fine-tuning code is DeiT's, which is not in this repository §6.3.

Compute tab. 3: per 100 ImageNet epochs on TPU v3 in the authors' TensorFlow code, ViT-S/16 takes 1.2 h on 256 TPUs, ViT-B/16 2.1 h on 256 TPUs, ViT-L/16 6.1 h on 256 TPUs, and ViT-H/14 9.8 h on 512 TPUs. ViT-B on 128 GPUs in PyTorch took about 24 h per 100 epochs §6.

Paper vs code

Compared with the official repository (github.com/facebookresearch/moco-v3) at commit c349e6e. Each difference cites both sources.

• Momentum schedule. The paper's ViT default is a fixed m = 0.99 §6.2. The released ViT commands add a cosine ramp from 0.99 to 1 CONFIG.md main_moco.py main_moco.py. The paper describes the ramp only for its ResNet-50 setting, with m = 0.996 app. A.
• Temperature. The paper uses tau = 0.2 for ViT §5; the released code default is 1.0 main_moco.py and the ViT commands pass 0.2 CONFIG.md.
• Learning-rate values. The paper's ablation curves use base lr 1.0e-4 fig. 1 fig. 5, with batch 4096 in fig. 5, and says lr and wd were searched §5. The released ViT commands use base lr 1.5e-4 and wd 0.1 CONFIG.md. The paper's text does not state these final values.
• Linear-probing batch size. The paper states batch 4096 §5. The released linear command does not set a batch size, so the code default of 1024 applies main_lincls.py CONFIG.md.
• Reported ViT accuracy. The paper gives 76.5 at 300 epochs for ViT-B and 76.7 at 600 epochs §6.2. The released README table lists 76.7 for the ViT-B 300-epoch checkpoint and 73.2 for ViT-S at 300 epochs CONFIG.md. The paper gives 72.5 for ViT-S at 300 epochs §6.2. These are not reconciled.
• Output-layer BN. The projection and prediction heads end in BatchNorm with no affine parameters, and all Linear layers in the heads have no bias. The code comment says this follows SimCLR builder.py. The paper says only that all layers have BN §5.
• Momentum update order. The code updates f_k at the start of the forward pass, before computing keys builder.py. Alg. 1 updates f_k after the optimiser step alg. 1. The two orders give the same keys.
• Patch projection. The paper describes a stop-gradient after the projection §4.2. The code sets the convolution's weight and bias to not require gradients vits.py. This also removes them from weight decay. The paper does not state the initialisation; the code uses a Xavier-style uniform bound vits.py.
• Initialisation. The paper does not describe it. The code initialises the query/key/value weights with a separate bound, the class token with standard deviation 1e-6, and other linear layers with Xavier uniform vits.py.
• Augmentation details. The paper lists operations only app. A. The code supplies strengths and probabilities main_moco.py loader.py loader.py.
• Weight decay scope. The code applies AdamW weight decay to all trainable parameters, including normalisation parameters and the class token, with no exclusion list main_moco.py. The paper is silent.
• ViT variants. The paper reports a BN variant of the ViT MLP blocks and patch size 7 §6.3. The released model definitions fix patch size 16 and use LayerNorm vits.py. The released file also includes a convolutional stem variant that the paper does not describe vits.py vits.py. The ViT-BN-L/7 result cannot be reproduced from this repository as released.
• Implementation. The paper's timings are on TPU in TensorFlow tab. 3. The repository is PyTorch, with SyncBN and fp16 autocast main_moco.py main_moco.py main_moco.py. The paper describes its TensorFlow AdamW weight-decay convention app. A. The repository uses the PyTorch convention.
• Other frameworks. The repository implements MoCo v3 only. It has no SimCLR, BYOL or SwAV ViT code, and no transfer or fine-tuning code repository tree.
• ResNet appendix. The paper gives lr 0.3 and wd 1.5e-6 for LARS app. A. The default lr in code is 0.6 main_moco.py, and the 100-epoch ResNet command does not override it CONFIG.md.
• Verified: the InfoNCE loss, the 2 tau scale, the symmetric sum, the projection and prediction head sizes, and the warm-up and cosine schedule match between the paper and the code alg. 1 §5 builder.py main_moco.py.

Pitfalls

• Base versus effective learning rate. Figures and text use the base lr, and the effective value is base x batch / 256 §4.1 main_moco.py. For the ViT-B command (base 1.5e-4, batch 4096) that gives 2.4e-3. Mixing the two is the most likely way to get a wrong result.
• Instability hides in the numbers. Dips are small, a second run often differs by only 0.1 to 0.3 points, and sparse kNN checks can miss the dips §4.1 app. A. Watch training curves, not only final accuracy.
• The random-projection trick reduces instability but does not remove it. Training can still become unstable if lr is too large §4.2, discussion.
• Batch size and BN. Without BN in the heads the authors needed batch 2048; otherwise training did not converge §6.2. A batch of 6144 was clearly unstable fig. 1.
• BN batching choice. Two layouts are possible: all views in one BN batch, or each view forwarded separately. The paper's Alg. 1 implies the second, and the code does that app. A builder.py. Using the other layout changes gradient variance app. A.
• Negatives across devices. The released code gathers keys from all devices, so each query's negatives span the global batch. This adds memory and communication builder.py builder.py.
• Mixed precision is not mentioned in the paper. The code uses fp16 autocast with a gradient scaler main_moco.py main_moco.py.
• Momentum coefficient. m = 0.99 is best in the table; 0.999 is worse, and m = 0 is 2.2 points lower §6.2. The released ViT commands also ramp m to 1 CONFIG.md. Pick one and record it.
• Position embedding. The learned variant is slightly worse than sin-cos, and no position embedding costs 1.6 points §6.2. The learned option is not in the released ViT models vits.py.
• Class token and LayerNorm. Without a class token, pooling works only if the final LayerNorm is also removed; keeping it costs about 7 points §6.2.
• Scaling saturates. ViT-L and ViT-H gain little over ViT-B under this recipe §6.2. ViT-BN-H/7 ran out of memory in the authors' implementation fig. 8 footnote.
• Hidden compute. The 300-epoch ViT-H run amounts to about 625 TPU-days §6. The authors' PyTorch run of ViT-B took about 24 h on 128 GPUs, against 2.1 h per 100 epochs on 256 TPUs with their TensorFlow code. The two are not like-for-like §6 tab. 3.
• Fine-tuning instability. The authors report that DeiT-L diverges under its default settings, so a different recipe may be needed §6.3.
• Linear probing protocol. The classifier is trained with cosine decay per epoch and no warm-up, and BN stays in eval mode main_lincls.py main_lincls.py. Keep these identical when comparing runs.
• Some numbers are reported only in text or figures. The SwAV, SimCLR and BYOL ViT rows, and the transfer and fine-tuning results, cannot be reproduced from this repository repository tree.

Lineage

• Builds on: MoCo v1 and v2, which MoCo v3 improves on, with the memory queue removed §3. The released code links the MoCo paper as arXiv:1911.05722 builder.py. Uses the Vision Transformer design §1 §5. The predictor head follows BYOL §3.
• Compared against: SimCLR, BYOL and SwAV in ViT form §6.1 tab. 4; iGPT and masked patch prediction tab. 1; DeiT for fine-tuning tab. 5; SimCLR v2 and BYOL with big ResNets §6.3.