PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Image Tokenization with Index Backpropagation Quantization

Fengyuan Shi, Zhuoyan Luo, Yixiao Ge, Yujiu Yang, Ying Shan, Limin Wang

arXiv:2412.02692v2cs.CVcs.AI

TL;DR

Existing VQ methods struggle to scale because partial codebook updates widen the gap between inactive codes and visual features. IBQ updates all code embeddings through a straight-through estimator on their categorical distribution. On ImageNet, it supports large, highly utilized codebooks with strong reconstruction results and competitive autoregressive generation.

  • Problem

    Partial updates in existing VQ methods widen the gap between visual features and inactive codes, contributing to codebook collapse and limiting tokenizer scalability.

    Source: §1 · p. 2, §3.1 · p. 4

  • Method

    IBQ passes gradients to all codebook embeddings through a straight-through estimator on the categorical distribution over visual features and codes.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    1.00 rFID with a 262,144-entry codebook and 256-dimensional codes, alongside competitive autoregressive generation with models ranging from 300M to 2.1B parameters.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The reported results support scaling visual tokenizers across codebook size, code dimension, and model size, including use in autoregressive image generation.

    Source: §1 · pp. 2–3

    Generation validation uses vanilla autoregressive models, excluding masked modeling and multiscale structures for simplicity.

    Source: §2.2 · p. 3

Abstract

from arXiv · show

Existing vector quantization (VQ) methods struggle with scalability, largely attributed to the instability of the codebook that undergoes partial updates during training. The codebook is prone to collapse as utilization decreases, due to the progressively widening distribution gap between non-activated codes and visual features. To solve the problem, we propose Index Backpropagation Quantization (IBQ), a new VQ method for the joint optimization of all codebook embeddings and the visual encoder. Applying a straight-through estimator on the one-hot categorical distribution between the encoded feature and codebook, all codes are differentiable and maintain a consistent latent space with the visual encoder. IBQ enables scalable training of visual tokenizers and, for the first time, achieves a large-scale codebook ($2^{18}$) with high dimension ($256$) and high utilization. Experiments on the standard ImageNet benchmark demonstrate the scalability and superiority of IBQ, achieving competitive results on reconstruction and the application of autoregressive visual generation. The code and models are available at https://github.com/TencentARC/SEED-Voken.

1. Introduction

The introduction frames scaling visual tokenizers as a response to quantization-related reconstruction loss, while existing VQ methods struggle to scale. IBQ updates the full codebook and is reported to support larger, highly utilized codebooks with strong reconstruction and generation results.

1. Introduction

  • VQ-based image tokenizers lose information during quantization relative to continuous VAEs [17] [31], and current visual tokenizers [37] [46] have not demonstrated scaling.The introduction presents larger codebooks and embedding dimensions as a way to mitigate quantization information loss.Source: §1 · p. 2
  • 68% to 0.002% codebook usage after one epoch illustrates collapse as partial updates widen the gap between inactive codes and visual features [11] [37].The introduction attributes the widening gap to partial updates of selected codes during backpropagation.Source: §1 · p. 2
  • IBQ globally updates the codebook using a straight-through estimator on the categorical distribution over visual features and all embeddings, making every code differentiable.The introduction says this keeps the codebook aligned with the visual encoder and reports about 96% codebook usage throughout training.Source: §1 · p. 2
  • 1.00 rFID with a 262,144-entry codebook and 256 dimensions accompanies autoregressive generation that outperforms competing methods.The reported generation models range from 300M to 2.1B parameters.Source: §1 · pp. 2–3

2. Related Work

Related work covers advances in vector quantization and approaches to tokenized visual generation. The paper positions IBQ as a scalable alternative and validates its tokenizer using vanilla autoregressive models.

2.1. Vector Quantization

  • VQ work has used hierarchical representations, adversarial and perceptual losses, residual or dynamic quantization, and large-codebook alternatives based on reduced dimensions or frozen codebooks.The frozen-codebook approach is described as introducing a bottleneck that constrains tokenizer capacity.Source: §2.1 · p. 3
  • IBQ is presented as contrasting with VQ methods that struggle to scale, showing consistent improvements as codebook size, code dimension, and model size increase.The related-work discussion says prior methods typically use small codebooks or low-dimensional embeddings, limiting representational capacity.Source: §2.1 · p. 3

