PaperlayerView exact revision on arXiv

Source-linked AI summary

ADDP: Learning General Representations for Image Recognition and Generation with Alternating Denoising Diffusion Process

Changyao Tian, Chenxin Tao, Jifeng Dai, Hao Li, Ziheng Li, Lewei Lu, Xiaogang Wang, Hongsheng Li, Gao Huang, Xizhou Zhu

arXiv:2306.05423v2cs.CV

TL;DR

Recognition and generation have developed separately, and prior unified approaches have remained weaker on recognition while operating mainly in VQ-token space. ADDP alternates pixel decoding and VQ-token generation within one diffusion-based representation-learning framework, achieving competitive generation, classification, detection, and segmentation performance. Its main scope boundary is reliance on a pre-trained VQ encoder-decoder, which may constrain generation diversity.

  • Problem

    Prior unified methods mainly emphasize generation, remain inferior on recognition, and use VQ-token inputs despite recognition methods benefiting from raw pixels.

  • Method

    ADDP alternates decoding pixels from previous VQ tokens with generating new VQ tokens from decoded pixels while progressively masking VQ tokens during diffusion.

  • Results

    ADDP achieves competitive performance on unconditional generation, ImageNet classification, COCO detection, and ADE20k segmentation, including a 1.5 FID improvement over the previous generation state of the art.

  • Takeaways & Limitations

    ADDP demonstrates general representations applicable to both image generation and dense recognition tasks.

  • Takeaways & Limitations

    ADDP relies on a pre-trained VQ encoder-decoder, which may constrain generation diversity; future work includes continuous diffusion and higher resolutions.

Abstract

from arXiv · show

