Contrastive Masked Autoencoders are Stronger Vision Learners

In plain words

A masked image model learns by hiding part of a picture and guessing what is missing. This teaches it good local detail, but the features it learns are weaker at telling one picture apart from another. Contrastive learning does the opposite: it teaches that two views of the same picture belong together and that views of different pictures do not.

CMAE runs both at once. The main branch hides most of the picture and rebuilds the missing pixels, as in a masked autoencoder. A second copy of the encoder, updated slowly, sees the whole picture instead. Each branch gets a slightly different crop of the same photo, so the two views are close but not identical. A small extra decoder on the main branch predicts what the whole-picture copy produces, and training pulls these two outputs together and pushes them away from the other pictures in the batch.

The two views are kept close on purpose. Crops that are far apart can show unrelated content, so here the crop positions move only a few pixels. After pretraining only the main encoder is kept for downstream tasks.

TL;DR

CMAE pretrains a ViT with two branches: an online masked autoencoder (pixel reconstruction plus a feature decoder) and a momentum encoder that sees the full image, joined by an InfoNCE loss on pixel-shifted views. The headline is 84.7% ImageNet top-1 for ViT-B/16 after 1600 pretraining epochs, against 83.6% for MAE. The authors also report gains on ADE20K and COCO. All results are author-reported. The released code differs from the paper in several places, notably the scale of the contrastive loss and the masking pattern (see Paper vs code). The best numbers use a hybrid convolutional ViT, which the release does not include.

Problem

Masked image modeling (MIM) gives strong pretraining, but the authors argue that its features discriminate between images less well than contrastive features do §1. They attribute this to MIM modelling relations among the tokens of one image rather than across images §2. Naive combinations of the two gave only marginal gains §1. The two families differ in input augmentation, objective and architecture, which the authors say makes combining them non-trivial §1.

The obvious recipe, two random crops with masking, adds 0.2 points over the MAE baseline in their ablation tab. 4a §4.4. The authors argue that random crops of a masked input create false positive pairs, because the masking already removes much of the content §3.2 §3.4.

Assumptions: pretraining uses only the ImageNet-1K training split §4.2. The default encoder is ViT-B with 16-pixel patches §4.1.

Method

Overview. Both branches receive views cut from one master crop §3.2. The online branch is an asymmetric encoder-decoder: a ViT-B encoder sees only the visible patches, a pixel decoder reconstructs all patches, and a feature decoder predicts features of the masked patches §3.1. The momentum branch is an exponential-moving-average (EMA) copy of the encoder and of a projection head. It sees the whole view and provides the contrastive targets §3.1. Only the online encoder is kept afterwards fig. 3.

1. Two views from one master crop §3.2. A master image is cut from the source image by resized random cropping. The online view is a fixed crop at the top-left corner of the master. The momentum view is a crop at a random offset of up to p = 31 pixels, drawn independently on each axis. The paper writes the offsets as drawn from [0, p) §3.2. Both views therefore overlap heavily. The online view gets no colour change. The momentum view gets colour augmentation §3.2.
2. Online encoder §3.1 eq. 1. A ViT-B encoder, following MAE §3.1. It embeds the visible tokens, adds positional embeddings and processes them with transformer blocks. The paper says only that a large ratio of patches is masked, and cites MAE's 75% as an example §3.4. The released configuration uses a masking ratio of 0.65, applied to 2 x 2 blocks of patches ...1600e_in1k.py cmae_vit.py.
3. Pixel decoder §3.1 eq. 3. It receives the encoded visible tokens together with mask tokens for the hidden positions. It predicts the pixels of every patch, and only the hidden positions enter the loss. The paper gives no decoder depth or width.
4. Feature decoder §3.1 eq. 4. It has the same structure as the pixel decoder but separate weights. The authors use a two-layer version and found that sharing weights with the pixel decoder performs worst §4.4. Its output is averaged over all token positions to give one vector per image eq. 4.
5. Momentum encoder §3.1 eq. 2. It has the same architecture as the online encoder but takes the full view, with no masking. Its output tokens are mean-pooled over the input patch positions eq. 2. The authors argue that a token-level target is ambiguous and that the global mean is simpler §3.1.
6. Projection and prediction §3.3. The feature decoder output goes through a projection head and then a prediction head. The momentum output goes through a projection head only. Both heads are updated as described in step 9 §3.3. The paper gives no widths for the heads. The released configuration uses 768 to 1536 to 256 with a prediction head of 256 to 1536 to 256 ...1600e_in1k.py.
7. Contrastive loss §3.3 eq. 6 eq. 7. The cosine similarity between the predicted and momentum features is computed. The InfoNCE loss takes the matching image as the positive and the other images in the batch as negatives. The temperature is tau = 0.07 eq. 7. The paper chose InfoNCE over a BYOL-style loss because it performed better (83.8% vs 83.4%) §4.4.
8. Reconstruction loss §3.3 eq. 5. The targets are normalised pixels, following MAE §3.3. The loss is the mean squared error over masked patches only eq. 5.
9. Total loss and EMA update §3.3 eq. 8 §3.1. The loss is L = L_r + lambda_c * L_c with lambda_c = 1 in the experiments eq. 8 §4.4. After each step the momentum encoder and projection head are updated by theta_t <- mu * theta_t + (1 - mu) * theta_s with mu = 0.996 §3.1.

