摘要

现有的向量量化（Vector Quantization, VQ）方法在可扩展性方面面临挑战，这主要归因于训练过程中代码本（codebook）仅进行局部更新所带来的不稳定性。随着部分代码未被激活，代码本中的非激活代码与视觉特征之间的分布差距逐渐扩大，导致代码本容易发生坍缩（collapse）。为解决这一问题，我们提出了一种新型向量量化方法——索引反向传播量化（Index Backpropagation Quantization, IBQ），该方法实现了代码本嵌入与视觉编码器的联合优化。通过在编码特征与代码本之间的独热（one-hot）类别分布上应用直通估计器（straight-through estimator），所有代码均具备可微性，并与视觉编码器保持一致的潜在空间表示。IBQ 支持视觉标记化器的可扩展训练，首次实现了大规模代码本（$2^{18}$）、高维（256维）且高利用率的向量量化结构。在标准 ImageNet 基准测试上的实验表明，IBQ 具有出色的可扩展性与性能优势，在图像重建任务以及自回归视觉生成应用中均取得了具有竞争力的结果。相关代码与模型已开源，地址为：https://github.com/TencentARC/SEED-Voken。