StableVQ: Turning Vector-Quantized Tokenizer Training Stability from a Fragile Outcome into a Guaranteed Property

The paper locates the root cause of VQ tokenizer training instability in the entanglement of Encoder–Decoder and Codebook training: neither subsystem can reliably fulfill its own responsibility in isolation, so the system works only when the two happen to cooperate — a condition that breaks down precisely when training is most stressed. Prior shared-projection methods (VQ-STE++, SimVQ, FVQ) addressed gradient sparsity and utilization mainly from the codebook side; this work shifts to a division of responsibilities, reframing instability as a failure of modular responsibility rather than a fundamental limitation of the quantization paradigm. A conceptual analysis and characterization of failure modes, organized around three cases of code-scale versus token-scale relationships (code scale below token scale, code scale above token scale, and scale divergence), supported by illustrative figures.

Dynamic STE targets the Encoder's learning objective: when tokens are assigned to distant codes, the reconstruction gradient passed through the straight-through estimator becomes an unreliable estimate, can conflict with the commitment loss, and may escalate to NaN; the method weights each token's gradient by its distance to the assigned code relative to the best-matched token for that code, attenuating farther tokens while preserving the full gradient for the best assignment. Unlike VQ-STE++'s alternating optimization or the Rotation Trick's rotation and rescaling, Dynamic STE only attenuates and never amplifies, and reduces to the standard STE when all tokens in a batch are well matched, requiring no threshold hyperparameter. A controlled experiment freezing the Codebook and training only the Encoder–Decoder shows standard STE producing severe commitment-loss spikes and reconstruction-loss oscillations that diverge to NaN, while Dynamic STE keeps losses stable; a further controlled 15-epoch comparison is run against the Rotation Trick and VQ-STE++.

Region VQ Loss targets the Codebook's learning objective: the standard VQ loss is asymmetric — every token receives an explicit target through the commitment loss, whereas only selected codes receive meaningful objectives; the method propagates the targets received by active codes to nearby inactive ones, using a FIFO queue to distinguish window-active from persistently inactive codes, so the Codebook can independently guarantee full tracking of the encoder output distribution without relying on encoder oscillations. Shared projection lets gradients reach all codes through the shared function, but the signal is indirect, undirected, and subject to attenuation; Region VQ shifts from point-wise to distribution-wise alignment, assumes no parametric form for the token distribution, and avoids the costly pairwise kernel computation of MMD VQ. In a frozen-Encoder, Codebook-only experiment, the standard VQ loss stagnates at around 12.5% utilization even after 5000 steps, whereas Region VQ Loss reaches full utilization by step 500 and maintains it; on a non-Gaussian mixture benchmark, Wasserstein VQ and MMD VQ fall to 34.8% and 75.6% utilization under strongly non-Gaussian targets while Region VQ maintains 99.4%, with training time close to Wasserstein VQ and faster than MMD VQ.

Decoupled Schedule holds that the Encoder–Decoder and the Codebook require distinct optimization dynamics: the former handles multi-objective reconstruction under discrete regularization and benefits from warmup-plus-annealing, while the latter must continuously track an evolving encoder distribution and benefits from a constant high learning rate with no warmup. Prior work observed that warmup-plus-annealing benefits training quality yet degrades codebook utilization, and FVQ compensated with a more expressive shared projector; this work attributes the tension to the two modules having inherently different objectives and therefore requiring independent schedules. Two controlled experiments show that, on the FVQ architecture, setting the Codebook to a constant learning rate incurs no performance degradation whereas applying a constant rate to the Encoder–Decoder causes a clear quality drop; on the SimVQ architecture, codebook utilization is not improved by complex schedules but benefits from a stable and sufficiently high constant learning rate.

The result targets discrete visual tokenizers built on shared-projection codebooks; experiments are conducted on ImageNet with a VQGAN-style encoder–decoder, covering diverse codebook sizes and initialization settings, with Region VQ additionally tested on a synthetic non-Gaussian mixture benchmark. All three components are training-only and leave tokenizer inference unchanged, so they can be layered onto existing shared-projection methods. The paper positions itself as providing a stable and reliable training foundation on which stronger training objectives and downstream modeling choices can be explored, and suggests the same perspective may extend to video, audio, multimodal representation learning, and compression.

The failure-mode analysis is organized around code-scale versus token-scale relationships, and readers may want to check how those cases map onto their own data and resolution; Dynamic STE's weights depend on within-batch relative quantization error, so its behavior under different batch compositions is worth observing; Region VQ's propagation quotas are automatically enlarged on the synthetic benchmark when recipient collisions are severe, and how often that mechanism triggers on real data remains to be understood; additionally, although the loaded text is the full paper, some table values appear in placeholder form, so specific numbers should be checked against the original.