Context Autoencoder for Self-Supervised Representation Learning

In plain words

Self-supervised pretraining for images often works like a fill-in-the-blank game. Parts of a picture are hidden, and the network learns to guess them from the parts it can see. The network that learns this way is then reused for other vision jobs. This paper keeps that game but changes how the guessing is organised.

The main idea is to split the network into three jobs: an encoder that only looks at the visible pieces and builds features, a small module that predicts the features of the hidden pieces from the visible features, and a decoder that turns those predicted features into a description of the hidden pieces. Only the encoder is kept afterwards. Because the predictions are made in the feature space itself, the encoder is not forced to do the decoding work.

The authors argue that this separation gives better features for tasks such as segmentation and detection than older masked-image methods, which mix the two jobs inside one network.

TL;DR

Context autoencoder (CAE) is a masked-image-modeling method for ViT pretraining without labels. It hides about half the patches, encodes only the visible ones, and regresses the hidden patches' encoder representations with a cross-attention regressor. A decoder then predicts discrete tokens for the hidden patches, and the two losses are summed. Only the encoder is kept for downstream tasks. Reported gains over MAE and BEiT are author-reported and depend on the probe and the tokenizer.

Problem

Masked image modeling pretrains a ViT by predicting hidden patches. In BEiT, the ViT paper and iBOT, one ViT both builds the representation and completes the prediction, so the two roles are mixed §1. MAE separates encoder and decoder, but the authors say the decoder can still learn representations for the visible patches, which limits representation quality §1 §4.2. Most earlier methods also lack an explicit model that maps visible-patch representations to masked-patch representations §1.

The authors assume that predicting masked patches is easier when the encoded representations carry more semantics, and that accurate prediction pushes the encoder toward semantics §1. They also argue that contrastive pretraining with random crops mostly learns from central image regions, whereas masked modeling covers all patches §4.2. Both points are arguments from the authors' analysis and visualisations, not separate proofs §4.1 §4.2 §4.3.

Method

CAE has three parts: an encoder F, a latent contextual regressor H and a decoder G §3.1 fig. 1. Each image is split into visible patches X_v and masked patches X_m, using random block-wise masking in the style of BEiT. For each 224 by 224 image, 98 of the 196 patches (14 by 14 grid) are hidden §3.2.

1. Mask. Hide 98 of 196 patches in random blocks §3.2. The block sampler is defined in the code, not in the paper's text masking_generator.py.
2. Encode visible patches. F is a ViT that embeds only the visible patches, adds position embeddings and outputs Z_v, with a class token kept. The masked patches never enter F during the prediction path §3.1 modeling_cae_helper.py.
3. Regress hidden representations. H is a stack of cross-attention blocks. Its queries are one learned mask token, repeated for every hidden patch, plus the hidden patches' positions. Its keys and values are Z_v together with the output of the previous regressor layer. Visible representations are not updated §3.1 modeling_cae_helper.py.
4. Targets for the hidden representations. The target Z̄_m is the encoder's output for the hidden patches, with gradient stopped §3.2 modeling_cae.py.
5. Decode. G is a stack of self-attention blocks, followed by a linear layer. It receives only the regressor output for the hidden patches plus their positions, not the visible representations §3.1 modeling_cae_helper.py.
6. Reconstruction targets. A discrete tokenizer assigns one code to each patch of the full image. The reported main runs use a d-VAE trained on ImageNet-1K without labels, and a second setting uses the DALL-E tokenizer, trained on 400M images §3.2 tab. 1 §5.5.
7. Loss. The total loss is a cross-entropy term on the decoder's token logits plus λ times a mean-squared-error term between the regressor output and Z̄_m, with λ = 2 eq. 1 §3.2 §5.2.
8. Downstream use. After pretraining, the encoder is kept and the pretext parts are replaced by the downstream head §3.1 fig. 1.

The paper also gives a probabilistic factorisation, read as encoder, then regressor, then decoder §4.3. It is an interpretation, not a training change. The probabilistic equations are not numbered in the source as checked here, so they are cited by section.

