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
华科与可灵给共享投影VQ加上动态STE、区域VQ损失和独立学习率，默认仍用单层线性投影，ImageNet上262144码本训练120轮后占用100%、rFID降到0.92。
VQ tokenizer 把图像压成离散码，给自回归和 masked 生成当词表。码本崩塌一直是硬伤：码很多，最近邻却总打在少数几个上。SimVQ、FVQ 这类 shared-projection 方法给码向量加一层共享可微映射，梯度能间接碰到未选中的码，占用率经常报到 100%。
训练照样不稳。初始化差一点，占用长时间趴在低位；commitment loss 冲到 NaN；训到一半占用突然塌回去。华中科技大学与可灵（KlingAI Research）的判断是：编码器-解码器和码本绑在同一套优化里，谁都完不成自己的职责，只能靠两边碰巧配合。
StableVQ 叠在 shared-projection 上，三条干预都不加可学习参数，默认投影就是一层线性。
Dynamic STE 针对编码器。量化不可导，训练用 Straight-Through Estimator 把重建梯度原样传回去。token 离自己分到的码很远时，这条梯度方向不可信，会跟 commitment loss 对着干，把两个分布越推越开。权重按「相对该码在本批次最近 token 的距离」来给：最近的权为 1，越远越压。占用健康时权重接近 1，退回普通 STE，没有阈值超参。
Region VQ Loss 针对码本。标准 VQ loss 只监督当前被选中的码。共享投影能把梯度渗到死码，但方向散、后期还会衰减。Region VQ 用 FIFO 窗口分出近期活跃码和长期死码，活跃码按被选次数把目标传给最近的死码。码本可以自己追编码器分布，不必等编码器抖一抖才激活死码。
Decoupled Schedule 把学习率拆开。编码器走 warmup 再退火；码本用恒定高学习率，消融里是 1e-3，编码器峰值 1e-4。码本的活是持续跟踪，早期分布变得最快，不该跟编码器一起从零爬升。
设置：ImageNet 256×256，VQGAN 风格编码器，下采样 16，每图 256 个 token。
| 方法 | 投影 | 码本 | epoch | rFID↓ | 占用 | 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 测的是码本和 token 分布对不齐时，占用率恢复曲线的面积。终点占用大家都是 100%，错位之后谁能爬回来差一个数量级。
小范围初始化更狠：三条全关占用 49.13%、rFID 2.06；只开 Region VQ 会 NaN；Dynamic STE 单独把占用压到 1.27%；三条一起占用 100%、rFID 1.70。ℓ2 归一化后码本铺得太开时，FVQ 的 ViT 投影在 uniform 初始化下只到 18.75%，加上 Region VQ 回到 100%。
下游按 IBQ 设置训类条件自回归：StableVQ + IBQ-B（342M）FID 2.35，对照 IBQ tokenizer 的 2.88；IBQ-L（649M）FID 2.18 对 2.45。IS 没有跟着涨，IBQ-L 上 250.4 对 267.5。
冻结码本只训编码器，普通 STE 会 commitment 尖峰到 NaN，Dynamic STE 能撑完。冻结编码器只训码本，标准 VQ loss 5000 step 还在约 12.5% 占用，Region VQ 500 step 就满占用。
做离散视觉 tokenizer 的人，现在默认会叠 shared-projection、k-means 初始化、复杂投影、码重置。这篇把稳定性拆成三个可单独关掉的旋钮，默认一层线性就够。UR-AUC 比终点占用更有用：终点 100% 会掩盖训练过程里的死码和塌缩。
这是工程向的渐进改进，不是新量化范式。适用范围停在已经用 shared-projection 的 VQ。
实验几乎全在 ImageNet 256、f=16、256 token，没有视频、多模态、更高分辨率。生成侧 FID 降了，IBQ-L 的 IS 从 267.5 掉到 250.4，重建变好不等于生成全面变好。Region VQ 单独用会 NaN，三条必须搭配。FIFO 窗口长度正文没给消融。对抗判别器仍按前人配置在用，重建数字里有多少来自 GAN 损失，论文没拆开。作者把结论落在纠缠目标而不是量化本身的上限，判断主要靠消融，没有更形式的证明。