Image recognition and generation have long been developed independently of each other. With the recent trend towards general-purpose representation learning, the development of general representations for both recognition and generation tasks is also promoted. However, preliminary attempts mainly focus on generation performance, but are still inferior on recognition tasks. These methods are modeled in the vector-quantized (VQ) space, whereas leading recognition methods use pixels as inputs. Our key insights are twofold: (1) pixels as inputs are crucial for recognition tasks; (2) VQ tokens as reconstruction targets are beneficial for generation tasks. These observations motivate us to propose an Alternating Denoising Diffusion Process (ADDP) that integrates these two spaces within a single representation learning framework. In each denoising step, our method first decodes pixels from previous VQ tokens, then generates new VQ tokens from the decoded pixels. The diffusion process gradually masks out a portion of VQ tokens to construct the training samples. The learned representations can be used to generate diverse high-fidelity images and also demonstrate excellent transfer performance on recognition tasks. Extensive experiments show that our method achieves competitive performance on unconditional generation, ImageNet classification, COCO detection, and ADE20k segmentation. Importantly, our method represents the first successful development of general representations applicable to both generation and dense recognition tasks. Code is released at \url{https://github.com/ChangyaoTian/ADDP}.

1 INTRODUCTION

Recognition and generation have developed separately, while unified methods have not matched recognition needs, especially for dense tasks. ADDP bridges pixel and VQ-token spaces through alternating denoising to support both capabilities.

  • Motivation: Unified representation learning remains challenging because prior approaches emphasize generation yet underperform on recognition, particularly dense recognition tasks.MIM-based unified methods such as ViT-VQGAN and MAGE show strong generation but remain limited to classification rather than dense recognition.
  • Motivation: Pixel inputs preserve spatially sensitive information better than VQ tokens, while VQ tokens are beneficial reconstruction targets for generation.Under matched training conditions, pixel inputs outperform VQ-token inputs on canonical recognition tasks, with a larger gap for dense recognition.
  • Method: ADDP addresses the central question of associating pixel and token spaces so models perceive raw pixels while generating latent visual tokens.Its alternating process decodes pixels from previous VQ tokens and then generates new VQ tokens from those decoded pixels.
  • Method: ADDP gradually masks VQ tokens during diffusion, uses an off-the-shelf VQ decoder for token-to-pixel decoding, and fine-tunes its pixel-input encoder for recognition.The learned encoder is applied to downstream recognition datasets after pretraining.
  • Results: ADDP achieves competitive performance across unconditional generation, ImageNet classification, COCO detection, and ADE20k segmentation, supporting general representations for generation and dense recognition.The paper identifies ADDP as the first approach to develop general representations applicable to both generation and dense recognition tasks.

2 RELATED WORK

Prior generative and representation-learning methods often operate in only one space or preserve recognition quality without directly supporting image generation. ADDP alternates refinement between raw pixels and latent representations to unify both objectives.

  • Deep Generative Models for Image Generation: Earlier generative models directly decode pixels, while VQ-VAE-style approaches use discrete latent representations to address the difficulty of generating raw pixels.Diffusion models generate images by iteratively refining results, including approaches operating directly in raw-pixel space or latent space.
  • Generative Modeling for Unifying Representation Learning and Image Generation: ADDP alternately refines raw pixels and latent representations, unlike prior methods that denoise entirely in continuous pixel space or discrete VQ-token space.This alternating design is intended to learn unified representations for both recognition and generation.
  • Generative Modeling for Unifying Representation Learning and Image Generation: During ADDP inference, masked tokens initialize the process, each step decodes a noisy image and predicts reliable and unreliable tokens, and the final output is the refined image x0.The alternating process is illustrated as a sequence from masked tokens through noisy images and token updates to a synthesized image.
  • Generative Pre-training for Image Representation Learning: Generative pre-training methods such as MIM learn expressive recognition representations but do not preserve reconstructed-image quality sufficiently for direct image generation.These methods focus on representation expressiveness for recognition rather than maintaining image reconstruction quality.
  • Generative Modeling for Unifying Representation Learning and Image Generation: Previous unified methods have recognition performance below state-of-the-art recognition methods because they learn entirely in raw-pixel space or entirely in latent space.ADDP instead exploits both raw-pixel and latent spaces to obtain competitive performance on recognition and generation tasks.

3 METHOD

ADDP alternates between decoding pixels from VQ tokens and predicting new VQ tokens, using masked-token diffusion to learn representations for both generation and recognition.

  • Alternating Denoising: ADDP alternates token-to-pixel decoding with pixel-to-token generation, allowing denoising to use both raw pixels and VQ tokens.Each step decodes an image from reliable and unreliable tokens, then predicts the next token distributions from that image.
  • Learning the Denoising Process: The pixel-to-token network masks unreliable regions in its encoded image representation and predicts reliable and unreliable VQ tokens with learnable encoder-decoder subnetworks.The network accepts full images, supports CNN or ViT encoders, and its encoder representations can transfer to recognition and generation tasks.
  • Diffusion Process: The diffusion process encodes images into VQ tokens and progressively masks token regions through a Markov chain until all tokens are masked.Reliable tokens are retained at each step, while complementary unreliable tokens are sampled for decoding and are not carried forward.
  • Diffusion Process: A token predictor estimates unreliable tokens from reliable tokens because the ideal conditional distribution q(z0|zt) is generally intractable.The predictor is used during training to approximate q(z0|zt) through q(¯zt−1|zt), while the learned model predicts unreliable-token distributions from decoded noisy images.
  • Learning the Denoising Process: Training minimizes an ELBO while optimizing only masked-token regions, using a pretrained VQ tokenizer and q(¯zt−1|zt) as the practical training target.The VQ reconstruction term is omitted during training because the VQ tokenizer is pretrained; decoded noisy images provide the model inputs.
  • Learning the Denoising Process: WeightedSum maps unreliable-token distributions to weighted codebook embeddings before VQ decoding, producing high-quality images and improving model performance empirically.Sampling, ArgMax, and WeightedSum are evaluated as candidate mappings for unreliable tokens.

4 EXPERIMENTS

ADDP is evaluated for unconditional generation, image classification, detection, segmentation, and design choices through ablations. It achieves strong joint performance while revealing effects of unreliable-token conditioning, mapping, token input, prediction target, and masking schedule.

  • 4.1 IMPLEMENTATION DETAILS: ADDP’s ViT-L implementation uses an off-the-shelf VQGAN tokenizer, MAGE-Large token predictor, and an eight-block Transformer decoder trained on ImageNet-1k.The reported comparison pre-trains the ViT-L backbone for 800 epochs; generated images and implementation details are provided in the appendix.
  • 4.2 MAIN RESULTS: 1.5 FID better than the previous state of the art demonstrates ADDP’s effectiveness for unconditional image generation.The comparison uses ImageNet-1k 256×256 validation images and reports FID and IS for unconditional generation.
  • 4.2 MAIN RESULTS: 2 points higher than MAGE on ImageNet fine-tuning shows ADDP improves recognition while retaining competitive generation performance.ADDP is comparable to recognition-specialized methods, and the paper reports transfer evaluation on classification, detection, and segmentation.
  • 4.3 ABLATION STUDY: Using reliable tokens z_t rather than z_t−1 to condition unreliable-token prediction improves performance, while WeightedSum mapping outperforms direct Sampling.The ablations evaluate generation metrics and ImageNet fine-tuning accuracy; feeding tokens directly can improve generation but rapidly degrade classification.
  • 4.3 ABLATION STUDY: Predicting q(¯z_t−1|z_t) outperforms estimating q(z_0|z_t), and the masking schedule is also evaluated for its effect on generation.The ablation settings report FID, IS, and ImageNet-1k fine-tuning accuracy using 20 inference steps.

5 CONCLUSIONS

The paper concludes that ADDP bridges pixel and VQ-token spaces to support both image generation and recognition. It reports broad applicability while identifying dependence on a pre-trained VQ encoder-decoder as a limitation.

  • 5 CONCLUSIONS: ADDP is presented as a general representation learning framework applicable to image generation, recognition, and dense recognition tasks.The conclusion states that it is the first demonstrated general representation supporting both generation and dense recognition.
  • 5 CONCLUSIONS: The framework combines pixel inputs for recognition with VQ tokens as reconstruction targets for generation and optimizes the network using the ELBO.This design bridges the two representational spaces within one alternating denoising diffusion process.
  • 5 CONCLUSIONS: ADDP currently relies on a pre-trained VQ Encoder-Decoder, which may constrain generation diversity.The paper identifies continuous diffusion and higher-resolution scaling as future directions.
  • 5 CONCLUSIONS: The authors release source code and checkpoints and provide algorithmic, implementation, and theoretical details in the appendix and main text.The reproducibility statement points to pseudocode, implementation sections, derivations, and the released repository.

A.1 PRE-TRAINING

The training algorithm samples a diffusion step and masking states, decodes noisy pixels from partially known VQ tokens, and trains the encoder-decoder to predict masked tokens.

  • A.1 PRE-TRAINING: Each training iteration samples a diffusion timestep, two random masks, and an image, then encodes the image into VQ tokens.The masked token states z_t and z_t+1 are formed from the encoded z_0.
  • A.1 PRE-TRAINING: The token predictor estimates unreliable tokens, while the VQ decoder combines reliable and predicted tokens to produce a noisy image x_t.The decoded image is then processed by the encoder-decoder network.
  • A.1 PRE-TRAINING: The encoder-decoder predicts masked-token distributions, and training minimizes cross-entropy together with the diffusion objective before taking a gradient step.The algorithm repeats these operations until convergence.
  • A.1 PRE-TRAINING: Unconditional generation starts from decoder predictions for fully masked tokens and iteratively alternates image decoding with token prediction until returning x_0.At each reverse step, the current image is encoded, masked token predictions are decoded, and the next image is reconstructed.

B DERIVATION FOR ALTERNATING DENOISING DIFFUSION PROCESS

The derivation expresses ADDP as an alternating process that decodes images from VQ tokens and predicts token states from decoded images. Its ELBO reduces to token-prediction objectives because the VQ decoder is deterministic.

  • Diffusion Process: The diffusion process masks VQ tokens over T + 1 steps, ending with all tokens masked, and distinguishes reliable tokens z_0:T from unreliable tokens ¯z_0:T.The notation shifts timestep subscripts for simplicity.
  • Alternating Denoising Process: The alternating process represents decoded images x_0:T between VQ-token states, with VQ decoding and the learned ADDP model forming successive transitions.The derivation treats the decoded image sequence as the intermediate visual states during denoising.
  • Evidence Lower Bound: The ELBO decomposes the training objective into diffusion, token-prediction, VQ-encoder, and VQ-decoder terms, with the VQ-decoder mapping deterministic.The pre-trained VQ tokenizer means the VQ training term is omitted from ADDP optimization.
  • Diffusion Process: The unreliable-token forward process conditions q(¯z_t−1|z_t) on reliable tokens z_t rather than z_t−1.This choice is explicitly identified as the process used for unreliable-token prediction.
  • Optimizing the Evidence Lower Bound: Because q(z_0|z_t+1) is generally intractable, ADDP uses q(¯z_t|z_t+1) as a practical approximation and computes the loss with cross entropy.The derivation identifies q(z_0|z_t+1) and q(¯z_t|z_t+1) as optimization targets before selecting the latter in practice.

C.1 RESULTS OF VIT-B

With ViT-B, ADDP delivers strong generation and recognition performance across unconditional generation, classification, and dense prediction tasks.

  • C.1 RESULTS OF VIT-B: ∼2 FID improvement over previous SOTA demonstrates ADDP’s strong unconditional generation capacity.
  • C.1 RESULTS OF VIT-B: ADDP achieves competitive performance across unconditional generation, ImageNet classification, COCO detection, and ADE20k segmentation with a ViT-B backbone.The reported evaluation uses the MAGE Base token predictor and 1600 training epochs.
  • C.1 RESULTS OF VIT-B: ADDP’s classification performance is comparable to methods specifically designed for recognition tasks.However, linear probing performs poorly, likely because training uses noisy synthetic images.
  • C.1 RESULTS OF VIT-B: ADDP achieves comparable performance to recognition-specialized methods on object detection and semantic segmentation, supporting its use for dense prediction.

C.2 RESULTS OF RESNET-50

ADDP remains effective with a ResNet50 image encoder, combining competitive generation with strong ImageNet classification transfer.

  • C.2 RESULTS OF RESNET-50: ADDP’s ResNet50 unconditional generation performance is comparable to methods specially designed for generation tasks.The authors identify this as the first use of ResNet as the image encoder for image generation.
  • C.2 RESULTS OF RESNET-50: ADDP’s ResNet50 finetuning performance on ImageNet-1k outperforms previous supervised and self-supervised methods using the same backbone.

C.3 ROBUSTNESS EVALUATION

ADDP shows robustness comparable to a leading robust representation method, while inference-token selection and sampling settings materially affect generation quality and diversity.

  • C.3 ROBUSTNESS EVALUATION: ADDP achieves on-par robustness with SiameseIM across ImageNet validation-set variants.The authors speculate that training with noisy synthetic images may enhance robustness.
  • C.3 ROBUSTNESS EVALUATION: The default reliable-token sampling setting yields the lowest FID, although slightly reducing top-p and τ improves IS.The explored factors are nucleus-sampling truncation and Gumbel-noise temperature.
  • C.3 ROBUSTNESS EVALUATION: Disregarding tokens with excessively low confidence can enhance synthesized-image quality.
  • C.3 ROBUSTNESS EVALUATION: ADDP achieves better FID, IS, and recall than ADM with classifier guidance and MAGE on unconditional generation.ADM without classifier guidance has high recall but significantly worse FID and IS than ADDP.

D.1 COMPARISON OF INPUTS FOR RECOGNITION TASKS

The evaluation compares ADDP across recognition settings and uses a shared transfer protocol spanning classification, detection, and segmentation, with multiple encoder architectures.

  • D.1 COMPARISON OF INPUTS FOR RECOGNITION TASKS: ADDP is evaluated on ImageNet classification, COCO detection, and ADE20k segmentation, using 256 × 256 inputs except additional 1024 × 1024 COCO detection experiments.
  • D.1 COMPARISON OF INPUTS FOR RECOGNITION TASKS: The framework uses ViT-B, ViT-L, and ResNet50 encoders, while the token predictor defaults to the MAGE ViT-Base model.
  • D.1 COMPARISON OF INPUTS FOR RECOGNITION TASKS: ADDP pretraining uses 100 denoising steps with cosine-derived mask ratios sampled from a truncated normal distribution between 0.5 and 1.0.
  • D.1 COMPARISON OF INPUTS FOR RECOGNITION TASKS: ADDP differs from MAGE by being raw-pixel based while using predicted token distributions as part of its training objective, and it performs better on both generation and recognition, especially dense recognition.
  • D.1 COMPARISON OF INPUTS FOR RECOGNITION TASKS: For transfer, pretrained encoders receive task-specific heads under a setting mainly following MAE, including Mask R-CNN for detection and UperNet for segmentation.

E VISUALIZATION

The visualizations show ADDP’s generated images, intermediate token refinement, synthetic training inputs, and applications to image completion. They emphasize how token reliability and the WeightedSum mapping affect visual quality and downstream performance.

  • Unconditional and Intermediate Generation: ADDP produces unconditional ImageNet-1k samples and progressive intermediate generations under a linear masking schedule with T = 256.The intermediate visualization distinguishes reliable and unreliable token regions and shows their corresponding token-space representations.
  • Synthetic Training Images: WeightedSum synthesizes higher-quality training images than the alternative mapping strategies and achieves better recognition and generation performance.The comparison is qualitative in the visualizations, while the associated passage points to Table 4 for the performance results.
  • Pre-training Distributions: Pre-training samples mask ratios and timesteps from discrete distributions visualized in Figure 12.The visualization reports the distributions used for pre-training rather than a performance comparison.
  • Image Inpainting and Outpainting: ADDP performs image inpainting and outpainting without further finetuning by initializing masked regions, encoding the image into VQ tokens, and continuing generation.The timestep and mask ratio are computed from the masked-area ratio before generation proceeds.
  • Intermediate Generated Results: Reliable tokens sharpen generated images and add fine-grained details, while unreliable tokens refine coarse spatial contours and preserve spatial consistency.The token embeddings are visualized by mapping their top three PCA components to RGB values at each intermediate step.
Loading 2306.05423v2…