画像を離散表現に量子化することは、統一された生成モデリングにおける基本的な問題でした。優勢なアプローチは、最も一致するトークンを選択することによる決定論的方法、または予測された分布からのサンプリングによる確率論的方法のいずれかで、離散表現を学習します。ただし、決定論的量子化は、深刻なコードブックの崩壊と推論段階との不整合に悩まされますが、確率論的量子化は、コードブックの利用率が低く、再構成目標が乱れることに悩まされます。この論文では、2 つの観点から正則化を適用することにより、上記の問題を効果的に軽減できる正則化されたベクトル量子化フレームワークを提示します。 1 つ目は、コードブックの崩壊と低いコードブックの使用率を回避するために、以前のトークン分布と予測されたトークン分布との間の不一致を測定する事前分布正則化です。 2 つ目は、量子化中に確率論を導入して、推論段階のミスアライメントと摂動のない再構成目的の間で適切なバランスを取る、確率論的マスク正則化です。さらに、摂動された再構成の目的をさらに軽減するための較正されたメトリックとして機能する確率的コントラスト損失を設計します。広範な実験により、提案された量子化フレームワークが、自己回帰モデルや拡散モデルを含むさまざまな生成モデル全体で一貫して一般的なベクトル量子化方法よりも優れていることが示されています。
Quantizing images into discrete representations has been a fundamental problem in unified generative modeling. Predominant approaches learn the discrete representation either in a deterministic manner by selecting the best-matching token or in a stochastic manner by sampling from a predicted distribution. However, deterministic quantization suffers from severe codebook collapse and misalignment with inference stage while stochastic quantization suffers from low codebook utilization and perturbed reconstruction objective. This paper presents a regularized vector quantization framework that allows to mitigate above issues effectively by applying regularization from two perspectives. The first is a prior distribution regularization which measures the discrepancy between a prior token distribution and the predicted token distribution to avoid codebook collapse and low codebook utilization. The second is a stochastic mask regularization that introduces stochasticity during quantization to strike a good balance between inference stage misalignment and unperturbed reconstruction objective. In addition, we design a probabilistic contrastive loss which serves as a calibrated metric to further mitigate the perturbed reconstruction objective. Extensive experiments show that the proposed quantization framework outperforms prevailing vector quantization methods consistently across different generative models including auto-regressive models and diffusion models.