2.2. Tokenized Visual Generation

  • Tokenized generation includes non-autoregressive masked-token prediction with iterative refinement and autoregressive next-token prediction; cited approaches also vary their architectures and tokenization designs.Examples include VQGAN with GPT2-medium, LlamaGen with Llama, VAR with next-scale prediction, and Open-MAGVIT2 with asymmetric token factorization.Source: §2.2 · p. 3
  • For generation validation, the paper uses vanilla autoregressive models, excluding masked modeling and multiscale structures for simplicity while leaving exploration of advanced models to future work.The text says IBQ is compatible with advanced generative models.Source: §2.2 · p. 3

3. Method

The method replaces selected-code updates with Index Backpropagation Quantization, which sends gradients through the categorical distribution to all codebook embeddings. It then applies the tokenizer to autoregressive generation by modeling flattened visual tokens with next-token prediction.

3.1. Preliminary: Vector Quantization

  • Standard VQ selects the nearest code and uses a straight-through estimator, while updating only selected codes can widen the gap to inactive codes and destabilize training.VQ maps encoder features to codebook representations for image reconstruction.Source: §3.1 · pp. 3–4

3.2. Index Backpropagation Quantization

  • IBQ adds a double quantization loss to pull selected code embeddings and encoded features toward each other, alongside reconstruction, quantization, perceptual, adversarial, and entropy losses.The entropy penalty encourages codebook utilization.Source: §3.2 · pp. 4–5
  • IBQ applies the straight-through estimator to the categorical distribution, sending gradients to all codebook embeddings instead of only selected codes.The method uses visual-feature/codebook dot products and softmax probabilities before transferring gradients to the hard one-hot index.Source: §3.2 · pp. 4–5

3.3. Vanilla Autoregressive Transformer

  • For generation, quantized visual tokens are flattened in raster-scan order, and an autoregressive transformer models their sequence with next-token prediction conditioned on context such as a class label.The described generation setup uses a vanilla autoregressive transformer architecture.Source: §3.3 · p. 5

4. Experiment

Experiments evaluate reconstruction and class-conditional generation on ImageNet. They find that IBQ scales across tokenizer and autoregressive-model sizes while maintaining strong results in both tasks.

4.1. Datasets and Metrics

  • Both tokenizers and autoregressive transformers are trained on 256 × 256 ImageNet; reconstruction is measured on the 50k validation set, with task-specific quality metrics.Generation uses the ADM evaluation protocol.Source: §4.1 · p. 5

4.2. Implementations Details

  • Reconstruction uses the VQGAN architecture and a default 16,384-code, 256-dimension, four-ResBlock tokenizer; generation tests vanilla Llama-based AR models from 300M to 2.1B.The reported evaluations use 256 × 256 ImageNet; reconstruction metrics include rFID, utilization, and LPIPS, while generation metrics include gFID, Inception Score, and Precision/Recall.Source: §4.2 · p. 5, §4.1 · p. 5

4.3. Main Results

  • IBQ reaches 1.00 rFID at 262,144 codes, surpasses Open-MAGVIT2; its 300M–2.1B vanilla AR models beat prior vanilla AR and DiT, matching AR variants.The reconstruction comparison is on 256 × 256 ImageNet 50k validation; the generation comparison is class-conditional ImageNet.Source: Table 1 · p. 6, §4.3 · p. 6, Table 7 · p. 8

4.4. Scaling Up IBQ Tokenizers

  • Scaling codebook size, code dimension, and ResBlocks improves reconstruction; larger codebooks and dimensions improve generation, especially with larger autoregressive models.The paper reports the reconstruction trends in its codebook-size, code-dimension, and model-size comparisons, and notes larger generation gains with larger AR models.Source: Table 3 · p. 6, Table 4 · p. 6, Table 5 · p. 6

