Source-linked AI summary

An Image is Worth 32 Tokens for Reconstruction and Generation

Qihang Yu, Mark Weber, Xueqing Deng, Xiaohui Shen, Daniel Cremers, Liang-Chieh Chen

arXiv:2406.07550v1cs.CV

TL;DR

Conventional 2D image tokenizers retain a grid structure that limits compactness despite image redundancy. TiTok encodes images as 1D discrete latent sequences with a transformer-based encoder, quantizer, and decoder. It achieves competitive ImageNet generation quality with 8× to 64× fewer tokens and substantially faster training and inference, while its scope includes vanilla VQ modules and ViT-based implementations.

  • Problem

    2D image tokenizers retain fixed spatial grids, limiting their ability to exploit image redundancy and produce compact latent representations.

  • Method

    TiTok uses a ViT encoder, vector quantizer, and ViT decoder to tokenize images into 1D discrete latent sequences that can be decoded back into images.

  • Results

    TiTok represents images with 8 to 64 times fewer tokens while achieving competitive ImageNet FID and improving generation-model training and inference throughput.

  • Takeaways & Limitations

    Compact 1D tokenization offers a more efficient representation for image reconstruction and generation without sacrificing competitive benchmark quality.

  • Takeaways & Limitations

    The experiments mainly use vanilla VQ modules, while more advanced quantization methods are outside the paper’s focus and may further benefit TiTok.

Abstract

from arXiv · show

Recent advancements in generative models have highlighted the crucial role of image tokenization in the efficient synthesis of high-resolution images. Tokenization, which transforms images into latent representations, reduces computational demands compared to directly processing pixels and enhances the effectiveness and efficiency of the generation process. Prior methods, such as VQGAN, typically utilize 2D latent grids with fixed downsampling factors. However, these 2D tokenizations face challenges in managing the inherent redundancies present in images, where adjacent regions frequently display similarities. To overcome this issue, we introduce Transformer-based 1-Dimensional Tokenizer (TiTok), an innovative approach that tokenizes images into 1D latent sequences. TiTok provides a more compact latent representation, yielding substantially more efficient and effective representations than conventional techniques. For example, a 256 x 256 x 3 image can be reduced to just 32 discrete tokens, a significant reduction from the 256 or 1024 tokens obtained by prior methods. Despite its compact nature, TiTok achieves competitive performance to state-of-the-art approaches. Specifically, using the same generator framework, TiTok attains 1.97 gFID, outperforming MaskGIT baseline significantly by 4.21 at ImageNet 256 x 256 benchmark. The advantages of TiTok become even more significant when it comes to higher resolution. At ImageNet 512 x 512 benchmark, TiTok not only outperforms state-of-the-art diffusion model DiT-XL/2 (gFID 2.74 vs. 3.04), but also reduces the image tokens by 64x, leading to 410x faster generation process. Our best-performing variant can significantly surpasses DiT-XL/2 (gFID 2.13 vs. 3.04) while still generating high-quality samples 74x faster.

1 Introduction

TiTok addresses the redundancy and rigidity of 2D image tokenizers by encoding images as compact 1D discrete latent sequences. Across reconstruction and generation, this design supports fewer tokens, competitive quality, and faster processing.

  • Motivation: 2D tokenizers preserve a spatial grid that limits their ability to exploit redundancy between similar image regions.The paper asks whether this 2D structure is necessary and draws inspiration from image-understanding systems that use fixed-length 1D token sequences.
  • Method: TiTok combines a ViT encoder, vector quantizer, and ViT decoder to convert image patches into a 1D discrete sequence and reconstruct the image.Latent tokens are concatenated with flattened image patches before encoding, then quantized and decoded with mask tokens.
  • Findings: 32 tokens are sufficient for reasonable image reconstruction, while reconstruction improvements become marginal beyond 128 tokens.The experiments examine the interaction between latent-space size, tokenizer size, reconstruction fidelity, and generative quality.
  • Findings: Scaling tokenizer model size improves reconstruction and generation most when the latent sequence is limited to 32 or 64 tokens.This supports compact representations without treating token count alone as the only determinant of quality.
  • Findings: 1D tokenization removes fixed grid constraints, allowing tokens to reconstruct regions beyond fixed positions and capture more semantic-rich information.The paper identifies this flexibility particularly at compact latent sizes.
  • Generation performance: 8× to 64× smaller latent spaces provide competitive image-generation quality while accelerating training and inference, with up to 410× faster generation than diffusion models such as DiT.Figure 2 compares speed and quality on ImageNet 256 × 256 and 512 × 512 benchmarks, measuring sampling speed including de-tokenization on an A100 GPU.

2 Related Work