# CAE pretraining loss  [§3.1] [§3.2] [eq. 1]
# image: (B, 3, H, W); masked: (B, 196) bool, True for 98 hidden patches

def cae_loss(image, masked):
    # Keep the visible patches as the only encoder input  [§3.1]
    visible = patches(image)[~masked]               # (B, N_v, d_in)

    # Encode visible patches; the class token is kept  [§3.1]
    visible_repr = encoder(visible, positions)      # (B, N_v + 1, d)

    # Target: the encoder on the hidden patches, no gradient  [§3.2]
    hidden_target = stop_gradient(
        target_encoder(patches(image)[masked])      # (B, N_m, d)
    )

    # Predict hidden representations from the visible ones  [§3.1]
    # Queries: learned mask token plus position of each hidden patch
    hidden_pred = regressor(
        mask_token, visible_repr, pos_hidden, pos_visible
    )                                               # (B, N_m, d)

    # Decode hidden representations into token logits  [§3.1]
    logits = decoder(hidden_pred + pos_hidden)      # (B, N_m, 8192)

    # Discrete targets from the tokenizer, no gradient  [§3.2]
    codes = tokenizer(image)[masked]                # (B, N_m)

    # Reconstruction plus alignment, lambda = 2  [eq. 1] [§3.2] [§5.2]
    loss_rec = cross_entropy(logits, codes)
    loss_align = mse(hidden_pred, hidden_target)
    return loss_rec + 2.0 * loss_align

The pseudocode above is a condensed, language-agnostic reading. In the code, target_encoder is a second module that shares the encoder's weights, with no gradient. See Paper vs code.

Hyperparameters

Pretraining is on ImageNet-1K. Values are the paper's unless the code column says otherwise.

Linear probing (frozen encoder, class token):

Attentive probing (frozen encoder, learned query token):

Fine-tuning on ImageNet-1K (ViT-B shown; ViT-S and ViT-L in the paper):

Downstream, from the text §5.2:

• ADE20K segmentation with UperNet, AdamW, input 512 by 512, batch 16, 160k steps, no multi-scale testing. For ViT-B, LR chosen from 1e-4, 2e-4, 3e-4 and 4e-4, decay 0.65, drop path 0.1. For ViT-L, LR from 3e-5, 4e-5 and 5e-5, decay 0.95, drop path 0.15.
• COCO Mask R-CNN 1× schedule with multi-scale training on short side 480 to 800, long side at most 1333, batch 32. ViT-S: LR 3e-4, decay 0.75, drop path 0.1. ViT-B: LR 3e-4, decay 0.75, drop path 0.2. ViT-L: LR 2e-4, decay 0.8, drop path 0.2.

Results

Author-reported. Setup in each row. All results are ImageNet-1K pretraining unless stated. CAE* uses the DALL-E tokenizer; CAE uses the d-VAE tokenizer trained on ImageNet-1K tab. 1 caption.

ImageNet-1K evaluation, ViT-B tab. 1:

ImageNet-1K evaluation, ViT-L, 1600 epochs tab. 1:

ADE20K semantic segmentation, mIoU, UperNet tab. 2:

COCO object detection and instance segmentation, Mask R-CNN 1× tab. 3:

Table 3's ViT-S block lists only DeiT, MoCo v3, BEiT and CAE*. The comparisons in the text use the ViT-B and ViT-L blocks.

Classification, fine-tuning on ViT-B, top-1 tab. 4:

Ablations, ViT-B, 300 epochs, ImageNet-1K tab. 5:

Reconstruction target, ViT-B, 1600 epochs tab. 8: DALL-E tokenizer 70.4 / 77.1 / 50.2 (linear / attentive / ADE20K); d-VAE tokenizer 71.4 / 77.4 / 50.1; raw RGB pixels 72.4 / 77.0 / 50.4. The authors conclude the target does not matter much.

Mask ratio, ViT-B, 300 epochs tab. 7: 40%: 63.1 / 73.0 / 47.2; 50%: 64.1 / 73.8 / 48.3; 60%: 64.8 / 74.2 / 48.1.

