PaperlayerView exact revision on arXiv

Source-linked AI summary

Autoregressive Image Generation without Vector Quantization

Tianhong Li, Yonglong Tian, He Li, Mingyang Deng, Kaiming He

arXiv:2406.11838v3cs.CV

TL;DR

Autoregressive image generation is commonly tied to vector-quantized tokens, but the paper argues that only per-token distribution modeling is essential. It introduces Diffusion Loss for continuous-valued tokens and reports fast, strong results across AR and MAR models, including FID < 2.0 at under 0.3 second per ImageNet image.

  • Problem

    The paper asks whether autoregressive models must use vector-quantized representations when modeling continuous-valued image tokens.

  • Method

    Diffusion Loss uses a denoising diffusion procedure conditioned on autoregressively predicted vectors to model continuous per-token probability distributions.

  • Results

    A MAR model achieves FID < 2.0 on ImageNet 256×256 while generating each image in < 0.3 second, and the best model approaches 1.55 FID.

  • Takeaways & Limitations

    The method removes the need for vector-quantized tokenizers and supports autoregressive image generation with continuous-valued representations.

  • Takeaways & Limitations

    The image generation system can produce noticeable artifacts, with a visual-quality gap relative to commercial models trained on massive data.

Abstract

from arXiv · show

Conventional wisdom holds that autoregressive models for image generation are typically accompanied by vector-quantized tokens. We observe that while a discrete-valued space can facilitate representing a categorical distribution, it is not a necessity for autoregressive modeling. In this work, we propose to model the per-token probability distribution using a diffusion procedure, which allows us to apply autoregressive models in a continuous-valued space. Rather than using categorical cross-entropy loss, we define a Diffusion Loss function to model the per-token probability. This approach eliminates the need for discrete-valued tokenizers. We evaluate its effectiveness across a wide range of cases, including standard autoregressive models and generalized masked autoregressive (MAR) variants. By removing vector quantization, our image generator achieves strong results while enjoying the speed advantage of sequence modeling. We hope this work will motivate the use of autoregressive generation in other continuous-valued domains and applications. Code is available at: https://github.com/LTH14/mar.

1 Introduction

The paper argues that autoregressive image generation does not require vector-quantized tokens, proposing diffusion-based modeling of per-token distributions in continuous space. Diffusion Loss supports AR and MAR models, removes discrete tokenizers, and achieves fast, strong ImageNet results.

  • Autoregressive modeling only requires a per-token probability distribution and sampler, so discrete-valued representations are not conceptually necessary.The paper contrasts categorical modeling for discrete tokens with alternative distribution models for continuous tokens.
  • Diffusion Loss models each continuous token's distribution by conditioning a denoising network on an autoregressively predicted vector z.The denoising network is trained jointly with the autoregressive model using continuous-valued inputs and targets.
  • The approach eliminates discrete-valued tokenizers, allowing autoregressive models to use higher-quality non-quantized tokenizers.The paper notes that vector-quantized tokenizers are difficult to train, gradient-sensitive, and often reconstruct less accurately than continuous-valued alternatives.
  • Masked generative models can be unified with standard autoregressive models as masked autoregressive models compatible with Diffusion Loss.MAR models predict multiple output tokens simultaneously in randomized order while conditioning on known tokens.
  • A MAR model with Diffusion Loss generates ImageNet 256×256 images in < 0.3 second per image with FID < 2.0, while the best model approaches 1.55 FID.The experiments cover both standard AR and MAR models and report consistently improved generation quality with sequence-model speed.

2 Related Work

Prior image-generation sequence models commonly use pixels or vector-quantized discrete tokens, while related continuous-token approaches use alternative distribution models. This work is distinguished by applying diffusion to model arbitrary per-token distributions for image generation.

  • Earlier autoregressive image models operated on pixel sequences using RNNs, CNNs, or Transformers, while later approaches modeled images as discrete-valued tokens.Both autoregressive and masked generative models have operated in discrete-valued token spaces.
  • GIVT and this work both study continuous-valued tokens, but GIVT represents token distributions with a fixed number of Gaussian mixtures.The predefined mixture count can limit the types of distributions GIVT represents.
  • Diffusion-based decoders have also been used for visual representation learning, where diverse image generation is not the primary goal.DiffMAE and DARL apply denoising diffusion decoders in representation-learning settings rather than image generation.
  • Diffusion Policy provides a conceptual connection by modeling action distributions as denoising processes, but sample diversity is less central in robotics than in image generation.In the paper's analogy, generating an image token corresponds to taking an action.

3 Method