Pseudocode, view generation and masking §3.1 §3.2:

# Two views of one image  [§3.2]
# Input:  image (3, H, W), an ImageNet picture
# Output: online view (3, 224, 224), momentum view (3, 224, 224)

def make_views(image):
    # 256 px resized random crop, then a horizontal flip  [§3.2]
    master = random_flip(random_resized_crop(image, size=256))
    # Online view: fixed top-left 224 px crop, no colour change  [§3.2]
    online_view = master[:, 0:224, 0:224]
    # Momentum view: random offset in 0..31 on each axis  [§3.2]
    offset_h = random_int(0, 31)
    offset_w = random_int(0, 31)
    momentum_view = master[:, offset_h:offset_h + 224, offset_w:offset_w + 224]
    # Colour and blur on the momentum view only  [§3.2]
    momentum_view = color_augment(momentum_view)
    return online_view, momentum_view

# Block masking of the online view  [§3.1]
# Input:  patch tokens (B, 196, D) for a 224 px view (14 x 14 patches)
# Output: visible tokens (B, 68, D), loss mask (B, 196)

def block_mask(tokens, mask_ratio=0.65):
    # 2 x 2 patch cells: 196 / 4 = 49 per image  [code: cmae_vit.py#L121]
    cells = tokens.shape[1] // 4
    # Visible cells: int(49 * (1 - 0.65)) = 17  [code: cmae_vit.py#L122]
    kept = int(cells * (1 - mask_ratio))
    # Random cell order; keep the first `kept`  [code: cmae_vit.py#L123-L126]
    cell_noise = random_uniform(batch_size, cells)
    kept_cells = argsort(cell_noise)[:, :kept]
    # Each kept cell gives 4 patches: 17 x 4 = 68  [code: cmae_vit.py#L130-L135]
    visible = cells_to_patches(kept_cells)
    # Loss mask: 1 on the 128 hidden patches  [code: cmae_vit.py#L141-L142]
    loss_mask = 1 - patch_mask(visible)
    return gather_patches(tokens, visible), loss_mask

One pretraining step, following the released loss §3.3 eq. 5-8:

# One pretraining step  [§3.1] [§3.3]
# Input:  images (B, 3, H, W); momentum encoder and projector are EMA copies
# Output: scalar loss

def cmae_loss(images, tau=0.07, lambda_c=1.0):
    # Two views of each image  [§3.2]
    online_view, momentum_view = make_views(images)
    # Online encoder sees only the visible patches  [eq. 1]
    visible, loss_mask = block_mask(patchify(online_view))
    latent_s = online_encoder(visible)  # (B, 1 + 68, 768)
    # Pixel decoder rebuilds all patches from visible and mask tokens  [eq. 3]
    pred_pixels = pixel_decoder(latent_s, visible_positions)  # (B, 196, 768)
    # Feature decoder: same design, separate weights, depth 2  [eq. 4] [§4.4]
    pred_feat = feature_decoder(latent_s, visible_positions)  # (B, 196, 768)
    # Momentum encoder sees the whole view with no masking  [§3.1]
    latent_t = momentum_encoder(momentum_view)  # (B, 1 + 196, 768)
    # Reconstruction: normalised targets, MSE on hidden patches  [eq. 5]
    target = normalise_per_patch(patchify(images))  # (B, 196, 768)
    loss_rc = masked_mean((pred_pixels - target) ** 2, loss_mask)
    # Query: projection then prediction on the mean feature  [§3.3]
    query = predictor(projector(mean_over_tokens(pred_feat)))  # (B, 256)
    # Key: projected mean of patch tokens, CLS excluded  [eq. 2]
    key = momentum_projector(mean_over_tokens(latent_t[:, 1:]))  # (B, 256)
    # Keys carry no gradient; the EMA update drives them  [§3.1]
    key = stop_gradient(gather_all_ranks(key))
    # Cosine similarity to every key in the global batch, scaled by tau  [eq. 6]
    logits = normalise(query) @ normalise(key).T / tau
    # InfoNCE: the matching key is the positive  [eq. 7]
    loss_ct = cross_entropy(logits, targets=arange(B) + rank * B)
    # Released code scales InfoNCE by 2 * tau  [code: cmae_head.py#L82]
    loss_ct = 2 * tau * loss_ct
    return loss_rc + lambda_c * loss_ct  # total loss  [eq. 8]