Other numbers in the text §5.4 §5.5:

• Layer count of regressor and decoder, linear probe for 1, 2, 4 and 5 layers: 58.7, 62.1, 64.1, 64.2. Attentive: 67.5, 71.1, 73.8, 73.7. Four layers were chosen §5.5.
• Cascade Mask R-CNN, ViT-B, 1600 epochs: AP^b 52.9 and AP^m 45.5 for CAE*, against MAE at 51.3 and 44.3 tab. 6.
• Detection with a ViT-Huge encoder pretrained on ImageNet-22K: 64.6 mAP on COCO test-dev, given in the text only §5.4. The paper gives no table for it, and this card did not check the repository for its configuration.

Paper vs code

Official repository: https://github.com/Atten4Vis/CAE at commit a7fd1628176358e3e76b0b042b0f6e9f3cd7c76c. The paper links it in the abstract and in its unnumbered Declarations section. The repository has no LICENSE file, so code_license is "not stated". No code is quoted in this card.

1. Alignment target is a second module that shares the encoder's weights. The paper defines the target as the same encoder applied to the hidden patches, with gradient stopped §3.2. The code builds a second ViT, the alignment encoder, copies the weights once at construction, and freezes it modeling_cae.py. After each target forward pass, it points each of its parameters at the encoder's tensor storage modeling_cae.py. From the first update on, the target therefore follows the encoder's in-place optimizer updates, as the paper describes. The paper does not mention the aliasing.
2. Regressor keys and values include hidden tokens. The paper's §3.1 names visible representations and the previous regressor output as keys and values. Section 5.1 calls self-attention over masked tokens together with visible representations a choice that costs more computation and slightly reduces performance. The code always concatenates visible and hidden tokens into the keys and values, with no switch modeling_cae_helper.py. The paper does not say which setting produced the reported numbers.
3. Default tokenizer is DALL-E. The Table 1 rows named CAE use the d-VAE tokenizer trained on ImageNet-1K. In the code, the default is dall-e run_pretraining.py, the README downloads the DALL-E encoder and decoder README.md, and all scripts use the default. The d-VAE setting appears to need --discrete_vae_type customized utils.py, and the repository does not show where its weights come from. The DALL-E weights are not included in this card.
4. Large-model script has a line-continuation bug. scripts/cae_large_1600e.sh ends the pretraining command at line 37 without a backslash. Lines 38 to 40 (--decoder_embed_dim 1024, --decoder_num_heads 16, --fix_init_weight) are therefore not passed to run_pretraining.py cae_large_1600e.sh. The script's defaults then apply: decoder width 768, 12 heads, no fixed init run_pretraining.py. Whether the reported ViT-L numbers used the intended 1024 width is not stated.
5. README numbers differ slightly from the paper. The README's CAE-base rows at 300 and 1600 epochs read linear 64.5 and 70.3, and attentive 74.0 and 77.2 README.md, README.md. The paper reports 64.1 and 70.4 for linear, and 73.8 and 77.1 for attentive tab. 1. The ViT-L row reads linear 77.8, against 78.1 in the paper README.md tab. 1. The README says the paper's results are reproduced README.md. The differences are small but are not the paper's numbers.
6. Code-only pretraining settings. These are not in the paper's text: gradient clipping at 3.0 cae_base_800e.sh; layer-scale initial values 0.1 for base and 1e-5 for large cae_base_800e.sh cae_large_1600e.sh; minimum block size 16 patches run_pretraining.py; cosine decay to 1e-5 run_pretraining.py; mixed-precision autocast engine_for_pretraining.py. Each one changes training, so the card lists them rather than assuming the paper's defaults.
7. Fine-tuning script adds mixup and cutmix. The paper's fine-tuning section lists layer-wise decay, weight decay, AdamW and the learning rates §5.2, but no mixup or cutmix. The base script uses mixup 0.8 and cutmix 1.0 cae_base_finetune.sh.
8. Regressor normalisation is not in the paper. Each regressor cross-attention block applies a LayerNorm after its residual, before the MLP, and the regressor ends with another LayerNorm modeling_cae_helper.py modeling_cae_helper.py.
9. Projection layer. An encoder-to-regressor linear layer with LayerNorm is added only when the encoder and decoder widths differ. The same projection is applied to the alignment targets modeling_cae.py modeling_cae.py. For ViT-B, both widths are 768, so no projection is used.
10. Settings that match. Pretraining learning rate 1.5e-3, batch 2048, warmup 10 / 20 / 40, weight decay 0.05, mask 98, λ = 2 and four regressor and decoder blocks agree between paper and scripts §5.1 §5.2 cae_base_800e.sh. Linear probing LR 6.4 equals 0.1 times 16384 over 256 run_linear.py. Attentive probing LR 0.4 and batch 8192 (256 per GPU on 32 GPUs) agree cae_base_800e.sh.
11. Not compared. Segmentation and detection configurations under downstream_tasks/ were not checked against the paper's per-run settings. ViT-S and ViT-L fine-tuning settings were not compared.

