Context Autoencoder for Self-Supervised Representation Learning

Xiaokang Chen; Mingyu Ding; Xiaodi Wang; Ying Xin; Shentong Mo; Yunhao Wang; Shumin Han; Ping Luo; Gang Zeng; Jingdong Wang

Source-linked AI summary

Context Autoencoder for Self-Supervised Representation Learning

Xiaokang Chen, Mingyu Ding, Xiaodi Wang, Ying Xin, Shentong Mo, Yunhao Wang, Shumin Han, Ping Luo, Gang Zeng, Jingdong Wang

arXiv:2202.03026v3Feb 7, 2022cs.CV

TL;DR

Masked image modeling seeks representations that transfer to downstream tasks, but prior methods may mix encoder learning with masked-patch reconstruction. CAE separates these roles with an encoder-regressor-decoder that predicts masked representations in encoded space, and it reports superior downstream performance across segmentation, detection, instance segmentation, and classification.

Problem
Prior MIM methods mix representation learning with masked-patch reconstruction, while MAE only partially separates these roles and often lacks explicit visible-to-masked representation modeling.
Method
CAE uses an encoder-regressor-decoder to predict masked-patch representations from visible patches in encoded space and reconstruct masked patches from those predictions.
Results
CAE outperforms supervised pretraining, contrastive self-supervised pretraining, and other MIM methods on semantic segmentation, object detection, instance segmentation, and classification.
Takeaways & Limitations
CAE’s encoded-space prediction and role separation empirically improve representation learning and transfer performance across the reported downstream tasks.
Takeaways & Limitations
CAE may struggle with large contiguous masked regions because high-quality reconstruction must infer missing information from limited context.

Abstract

from arXiv · show

We present a novel masked image modeling (MIM) approach, context autoencoder (CAE), for self-supervised representation pretraining. We pretrain an encoder by making predictions in the encoded representation space. The pretraining tasks include two tasks: masked representation prediction - predict the representations for the masked patches, and masked patch reconstruction - reconstruct the masked patches. The network is an encoder-regressor-decoder architecture: the encoder takes the visible patches as input; the regressor predicts the representations of the masked patches, which are expected to be aligned with the representations computed from the encoder, using the representations of visible patches and the positions of visible and masked patches; the decoder reconstructs the masked patches from the predicted encoded representations. The CAE design encourages the separation of learning the encoder (representation) from completing the pertaining tasks: masked representation prediction and masked patch reconstruction tasks, and making predictions in the encoded representation space empirically shows the benefit to representation learning. We demonstrate the effectiveness of our CAE through superior transfer performance in downstream tasks: semantic segmentation, object detection and instance segmentation, and classification. The code will be available at https://github.com/Atten4Vis/CAE.

1 Introduction

The paper identifies limits in prior masked image modeling methods and introduces CAE to improve representation learning by predicting masked-patch representations in encoded space. CAE separates representation learning from pretraining-task completion and reports stronger downstream performance across several tasks.

Prior MIM methods often mix encoder learning with masked-patch reconstruction, while MAE only partially separates them and lacks explicit masked-representation prediction.
CAE uses an encoder-regressor-decoder architecture combining masked representation prediction with masked patch reconstruction.The encoder processes visible patches, the regressor predicts masked representations, and the decoder reconstructs masked patches from those predictions.
CAE predicts from visible-patch representations to masked-patch representations in encoded space, with targets expected to align with encoder-computed masked representations.
CAE separates representation learning from task completion by restricting the encoder to visible patches and assigning regression and reconstruction to separate components.
CAE outperforms supervised pretraining, contrastive self-supervised pretraining, and other MIM methods on semantic segmentation, detection, instance segmentation, and classification.

2 Related Work

The related work spans self-supervised learning, autoencoding, contrastive learning, and masked image modeling. CAE differs from earlier MIM approaches by explicitly separating encoder learning from masked-representation estimation and predicting in encoded representation space.

Self-supervised representation learning includes context prediction, clustering, contrastive learning, instance discrimination, image discretization, masked image modeling, and information maximization.
Autoencoding: Traditional autoencoders support dimensionality reduction or feature learning, while denoising autoencoders reconstruct uncorrupted data from corrupted inputs.
Contrastive self-supervised learning: Contrastive self-supervised learning compares augmented views, with random cropping often emphasizing central image regions; dense variants partially reduce this tendency.
Masked image modeling: Earlier MIM methods estimate pixels, discrete tokens, or other targets, but generally lack explicit separation between encoder representation learning and pretraining-task completion.
Masked image modeling: CAE jointly predicts masked representations and reconstructs masked patches in the encoder’s representation space.
Masked image modeling: Unlike MAE, CAE adds masked representation prediction and explicitly separates learning the encoder from completing pretraining tasks.
Masked image modeling: Unlike data2vec and iBoT, CAE forms target representations with a separate encoder rather than using the same network as teacher and student.

3 Approach

CAE is an encoder-regressor-decoder MIM architecture that predicts masked-patch representations from visible patches and reconstructs masked patches from those predictions. Its objective combines representation alignment and patch reconstruction under random block-wise masking.

