PaperlayerView exact revision on arXiv

Source-linked AI summary

MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis

Tianhong Li, Huiwen Chang, Shlok Kumar Mishra, Han Zhang, Dina Katabi, Dilip Krishnan

arXiv:2211.09117v2cs.CV

TL;DR

Computer-vision generation and representation models are typically trained independently, leaving their potential complementarity unused. MAGE unifies them with variable-masking masked image modeling over semantic tokens, achieving strong ImageNet-1K results including 9.10 FID and 78.9% linear-probing accuracy.

  • Problem

    Image generation and representation learning are typically trained independently, and computer vision lacks widely adopted models that unify them despite their complementary visual and semantic requirements.

  • Method

    MAGE uses variable masking ratios in a token-based masked image modeling framework, with semantic tokens as inputs and targets and an optional contrastive loss for representations.

  • Results

    MAGE ViT-L achieves 9.10 FID for class-unconditional ImageNet-1K generation and 78.9% top-1 accuracy for linear probing, with MAGE-C reaching 80.9% accuracy.

  • Takeaways & Limitations

    A single MAGE model achieves close to state-of-the-art results for generation and representation learning using the same data and training paradigm.

  • Takeaways & Limitations

    The authors report that quantized tokens may lead to subpar from-scratch and fine-tune performance, leaving this phenomenon for future investigation.

Abstract

from arXiv · show

Generative modeling and representation learning are two key tasks in computer vision. However, these models are typically trained independently, which ignores the potential for each task to help the other, and leads to training and model maintenance overheads. In this work, we propose MAsked Generative Encoder (MAGE), the first framework to unify SOTA image generation and self-supervised representation learning. Our key insight is that using variable masking ratios in masked image modeling pre-training can allow generative training (very high masking ratio) and representation learning (lower masking ratio) under the same training framework. Inspired by previous generative models, MAGE uses semantic tokens learned by a vector-quantized GAN at inputs and outputs, combining this with masking. We can further improve the representation by adding a contrastive loss to the encoder output. We extensively evaluate the generation and representation learning capabilities of MAGE. On ImageNet-1K, a single MAGE ViT-L model obtains 9.10 FID in the task of class-unconditional image generation and 78.9% top-1 accuracy for linear probing, achieving state-of-the-art performance in both image generation and representation learning. Code is available at https://github.com/LTH14/mage.

1. Introduction

MAGE unifies image generation and representation learning through variable-masking masked image modeling with semantic tokens. It achieves strong ImageNet-1K results across both tasks, including state-of-the-art generation and representation performance.

  • Motivation: Generation and representation learning are usually trained independently, despite requiring both visual and semantic understanding.The paper identifies a lack of widely adopted computer-vision models unifying both capabilities.
  • Key insight: MAGE views generation as reconstructing 100%-masked images and representation learning as encoding 0%-masked images.Variable masking ratios support both regimes under the same architecture, training scheme, and loss function.
  • Key insight: Variable masking alone is insufficient for high-quality generation because pixel-level reconstruction losses produce blurry outputs lacking fine details and textures.The paper uses MAE as a representative example of this limitation.
  • Method: MAGE uses semantic tokens as both inputs and reconstruction targets, enabling iterative generation, diverse outputs, and high-semantic-level representations without losing low-level details.Token-space prediction models a probability distribution over masked tokens rather than an average over possible masked pixels.
  • Results: 9.10 FID is achieved by ViT-L for class-unconditional ImageNet-1K generation, compared with 20.68 FID for MaskGIT.ViT-B achieves 11.11 FID in the same task.
  • Results: 78.9% top-1 accuracy is obtained by MAGE ViT-L for ImageNet-1K linear probing, rising to 80.9% with the added contrastive loss in MAGE-C.The paper reports state-of-the-art self-supervised representation learning performance for MAGE-C.

2. Related Work