Pitfalls

• Aliasing. The target follows the encoder only because the copy shares tensor storage and is re-pointed after each target pass modeling_cae.py. If a change replaces encoder tensors rather than updating them in place, the target lags by one step, because the re-pointing runs after the target forward pass. Check this when changing checkpoint loading or optimizer code.
• Extra cost. Each step runs a no-gradient forward pass through the alignment encoder on the hidden patches modeling_cae.py. The speed cost of this is not reported separately in the paper; tab. 5 gives one combined training-time ratio.
• Parameter count. The script's printed count includes only parameters with gradients, so it excludes the frozen alignment module run_pretraining.py. The paper's parameter column in tab. 5 does not say whether that module is counted.
• Do not drop the reconstruction term. The authors say training collapses to a trivial solution without masked patch reconstruction §5.5. Dropping the decoder therefore changes what the encoder learns.
• Decoder alone helps little. tab. 5 and §5.5: adding the decoder alone leaves downstream results almost unchanged, while adding both the decoder and the alignment term gives the best linear and attentive scores. Without the decoder, a linear output layer must map semantic features to low-level targets, which the authors say performs worse §5.5.
• Block masking details. The block sampler adds blocks with random size and aspect ratio, in up to 10 attempts per block, until exactly 98 patches are hidden masking_generator.py. A block that does not fit is retried. Changing the number of hidden patches also changes the block sizes, since the maximum block size equals the total hidden count by default masking_generator.py.
• Learning rate depends on batch size for linear probing. The learning rate is base LR times total batch over 256 run_linear.py. The attentive probe uses its LR as given cae_base_800e.sh. Changing batch size for one probe changes its effective schedule.
• Attentive probe baselines. The paper says scores for other models come from the authors' own implementations unless noted tab. 1 caption. Attentive-probe comparisons therefore depend on one implementation.
• Tokenizer and weights. The DALL-E tokenizer is downloaded from OpenAI's CDN README.md. No tokenizer weights are included in this card, and their license is not checked here.
• Code license not established. The repository has no LICENSE file. One source header names the MIT license modeling_discrete_vae.py, but that does not establish a repository-wide license. Quoting code is avoided for that reason.
• Mixed precision. Losses are computed under autocast and the MSE on the alignment is taken in float engine_for_pretraining.py. Numerical changes are possible if you change precision.

Lineage

• Builds on: BEiT for block-wise masking, tokenizer targets and pretraining settings §3.2 §5.2; BERT's masked language modeling §2; iBOT's combination of masked modeling with DINO-style self-distillation §2 §4.2; denoising autoencoders §2.
• Compared against: MAE, a concurrent masked-autoencoder method, which has a card §2 tab. 1; BEiT, iBOT, MoCo v3, DINO, DeiT, SimMIM, SplitMask, Ge²-AE and A²MIM in the tables §2 tab. 1 tab. 2.
• Superseded by: not stated in the paper. The repository points to CAE v2, which adds CLIP latent alignment README.md. CAE v2 is not covered by this card.