Prior work uses encoder–decoder tokenization for image compression, semantic encoders for image understanding, and sequential or masked-token models for generation. TiTok applies a related encoder–decoder pipeline while replacing the conventional 2D latent grid with a compact 1D sequence.

  • Image Tokenization: Autoencoder-based image tokenization compresses high-dimensional images into low-dimensional latent representations and decodes them back into images.
  • Tokenization for Image Understanding: Image-understanding systems commonly use feature encoders to produce highly semantic tokens for classification, detection, segmentation, captioning, and visual question answering.
  • Image Generation: Autoregressive image-generation transformers model patches step by step, requiring as many generation steps as tokens, such as 256 or 1024.
  • TiTok framework: TiTok reconstructs images with a ViT encoder, vector quantizer, and ViT decoder after representing them with a small set of 1D latent tokens.
  • Image Generation: MaskGIT-style bidirectional transformers predict multiple masked tokens per step, reducing the number of steps needed to generate a complete image.TiTok uses this non-autoregressive sampling scheme after tokenization.

3 Method

TiTok replaces the fixed 2D latent grid of conventional VQ tokenizers with a compact 1D discrete sequence encoded and decoded by Vision Transformers. Its two-stage training uses proxy codes and optional decoder fine-tuning to improve stability and reconstruction quality.

  • VQ background: A typical VQ tokenizer encodes an image into a 2D latent representation, quantizes embeddings to codebook entries, and decodes the quantized representation back into an image.This workflow preserves a fixed spatial correspondence between latent tokens and image patches.
  • VQ background: The fixed 2D grid limits compression because it constrains tokens to correspond one-to-one with image patches despite redundancies among adjacent regions.The same design also limits flexibility in selecting latent sizes.
  • TiTok architecture: TiTok patchifies an image, concatenates the patches with K learnable latent tokens, and retains only those latent tokens as a 1D representation.This decouples latent sequence length from image resolution and permits more flexible design choices.
  • TiTok architecture: Unlike flattening a 2D grid, TiTok removes the implicit grid correspondence, allowing tokens to reconstruct regions beyond fixed patches and capture more high-level information.The paper presents this distinction as the basis for exploring compact 1D tokenization and de-tokenization.
  • Image generation: TiTok can replace VQGAN in MaskGIT without other specific generator modifications, using masked-token prediction during training and iterative confidence-based sampling at inference.This evaluates the tokenizer within a standard image-generation pipeline.
  • Training: The two-stage training paradigm first learns from proxy codes, then optionally fine-tunes only the decoder while keeping the encoder and quantizer frozen.The authors report improved training stability and reconstructed image quality from this strategy.

4 Experimental Results

Experiments evaluate TiTok across reconstruction, semantic representation, generation, throughput, and ablations. The results show that compact 1D token sequences can preserve quality while improving efficiency, especially with larger tokenizer models and suitable training designs.

  • Preliminary Experiments: 32 latent tokens suffice for reasonable reconstruction, while gains from increasing token count become marginal beyond 128 tokens.TiTok-L with 32 tokens performs better than a 2D VQGAN using 256 tokens.
  • Preliminary Experiments: Larger TiTok models enable smaller latent representations without compromising reconstruction performance.TiTok-B with 64 tokens matches TiTok-S with 128, and TiTok-L with 32 matches TiTok-B with 64.
  • Preliminary Experiments: Compact latent representations increasingly produce semantically rich encoder features, as shown by improved linear-probing accuracy when token count decreases.
  • Preliminary Experiments: 12.8× faster training and 4.5× faster sampling result when using 32 rather than 256 latent tokens.The reported throughput changes are 2815.2 vs. 219.7 samples/s/gpu for training and 123.1 vs. 27.5 samples/s/gpu for sampling.
  • Main Results: 2.21 rFID for TiTok-L-32 matches MaskGIT’s VQGAN at 2.28 rFID while using 8× fewer latent tokens.On ImageNet 256 × 256, TiTok-L-32 also improves generation over the MaskGIT baseline under the same generator framework and sampling steps.
  • Main Results: 2.74 vs. 3.04 gFID gives TiTok-L-64 better ImageNet 512 × 512 generation performance than DiT-XL/2 while running 410× faster.TiTok-B-128 reaches 2.13 vs. 3.04 gFID and generates high-quality samples 74× faster.
  • Ablation Studies: 2.21 rFID is achieved after decoder fine-tuning, improving the TiTok-L-32 baseline of 6.59 rFID.A larger codebook improves rFID by 0.74 and longer training adds another 0.37 improvement before decoder fine-tuning.
  • Ablation Studies: 15.58 vs. 5.15 rFID shows that the 2D TiTok variant performs much worse than 1D TiTok-B-64 under the same compact representation.The ablation attributes the gap to fixed 2D correspondences limiting reconstruction in a compact latent space.

5 Conclusion

TiTok represents images as compact 1D latent sequences rather than 2D grids. It uses substantially fewer tokens while retaining competitive ImageNet generation quality and improving generative throughput.

  • TiTok represents images with 8 to 64 times fewer tokens than commonly used 2D tokenizers.
  • Compact 1D tokens improve generation training and inference throughput while achieving competitive FID on ImageNet benchmarks.

