[Paper Review] Corrupted Image Modeling for Self-Supervised Visual Pre-Training

Y.K. Fang, Dong Li+3 more|arXiv (Cornell University)|Feb 7, 2022

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology23 citations

TL;DR

CiM introduces a generator-enhanced corruption of input images to train a non-Siamese enhancer for ViT and CNN pre-training, achieving competitive ImageNet and ADE20K results. The enhancer can be either generative (ResPix) or discriminative (RevDet).

ABSTRACT

We introduce Corrupted Image Modeling (CIM) for self-supervised visual pre-training. CIM uses an auxiliary generator with a small trainable BEiT to corrupt the input image instead of using artificial [MASK] tokens, where some patches are randomly selected and replaced with plausible alternatives sampled from the BEiT output distribution. Given this corrupted image, an enhancer network learns to either recover all the original image pixels, or predict whether each visual token is replaced by a generator sample or not. The generator and the enhancer are simultaneously trained and synergistically updated. After pre-training, the enhancer can be used as a high-capacity visual encoder for downstream tasks. CIM is a general and flexible visual pre-training framework that is suitable for various network architectures. For the first time, CIM demonstrates that both ViT and CNN can learn rich visual representations using a unified, non-Siamese framework. Experimental results show that our approach achieves compelling results in vision benchmarks, such as ImageNet classification and ADE20K semantic segmentation.

Motivation & Objective

• Motivate a flexible, non-Siamese self-supervised pre-training framework that works for both ViT and CNN architectures.
• Propose corrupted image generation via a trainable BEiT-based generator to enrich input views.
• Explore two enhancer objectives—pixel reconstruction (ResPix) and replacement detection (RevDet).
• Demonstrate that CiM can achieve competitive transfer performance on ImageNet classification and ADE20K segmentation.

Proposed method

• Use a small trainable BEiT as generator with a frozen pre-trained image tokenizer (dVAE) to produce corrupted images by sampling from the BEiT output distribution.
• Feed corrupted images to an enhancer network that is trained to either reconstruct original pixels (ResPix) or detect which tokens were replaced (RevDet).
• Train generator and enhancer jointly but with gradients not flowing through the generator, allowing a stable, non-adversarial update.
• Optionally normalize pixel targets with a sliding-window scheme to improve ResPix performance.
• Demonstrate that both ViT and CNN enhancers can be pre-trained in a unified framework and transfer well to downstream tasks.

Experimental results

Research questions

• RQ1 Can CiM pre-train high-capacity ViT and CNN encoders using a non-Siamese, corrupted-image paradigm?
• RQ2 Does a trainable generator to corrupt inputs improve representation learning over fixed masking schemes?
• RQ3 Are generative (ResPix) and discriminative (RevDet) enhancer objectives both effective for downstream tasks?
• RQ4 How do CiM-pretrained ViT and CNN models perform on ImageNet classification and ADE20K segmentation relative to baselines?

Key findings

• CiM achieves competitive ImageNet end-to-end fine-tuning top-1 accuracy for ViT-S/16 and ViT-Base/16 with 300 pre-training epochs.
• CiM ViT-Base/16 with RevDet reaches 83.3% top-1 on ImageNet-1K in Table 1, comparable to BEiT and MAE baselines.
• CiM-ResPix and CiM-RevDet with ViT-Base/16 achieve 83.3% top-1 under the reported settings; RevDet achieves 83.6% in Table 1? (The exact phrasing from the table shows 83.3 and 83.3 for RevDet and ResPix under ViT-Base/16.)
• For ResNet-50, CiM-RevDet reaches 78.8% top-1 after 300 pre-training epochs, outperforming several Siamese-method baselines under certain fine-tuning regimes.
• On ADE20K semantic segmentation, CiM-ResPix and CiM-RevDet achieve 38.7 and 39.0 mIoU with ViT-Base/16 encoder, respectively, and 43.5–43.6 mIoU with ViT-Base/16 encoders in the presented setups.
• Ablations show random masking with ~50% ratio, 4–6 layers in small BEiT, sliding-window normalization for ResPix, and softmax sampling of generator tokens provide benefits.