Efficient Visual Pretraining with Contrastive Detection

In plain words

Learning good image features without labels usually takes a lot of computing. Methods that compare whole pictures learn one signal per picture, and they need huge numbers of passes over the data to get good. DetCon asks a simpler question: can the network learn from the objects inside each picture instead?

The method cuts each image into rough regions, such as a dog, a sofa or the sky, using a cheap classical segmentation algorithm that needs no labels. The image is then shown twice, each time with random crops, flips and colour changes. The network has to recognise that a region in one version is the same region in the other version, and that it differs from the regions in other images.

Each picture therefore gives many learning signals instead of one. The regions do not have to be accurate, but the better they line up with real objects, the more the network learns. The result is features that transfer well to detection, segmentation and depth tasks while using less pretraining time than comparable methods.

TL;DR

• DetCon is a contrastive pretraining objective. It matches object-level features across two augmented views of an image, using unsupervised segments as the regions.
• Two variants are described: DetCon-S, built on SimCLR, and DetCon-B, built on BYOL. The paper reports both §3.1.
• Author-reported: DetCon-B on ImageNet beats BYOL on COCO, PASCAL, Cityscapes and NYU, with about 3x to 10x less pretraining tab. 1.
• The official code releases only the DetCon-B training loop, the mask generator and the ImageNet pipeline. DetCon-S, the COCO pretraining and all transfer experiments are not in the repository. See "Paper vs code".

Problem

Self-supervised pretraining reaches strong transfer accuracy, but state-of-the-art methods need roughly an order of magnitude more compute than supervised pretraining §1. The paper targets this cost. It also notes that instance-discrimination methods have mostly been shown on simple datasets such as ImageNet, and that complex scenes have been a hard domain for self-supervised methods §2 §1.

The method assumes that cheap, image-computable segmentations are available for every training image. These masks are computed once and reused during training §3.1. It also assumes access to a large batch (4096 images) spread across many accelerators app. 6.3.

Method

Framework. Each image is augmented twice, giving two views. The method reuses the augmentation pipeline and the network design of SimCLR (DetCon-S) or BYOL (DetCon-B). The new parts are the mask pooling and the contrastive detection loss §3.1 app. 6.1 app. 6.2. Images are resized to 224x224 §3.1.

Masks. Every image has a set of masks, one per segment. The masks are transformed with the same crop, resize and flip as their image, so they stay aligned with both views §3.1.

Features. A ResNet-50 gives a 7x7 grid of 2048-dimensional vectors §3.1. Each mask is downsampled to the 7x7 grid and used to average the grid vectors it covers §3.1. A two-layer projection MLP maps each pooled vector to the latent used by the loss §3.1 app. 6.2. In DetCon-B, a prediction MLP is applied on top of the online projection, and the target projection comes from an exponential moving average (EMA) of the online weights §3.1 app. 6.2.

Sampling. Each image has a variable number of masks. At each iteration, 16 masks are drawn per image, with replacement, so the same object can be drawn twice §3.1.

Loss. Latents are scaled to unit norm and divided by a temperature of 0.1 §3.1. For a mask in one view, the positive is the mask with the same region in the other view, and the negatives are all other masks in the batch eq. 1 eq. 3. The paper derives the objective from a paired-mask expectation eq. 2, then uses the batched version eq. 3. The sum runs only over pairs that belong to the same underlying region §3.1.

Mask generators. The default is the Felzenszwalb-Huttenlocher algorithm, with scale s = 1000 and minimum size 1000 on ImageNet §3.2. The paper also tests spatial grids (2x2, 5x5, 10x10), MCG masks and ground-truth COCO masks §3.2 §4.3 §4.4.

Transfer. Only the encoder f is kept for downstream tasks. The projection and prediction heads and the EMA copy are discarded §3.1. Downstream tasks are COCO detection and instance segmentation, PASCAL and Cityscapes segmentation, and NYU depth §3.3.

Equations.

\mathbf{h}_m = \frac{1}{\sum_{i,j} m_{ij}} \sum_{i,j} m_{ij}\, \mathbf{h}[i,j] §3.1