Hyperparameters

Pretraining, unless stated. ...1600e_in1k.py is the 1600-epoch pretraining config in configs/pre-training/cmae/. Values marked "computed" are results of the released code's formula, not numbers printed in the paper.

Fine-tuning on ImageNet-1K, for the 1600-epoch checkpoint unless stated:

Detection and segmentation (paper only, not in the release): ADE20K with UperNet, the default segmentation framework §4.3; COCO with Mask R-CNN, 100 epochs, base learning rate 1e-4 with cosine annealing and weight decay 0.1 §4.3.

Results

Author-reported. ImageNet-1K top-1 after fine-tuning, ViT-B/16 unless marked hybrid §4.2 tab. 2:

The 300-epoch CMAE model reaches 83.8 in the ablation table tab. 4a tab. 4c. The README lists the same fine-tuned value for it README.md.

Downstream tasks §4.3 tab. 3 tab. 5:

Other reported results, all author-reported and given only in the text: linear probing and partial fine-tuning with the 1600-epoch ViT-B: one fine-tuned block gives +2.5 over MAE, and linear probing gives +5.9 over MAE §4.5. Convergence: CMAE surpasses MAE's final fine-tuned accuracy after 55 fine-tuning epochs §4.4 fig. 5. Feature analysis on ImageNet-1K validation (pretraining epochs not stated): intra-class distance 0.0377 vs 0.0380 for MAE, and the standard deviation of intra-class distance 0.0189 vs 0.0371 §4.7 fig. 6. Model scaling over ViT-S, ViT-B and ViT-L is shown only in fig. 4; no values are given in the text §4.6.

Ablations, ViT-B, 300 pretraining epochs, top-1 on ImageNet-1K tab. 4 tab. 6:

Paper vs code

Code: ZhichengHuang/CMAE at e786a5253504b1c5d5ac4027c485711b58ba4ddd. The repository is linked from the abstract, and its owner is the paper's first author. It contains the pretraining and ImageNet fine-tuning configurations for ViT-B, for both CMAE and MAE. The pixel-shift views follow the paper. The masking pattern, the loss scale and the items below differ from the paper.

• Contrastive weight. The paper's total loss is L_r + lambda_c * L_c with lambda_c = 1 §3.3 eq. 8. The code multiplies the InfoNCE term by 2 * tau = 0.14 before adding it cmae_head.py. The lambda_c sweep in the paper tab. 4d does not say which scale it uses.
• Masking pattern. The paper describes random masking of patches, but the code masks 2 x 2 blocks of patches. At 224 px, 17 of 49 cells stay visible, which is 68 of 196 patches cmae_vit.py. The paper does not mention blocks §3.1.
• Masking ratio. The paper says only "large" §3.1. The released configuration sets 0.65 ...1600e_in1k.py. The paper's own momentum-branch sweep also includes 0.65 tab. 4e.
• Momentum branch normalisation. The momentum encoder is built without its final LayerNorm ...1600e_in1k.py, while the online encoder keeps it cmae_vit.py. The paper does not mention this.
• Extra augmentation. The momentum view also gets Gaussian blur with p = 0.5 cmae_dataset.py. The paper names colour jitter and random grayscale only §3.2.
• Pixel offsets. The paper draws the offsets from [0, p) §3.2. The code draws them with random.randint(0, 31), which includes 31 processing.py.
• Momentum mean. The code averages momentum tokens over patch positions and drops the CLS token cmae.py. This matches the indexing of eq. 2 eq. 2.
• Feature decoder output. The feature decoder ends in the same linear layer as the pixel decoder, which outputs 16 x 16 x 3 = 768 values per token mae_neck.py. Its features are therefore pixel-shaped vectors of width 768, not separately defined encoder features. The paper says it has the same structure §3.1.
• Projector normalisation. The projection head uses SyncBN in its hidden layers nonlinear_neck.py. The paper does not state a normalisation layer for the heads.
• Negatives. The code gathers the momentum keys from all 32 GPUs, so the negatives are the full global batch of 4096 cmae_head.py. This matches "K is the batch size" eq. 7, but the loss depends on the global batch, not the per-GPU batch.
• Not in the release. The configs and code cover only ViT-B pretraining and ImageNet fine-tuning. Missing: the ADE20K and COCO setups §4.3, the hybrid convolutional ViT used for the best numbers §4.2, ViT-S and ViT-L scaling §4.6, linear probing and partial fine-tuning §4.5, and the feature analysis §4.7. A HybridEmbed utility exists embed.py but no CMAE config uses it.
• Licence file. The repository has no LICENSE file. setup.py declares "Apache License 2.0" setup.py. Treat the code licence as not confirmed until the authors add one. No code is quoted here.

