StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training
StableVQ provides practical guidelines for stable Vector-Quantized tokenizer training, addressing instability between encoder and codebook optimization.
StableVQ is a new method to improve the training stability of Vector Quantization (VQ) for discrete visual tokenizers used in autoregressive image generation. The approach identifies the root cause as the entanglement of Encoder-Decoder and Codebook training objectives. It proposes Dynamic STE, Region VQ Loss, and a Decoupled Schedule to allow each component to learn independently and robustly. Experiments on ImageNet demonstrate consistent improvements in stability, codebook utilization, and reconstruction quality across various configurations.
- StableVQ addresses training instability in Vector Quantization for visual tokenizers.
- It introduces Dynamic STE, Region VQ Loss, and Decoupled Schedule.
- The method stabilizes encoder and codebook training under discrete regularization.
- Experiments on ImageNet show improved codebook utilization and reconstruction quality.
Full article239 words · extracted from huggingface.co · click to collapse
Vector Quantization (VQ) is fundamental to discrete visual tokenizers that power modern autoregressive and masked image generation models. While recent shared-projection codebook methods have substantially advanced codebook utilization, training stability remains a critical and underexplored challenge. We argue that the root cause lies in the entanglement of the Encoder--Decoder and Codebook training: because neither module can reliably fulfill its own responsibility in isolation, the system can only function when the two subsystems happen to cooperate---a fragile condition that breaks down precisely when training is most stressed. We propose StableVQ, which revisits the proper learning objective of each module and resolves the problems that arise when each is trained to fulfill its own role independently. Concretely, (1) Dynamic STE corrects the instability in the Encoder's learning objective, enabling it to robustly optimize the reconstruction space under discrete regularization even when codebook utilization is low. (2) Region VQ Loss reconceives the Codebook's learning objective so that it can independently guarantee full tracking of the encoder output distribution, without relying on encoder oscillations to drive activation. (3) Decoupled Schedule recognizes that the distinct responsibilities of the Encoder--Decoder and the Codebook demand distinct optimization dynamics, and assigns each an independent learning rate schedule to ensure robust system-level behavior. Built on top of shared-projection codebooks, StableVQ is lightweight and introduces no learnable parameters. Experiments on ImageNet demonstrate consistent improvements in training stability, codebook utilization, and reconstruction quality across diverse codebook sizes and initialization settings.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2609.26774