Scalable Image Tokenization with Index Backpropagation Quantization

TL;DR

This paper tackles the scalability bottleneck of visual tokenizers by addressing codebook collapse through Index Backpropagation Quantization (IBQ), which allows gradients to flow to all codebook embeddings via a soft-to-hard categorical index. By jointly optimizing the encoder and the entire codebook, IBQ achieves high codebook utilization and enables training of extremely large codebooks (up to 2^18 entries) with strong reconstruction (rFID as low as 1.00) and competitive autoregressive generation when paired with vanilla transformers. Key contributions include introducing a double quantization loss, demonstrating scalable codebook growth across size and dimensionality, and validating IBQ with AR generation across model scales from hundreds of millions to billions of parameters. The results suggest IBQ's potential to advance discrete visual tokenization, enabling scalable, high-quality image reconstruction and generation for large-scale multimodal models.

Abstract

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.

Figures (9)

• Figure 1: Reconstruction and generation samples of IBQ. We show 1024 $\times$ 1024 reconstructed samples (top) and 256 $\times$ 256 generated samples (middle and bottom).
• Figure 3: Gradient flow of different VQ methods. VQGAN/VQVAE only update the selected codes in each backward process. IBQ updates all codes simultaneously by transferring the gradients of soft one-hot categorical distribution to hard one-hot index.
• Figure 4: Qualitative Reconstruction Comparison. We compare IBQ with LlamaGen and Open-MAGVIT2 tokenizer.
• Figure 5: Scaling up visual tokenizers (e.g., codebook size, code dimension and model size) improves visual soundness of reconstruction.
• Figure 6: Scaling up IBQ tokenizers enables better generation, especially with larger autoregressive models (e.g., 1B param.).