Related work spans self-supervised representation learning, masked image modeling, and generative image synthesis. MAGE addresses the remaining gap between representation-focused MIM methods and generation-focused models by targeting both capabilities in one framework.

  • Self-supervised learning: Early unsupervised representation methods used pretext tasks such as jigsaw solving, missing-patch restoration, and rotation prediction.These representations substantially trailed supervised training.
  • Self-supervised learning: Contrastive learning became a competitive systematic approach, with methods including SimCLR and MoCo using negative samples and momentum encoders.The passage describes representative design differences among contrastive-learning variants.
  • Masked image modeling: Masked image modeling adapts NLP-style masking to vision, with methods such as BEiT recovering discrete visual tokens from masked inputs.PeCo improves tokenization by using MoCo-v3 as the perceptual model in VQGAN training.
  • Masked image modeling: MIM variants differ in prediction targets: MAE reconstructs pixels, CMAE adds contrastive loss, and MaskFeat and MVP predict teacher-generated features.These approaches use masked image modeling for self-supervised representation learning.
  • Research gap: Existing MIM methods favor downstream representations over reconstructed-image quality, whereas MAGE reports both fine-grained representations and high-visual-fidelity image generation from one model.The paper presents this as a first demonstration within its stated scope.
  • Generative models: Generative image models include GAN-based systems and two-stage approaches that tokenize images into a latent space before generation.GAN-based models can generate realistic images but may suffer training instability and mode collapse.
  • Research gap: Previous generative models lacked high-quality semantic representations, while MAGE reports strong performance on both class-unconditional generation and representation learning.The comparison supports the feasibility of a unified high-performance framework.

3. Method

MAGE unifies image generation and representation learning by masking semantic tokens at variable ratios and reconstructing the masked tokens with a ViT encoder-decoder. An optional contrastive objective further improves the separability of encoder representations.

  • Tokenization and masking: MAGE tokenizes images into semantic tokens, randomly masks them with a variable ratio, and uses a ViT encoder-decoder to reconstruct the masked tokens.The tokenizer operates before masking, while the decoder predicts the original token sequence from encoded unmasked information.
  • Masking strategy: A truncated-Gaussian masking strategy samples ratios from 0.5 to 1, then replaces masked tokens with a learnable mask token and randomly drops half the sequence from those masked tokens.The ratio is centered at 0.55 and bounded between 0.5 and 1.
  • Reconstructive training: The training objective reconstructs masked tokens using cross-entropy between ground-truth one-hot tokens and decoder predictions conditioned on unmasked tokens.Following MAE, the loss is optimized only on masked tokens.
  • Contrastive co-training: MAGE adds an optional SimCLR-like contrastive loss to globally pooled encoder features, using a two-layer MLP and InfoNCE to improve linear separability.Positive pairs come from two augmented views of the same image, while other batch samples provide negatives.
  • Training objective: The final objective combines reconstructive and contrastive losses, with λ = 0.1 balancing their scales.The reconstructive loss also acts as a regularizer, allowing strong performance without extensive contrastive augmentations.
  • Image generation: Images are generated through iterative decoding that begins with all tokens masked and progressively replaces selected predictions with sampled tokens.Tokens with higher predicted probability are more likely to be sampled during decoding.

4. Results

MAGE is evaluated as a unified framework for image generation and visual representation learning across ImageNet-1K and transfer benchmarks. It achieves strong generation, linear-probing, few-shot, transfer-learning, and fine-tuning results, while augmentation and tokenization choices affect performance.

  • Image Generation: MAGE generates realistic and diverse images with both strong and weak augmentations, although weak augmentation produces better reported evaluation scores.Strong augmentation crops a larger portion of the image, while weak augmentation crops a smaller portion.
  • Image Classification: 78.9% top-1 accuracy is achieved by MAGE ViT-L for ImageNet-1K linear probing, surpassing previous MIM methods and strong contrastive baselines.MAGE-C reaches 78.2% with ViT-B and 80.9% with ViT-L without multi-crop augmentation; MAGE exceeds MAE by 6.7% for ViT-B and 3.1% for ViT-L.
  • Image Classification: MAGE ViT-B outperforms MAE by a very large margin in few-shot ImageNet-1K learning, while MAGE-C ViT-L surpasses MSN using 13 images per class.The few-shot protocol freezes the pretrained model and trains a linear classifier using limited labeled samples.
  • Image Classification: MAGE outperforms SimCLR and MAE on 6 of 8 transfer-learning datasets under a 25-samples-per-class setting.The models are pretrained on ImageNet-1K and evaluated on other datasets.
  • Image Classification: Fine-tuning performance is at par with DINO and slightly under MoCo-v3, while still improving over supervised training from scratch by a similarly large margin.The authors attribute subpar from-scratch and fine-tuning performance to quantized tokens and leave further investigation for future work.
  • Analysis: Variable masking-ratio distributions affect both linear probing and class-unconditional generation, while fixed masking can yield generation FID above 50.The masking-ratio analysis varies the mode and standard deviation of a truncated Gaussian distribution.
  • Analysis: Using unquantized VQGAN features lowers reconstruction loss from 5.76 to 3.31 but lowers linear-probe accuracy from 74.7% to 49.5%.The result suggests that unquantized features make the pretraining task too easy.

5. Discussion

