Scalable Image Tokenization with Index Backpropagation Quantization

2024-12-03

AI Summary

Existing vector quantization (VQ) methods suffer from codebook instability and poor scalability: certain codewords remain inactive for extended periods, causing their embeddings to drift from the visual feature distribution and leading to codebook collapse. To address this, we propose Index Backpropagation Quantization (IBQ), the first VQ framework enabling end-to-end differentiable joint optimization of the codebook embeddings and visual encoder. IBQ leverages the straight-through estimator to construct a one-shot differentiable categorical distribution, ensuring full codeword differentiability and uniform activation throughout training. This design supports stable training with ultra-large codebooks (up to $2^{18}$ entries) and high-dimensional representations (256-D). On ImageNet, IBQ significantly improves image reconstruction fidelity and autoregressive generation performance, achieving state-of-the-art scalability among VQ-based tokenization methods.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Querying, indexing, and retrieval in Web-related graphs

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.

Problem

• Addresses scalability issues in vector quantization methods
• Proposes Index Backpropagation Quantization for joint optimization
• Enables large-scale codebook with high utilization and dimension

Innovation

• Index Backpropagation Quantization for scalable VQ
• Joint optimization of codebook and visual encoder
• Large-scale codebook with high utilization and dimension