PaperlayerView exact revision on arXiv

Source-linked AI summary

StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training

Bao Tang, Jiahao Guo, Haoxiang Cao, Wenyu Liu, Changqian Yu, Kun Gai, Xinggang Wang

arXiv:2609.26774v1cs.CV

TL;DR

VQ training stability remains underexplored, with entangled Encoder–Decoder and Codebook objectives preventing reliable independent responsibility. StableVQ targets that problem with lightweight interventions, and the paper reports improved stability, utilization, and reconstruction quality across settings.

  • Problem

    Training stability remains underexplored, while entangled Encoder–Decoder and Codebook objectives keep the modules from reliably fulfilling their responsibilities independently.

    Source: §1 · p. 2, §4.1 · p. 4

  • Method

    StableVQ uses parameter-free interventions to restore separate responsibilities for the Encoder–Decoder and Codebook, including region-based targets for codes.

    Source: §1 · p. 2, §4.1 · p. 4, §4.3 · p. 5

  • Results

    StableVQ rapidly restores full codebook usage across diverse codebook-token distribution relationships, and combining its components recovers full utilization with strong reconstruction quality.

    Source: §5.1 · p. 8, §5.2 · p. 9

  • Takeaways & Limitations

    StableVQ frames full codebook utilization as a property that can be directly guaranteed through separation of concerns, providing a foundation for broader and more challenging scenarios.

    Source: §6 · p. 9

    Experiments focus on image tokenizers; usefulness of the same perspective in video, audio, multimodal representation learning, or compression remains prospective.

    Source: B Limitations · pp. 13–14

Abstract

from arXiv · show

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.

1 Introduction

Discrete visual tokenizers support modern generative vision systems, but VQ training stability remains underexplored despite improved codebook utilization from shared-projection methods. StableVQ addresses instability as an entanglement of module responsibilities and proposes lightweight interventions for stable training.

1 Introduction

  • VQ tokenizers map continuous image features to discrete tokens, and vocabulary expressiveness sets an upper bound on downstream generation quality.They support autoregressive transformers, masked generative models, and multimodal language models.Source: §1 · p. 1
  • Shared-projection methods improve codebook utilization, but training stability remains a critical, underexplored challenge.The methods reparameterize codes through a shared differentiable function so gradients can propagate across the code distribution.Source: §1 · pp. 1–2
  • Initialization sensitivity, prolonged low utilization, and abrupt utilization collapse expose fragile cooperation between the Encoder–Decoder and Codebook.The paper argues that entangled training prevents either module from reliably fulfilling its responsibility in isolation.Source: §1 · p. 2
  • StableVQ proposes lightweight interventions that let each module fulfill its responsibility independently, enabling principled and stable VQ training.The interventions are parameter-free and target stable training across diverse codebook sizes and initialization settings.Source: §1 · p. 2
  • StableVQ aims to realize VQ tokenizers’ expressive potential without heuristic initialization or complex projector design.The introduction reports that a single linear projection reaches state-of-the-art quality across diverse settings.Source: §1 · p. 2

2 Related Work

Related work develops VQ tokenizers through changes to latent structure, objectives, quantization stages, and architecture, while codebook collapse remains a central challenge. Shared-projection methods improve utilization, and StableVQ studies the remaining instability problem.

2 Related Work

  • VQ-VAE introduced mapping continuous encoder features to discrete code indices through nearest-neighbor lookup in a learned codebook [30].Later tokenizers pursued reconstruction quality and capacity through hierarchical latents [22], perceptual and adversarial objectives [6], residual or multi-stage quantization [13], and ViT-based architectures [31].Source: §2 · p. 2
  • Codebook collapse and low utilization persist, especially as codebook size or embedding dimension increases.Remedies include code reset or replacement [34], low-dimensional embeddings and normalization [31], relaxed or soft-assignment paths [21] [12] [25], and scalar or binary quantization [18] [32].Source: §2 · p. 2
  • Shared-projection methods such as VQ-STE++, SimVQ, and FVQ substantially improve utilization with little architectural overhead, while StableVQ addresses remaining training instability [11] [36] [2].StableVQ builds on this line of work.Source: §2 · p. 2

