StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training作者Bao Tang, Jiahao Guo, Haoxiang Cao, Wenyu Liu, Changqian Yu, Kun Gai, Xinggang Wang核心发表机构Huazhong University of Science and Technology、KlingAI Research、South China Normal University论文链接arXiv:2609.26774v1发布于arXiv 预印本cs.CV|—|—|—|—|—|—|—|—|| | | | 53.11 | 49.13% | 2.06 | 0.2297 | 21.65 | 0.5730 || ✓ | | | $$200 (NaN) | — | — | — | — | — || | ✓ | | $$0.1 | 1.27% | 6.68 | 0.3100 | 19.89 | 0.5026 || | | ✓ | $$0.1 | 83.77% | 1.71 | 0.2209 | 21.79 | 0.5864 || | ✓ | ✓ | $$0.1 | 21.66% | 2.07 | 0.2354 | 21.36 | 0.5665 || ✓ | ✓ | | $$0.1 | 100% | 1.72 | 0.2200 | 21.84 | 0.5854 || ✓ | | ✓ | $$0.1 | 100% | 1.75 | 0.2206 | 21.85 | 0.5860 || ✓ | ✓ | ✓ | $$0.1 | 100% | 1.70 | 0.2208 | 21.81 | 0.5879 |这张表格是理解三个组件分工的核心证据。第一行基线表现出典型的 commitment loss 尖峰53.11与中等利用率49.13%。Region VQ Loss 单独使用会导致 NaN collapsepeak commit $$200原因是 Codebook 在共享的 warmup 调度下更新太慢——这恰好说明只有正确的码本目标还不够优化动态也必须匹配。Dynamic STE 单独使用把 peak commitment loss 压到 $$0.1却把利用率压到仅 1.27%、rFID 恶化到 6.68这直接证明了传统训练在相当程度上是依赖不稳定的 Encoder 振荡来激活码本的一旦把这种振荡压住码本就没有别的激活来源。Decoupled Schedule 单独使用把利用率提高到 83.77%、rFID 改善到 1.71但在标准 VQ loss 下仍达不到 full utilization。Region VQ Loss 与二者中任一结合Region VQ Dyn. STE或 Region VQ De. Sch.即可恢复 100% 利用率并保持强重建质量三者同时使用最为完整联合保证 Encoder 稳定性、稠密的 Codebook 监督与快速的分布追踪。消融二Codebook shrinkage对应失败模式中的 code scale≫ \gg≫token scale。这一情形常出现在空间被ℓ 2 \ell_2ℓ2​归一化约束时归一化后uniform 或 gaussian 初始化会把码广泛散布在整个空间而 Encoder 输出反映自然图像统计、集中在更小区域于是自然形成 codebook shrinkage。实验使用 codebook dimension 为 1024、更强的 ViTBlock shared projector遵循 FVQ 配置ProjectorRegion VQInitUtil.rFIDPSNRSSIMLPIPSViTBlock-2uniform18.75%2.435520.97890.56770.2414ViTBlock-2gaussian62.5%1.953021.37740.58360.2273ViTBlock-2✓uniform100%1.817621.54970.58960.2210ViTBlock-2✓gaussian100%1.896621.46030.59310.2255FVQ 在两种初始化下均未达到 full utilization且利用率强烈受初始化方式影响uniform 仅 18.75%gaussian 62.5%加入 Region VQ 后无论 uniform 还是 gaussian 都一致达到 100% 并改善重建质量。这支持了论文更原则化方案的主张稳定的利用率不应取决于初始化这种偶然条件。额外分析。论文附录进一步考察了两个实践变量。其一是Codebook 初始化的影响pilot 设置冻结 Encoder、仅优化 Codebook 以拟合固定目标分布变化其 Gaussian 初始化尺度并报告 5k steps 时、65,536 token 窗口内的利用率。结果显示初始化尺度太小时 SimVQ 风格的线性 shared projector 恢复利用率很慢尺度变大时 FVQ 风格的 ViT-based shared projector 则面临明显的 dead-code 风险而用 Region VQ Loss 替换标准 VQ loss 后两种情况都大幅改善、5k steps 时接近 full usage且跨初始化尺度均成立。其二是shared projector 的影响同样冻结 Encoder测量 5k steps 内不同 learning rate 下的利用率窗口 65,536 tokensCodebook 初始化为小 Gaussian 标准差10 − 4 10^{-4}10−4以更直接地反映 projector 的学习容量并减少 dead-code 干扰。结果显示两层线性 projector 与两层 MLP 即使把学习率增至5 × 10 − 3 5\times10^{-3}5×10−3也无法在 5k steps 内把利用率提到高水平两层 ViT block 在中等学习率下更好但学习率过大时利用率急剧下降使用 Region VQ Loss 后不同 projector 结构都能在 5k steps 内达到强利用率可用学习率范围大幅拓宽。两组分析共同说明仅靠改变 projector 容量不足以保证稳定快速的利用率增长更关键的是给 Codebook 提供有原则的学习目标。与先前 STE 修正方法的对比。在一个 15 epoch 的受控实验中Region VQ 与其他设置固定论文把 Dynamic STE 与 VQ-STE、Rotation Trick 比较。VQ-STE 采用 alternating optimization动机与 Dynamic STE 类似量化误差大时抑制不可靠任务梯度但它依赖交替的 inner/outer 更新、引入多个敏感超参数与额外训练开销移除了 encoder commitment loss并假设通过k kk-means 初始化建立初始的 code–token 对齐在未联合使用k kk-means 初始化与 norm bottleneck 时该实验中持续出现明显的利用率崩溃无法给出有效指标。Rotation Trick 通过旋转矩阵R RR与 rescaling factor∥ q ∥ / ∥ z ∥ \|q\|/\|z\|∥q∥/∥z∥变换 Encoder 梯度但论文指出其旋转并无严格保证、可能把梯度重定向到远离期望的方向且 rescaling factor 在∥ q ∥ ∥ z ∥ \|q\|\|z\|∥q∥∥z∥时会放缩放大梯度。对比结果是Rotation Trick 与 Dynamic STE 在 Region VQ 支持下都达到 100% 利用率但 Dynamic STE 的 rFID2.7942、PSNR21.1882、LPIPS0.2359显著优于 Rotation Trick 的 3.6437、20.4024、0.2568该表未在片段中显式标注数据集摘要指出论文实验在 ImageNet 上进行。与显式分布对齐方法的对比。Wasserstein VQ 与 MMD VQ 把全局分布匹配作为训练损失分别依赖 Gaussian moments 与 kernel statisticsRegion VQ 则不做全局参数化分布匹配、不对 token 分布作参数假设而是针对标准 VQ 的不对称监督从统计上得到支持的活跃区域向持续不活跃码传播局部目标从而避免 MMD VQ 的昂贵 pairwise kernel 计算。在 VQ-Transplant 提出的 non-Gaussian mixture benchmark 上参数ζ \zetaζ控制非高斯程度ζ 0 \zeta0ζ0退化为单个高斯ζ \zetaζ增大则两个混合模态分离、非高斯结构增强16,384 个码、code dimension 8、每步采样 20K tokens报告 10K steps 时的利用率Wasserstein VQ 与 MMD VQ 的结果引自 VQ-Transplant 的 Table 13Region VQ 在相同ζ \zetaζ网格上以线性 shared projector 评估并在 recipient collisions 严重时自动增大传播配额结果如下方法ζ 0 \zeta0ζ0ζ 1 \zeta1ζ1ζ 2 \zeta2ζ2ζ 3 \zeta3ζ3ζ 4 \zeta4ζ4Time (ms)Relative TimeWasserstein VQ99.9%97.0%62.7%44.8%34.8%6.981.00 × 1.00\times1.00×MMD VQ99.9%99.8%92.5%85.7%75.6%295.5742.36 × 42.36\times42.36×Region VQ (ours)99.8%99.3%93.2%99.4%99.4%8.641.24 × 1.24\times1.24×在最强非高斯条件ζ 4 \zeta4ζ4下Wasserstein VQ 降至 34.8%、MMD VQ 降至 75.6%Region VQ 仍保持 99.4%训练时间上 Region VQ8.64 ms1.24 × 1.24\times1.24×接近 Wasserstein VQ6.98 ms1.00 × 1.00\times1.00×比 MMD VQ295.57 ms42.36 × 42.36\times42.36×快约34.2 倍。这说明其局部传播策略在非高斯分布关系下兼具稳健高利用率与低开销。计算与内存开销。论文在同一机器上对 16K 与 262K 码本的匹配基线做了 profiling16K 时基线为 438.1 ms、34.06 GiBDynamic STE 为 437.2 ms差异在噪声内Region VQ 90% 为 440.2 ms 0.46 % 0.46\%0.46%内存 34.07 GiB 0.04 % 0.04\%0.04%262K 时基线为 462.2 ms、35.12 GiBDynamic STE 为 462.7 msRegion VQ 90% 为 466.2 ms 0.85 % 0.85\%0.85%内存 35.55 GiB 1.24 % 1.24\%1.24%。Region VQ 的成本随利用率升高而降低因为更少的不活跃码需要传播目标Decoupled Schedule 不增加任何前向或反向操作所有组件都是 training-onlytokenizer 推理不变。五、相关工作 / Related Work向量量化 tokenizer 线。VQ-VAE 首次通过在学习到的码本上做最近邻查找把连续 encoder 特征映射为离散码索引奠定了离散表示学习的基础。VQ-VAE-2 引入层次化潜变量图VQGAN 结合感知损失与对抗目标使离散视觉 tokenizer 成为高保真图像合成的标准接口。后续沿着量化结构改进的工作包括 RQ-VAE 等残差/多阶段量化与 ViT-VQGAN 等 ViT 式 tokenizer。针对 codebook collapse / 低利用率已有缓解手段可分为若干类码本重置与替换CVQ、低维码嵌入与归一化ViT-VQGAN、松弛或软分配DALL-E、Gumbel-Softmax、IBQ后者通过 relaxed 或 soft-to-hard 的 categorical 路径改善码本梯度以及用标量/二值量化替代可学习向量码本查表FSQ、LFQ简化优化并改善大规模使用但引入不同的容量权衡。此外把预训练视觉基础模型用于改进 visual tokenizer如 VQGAN-LC、VQRAE可以增强语义表示质量但与本文关注的训练稳定性问题互补。Shared-projection 方法。这是与 StableVQ 最直接相关的一支。VQ-STEStraighten通过仿射重参数化、以码本级的 scale/shift 变换码向量SimVQ 用一个作用于固定 latent basis 的可学习线性层W WW生成码c k q k W \mathbf{c}_k \mathbf{q}_k Wck​qk​W从而优化码本张成的整个线性空间而非单个码向量FVQ 使用更具表达力的非线性 projector 重映射码向量以实现全码本利用。它们的共同点是把码向量重参数化为共享可微函数的输出以较小的架构开销改善利用率。StableVQ 建立在这条线上但研究的是它们未解决的剩余训练不稳定问题并明确指出这些方法仍主要操作码本侧没有同时处理 Encoder 优化自身的不可靠性与模块训练的纠缠。论文还具体指出两个实践限制SimVQ 采用常数学习率维持满利用率但缺少 annealing 导致重建质量大幅下降FVQ 依赖精心设计的 projector 架构包括 ViT block 深度与 patch size才能在 warmup annealing 下维持利用率且 patch embedding 操作把码本大小约束为完全平方数。StableVQ 用单个线性投影层即达到优于或匹配复杂 projector 的重建质量且不施加结构约束。分布对齐与 STE 修正。Wasserstein VQ 与 MMD VQ 把全局分布匹配作为训练损失分别使用 Gaussian moments 与 kernel statisticsWasserstein VQ 分析的是 token 与码分布之间的静态关系而 StableVQ 分析的是训练过程中这种关系如何演化并据此归纳三种失败模式。Region VQ Loss 与它们的关键区别在于不做全局参数化分布匹配、不假设 token 分布形式、使用局部传播目标因而避免了 MMD VQ 的昂贵 pairwise kernel 计算。在 Encoder 梯度侧VQ-STE 采用 alternating optimization与 Dynamic STE 的动机类似但依赖交替的 inner/outer 更新、多个敏感超参数与额外训练开销并假设通过k kk-means 初始化建立初始 code–token 对齐Rotation Trick 通过旋转矩阵与 rescaling factor 变换梯度但论文认为它没有严格保证梯度更准确、可能重定向到远离期望的方向且在∥ q ∥ ∥ z ∥ \|q\|\|z\|∥q∥∥z∥时会放大梯度、在ℓ 2 \ell_2ℓ2​归一化下失去衰减效果。Dynamic STE 则使用 batch 内相对量化距离衰减不可靠梯度且从不放大梯度。图像生成线。PixelCNN、iGPT 等在像素空间直接建模但顺序生成成本高、压缩能力弱。离散 tokenizer 路线由 VQGAN 在 VQ latent space 上用 transformer 做自回归VQ-VAE2、RQ-Transformer 等探索多级离散表示。非自回归与 masked-token 生成器MaskGIT、MAGE、MAGVIT-v2预测缺失的 visual tokens 并迭代细化大规模自回归图像生成器LlamaGen、VAR、RandAR、Open-MAGVIT2使用 next-token、next-scale 或随机顺序预测。扩散 / score-based / flow-matching 模型DDPM、score SDE、flow matching、latent diffusion、DiT以及混合替代方案 MAR 则减少或移除对硬 VQ 的依赖。论文明确把自身定位为与 generator design正交目标不是提出新的生成器而是提供更可靠的离散表示基板。六、局限性与展望 / Limitations Future Work论文作者自述的定位是StableVQ 聚焦于让 VQ tokenizer 训练变得稳定可靠从而为更强的训练目标与下游建模选择提供基础。基于这一判断论文给出了三个自然方向。第一当 codebook utilization 与优化稳定性不再是主要瓶颈之后研究如何更好地平衡重建质量、语义结构与 generation-friendliness 的训练配方——此前的很多设计选择如码本初始化范围、projector 复杂度之所以被采用本质上是在与不稳定作斗争而当这一约束被解除后优化空间才能被真正打开。第二把稳定的视觉 tokenization 与统一的生成和理解目标连接起来使离散 tokens 同时保留低层保真与高层语义。第三虽然实验聚焦图像 tokenizer但同样的责任分离视角可能适用于其他依赖 VQ 的领域例如视频、音频、多模态表示学习或压缩。需要如实指出的局限还包括论文片段中没有独立的 Limitations 小节因此以上方向来自结论中的展望而不是作者对具体失败案例的自述此外从实验设置本身也能看到若干边界。其一是鲁棒性测试主要通过 codebook initialization scale 来构造分布失配属于受控的、由初始化诱导的 mismatch未必覆盖训练中后期由数据分布或多任务变化引起的分布移动。其二是非高斯混合基准是合成 benchmark且当 recipient collisions 严重、导致覆盖不足时需要自动增大传播配额说明 Region VQ 的传播机制在极端重叠情形下仍需要启发式调节。其三是本精读所依据的源码片段未包含若干关键实现细节例如 FIFO 窗口长度W WW的具体取值、损失权重β \betaβ与对抗损失的权重、UR-AUC 主表的完整数值、16k / 262k 定性重建的定量对照以及下游生成的详细数据配置同时论文也未展示完整的参数量比较来量化轻量的程度。这些空白不影响核心结论但在复现与方法推广时需要额外确认。七、总结 / ConclusionStableVQ 的核心贡献不在于提出一个新的量化结构而在于重新表述了 VQ tokenizer 训练不稳定的成因。论文论证长期被视为量化范式固有缺陷的训练崩溃与低利用率其实是 Encoder–Decoder 与 Codebook 优化目标相互纠缠的后果——两个子系统都不能独立履行自己的职责系统只能在二者恰好合作时运转而这种合作恰好最容易在训练压力最大的时刻断裂。围绕这一诊断论文给出三个彼此独立又相互补充的干预Dynamic STE 按 batch 内相对量化距离衰减不可靠的 Encoder 梯度且具有自失活性质、不需要阈值超参数Region VQ Loss 把活跃码的目标按配额局部传播给持续不活跃码使 Codebook 能够独立保证对 Encoder 输出分布的完整追踪Decoupled Schedule 承认两个子系统的优化特性不同为 Encoder–Decoder 保留 warmup annealing、为 Codebook 使用常数高学习率。实验证据在三个层面互相印证pilot study 分别隔离验证了每个组件的必要性标准 STE 在冻结码本下发散到 NaN标准 VQ loss 在冻结 Encoder 下 5000 步仍只有约 12.5% 利用率主实验在 ImageNet256 × 256 256\times256256×256上跨16,384 × 256 16{,}384\times25616,384×256与262,144 × 256 262{,}144\times256262,144×256两种码本规模取得 100% 利用率与领先的 rFID120 epoch 下分别为 1.13 与 0.92鲁棒性测试中 UR-AUC 从 SimVQ 的 2.17±0.32、FVQ 的 8.08±0.24 提升到 60.59±1.52说明收益不只是最终利用率更高而是分布失配下的恢复速度与一致性发生了质变。消融实验进一步揭示了组件间的依赖关系Dynamic STE 单独使用会把利用率压到 1.27%说明传统训练确实依赖 Encoder 的不稳定振荡来激活码本Region VQ Loss 单独使用会因码本更新过慢而 NaN 崩溃说明好的目标仍需匹配的优化动态。所有这些组件都是 training-only、不引入可学习参数计算开销在 262K 码本下最高仅约 0.85 % 0.85\%0.85%时间与 1.24 % 1.24\%1.24%内存。如果用一句话概括这篇论文的立场它把 full codebook utilization 从一个依赖初始化技巧与复杂 projector 架构的脆弱启发式结果转变为一个可以通过目标设计直接保证的有原则属性。这一转变的意义超出重建质量本身——当稳定性不再是瓶颈VQ tokenizer 的研究重心才能从如何让它不崩转向如何让它更好这也正是论文展望中反复强调的方向。原文摘要: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.PDF链接:https://arxiv.org/pdf/2609.26774v1部分平台可能图片显示异常请以我的博客内容为准