MAGE unifies image generation and representation learning through quantized tokens and variable masking ratios, achieving close to state-of-the-art results across both tasks with one model and training paradigm.

  • MAGE uses quantized tokens and variable masking ratios to adapt one masking-based framework to generation and representation learning.The paper reports evaluations on linear probing, few-shot transfer learning, and unconditional image generation.

A.1. Qualitative Results

MAGE produces qualitative results for class-unconditional generation and supports image editing applications including inpainting, outpainting, and uncropping.

  • MAGE reconstructs realistic, high-quality images for inpainting, outpainting, and uncropping.These applications are illustrated in Figure 7 and additional figures.
  • Additional class-unconditional generation results are shown under both strong and weak augmentation settings.Strong augmentation uses random crop and resize from 0.2 to 1, while weak augmentation uses 0.8 to 1; both include random flipping.

A.2. Quantitative Results

MAGE is applied to class-conditional generation and few-shot transfer learning, with competitive generation performance and consistent transfer advantages over MAE and SimCLR.

  • MAGE supports class-conditional image generation by pairing its pre-trained encoder with a class-conditional decoder.The decoder receives the class label alongside padded encoder features, while the encoder remains frozen during downstream training.
  • Comparable performance to state-of-the-art image generation methods is achieved on ImageNet-1K class-conditional generation.Only the 113M-parameter decoder is trained for this downstream task.
  • MAGE is consistently better than MAE and SimCLR on most few-shot transfer datasets across different numbers of samples per class.The additional evaluations use 1, 5, and 10 samples per class.

B. Ablation Studies

Ablations examine losses, tokens, augmentations, training duration, decoder design, and masking distributions, identifying choices that balance generation and representation learning.

  • MAE with GAN loss: A GAN loss makes MAE reconstructions more realistic but decreases ViT-L linear-probing performance by 11%.MAGE instead achieves strong performance on both generation and representation learning with one model.
  • Pad with [CLS] token: Using the image-specific [CLS] feature for decoder padding improves both class-unconditional generation and linear-probing results.Unlike MAE’s shared learnable mask token, this padding supplies global features from each image.
  • Augmentations: Strong augmentations favor representation learning, whereas weak augmentations favor generation quality.The two settings differ in the zoom-in range of random crop and resize: 0.2–1 for strong and 0.8–1 for weak.
  • Pre-training Epochs: 400 pre-training epochs already provide good performance on both tasks, while longer training consistently improves MAGE.The comparison uses 400, 800, and 1600 epochs.
  • Decoder Design: An 8-block decoder with 768-dimensional features gives ViT-B the best balance between computation cost and performance.The paper uses 8 blocks with 768 dimensions for ViT-B and 1024 dimensions for ViT-L.
  • Complement MIM with Contrastive Loss: Combining reconstructive and contrastive losses outperforms applying contrastive loss alone, whose performance is worse than the SimCLR baseline.The reconstructive loss helps prevent shortcut solutions associated with augmentation-sensitive cues such as color distribution.
  • Masking Design: Variable masking ratios are necessary for generation and also improve representation learning and linear-probe performance.For MAGE-C, smaller maximum masking ratios improve linear probing but worsen FID because generation requires exposure to relatively high masking ratios.

C. Implementation Details

MAGE uses VQGAN-derived semantic tokens with a ViT encoder-decoder, supporting both representation learning and iterative masked-token image generation. The implementation specifies standard augmentation, encoder feature pooling, and cosine-schedule decoding details.

  • Tokenizer and Detokenizer: 256×256 images are tokenized into 16×16 discrete tokens by a CNN-based VQGAN encoder and quantizer, then reconstructed by a detokenizer.The tokenizer uses residual blocks, average-pooling downsampling, and a 1024-entry codebook.
  • ViT architecture: The Transformer follows a standard ViT encoder-decoder architecture with separate learnable positional embeddings for the encoder and decoder.The latent sequence has length 256 plus one fake class token.
  • Downstream evaluation: Encoder outputs are average-pooled without the class token for linear probing, few-shot transfer learning, and fine-tuning.The resulting representation is passed to a linear classifier for classification tasks.
  • Pre-training: Pre-training uses random crop-and-resize augmentation with scale 0.2–1 and random horizontal flips as the default augmentation scheme.The detailed default configuration is provided in Table 15.
  • Generation: Image generation starts from fully masked tokens and iteratively fills them using prediction, sampling, and remasking steps.At each iteration, sampled tokens receive confidence scores, and the lowest-confidence locations are remasked according to a cosine schedule.
  • Generation: For class-unconditional generation, MAGE uses temperature τ = 6.0 and T = 20 decoding iterations.The implementation also specifies that the number of masked tokens follows a cosine masking schedule.
Loading 2211.09117v2…