3 Background

The background describes VQ’s training objective and the gradient sparsity of independent code vectors. Shared projections spread gradient influence across codes, motivating distribution alignment as a central training concern.

3.1 Vector Quantization

  • The VQ training objective includes commitment and VQ losses, while the reconstruction gradient reaches the Encoder through the Straight-Through Estimator.Commitment loss keeps encoder outputs near assigned codes; VQ loss moves codebook entries toward the encoder output distribution.Source: §3.1 · pp. 2–3

3.2 Codebook with a Shared Projection

  • With independent code parameters, only selected codes receive nonzero VQ-loss gradients, a property termed gradient sparsity.As codebook size grows, the fraction of codes updated per step diminishes, increasing collapse risk.Source: §3.2 · p. 3
  • Shared projection reparameterizes code vectors through a differentiable function shared across all codes.Gradients from selected codes can then influence the full code distribution.Source: §3.2 · p. 3
  • Affine reparameterization, SimVQ’s learnable linear layer, and FVQ’s nonlinear projector are examples of shared-projection methods that improve codebook utilization [11] [36] [2].The approaches differ in how they transform the code vectors, and the passage notes that training stability remains unresolved.Source: §3.2 · p. 3

4 Method

StableVQ treats training failures as a consequence of Encoder–Decoder and Codebook entanglement, which prevents either module from reliably fulfilling its responsibility independently. It targets the Encoder’s gradient objective, Codebook distribution alignment, and distinct scheduling needs.

4.1 Failure Modes of VQ Training

  • Small code scale yields sparse targets and prolonged low utilization, oversized code scale leaves dead codes, and tracking lag can trigger positive-feedback collapse.The paper identifies these failures as arising when the Encoder and Codebook cannot independently fulfill their responsibilities.Source: §4.1 · pp. 3–4

4.2 Dynamic Straight-Through Estimator

  • Dynamic STE weights tokens by relative quantization error, suppressing unreliable distant-token gradients; with the Codebook frozen, it avoids standard STE’s diverging losses.For well-matched codes, the intervention reduces to standard STE; the pilot study reports stable losses under Dynamic STE.Source: §4.2 · pp. 4–5, Figure 2 · p. 5

4.3 Region VQ Loss

  • 12.5% utilization after 5,000 steps with standard VQ loss contrasts with full utilization by Step 500 under Region VQ Loss when the Encoder is frozen.Figure 3 visualizes the encoder outputs and codebook entries at Steps 0, 500, and 5000.Source: §4.3 · p. 6, Figure 3 · p. 6

4.4 Decoupled Schedule

  • Warmup-plus-annealing suits the Encoder–Decoder, while the Codebook benefits from a high constant rate without warmup; a constant Encoder rate drops quality in the pilot study.The pilot also reports that Codebook utilization benefits from a stable, sufficiently high constant learning rate.Source: §4.4 · pp. 6–7, Figure 4 · p. 7

5 Experiments

StableVQ pairs a simple linear projector with changes to improve reconstruction quality and codebook-usage recovery. Ablations show that full utilization across expansion and shrinkage settings requires Region VQ Loss combined with other proposed strategies.

5.1 Main Results

  • The ImageNet evaluation uses 256 × 256 images represented by 256 tokens and measures reconstruction, perceptual quality, and codebook utilization.The encoder–decoder downsamples by a factor of 16; utilization is the fraction of activated codes.Source: §5.1 · p. 7
  • StableVQ delivers superior reconstruction with a single linear projection, while SimVQ and FVQ reach full utilization with notable quality or design limitations.FVQ’s patch embedding constrains codebook sizes to perfect squares; StableVQ requires no projector-specific design or structural constraints.Source: §5.1 · p. 7, Table 1 · p. 8
  • Across diverse codebook-token distribution mismatches, StableVQ rapidly restores full usage, while SimVQ and FVQ recover slowly under limited conditions.UR-AUC measures usage recovery under mismatch settings.Source: §5.1 · pp. 7–8, Figure 5 · p. 8
  • StableVQ achieves competitive class-conditional ImageNet 256 × 256 generation results using its tokens.The evaluation follows the IBQ generation setup.Source: §5.1 · p. 8, Table 2 · p. 8

