Stable training method improves image token encoding and reconstruction
StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training
Computer Vision and Pattern Recognition
Summary
Training image compressors that use discrete building blocks can be unstable, causing poor performance. The authors found this instability comes from the encoder and codebook modules not working well independently during training. They designed StableVQ, which fixes this by giving each part its own clear training goal and learning schedule. This method helps the system train more stably and perform better when recreating images.
What this means in practice
- •For image generation engineers: Train more stable and efficient discrete image encoders for improved generative model tokenization.
- •For compression software developers: Build image compression tools that use discrete vector quantization more reliably under varied conditions.
Authors
Bao Tang, Jiahao Guo, Haoxiang Cao, Wenyu Liu, Changqian Yu, Kun Gai, Xinggang Wang
Abstract
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.