StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training
Bao Tang, Jiahao Guo, Haoxiang Cao, Wenyu Liu, Changqian Yu, Kun Gai, Xinggang Wang
cs.CV
2026-09-23
HUST and KlingAI add Dynamic STE, Region VQ Loss, and a decoupled LR. On ImageNet 256 with a 262,144-code book at 120 epochs, rFID is 0.92 at 100% usage.
Discrete visual tokenizers map images to codebook indices so autoregressive and masked generators can treat pixels as tokens. Codebook collapse is the chronic failure: most entries never win a nearest-neighbor lookup. Shared-projection methods such as SimVQ and FVQ send every code through one differentiable map, so gradients reach unused entries, and utilization often prints as 100%.
Training still blows up. A bad init leaves utilization stuck low. Commitment loss spikes to NaN. A well-used codebook can collapse mid-run. The HUST and KlingAI Research team treat this as entangled optimization: the encoder-decoder and the codebook cannot each do their job unless the other happens to cooperate.
StableVQ sits on a shared projection and adds three parameter-free changes. The default projector is a single linear layer.
Dynamic STE is for the encoder. Quantization blocks gradients, so training copies decoder gradients through a straight-through estimator. When a token sits far from its assigned code, that copied gradient is a bad reconstruction direction and fights the commitment loss. Each token is weighted by how close it is relative to the best-matched token for the same code in the batch. The nearest token keeps weight 1; farther ones are down-weighted. When utilization is healthy the weights go to 1 and the rule becomes ordinary STE. No threshold hyperparameter.
Region VQ Loss is for the codebook. Standard VQ loss only supervises codes selected in the current step. Shared projection leaks a diffuse signal to dead codes, and that signal fades once a subset already covers the tokens. Region VQ keeps a FIFO window, splits recently active codes from persistently inactive ones, and lets each active code hand its target to nearby dead codes in proportion to how often it was selected. The codebook can track the encoder distribution on its own.
Decoupled Schedule splits the learning rates. The encoder uses warmup then annealing. The codebook uses a constant high rate, 1e-3 in the ablation, against a 1e-4 encoder peak. Tracking a moving distribution wants a large step from the first iteration, not a shared warmup.
ImageNet 256×256, VQGAN-style encoder, downsample 16, 256 tokens per image.
| Method | Projector | Codebook | Epochs | rFID↓ | Usage | UR-AUC↑ |
| SimVQ | Linear-1 | 16384×256 | 40 | 2.89 | 100% | 2.17 |
| FVQ | ViTBlock-2 | 16384×256 | 40 | 1.70 | 100% | 8.08 |
| StableVQ | Linear-1 | 16384×256 | 40 | 1.22 | 100% | 60.59 |
| StableVQ | Linear-1 | 262144×256 | 120 | 0.92 | 100% | — |
UR-AUC is the area under utilization-recovery curves when the codebook and token distributions start mismatched. Terminal usage is 100% for SimVQ, FVQ, and StableVQ. Recovery is not: 60.59 versus 8.08 and 2.17.
Under a tiny-range codebook init, turning all three knobs off leaves 49.13% usage and rFID 2.06. Region VQ alone NaNs. Dynamic STE alone drops usage to 1.27%. All three reach 100% usage and rFID 1.70. When l2 normalization spreads codes across a large region, FVQ with a ViT projector reaches 18.75% usage from uniform init; adding Region VQ restores 100%.
Class-conditional AR following the IBQ recipe: StableVQ plus IBQ-B (342M) scores FID 2.35 against 2.88 for the IBQ tokenizer; IBQ-L (649M) scores 2.18 against 2.45. IS does not follow. On IBQ-L it falls from 267.5 to 250.4.
With the codebook frozen, standard STE spikes commitment to NaN; Dynamic STE stays finite. With the encoder frozen, standard VQ loss is still around 12.5% usage at 5,000 steps; Region VQ fills the codebook by step 500.
Tokenizer training still piles on shared projections, k-means inits, deep projectors, and code resets. This paper turns stability into three switches you can ablate, and a linear projector is enough in the standard setting. UR-AUC is the number to copy into a training report. Terminal 100% usage hides dead codes and mid-run collapse.
This is incremental systems work, not a new quantization primitive. The claim is scoped to VQ that already uses a shared projection.
Almost every number is ImageNet 256, f=16, 256 tokens. No video, no multimodal, no higher resolution. FID improved while IBQ-L IS dropped from 267.5 to 250.4, so better reconstruction is not a full generation win. Region VQ alone NaNs; the three pieces have to travel together. The FIFO window length is not ablated in the main text. A discriminator is still in the recipe, and the paper never splits GAN loss from the VQ terms. The entangled-objectives story is supported by ablations, not a proof.