5.2 Ablation Studies

  • Under codebook expansion, Region VQ Loss combined with Dynamic STE or Decoupled Schedule recovers full utilization with strong reconstruction quality.Region VQ Loss alone causes NaN collapse; all three components together provide encoder stability, dense codebook supervision, and fast distribution tracking.Source: §5.2 · p. 9, Table 3 · p. 9
  • Under codebook shrinkage, Region VQ consistently reaches full utilization and improves reconstruction regardless of initialization, unlike FVQ.The test uses a broad codebook range and finds FVQ’s utilization strongly affected by initialization.Source: §5.2 · p. 9, Table 4 · p. 9

6 Conclusion

StableVQ attributes VQ training instability to entangled optimization objectives, not to vector quantization itself. It builds on shared-projection methods and targets a more reliable discrete representation substrate for token-based image generation.

6 Conclusion

  • StableVQ frames VQ instability as a consequence of entangled optimization objectives and restores separation of responsibilities between the Encoder–Decoder and Codebook.The conclusion presents full codebook utilization as a property that can be directly guaranteed rather than a fragile heuristic outcome.Source: §6 · p. 9

A Related Work

  • StableVQ builds on shared-projection codebooks while addressing optimization instability that remains after utilization improvements.Related work describes shared projection as one approach to low utilization, alongside alternatives such as reset strategies, soft assignments, and scalar quantization.Source: A Related Work · p. 13
  • Because tokenizers are upstream of autoregressive, masked, multi-scale, and hybrid generators, unstable VQ training can limit reconstruction fidelity, code utilization, and generation quality.StableVQ is presented as orthogonal to generator design and aims to provide a more reliable discrete representation substrate.Source: A Related Work · p. 13

B Limitations

StableVQ’s experiments focus on image tokenizers, while the authors identify broader training objectives and other vector-quantization domains as directions for future work.

B Limitations

  • A future direction is to balance reconstruction quality, semantic structure, and generation-friendliness after utilization and optimization stability cease to be the main bottlenecks.The paper positions StableVQ as a foundation for exploring these stronger objectives and downstream choices.Source: B Limitations · p. 13
  • The paper’s stated focus is making VQ tokenizer training stable and reliable.It does not present broader training objectives as part of StableVQ’s focus.Source: B Limitations · p. 13
  • The experiments focus on image tokenizers; applying the separation-of-concerns perspective to video, audio, multimodal representation learning, or compression is proposed as a possibility.The authors frame these domains as potential extensions, not as evaluated settings.Source: B Limitations · p. 14

C Additional Experimental Results

Additional experiments report reconstruction and generation results, then examine how initialization, projector design, computational overhead, and STE corrections affect training.

C Additional Experimental Results

D.1 Effect of Codebook Initialization

  • At 5k steps, Region VQ brings utilization near full across initialization scales, whereas small scales slow recovery and large scales pose dead-code risk.Under ℓ2 normalization, shrinking raw initialization cannot control the effective code distribution; overly small scales may prolong recovery.Source: §D.1 · p. 15, Figure 6 · p. 15

D.2 Effect of Shared Projector

  • At 5k steps, Region VQ yields strong utilization across projector structures and widens usable learning rates; linear and MLP projectors lag, while ViT utilization drops at excessive rates.The passage concludes that changing projector capacity alone is insufficient for stable, rapid utilization growth.Source: Figure 6 · p. 15, §D.2 · p. 16

D.3 Computational and Memory Overhead

  • StableVQ components are training-only: Decoupled Schedule adds no forward or backward operation, Dynamic STE timing differs within noise, and Region VQ cost falls as utilization rises.The table reports matched baseline and StableVQ profiles for 16K and 262K codebooks on the same machine.Source: §D.3 · p. 16, Table 8 · p. 16