Appendix

The appendix provides supplementary implementation details, expanded experiments, visualizations, limitation discussion, broader-impacts discussion, and dataset licensing information.

  • Supplementary material includes TiTok training and testing hyperparameters and training costs.
  • It reports detailed preliminary-experiment results and broader comparisons using additional metrics and baselines.
  • It provides qualitative visualizations of the results.
  • The appendix contains discussions of limitations, broader impacts, and dataset licenses.

A Training and Testing Protocols

The protocols specify training data augmentation, optimization, model configurations, generation settings, and compute requirements for TiTok and its MaskGIT generator.

  • Tokenizer Training: Tokenizer training uses random cropping and flipping, batch size 256, 500k iterations, AdamW, learning rate 1 × 10−4, cosine decay, and weight decay 1 × 10−4.
  • Tokenizer Configurations: TiTok-L, TiTok-B, and TiTok-S denote ViT-L, ViT-B, and ViT-S encoder-decoder configurations.
  • Compute: Tokenizer training requires up to 64 A100-40G GPUs for 91 hours, while generator training uses 32 or 16 A100-40G GPUs depending on the variant.
  • Generator Training: Generation uses pre-tokenized images and MaskGIT masked-image modeling with a batch size of 2048 and 500k iterations.
  • Generation Settings: Inference uses variant-specific guidance scales and temperatures at ImageNet 256 × 256 and 512 × 512 resolutions.

B Detailed Results of Preliminary Experiments

This section provides detailed preliminary-experiment results in Table 4 and identifies the corresponding main-paper figure. It also presents the ImageNet-1K 256 × 256 generation-results table.

  • The preliminary-experiment results summarized in Table 4 correspond to Fig. 4 of the main paper.
  • Table 4 contains detailed results for the preliminary experiments reported in the main paper.
  • Table 5 reports ImageNet-1K 256 × 256 generation results evaluated with ADM.

C Additional Results

Additional ImageNet generation results show that TiTok variants achieve strong performance and high throughput, including against methods using the same generator and at higher resolution. The experiments also identify limited gains from classifier-free guidance for compact-token variants.

  • ImageNet generation: 3.15 gFID lets TiTok-L-32 outperform MaskGIT’s 6.18 with the same generator model at 256 × 256 resolution.TiTok-L-32 also uses 8 sampling steps versus 64 for MAGVIT-v2 and achieves 92.4× faster sampling than that baseline.
  • ImageNet generation: 3.64 vs. 7.32 gFID shows TiTok-L-64 outperforming MaskGIT at 512 × 512 while using 16× fewer tokens.
  • ImageNet generation: 2.74 vs. 3.04 gFID and 410× higher throughput distinguish TiTok-L-64 from guided DiT-XL/2 at 512 × 512.
  • Classifier-free guidance: TiTok-L-32 and TiTok-L-64 outperform most other methods without guidance but benefit relatively less from classifier-free guidance.The passage identifies adapting inference-time improvements to compact 1D tokens as a promising future direction.
  • Reconstruction: Limited latent representations preserve high-level layouts or salient objects, while larger models reconstruct more image details.

E Limitations

The paper validates its 1D tokenization formulation with a VQ tokenizer and Masked Transformer generator, leaving broader extensions for future work.

  • The study evaluates 1D tokenization only with a VQ tokenizer and a Masked Transformer generator.
  • Extensions to 1D-VAE, diffusion models, and video remain unexplored because of limited computational resources.
  • The proposed formulation is described as potentially generalizable beyond the evaluated tokenizer, generation framework, and image modality.

F Broader Impacts

The paper acknowledges that generative models can create social risks, including misinformation, harassment, and reinforcement of dataset-driven stereotypes. Its evaluation is limited to class-conditional generation over a fixed, public, controlled class set, while visualizations show generated and reconstructed images.

  • Risks: Generative models can enhance creativity but may also be misused for misinformation and harassment.
  • Risks: Dataset biases can reinforce negative social stereotypes and viewpoints in generative models.
  • Visualizations: Figure 5 visualizes generated images from TiTok variants with MaskGIT and displays corresponding ImageNet class names.
  • Scope boundary: The paper limits its study to class-conditional generation using a fixed, public, controlled set of classes.
  • Visualizations: Figure 7 compares reconstruction results across token counts and model sizes, showing retained salient regions and improved quality with scaling.

G Dataset Licenses

The paper describes the datasets and evaluation protocol used to train and test TiTok, including separate reconstruction and generation FID evaluations.

  • ImageNet-1K is used for the ImageNet generation benchmark and for training and evaluating TiTok.The dataset contains 1,000 object classes, 1,281,167 training images, 50,000 validation images, and 100,000 test images.
  • The training set supports tokenizer and generator training, while the validation set is used to compute reconstruction FID.
  • Generation results are evaluated with generation FID using pre-computed statistics and scripts from ADM.
Loading 2406.07550v1…