Regularized Vector Quantization for Tokenized Image Synthesis
This work addresses a fundamental bottleneck in unified generative modeling for image synthesis, offering an incremental improvement over existing vector quantization techniques.
The paper tackles the problem of deterministic and stochastic vector quantization in tokenized image synthesis, which suffer from codebook collapse, misalignment, low utilization, and perturbed reconstruction. The proposed regularized vector quantization framework, using prior distribution and stochastic mask regularization, outperforms prevailing methods across auto-regressive and diffusion models.
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.