\ell_{m,m'} = -\log \frac{\exp(\mathbf{v}_m \cdot \mathbf{v}'_{m'})}{\exp(\mathbf{v}_m \cdot \mathbf{v}'_{m'}) + \sum_n \exp(\mathbf{v}_m \cdot \mathbf{v}_n)} eq. 1

\mathcal{L} = \mathbb{E}_{(m,m') \sim \mathcal{M}}\, \ell_{m,m'} eq. 2

\mathcal{L} = \sum_m \sum_{m'} \mathbb{1}_{m,m'}\, \ell_{m,m'} eq. 3

Pseudocode. The paper has no algorithm box. This sketch follows §3.1 and eq. 3, with the detail that the released code adds (marked "code"). Shapes are for ImageNet with 224 pixel inputs.

# One DetCon-B training step, reconstructed from [§3.1] and [eq. 3]
# Input:  views x1, x2 (B, 224, 224, 3); segment maps s1, s2 (B, 224, 224)
# Output: scalar loss
#   [code: utils/losses.py#L30-L153] [code: detcon_b.py#L111-L154]

def detcon_step(x1, x2, s1, s2):
    # Encode both views with the online encoder  [§3.1]
    h1 = encoder(x1)                            # (B, 7, 7, 2048)
    h2 = encoder(x2)                            # (B, 7, 7, 2048)

    # Give each 7x7 cell the one segment that covers most of it (code)
    # [code: detcon_b.py#L86-L93]; the paper averages membership instead
    g1 = majority_segment_per_cell(s1)          # (B, 7, 7) ints
    g2 = majority_segment_per_cell(s2)

    # Draw 16 segment ids per image, with replacement, uniformly among
    # segments that own at least one cell  [§3.1] [code: detcon_b.py#L95-L109]
    ids1 = sample_segment_ids(g1, count=16)     # (B, 16)
    ids2 = sample_segment_ids(g2, count=16)

    # Mask-pool each sampled segment, then project and predict  [§3.1]
    z1 = predict(project(mask_pool(h1, g1, ids1)))   # (B, 16, 256)
    z2 = predict(project(mask_pool(h2, g2, ids2)))   # (B, 16, 256)

    # Targets from the EMA network, with no gradient  [§3.1] [app. 6.2]
    # The target pass draws the same ids because it reuses the same rng
    # [code: detcon_b.py#L241-L263] (inferred from code)
    t1 = stop_gradient(target_mask_pool_project(x1, ids1))
    t2 = stop_gradient(target_mask_pool_project(x2, ids2))

    # Cosine similarity divided by tau = 0.1  [§3.1]
    # [code: utils/losses.py#L71-L106]
    # Positives: same image, same segment id, other view.
    # Negatives: every other target vector in the global batch, from
    # both views, except same-image same-segment pairs (code)
    # [code: utils/losses.py#L108-L121]
    loss1 = weighted_nce(z1, t1, t2, ids1, ids2)
    loss2 = weighted_nce(z2, t2, t1, ids2, ids1)

    # Sum the two directions, with no factor of 1/2
    # [code: utils/losses.py#L149-L151]
    return loss1 + loss2

Compute. The projection head runs on 16 vectors per image instead of one, which adds a small cost. The paper gives a 5.3% overhead over SimCLR for DetCon-S and 11.6% over BYOL for DetCon-B, counting the loss §3.1 app. 6.2. Computing the masks costs nothing per step because they are computed once §3.1.

Hyperparameters

Pretraining on ImageNet, DetCon-B. The code column reports the default in detcon_b_config.py, which is 100 epochs and the 1000-epoch values (see "Paper vs code").

Augmentation parameters. The paper gives a table for both variants. Column T is the first view and T' the second. Probabilities are per image.

Source: app. 6.1. Crop area is drawn uniformly from 8% to 100% of the original area, aspect ratio log-uniformly from 3/4 to 4/3, then resized with bicubic interpolation to 224x224 app. 6.1. Gaussian blur uses a 23x23 kernel and sigma uniform in [0.1, 2.0] app. 6.1. Solarization keeps a pixel value x when x is below 0.5 and maps it to 1 - x otherwise app. 6.1.

Transfer fine-tuning.

Input size for transfer: 1024 pixels on the longest side for COCO, with random scaling in [0.8, 1.25] during training app. 6.1. PASCAL uses 513x513 crops and scale jitter in [0.5, 2.0] app. 6.1. Cityscapes uses 769x769 training crops and 1025x2049 test images app. 6.1. NYU uses 304x228 images app. 6.1. COCO pretraining uses a piecewise-constant schedule, dropping the learning rate x0.1 at the 96th and 98th percentiles, instead of cosine app. 6.3.

Results

Author-reported. Setup: ImageNet (1.3M training images) or COCO for pretraining; ResNet-50 unless stated; COCO Mask-RCNN, FCN for PASCAL and Cityscapes, and a depth head for NYU v2 §3.3 tab. 3 app. 6.3. Baselines are SimCLR, BYOL, supervised pretraining and, in one case, SEER tab. 2 tab. 3.

Table 1. DetCon-B against BYOL, ResNet-50, pretrained on ImageNet tab. 1. Each cell gives the two pretraining budgets as "300 / 1000 epochs"; NYU uses "100 / 1000".

The efficiency gains are read from training curves. The paper does not state how they were computed tab. 1.

Table 2. Comparison to prior art, ResNet-50, ImageNet pretraining, COCO fine-tuning tab. 2. Columns are box AP and mask AP after 12 epochs (1x) and 24 epochs (2x). "-" marks a value the paper did not report.

Table 3. Comparison to large-scale transfer tab. 3. Data and architecture differ between rows, so the comparison is not like-for-like. The paper says so. SEER and the supervised baseline use RegNet, while DetCon uses a plain ResNet-200 with 2x width.

Table 4. From SimCLR to DetCon-S, COCO pretraining, frozen-feature COCO accuracy tab. 4. Columns are box AP and mask AP, as the table header names them. "Masks" is the pooling region. "Sampled segments" is how many segments per image enter the loss.

Table 5. Simplifying the contrastive framework, DetCon-S on COCO tab. 5. "All negatives" means negatives are drawn from the whole batch rather than one worker out of 128. "Two views" means positives are matched across augmentations rather than within one view.

Larger encoders, COCO fine-tuning app. 6.4. The paper's table is numbered inside the appendix; its numbering is not checked against the compiled PDF. Columns are box AP and mask AP after 1x and 2x schedules.

Results described in the figures (values not transcribed):

• Figure 1: SimCLR matches supervised pretraining only with about 10x more iterations. DetCon surpasses both with about 5x less computation than SimCLR fig. 1.
• Figure 4: on ImageNet, DetCon-S beats SimCLR at every pretraining budget. The paper reports 5x efficiency on COCO detection, 2x on PASCAL, and about 10x on Cityscapes and NYU fig. 4 §4.1.
• Figure 5: scaling to ResNet-101, -152 and -200 keeps the DetCon-B lead fig. 5.
• Figure 6: pretraining on COCO, DetCon-S beats SimCLR at every schedule from 324 to 5184 epochs. SimCLR reaches its best accuracy with about 4x more pretraining time. DetCon also beats supervised COCO pretraining fig. 6 §4.2.
• Figure 7: frozen-feature COCO accuracy rises with the overlap between masks and ground-truth objects, so better masks give better features fig. 7 §4.3.
• Figure 8: with ground-truth masks, higher input resolution helps DetCon. With FH masks it helps only a little, and SimCLR gets worse fig. 8 §4.4.

Paper vs code

The code is the google-deepmind/detcon repository (Apache-2.0), commit dda05678d9e339f429f5e0fde1624bcbce56fc68. The paper links no DetCon code; its only code link is the TPU detection implementation it uses for COCO §3.3. The repository was found by GitHub search for "detcon" and belongs to google-deepmind, the organisation of the authors' lab. The README calls it "not an officially supported Google product" README.md. File paths below are relative to the repository root.

Scope. The release contains the DetCon-B (BYOL-based) loss and training loop, the ImageNet Felzenszwalb mask generator, and the data pipeline, plus a mock test README.md. It does not contain DetCon-S, COCO pretraining, or any transfer or evaluation code. Every transfer number in the paper rests on code that is not here.

• Paper trains DetCon-S and DetCon-B §3.1; code contains only DetCon-B. DetCon-S projection and augmentation settings come only from the paper app. 6.1 app. 6.2.
• Paper sets the 300-epoch recipe to base learning rate 0.3, weight decay 1e-6 and lambda_0 = 0.99 app. 6.2 app. 6.3. The config defaults to 100 epochs with base learning rate 0.2, weight decay 1.5e-6 and lambda_0 = 0.996 detcon_b_config.py detcon_b_config.py detcon_b_config.py detcon_b_config.py. The config's values are the paper's 1000-epoch values.
• Paper gives no warmup length app. 6.3. Code uses num_epochs // 100 epochs of linear warmup detcon_b_config.py. The optimizer block also has warmup_epochs=10 detcon_b_config.py, which the training code never reads. Only warmup_steps is used schedules.py.
• Paper states Gaussian blur uses a 23x23 kernel app. 6.1. Code uses kernel size equal to image width / 10, which is 22.4 augmentations.py. How the library rounds 22.4 is not checked here.
• Paper averages the downsampled mask, so each cell carries a fractional weight in the pooled mean §3.1. Code assigns each 7x7 cell to a single segment, the one that covers most of it, and then pools with weight 0 or 1 detcon_b.py. Small segments can disappear completely in this scheme.
• Paper applies the prediction head q to the first view only, so the loss is asymmetric §3.1. Code applies the online predictor to both views and sums two symmetric terms detcon_b.py losses.py. The paper's eq. 3 has no normaliser, while code divides each direction by the number of same-segment duplicates and by the number of positives losses.py.
• Paper says masks are sampled "randomly" without giving the distribution §3.1. Code samples segments uniformly among those that own at least one cell detcon_b.py.
• Paper's denominator sums over negative latents v_n eq. 1. Code uses target projections from both views, across the whole global batch (gathered from all devices) losses.py losses.py. It excludes only pairs from the same image and the same segment losses.py.
• Code adds a linear classifier on stop-gradient pooled features and includes its cross-entropy in the reported loss detcon_b.py detcon_b.py pretrain_common.py. The paper does not mention this term.
• Code stores segment ids as uint8 and sizes the one-hot mask to 256 ids generate_fh_masks_for_imagenet.py detcon_b.py. It does not check for images with more than 256 segments.
• Code normalises images with the ImageNet mean and standard deviation after augmentation detcon_b.py image_dataset.py. The paper does not state this.
• Code uses cross-replica batch norm over the device axis i detcon_b_config.py. The paper does not state the BN statistics mode.

No discrepancy was found in these pieces: temperature 0.1 detcon_b_config.py; projection sizes 4096 and 256 for DetCon-B detcon_b_config.py; cosine annealing of the EMA decay schedules.py; crop ranges and bicubic resize tf_image_ops.py tf_image_ops.py; the FH defaults of scale and minimum size 1000 generate_fh_masks_for_imagenet.py; and the DetCon-B augmentation probabilities and magnitudes in app. 6.1 augmentations.py.

Pitfalls

• The masks decide the result. With the same objective, ground-truth masks score 37.0 against 33.6 for FH masks in the full model, on the frozen-feature COCO measure tab. 5. Mask quality correlates with the features §4.3 fig. 7. Budget for the segmenter: the code's README says the ImageNet FH pass takes 2 to 3 days on a 16-core CPU README.md.
• Segments below one cell vanish. Cells are 32 pixels wide, so a segment that never wins a majority of any cell is never sampled. The code does not warn about this detcon_b.py.
• Some numbers are not reproducible from the release. DetCon-S, COCO pretraining and all transfer tasks have no code here README.md. The paper also reports PASCAL, Cityscapes and NYU results with settings that the release does not contain app. 6.3. A reader can reproduce only the DetCon-B ImageNet pretraining.
• The loss has hidden normalisation. Positives are split evenly, duplicates are down-weighted, and both directions are summed with no 1/2 factor losses.py. Changing any of these changes the effective learning rate.
• Duplicate sampling is a design choice. 16 draws with replacement mean that the same segment can appear more than once. The code handles this by weighting, so it is not just a sampling detail §3.1 losses.py.
• Memory grows with the global batch. Each image contributes 16 vectors, so each similarity matrix has size (local batch x 16) by (global batch x 16). With a global batch of 4096 spread over 128 devices, that is 32 x 16 by 4096 x 16, about 33.6M values per matrix, and the code builds four of them app. 6.3 losses.py.
• Learning rate depends on batch size. The base 0.2 becomes 3.2 at batch 4096 through the scaling in the code pretrain_common.py. A smaller batch with the same base rate is much gentler.
• Presets are not defaults. Using the config as is gives the 1000-epoch learning rate and EMA, not the 300-epoch recipe detcon_b_config.py app. 6.2 app. 6.3.
• Speedups are read from curves. Efficiency claims such as "10x" are read from training curves, and the method is not stated tab. 1 fig. 4.
• Large-model comparisons mix data and architecture. SEER is trained on 1 billion images with a RegNet, while DetCon uses ImageNet and a ResNet-200. The paper says the comparison is imperfect tab. 3 §4.1.
• Classifier loss is in the total. The reported loss includes the monitor classifier's cross-entropy, so logs are not directly comparable to the paper's loss detcon_b.py.

Lineage

• Builds on: SimCLR (DetCon-S augmentations, encoder and projection head) and BYOL (DetCon-B framework, EMA target, prediction head) §3.1. Neither has a card in this wiki, so nothing from them is used here.
• Compared against: SwAV, MoCo, MoCo v2, InfoMin, PixPro, VADeR and SEER, on COCO transfer tab. 2 tab. 3.
• Superseded by: not stated in this wiki.