D.4 Comparison with Prior STE Corrections

  • VQ-STE++ requires k-means pre-alignment and alternating inner and outer updates; Rotation Trick can amplify gradients for distant token–code pairs.VQ-STE++ also removes encoder commitment loss, while under ℓ2 normalization Rotation Trick’s rescaling loses its attenuation effect.Source: §D.4 · p. 16
  • Both Rotation Trick and Dynamic STE achieve full codebook utilization with Region VQ, while Dynamic STE delivers better reconstruction across metrics.In the controlled 15-epoch experiment, VQ-STE++ shows pronounced utilization collapse without both k-means pre-alignment and a norm bottleneck.Source: §D.4 · p. 17, Table 9 · p. 17

D.5 Comparison with Explicit Distribution-Alignment Methods

On a non-Gaussian distribution-fitting benchmark, Region VQ maintains high codebook utilization while keeping training time low relative to explicit distribution-alignment methods.

D.5 Comparison with Explicit Distribution-Alignment Methods

  • Region VQ propagates local targets from statistically supported active regions to persistently inactive codes, without assuming a parametric token distribution or requiring MMD’s pairwise kernel computation.Wasserstein VQ and MMD VQ instead use Gaussian moments and kernel statistics, respectively, for global distribution matching.Source: §D.5 · p. 17
  • The benchmark varies ζ to strengthen the target distribution’s non-Gaussian structure and measures utilization after 10K steps using 16,384 codes of dimension 8.Region VQ uses a linear shared projector; the comparison also measures per-step training time.Source: §D.5 · p. 17
  • At ζ = 4, Region VQ reaches 99.4% utilization, versus 34.8% for Wasserstein VQ and 75.6% for MMD VQ, and is 34.2× faster than MMD VQ.The table reports utilization at 10K steps and per-step training time.Source: §D.5 · p. 17, Table 10 · p. 17

E Robustness Test Illustration

The robustness test varies codebook initialization to probe recovery under codebook–token distribution mismatch; the accompanying algorithm describes Region VQ target assignment and loss computation.

E Robustness Test Illustration

  • The test varies Gaussian codebook initialization while holding data, architecture, optimizer, learning-rate schedule, and training budget fixed to create controlled mismatch settings.Small scales compact code vectors, while larger scales spread them relative to the token distribution.Source: E Robustness Test Illustration · pp. 17–18
  • Figure 5 heatmaps use rows for mismatch settings and brighter colors for higher codebook usage; robust recovery means high usage across most rows.Usage is measured within 65,536-token windows, so values are slightly lower than full-validation-set utilization.Source: E Robustness Test Illustration · p. 18
  • UR-AUC averages normalized areas under usage curves across mismatch settings, with higher values indicating faster and more consistent recovery.Table 1 error bars are sample means and standard deviations across three random-seed runs.Source: E Robustness Test Illustration · p. 18

F Region VQ Algorithm

  • In each Region VQ step, assignments identify active and persistently inactive codes; targets from reliable active codes are propagated to nearby inactive codes before codebook loss is computed.The procedure uses a FIFO queue to identify persistently inactive codes.Source: F Region VQ Algorithm · p. 18

G Experimental Details

The appendix gives the main experiment configurations and describes training settings for reconstruction and generation experiments.

G Experimental Details

H Visualization Results

The section provides qualitative comparisons of reconstruction at two codebook sizes and samples from downstream class-conditional image generation.

H Visualization Results

  • Qualitative reconstruction comparisons are provided for 16k- and 262k-code codebooks, alongside downstream class-conditional generation samples.Figures 7 and 8 compare reconstruction results; Figure 9 shows generation samples.Source: H Visualization Results · p. 19, Figure 7 · p. 19, Figure 8 · p. 20
  • The 16k-code reconstruction comparison labels the alternatives StableVQ-16k and FVQ-16k alongside the original.Source: Figure 7 · p. 19
  • The 262k-code reconstruction comparison labels the alternatives StableVQ-262k and FVQ-262k alongside the original.Source: Figure 8 · p. 20

Related papers

Loading 2609.26774v1…