Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models

TL;DR

This work tackles codebook collapse in vector-quantized VAEs by introducing Group-VQ, a group-wise optimization framework that partitions the codebook into multiple independent groups with within-group joint optimization, balancing codebook utilization and reconstruction quality. It further adds a training-free codebook resampling method enabling post-hoc adjustment of codebook size via a generative core and projector, allowing flexible scaling without retraining. Empirical results on image reconstruction with diverse backbones show Group-VQ achieving state-of-the-art metrics and revealing distinct, specialized distributions across groups; ablations guide group-number selection and projector design. Overall, Group-VQ provides a practical, scalable way to improve VQ-based tokenization while retaining adaptability for size changes post-training, with potential applicability to broader Joint VQ paradigms.

Abstract

Vector Quantized Variational Autoencoders (VQ-VAEs) leverage self-supervised learning through reconstruction tasks to represent continuous vectors using the closest vectors in a codebook. However, issues such as codebook collapse persist in the VQ model. To address these issues, existing approaches employ implicit static codebooks or jointly optimize the entire codebook, but these methods constrain the codebook's learning capability, leading to reduced reconstruction quality. In this paper, we propose Group-VQ, which performs group-wise optimization on the codebook. Each group is optimized independently, with joint optimization performed within groups. This approach improves the trade-off between codebook utilization and reconstruction performance. Additionally, we introduce a training-free codebook resampling method, allowing post-training adjustment of the codebook size. In image reconstruction experiments under various settings, Group-VQ demonstrates improved performance on reconstruction metrics. And the post-training codebook sampling method achieves the desired flexibility in adjusting the codebook size.

Figures (12)

• Figure 1: The differences among Vanilla VQ (a), Group-VQ (b), and Joint VQ (c) lie in their codebook update strategies: Vanilla VQ updates codes independently; Joint VQ optimizes the entire codebook jointly; and Group-VQ updates groups independently while optimizing jointly within each group.
• Figure 2: The codebook utilization rate and rFID for Joint VQ and Vanilla VQ at each epoch.
• Figure 3: Left: The reparameterization method of the codebook by Group-VQ. It partitions the codebook into multiple disjoint groups (sub-codebooks), each of which undergoes separate reparameterization. Right: Post-training sampling. Simply sampling a new codebook core and applying the trained projector yields new codes.
• Figure 4: Using resampling and self-expansion methods, the rFID and PSNR under different multiples of expanding the codebook size. The codebook size during training is 1024.
• Figure 5: Left: The average value and variance of the $l_2$-norm of vectors in each group (0 15) in the codebook; Right: The usage frequency of each group when tested on the ImageNet validation set.