4.5. Ablation Studies

  • Ablations raise codebook utilization from 5.3% to 98% and improve rFID from 3.98 to 1.67; double quantization encourages closer code-feature alignment, while larger models improve reconstruction.The baseline is the re-implemented VQGAN; the ablation text attributes more precise quantization to double quantization loss.Source: §4.5 · pp. 7–8

5. Conclusion

IBQ addresses codebook collapse linked to partial updates by jointly updating the codebook and visual encoder. On ImageNet, it reports 1.00 rFID for reconstruction and 2.05 gFID for generation.

5. Conclusion

  • IBQ identifies partial updates as a source of widening distribution gaps and codebook collapse, then updates all codes with the visual encoder.It applies a straight-through estimator to the categorical distribution over visual features and all codebook embeddings.Source: §5 · p. 8
  • On ImageNet, IBQ reports 1.00 rFID for reconstruction and 2.05 gFID for generation.Source: §5 · p. 8

A. Autoregressive Model Configurations

The appendix describes IBQ autoregressive model configurations spanning 300M to 2.1B parameters, scaled according to VAR's rules.

A. Autoregressive Model Configurations

B. Comparison with Soft Vector Quantization

Soft VQ trains with weighted-average code embeddings and uses hard quantization at inference, but it trails IBQ in reconstruction and codebook usage. The authors report training corruption and low codebook usage for Soft VQ.

B. Comparison with Soft Vector Quantization

  • Soft VQ trains with a weighted average of code embeddings and a cosine-decayed temperature from 0.9 to 1e-6, then uses highest-probability hard quantization at inference.Source: B. Comparison with Soft Vector… · p. 9
  • Soft VQ trails IBQ in reconstruction quality and codebook usage, and its training corrupts within fewer than 10 epochs.The table caption notes a significant rFID drop when hard quantization is used for inference.Source: B. Comparison with Soft Vector… · p. 9
  • Soft VQ's codebook visualization shows mistakenly clustered codes and 2.5% usage, while IBQ combines hard quantization with index backpropagation.The authors speculate that weighted averaging smooths codebook representations, making embeddings similar and less informative.Source: B. Comparison with Soft Vector… · p. 9, Figure 7 · p. 9

C. Training Costs

IBQ improves utilization with modest reported training overhead at codebook sizes up to 16,384, and trains successfully at 262,144 entries where VQGAN fails.

C. Training Costs

  • At 262,144 entries, VQGAN fails to train while IBQ reaches 84% utilization.Source: C. Training Costs · p. 9
  • At codebook size 16,384, IBQ adds 0.2 GB memory and 19 minutes of training time while increasing utilization from 5.3% to 96%.Source: C. Training Costs · p. 9
  • The authors characterize IBQ's additional training cost as marginal compared with VQGAN for codebooks up to 16,384 entries.Source: C. Training Costs · p. 9

D. Pretraining Tokenizer

The section evaluates IBQ tokenizers pretrained on large-scale general-domain and high-quality datasets for text-conditional image generation. It reports state-of-the-art performance against concurrent methods in zero-shot reconstruction comparisons.

D. Pretraining Tokenizer

  • IBQ is pretrained on large-scale general-domain and high-quality datasets, with training steps of approximately 800,000.The datasets are intended to serve text-conditional image generation.Source: D. Pretraining Tokenizer · pp. 9–10
  • IBQ achieves state-of-the-art performance against concurrent methods including Cosmos and Show-o in zero-shot reconstruction comparisons.Table 11 covers ImageNet 50k validation and MS-COCO val2017.Source: Table 11 · p. 10, D. Pretraining Tokenizer · p. 10
  • Some residual-tokenization efforts achieve better results, but are excluded because those techniques are orthogonal and compatible with IBQ.Source: D. Pretraining Tokenizer · p. 10

E. Improving Face Reconstrution

The section reports that scaling the codebook and incorporating face data can improve face reconstruction. It also points to further qualitative reconstruction and generation examples.

E. Improving Face Reconstrution

F. Additional Visualizations

Loading 2412.02692v2…