3.1 Architecture: CAE makes predictions from visible patches to masked patches in encoded representation space using masked representation prediction and masked patch reconstruction.
Encoder: The encoder maps visible patches Xv to latent representations Zv and processes only the visible patches.
Regressor: The contextual regressor predicts masked representations Zm from visible representations Zv while conditioning on masked-patch positions.
Decoder: The decoder maps predicted masked representations to reconstructed masked patches and receives masked representations and positions without directly using visible-patch information.
Masking: Random block-wise masking divides each image into visible and masked patches, masking 98 of 196 patches.
3.2 Objective Function: The objective combines reconstruction loss ℓy(Ym, ¯Ym) and alignment loss ℓz(Zm, ¯Zm) as a weighted sum.The reconstruction targets are discrete tokenizer outputs, while representation targets are computed by encoding masked patches.

4 Discussions

CAE separates representation learning from masked-image task completion by predicting masked-patch representations from visible-patch representations, then reconstructing masked patches from those predictions. Its alignment constraint supports representation-space prediction, while the design is contrasted with BEiT, MAE, and contrastive methods.

CAE design: CAE predicts masked-patch representations from visible-patch representations and reconstructs masked patches from the predicted representations.Its encoder processes visible patches, the regressor predicts masked representations, and the decoder predicts masked-patch targets.
Representation-space prediction: Explicit prediction alignment ensures that regressor outputs lie in the encoder’s representation space.With alignment, reconstruction preserves image semantics; without it, reconstructed images are noisy or meaningless.
Connections to prior methods: CAE explicitly separates encoding from masked-image task completion, unlike BEiT and MAE, which perform these roles implicitly or simultaneously.In MAE, the decoder also updates visible-patch representations, whereas CAE’s decoder processes only predicted masked representations.
Downstream transfer: Separating representation learning from pretext-task completion helps downstream applications use the pretrained encoder with task-specific layers.The paper gives segmentation and detection layers as examples of replacing the pretext-task completion component.

5 Experiments

CAE experiments evaluate representation quality through probing, segmentation, detection, instance segmentation, and classification. Across these evaluations, CAE generally outperforms supervised, contrastive self-supervised, and other MIM baselines, while ablations support encoded-space prediction and task separation.

ImageNet Evaluation: Attentive probing scores are much larger than linear probing for MIM methods because MIM extracts representations for all image patches.The authors attribute the difference to classification attending to the corresponding patch regions.
Semantic Segmentation: CAE* with ViT-B trained for 1600 epochs outperforms MAE, MoCo v3, and DeiT on ADE20K segmentation by 2.1, 3.0, and 3.2, respectively.With ViT-L, CAE* also outperforms BEiT and MAE by 1.4 and 1.1, respectively.
Object Detection and Instance Segmentation: CAE* with ViT-B trained for 1600 epochs outperforms MAE, MoCo v3, and DeiT on COCO by 1.6, 4.5, and 3.1, respectively.Using ViT-L, CAE achieves 54.6 box AP and outperforms MAE by 0.6; CAE also performs better under Cascaded Mask R-CNN.
Object Detection and Instance Segmentation: 64.6 mAP on COCO test-dev is obtained with CAE-pretrained ViT-Huge, exceeding BEiT-3 at 63.7 mAP and SwinV2-G at 63.1 mAP.The authors present this as a scaling experiment using ImageNet-22K pretraining.
Classification and Ablations: CAE outperforms prior supervised and self-supervised methods on Food-101, Clipart, and Sketch classification, and its performance is largely insensitive to the reconstruction target.RGB pixels, DALL-E tokenizer targets, and d-VAE tokenizer targets obtain similar overall performance, with differences across probing schemes.
Ablation Studies: Adding both decoder and alignment improves downstream performance, whereas removing the decoder reduces it; the added components increase parameters to 1.23× and training time to 1.24×.The ablation attributes the improvement to alignment keeping predicted masked representations in the encoded representation space.
Ablation Studies: A 60% mask ratio improves linear and attentive probing but reduces semantic segmentation performance by 0.2%, so the experiments generally use 50%.The choice of λ also matters modestly: λ = 1 performs slightly worse than the adopted λ = 2.

6 Conclusion

The conclusion identifies encoded-space prediction and the combination of masked representation prediction with masked patch reconstruction as CAE’s core design. It reports effective downstream performance while acknowledging difficulty with large contiguous masked regions.

Conclusion: CAE predicts masked-patch representations from visible patches in encoded representation space and combines masked representation prediction with masked patch reconstruction.The architecture is presented as the core design for masked image modeling.
Conclusion: Experiments demonstrate the effectiveness of the CAE design, including its downstream transfer performance.The conclusion summarizes the design and its empirical validation rather than reporting a single aggregate metric.
Conclusion: CAE may face challenges when large, contiguous masked regions leave limited context for reconstructing missing information.The conclusion describes this as a common limitation of MIM methods and states that CAE is not exempt.

Declarations

The paper reports funding from multiple research programs and laboratories, and states that its code and datasets are publicly available or will be made available.

Funding: The work is partially supported by national research programs, provincial and joint laboratories, and institutional grants.The listed support includes Chinese national programs, the Intelligent Terminal Key Laboratory of SiChuan Province, and the Beijing Advanced Innovation Center for Intelligent Robots and Systems.
Code availability: The authors state that the CAE code will be available on GitHub.The repository is identified as Atten4Vis/CAE.
Availability of data and materials: The datasets used are publicly available, including ImageNet, ADE20K, COCO, Food-101, Clipart, and Sketch.The passage provides links for each listed dataset.