StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training

StableVQ. Practical Guidelines for Stable Vector-Quantized Tokenizer Training. By Bao Tang and six others, from the Huazhong University of Science and Technology, Kling AI Research, and one other institution. Published on the twenty-second of September. No code repository is linked to it.

Vector Quantization, or VQ, sits at the heart of the discrete visual tokenizers that power modern autoregressive and masked image generation. Recent shared-projection codebook methods have made large gains in codebook utilization. Training stability, though, remains a critical challenge that has received little attention. Utilization asks how much of the codebook gets used. Stability asks whether training holds together at all, and the second can fail even when the first looks good.

The authors locate the root cause in entanglement between the encoder-decoder and the codebook during training. Neither module can reliably do its own job in isolation. So the system only works when the two happen to cooperate, and that is a fragile condition, one that breaks down precisely when training is under the most stress.

StableVQ revisits the proper learning objective of each module, so that each can be trained to fulfil its own role independently, and resolves the problems that arise when it is. The aim is to replace a fragile cooperation with two parts that each work on their own. If the encoder can do its job regardless of the codebook, and the codebook can do its job regardless of the encoder, then a bad moment for one no longer drags down the other.

It makes three changes. First, Dynamic STE corrects the instability in the encoder's learning objective. That lets the encoder robustly optimize the reconstruction space under discrete regularization, even when codebook utilization is low. Second, Region VQ Loss reconceives the codebook's learning objective, so that the codebook can independently guarantee full tracking of the encoder's output distribution, without relying on oscillations in the encoder to drive activation.

Third, Decoupled Schedule recognises that the encoder-decoder and the codebook have distinct responsibilities, and so need distinct optimization dynamics. Each is given its own learning rate schedule, to ensure robust behaviour at the level of the whole system. StableVQ is built on top of shared-projection codebooks, is lightweight, and adds no learnable parameters.

Experiments on ImageNet show consistent improvements in training stability, codebook utilization and reconstruction quality. Those improvements hold across diverse codebook sizes and initialization settings. Breadth is the relevant test for a stability method. A fix that worked at one codebook size, from one starting point, would say little about whether the underlying fragility had been removed.

The diagnosis shapes the cure. If instability comes from two modules that only work when they happen to cooperate, the fix is to give each one an objective, and a schedule, that it can satisfy on its own. Every one of the three changes is a way of cutting that dependency.