Scalable Image Tokenization with Index Backpropagation Quantization
This addresses scalability issues in visual tokenization for researchers and practitioners in computer vision, representing a novel method for a known bottleneck.
The paper tackles the scalability problem in vector quantization (VQ) methods by proposing Index Backpropagation Quantization (IBQ), which enables joint optimization of codebook embeddings and visual encoders, achieving a large-scale codebook of 2^18 with high utilization and competitive results on ImageNet benchmarks.
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.