The method replaces discrete categorical token modeling with diffusion-based per-token distribution modeling, enabling autoregressive generation over continuous-valued tokens. It also unifies standard and masked autoregression, using bidirectional attention to predict one or multiple tokens while preserving next-token prediction.

  • 3 Method: Continuous-valued tokenizers can support autoregressive image generation when the per-token distribution is modeled with a loss function and sampler rather than a categorical representation.The approach uses Diffusion Loss to make sequence models compatible with continuous-valued tokens.
  • 3.2 Diffusion Loss: Diffusion Loss models p(x|z) by denoising noisy continuous tokens conditioned on an autoregressive vector z, with reverse diffusion sampling producing token samples.The denoising network is trained jointly with the autoregressive model, and gradients from the loss update the network producing z.
  • 3.3 Diffusion Loss for Autoregressive Models: Autoregressive generation factorizes sequence modeling into conditioning vectors from previous tokens and per-token conditional distributions that Diffusion Loss can train.The same formulation permits continuous-valued next tokens instead of requiring categorical outputs.
  • 3.4 Unifying Autoregressive and Masked Generative Models: Masked autoregressive models predict randomized sets of tokens simultaneously, generalizing next-token prediction while remaining autoregressive.MAR uses random-order token sets whose union covers the sequence.
  • 3.4 Unifying Autoregressive and Masked Generative Models: Bidirectional attention lets known tokens communicate fully and enables one or more unknown tokens to be generated per step, trading away causal key-value caching for fewer generation steps.The paper states that full attention can improve quality and the speed/accuracy trade-off.

4 Implementation

The implementation combines Transformer-based AR/MAR sequence models with a cosine-schedule diffusion process and a small denoising MLP. Experiments use publicly available VQ-16 and KL-16 tokenizers and distinct causal or masked-attention inference procedures.

  • 4 Implementation: The implementation is presented as a general framework rather than a concept limited to the specific implementation details described in this section.Further specifics are deferred to Appendix B.
  • 4 Implementation: The diffusion process uses a 1000-step cosine noise schedule during training and fewer resampled steps at inference, with the denoising network predicting noise.The default inference schedule uses 100 steps, and classifier-free guidance is supported.
  • 4 Implementation: The denoising head is a small residual MLP whose conditioning vector z comes from the AR/MAR model and is combined with the diffusion timestep embedding.The default MLP uses three residual blocks with width 1024.
  • 4 Implementation: Experiments use LDM's publicly available VQ-16 tokenizer and its non-quantized KL-16 counterpart.VQ-16 is a VQ-GAN, whereas KL-16 is regularized by KL divergence without vector quantization.
  • 4 Implementation: The Transformer follows the ViT implementation, adding positional embeddings and class tokens before processing the token sequence.The default Large model has 32 blocks, width 1024, and approximately 400M parameters.
  • 4 Implementation: The causal baseline uses shifted inputs, triangular attention masking, temperature sampling, and key-value caching for efficient inference.MAR instead uses bidirectional attention, random masking during training, and a cosine schedule that reduces masking from 1.0 to 0 over 64 steps at inference.

5 Experiments

Experiments on ImageNet 256×256 show that Diffusion Loss improves continuous-token autoregressive generation across AR and MAR variants, supports diverse tokenizers, and preserves a favorable speed–quality trade-off.

  • Diffusion Loss properties: Diffusion Loss consistently outperforms cross-entropy across four AR/MAR variants, reducing MAR FID by approximately 50%-60%.The comparison uses discrete VQ-16 and continuous KL-16 tokenizers under matched settings.
  • Tokenizer flexibility: 7.82 FID with Diffusion Loss outperforms 8.79 FID from cross-entropy using the same VQ tokenizer.The method treats the continuous latent before the VQ layer as the modeled tokens.
  • Tokenizer flexibility: 2.05 FID for KL-8 is close to 1.98 FID for KL-16 despite grouping 2×2 tokens to match sequence length.Diffusion Loss also supports tokenizers with mismatched strides.
  • Denoising MLP: A 2M denoising MLP is competitive, while the default 21M MLP adds approximately 5% of MAR-L parameters and diffusion sampling costs approximately 10% of runtime.Increasing MLP width improves generation quality with negligible extra inference cost in the reported implementation.
  • Sampling: 100 diffusion sampling steps are sufficient for strong generation quality, while temperature τ influences both FID and IS.Temperature controls the diffusion sampler at inference time and plays a role analogous to temperature in cross-entropy autoregression.
  • Speed and accuracy: MAR improves the speed–accuracy trade-off over AR and achieves less than 0.3 second per image with FID below 2.0.The comparison notes that AR uses an efficient kv-cache, while the method’s diffusion process operates on a small MLP rather than the full Transformer.
  • Benchmarking: The best reported model reaches 1.55 FID, while the no-CFG result is 2.35 FID on ImageNet 256×256.The experiments compare multiple model sizes and train for 800 epochs.

6 Discussion and Conclusion