No other discrepancies found in the online encoder, the pixel decoder, the momentum EMA, the warmup and cosine schedule, the learning-rate rule, or the pixel-shift ranges.

Pitfalls

• Block masking. An implementation that masks single patches at random will not match the released model, because the code masks 2 x 2 blocks cmae_vit.py. The paper does not describe blocks.
• Contrastive scale. If you implement eq. 7 and 8 as written, the InfoNCE term is about 7 times larger than in the released code cmae_head.py. Pick one convention and report it.
• Feature target is 768-dimensional. The "feature" branch outputs pixel-sized vectors through the same linear head shape as the pixel decoder. The paper's word "features" can suggest encoder-width targets, but it says the feature decoder has the same structure as the pixel decoder §3.1. Implement the output head as the code does mae_neck.py.
• Momentum LayerNorm. Keeping the final LayerNorm on the momentum branch changes the contrastive target. The released code removes it ...1600e_in1k.py.
• Reconstruction loss. Targets are normalised per patch over the 768 values, with variance epsilon 1e-6 cmae_head.py. The loss is averaged over the 768 values of each masked patch, then over masked patches in the batch cmae_head.py. Every image has the same count of hidden patches (128), so this equals the per-image average in eq. 5.
• Pixel-shift geometry. The master crop is 256 px and the views are 224 px, so offsets 0 to 31 fit with one pixel to spare processing.py. The online view is always the top-left corner. The momentum view has independent offsets on each axis. Both views come from the same flipped master.
• Ablation tables disagree with each other. Pixel shift without colour is 83.6 in tab. 4a and 83.4 in tab. 4b. Random crop is 83.0 in tab. 4b and 83.29 in tab. 6. The text says random crops hurt §3.2, but the random-crop row of tab. 4b (83.0) is above the 82.9 baseline. Differences of 0.2 to 0.5 points are called significant §4.4, and no seeds or variance are reported.
• Depth-8 feature decoder. The text reports a "trivial solution" at depth 8 §4.4, but the depth-8 result is not in any table.
• Text vs table. The text gives a 0.8-point gain over SIM §4.2. The table gives 84.7 - 83.8 = 0.9 tab. 2. The COCO text compares against the authors' reproduced MAE row (51.7), not the reported MAE row (50.3) §4.3 tab. 3b.
• Headline numbers. The abstract's 85.3% ImageNet and 52.5 mIoU ADE20K use the hybrid convolutional ViT. For plain ViT-B/16 the best result is 84.7% tab. 2. The hybrid model is not in the release.
• Cost. The momentum branch runs a full ViT-B forward on all 196 patches at every step, and the online branch runs two decoders. The paper says CMAE is "less computationally expensive than iBOT" §3.4 but reports no wall-clock, memory or FLOP figures.
• Fine-tuning details beyond the paper. Layer-wise decay 0.65, label smoothing 0.1, RandAugment and random erasing are used in the release cmae_base_pre1600_8x128_100e_in1k.py but are not in the paper's text.

Lineage

• Builds on: MAE for the ViT encoder, masking default and normalised-pixel target §3.1 §3.3. MoCo v3 for the projection-prediction heads and the InfoNCE setup §3.3. BYOL, cited for the EMA momentum encoder and the prediction head, has no card in this wiki §3.1 §3.3.
• Compared against: iBOT and SIM, both mixing contrastive and masked objectives §1 tab. 2. SimMIM tab. 2. ConvMAE is the hybrid baseline for the best numbers tab. 2 and has no card.
• Superseded by: not stated in this wiki.