The paper concludes that autoregressive image generation can model token interdependence autoregressively while modeling each token’s distribution with diffusion, without vector-quantized representations. Its ImageNet results motivate exploring continuous-valued sequence models beyond language and image generation.

  • Discussion and Conclusion: Diffusion Loss separates autoregressive modeling of token interdependence from diffusion modeling of each token’s distribution.This differs from latent diffusion models that use diffusion to model the joint distribution of all tokens.
  • Discussion and Conclusion: The strong image-generation results indicate that autoregressive models need not be constrained by vector-quantized representations.The conclusion frames this as an opportunity for sequence models with continuous-valued representations in other domains.

A Limitations and Broader Impacts

The paper identifies artifacts, dependence on existing tokenizers, limited ImageNet-centered validation, and broader risks from dataset bias and misuse.

  • Limitations: The image generation system can produce noticeable artifacts, particularly on controlled academic datasets such as ImageNet.The paper notes a visual-quality gap between research-driven ImageNet models and commercial models trained on massive datasets.
  • Limitations: System quality can be limited by the quality of existing pre-trained tokenizers.Developing better tokenizers is outside the paper’s scope.
  • Limitations: The method was primarily tested on ImageNet because of limited computational resources.Further validation is needed for scalability and robustness in more diverse, real-world scenarios.
  • Broader Impacts: The system may reflect training-data bias and may be misused to generate disinformation.The paper presents these as broader-impact concerns requiring further consideration.

B Additional Implementation Details

The implementation trains a denoising network jointly with the autoregressive model and uses diffusion sampling, guidance, and specified optimization configurations.

  • Classifier-free guidance: Classifier-free guidance replaces class conditions with a dummy token for 10% of training samples and combines conditional and unconditional noise predictions at inference.The guidance strength is controlled by the scale ω and a scheduled inference procedure.
  • Training: Models use AdamW for 400 epochs with batch size 2048, learning rate 8e-4, and a 100-epoch linear warmup for Diffusion Loss models.Cross-entropy counterparts use a cosine learning-rate schedule.
  • Diffusion Loss: Diffusion Loss uses a SimpleMLP denoiser and GaussianDiffusion forward and backward processes conditioned on the AR/MAR output z.The pseudo-code defines the denoiser, diffusion process, and loss over condition z and ground-truth token x.
  • Diffusion Loss: Training samples noise and timesteps, forms x_t from x with q_sample, and predicts noise from x_t, the timestep, and z.The predicted noise is used in the denoising objective.
  • Sampling: Inference initializes x with noise and iterates timesteps in reverse using the p_sample reverse-diffusion process.The sampling procedure takes the conditioning vector z and initial noise as inputs.
  • Compute Resources: Training mainly uses 16 servers with 8 V100 GPUs each, and a 400-epoch MAR-L run takes approximately 2.6 days.The reported comparison is 4.6 days for DiT-XL/2 and 9.5 days for LDM-4 under the stated setup.

C Comparison between MAR and MAGE

The comparison uses MAR inference with MAGE’s confidence-sampling order and reports a 512×512 ImageNet system-level setup with CFG scale 4.0.

  • MAR and MAGE: MAR’s comparison with MAGE uses MAGE’s strategy of determining token order on the fly through confidence sampling across the spatial domain.The entries use tokenizers provided by the LDM codebase.
  • System-level comparison: The ImageNet 512×512 system-level comparison sets MAR’s CFG scale to 4.0 and otherwise follows the MAR-L configuration.These settings are specified in the Table 6 caption.

D.1 Autoregressive Image Generation in Pixel Space

MAR+DiffLoss can model RGB pixels directly without an image tokenizer, using grouped pixels as tokens, but higher-resolution pixel modeling is more computationally expensive.

  • Pixel-space generation: FID 2.93 was achieved by MAR-L+DiffLoss on ImageNet 64×64 when grouping every 4×4 pixels into one token.This experiment demonstrates direct RGB pixel-space modeling without an image tokenizer.
  • Pixel-space generation: Direct pixel-space modeling is significantly more computationally expensive than using an image tokenizer.At higher resolutions, it may require much longer autoregressive sequences or a substantially larger Diffusion Loss network.

D.2 ImageNet 512×512

At 512×512 ImageNet resolution, the method achieves competitive results against leading systems, while the L2 baseline performs disastrously.

  • FID 2.74 without CFG and 1.73 with CFG on ImageNet at 512×512 demonstrate competitive performance with previous systems.The experiment uses the KL-16 tokenizer and a 32×32 sequence length.
  • The 512×512 evaluation uses MAR-L settings, while limited resources prevented training the larger MAR-H configuration.
  • FID >100 for the MAR model trained with L2 loss shows that directly predicting continuous-valued tokens is a disastrous baseline.With L2 loss, predictions are deterministic for any given sequence order, limiting randomness and sample diversity.
Loading 2406.11838v3…