StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training¶

会议: NeurIPS 2026
arXiv: 2609.26774
OpenReview: https://openreview.net/forum?id=XrO6sGUEHB
领域: 模型压缩
关键词: 向量量化、视觉分词器、Straight-Through Estimator、码本利用率、训练稳定性

一句话总结¶

StableVQ 揭示了离散视觉分词器中编码器-解码器与码本之间的训练纠缠是导致失稳与坍塌的根因，通过提出动态直通估计器（Dynamic STE）、区域量化损失（Region VQ Loss）与解耦学习率调度（Decoupled Schedule），在无需额外可学习参数的情况下以单层线性投影实现了全量码本激活与顶尖重建性能。

研究背景与动机¶

离散视觉分词器（Visual Tokenizer）是将连续图像压缩表征为结构化离散词元序列的核心基础设施，直接决定了自回归大模型（如 LlamaGen、VAR）与掩码生成模型（如 MaskGIT）的生成质量上限。在经典的 VQ-VAE 框架中，码本崩溃（Codebook Collapse）与利用率低下是长期存在的瓶颈。近期以 SimVQ 和 FVQ 为代表的共享投影（Shared-Projection）方法，通过引入可微共享函数重参数化码本向量，将离散最近邻匹配转化为投影码本分布与编码器特征分布之间的对齐问题，显著提升了名义上的码本利用率。

然而，现有共享投影方法依然面临严重的训练不稳定性：训练过程对码本初始化极度敏感，极易陷入低利用率平台期，或在中后期突发码本崩溃甚至梯度爆炸（NaN）。深入剖析其失稳模式可以发现，根本症结并不在于向量量化范式本身的容量限制，而在于编码器-解码器与码本两套子系统在优化目标上的严重相互纠缠。在传统训练机制下，任何一个模块都无法独立达成自身职责——编码器无法在低码本利用率下稳定维持离散正则化重构，而码本也无法自主覆盖编码器分布，只能寄希望于编码器特征的随机震荡来被动激活未选中码字。这种脆弱的“协同假设”一旦遭遇尺度失配或剧烈特征漂移便会彻底瓦解。

本文从职责分离（Separation of Concerns）的系统设计原则出发，重新审视并确立了两个子系统应有的独立优化目标：编码器应在离散约束下稳健优化连续重构流形，码本应自主且确定性地追踪编码器输出分布，且二者异质的优化动力学必须配备独立的调度策略。核心 idea：通过 Dynamic STE 纠正编码器梯度估计误差以消除失真冲击，通过 Region VQ Loss 实现活跃码向邻近冷码的确定性目标传播，配合解耦学习率调度，彻底解绑编码器与码本的训练纠缠，使单层共享线性投影即可稳健达成 100% 码本利用率与最优重构质量。

方法详解¶

整体框架¶

StableVQ 构建在通用的共享投影向量量化架构之上。输入图像 $\mathbf{x}$ 首先由编码器 $E$ 映射为连续特征图 $\mathbf{z} = E(\mathbf{x}) \in \mathbb{R}^{H \times W \times d}$；码本 $\mathcal{C}=\{\mathbf{e}_k\}_{k=1}^K$ 经过共享线性变换层 $f_\theta$ 得到投影码本空间，特征点通过欧氏空间最近邻搜索完成离散量化 $\hat{\mathbf{z}}_{ij} = f_\theta(\mathbf{e}_{k^*_{ij}})$；最后由解码器 $D$ 根据量化特征重建图像。

为了彻底解除模块间的病态耦合，StableVQ 在整个前向与反向流水线中实施了三项轻量化解耦干预：在前向/反向连接处，Dynamic STE 根据相对量化误差加权调制反向传播给编码器的梯度；在码本更新端，Region VQ Loss 依据特征负载将活跃码本的目标均值定向传播给持续非活跃的码字；在优化器层面，编码器-解码器采用带预热和退火的优化策略，而码本使用固定高学习率持续追踪特征漂移。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 x"] --> B["编码器 E(x)<br/>连续潜变量特征空间"]
    B --> C["Dynamic STE<br/>按相对量化误差加权调制梯度"]
    C --> D["单层共享线性投影 fθ<br/>映射离散码本 entries"]
    D --> E["Region VQ Loss<br/>活跃码向非活跃码传播局部目标"]
    E --> F["Decoupled Schedule<br/>编码器退火 + 码本高常数学习率"]
    F --> G["解码器 D(z_hat)<br/>高质量重构图像与离散 Token"]

关键设计¶

1. Dynamic STE：抑制远距量化伪影带来的失真梯度 在标准直通估计器（Straight-Through Estimator, STE）中，假设量化前后梯度无损穿透 $\partial\mathcal{L}/\partial\mathbf{z}_{ij} \approx \partial\mathcal{L}/\partial\hat{\mathbf{z}}_{ij}$。然而，当某个特征 $\mathbf{z}_{ij}$ 与匹配到的码字距离较远时，该处真实梯度方向与近似方向存在严重偏差。这种失真梯度会与 commitment loss 发生冲突，将编码器特征推向不稳定区域，诱发 commitment loss 暴增甚至引发 NaN 异常。StableVQ 敏锐观察到：若多个特征同时量化到同一码字，与码字几何距离相对更远的特征其梯度最不可信。为此，设计了基于批次内相对量化误差的动态加权机制。定义码字 $k$ 在当前批次内的最小距离为： $$d^*_k = \min_{(i',j')} \|\mathbf{z}_{i'j'} - f_\theta(\mathbf{e}_k)\|_2^2$$ 每个空间位置的特征根据其相对最优距离获得截断梯度权重： $$w_{ij} = \text{sg}\left[ \frac{d^*_{k^*_{ij}}}{\|\mathbf{z}_{ij} - f_\theta(\mathbf{e}_{k^*_{ij}})\|_2^2} \right] \in (0, 1]$$ 在前向计算中注入加权恒等映射 $\hat{\mathbf{z}}_{ij} = w_{ij} \cdot \mathbf{z}_{ij} + \text{sg}[\hat{\mathbf{z}}_{ij} - w_{ij} \cdot \mathbf{z}_{ij}]$。最契合码字的特征获得全额梯度（$w_{ij}=1$），偏离较远的特征其穿透梯度被自适应衰减。在训练进入良好收敛状态时，$w_{ij} \to 1$，该机制平滑回退至标准 STE，全过程不引入任何超参数阈值。

2. Region VQ Loss：跨越活跃-非活跃区域的显式目标传播 传统 VQ loss 存在天然的监督不对称性：所有图像特征都有 commitment loss 约束，但只有在当前步被选中的码字才能获得显式学习目标，未选中码字仅能通过共享投影层的弥散梯度获得微弱且快速衰减的更新。这导致非活跃码本无法自主向数据流形靠拢，陷入永久死码。Region VQ 摒弃了点对点强制拉扯，改为局部区域分布传播。系统维护一个滑动窗口长度为 $W$ 的 FIFO 队列记录各步激活码字，从而精确分离出时间窗口内的活跃集合 $\mathcal{A}_t$ 与持续非活跃集合 $\mathcal{N}_t = [K] \setminus \mathcal{A}_t$。对于在当前步承载多个特征的活跃码字 $k \in \mathcal{S}_t^+$（其对应特征集合为 $\mathcal{U}_k$），先计算其特征均值中心 $\mathbf{t}_k = \frac{1}{|\mathcal{U}_k|}\sum_{u \in \mathcal{U}_k}\mathbf{z}_u$。接着根据其承担的样本配额： $$q_k = \left\lceil \frac{|\mathcal{U}_k| \cdot |\mathcal{N}_t|}{\sum_{\ell \in \mathcal{S}_t^+} |\mathcal{U}_\ell|} \right\rceil$$ 将中心目标 $\mathbf{t}_k$ 广播给在投影空间中与其欧氏距离最近的 $q_k$ 个冷码。非活跃码字汇总所有指向它的活跃源目标计算平均目标 $\mathbf{t}_j = \frac{1}{|\mathcal{K}_j|}\sum_{k \in \mathcal{K}_j} \mathbf{t}_k$。码本仅在拥有有效目标的集合 $\mathcal{M}_t$ 上计算均方误差损失： $$\mathcal{L}_{\text{code}} = \frac{1}{|\mathcal{M}_t|} \sum_{k \in \mathcal{M}_t} \|f_\theta(\mathbf{e}_k) - \text{sg}[\mathbf{t}_k]\|_2^2$$ 这一机制在不依赖数据高斯分布等先验假设的前提下，赋予非活跃码字确定性的牵引力，主动拉动冷码进入编码器活跃流形内部。

3. Decoupled Schedule：顺应模块异质特性的学习率解耦 编码器-解码器与码本在优化景观上具有本质异质性。编码器-解码器同时面临像素重构、感知对抗以及离散正则化等多目标权衡，地形极其复杂，必须依赖 Warmup-plus-Annealing（预热+余弦退火）来平抑初期波动并确保平滑收敛；相反，码本的职责仅仅是纯粹、及时地追踪编码器输出特征的分布漂移，其在训练初期编码器剧烈变化阶段反而最需要充沛的梯度步长。以往方法将两类参数绑定在统一的学习率曲线上，若采用退火策略会导致码本在后期更新停滞、利用率崩塌，逼得 FVQ 不得不通过堆叠复杂的 ViT Block 投影器来补偿容量。StableVQ 将二者优化体系完全解耦：编码器-解码器执行 10% 步数 Warmup 与 63% 步数 Cosine 退火（基础学习率 $1\times 10^{-4}$），而码本从第 0 步开始全程保持 $1\times 10^{-3}$ 的高恒定学习率，使码本从始至终具备充沛的流形追踪敏捷度。

损失函数 / 训练策略¶

系统端到端优化的总体损失函数由重构项、正则项与码本更新项构成： $$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{recon}}(\mathbf{x}, D(\hat{\mathbf{z}})) + \beta \|\mathbf{z} - \text{sg}[\hat{\mathbf{z}}]\|_2^2 + \mathcal{L}_{\text{code}}$$ 其中重构损失 $\mathcal{L}_{\text{recon}}$ 包含像素级 $\ell_1$ 损失、LPIPS 感知损失以及 StyleGAN-T 风格的鉴别器对抗损失；commitment loss 权重设为 $\beta = 0.25$。训练采用 Adam 优化器（$\beta_1=0.9, \beta_2=0.95$），基础批大小设为 128。共享投影器默认仅使用单一线性层 $f_\theta(\mathbf{e}_k) = \mathbf{e}_k \mathbf{W}$，将码本向量从潜变量基底直接线性映射至嵌入空间，零额外推理开销。

实验关键数据¶

主实验¶

在 ImageNet $256 \times 256$ 图像重构基准上，StableVQ 与代表性传统 VQ 方法及前沿共享投影方法展开对比。下采样率设为 $f=16$（每张图 $16 \times 16=256$ 个离散 tokens）。除标准重构指标外，测试引入了量化抗失配与恢复能力的综合指标 UR-AUC（Usage Recovery AUC）。

方法	投影器类型	训练轮数	码本规格 ($K \times d$)	rFID↓	LPIPS↓	码本利用率↑	UR-AUC↑
LlamaGen	—	40	16,384 × 8	2.19	0.2281	97.0%	—
LlamaGen	—	40	16,384 × 256	9.21	—	0.29%	—
IBQ†	—	330	16,384 × 256	1.37	0.2235	96.0%	—
IBQ†	—	330	262,144 × 256	1.00	0.2030	84.0%	—
SimVQ	Linear-1	40	16,384 × 256	2.89	0.2492	100.0%	2.17±0.32
SimVQ	Linear-1	40	262,144 × 256	3.16	0.2516	100.0%	—
FVQ	ViTBlock-2	40	16,384 × 256	1.70	0.2176	100.0%	8.08±0.24
FVQ	ViTBlock-2	40	262,144 × 256	1.29	0.2003	100.0%	—
StableVQ (本文)	Linear-1	40	16,384 × 256	1.22	0.2235	100.0%	60.59±1.52
StableVQ (本文)	Linear-1	40	262,144 × 256	1.05	0.1947	100.0%	—
StableVQ (本文)	Linear-1	120	16,384 × 256	1.13	0.2134	100.0%	—
StableVQ (本文)	Linear-1	120	262,144 × 256	0.92	0.1893	100.0%	—

注：† 表示使用了更庞大的编码器-解码器网络。

下游类条件图像生成评估（ImageNet $256 \times 256$，自回归架构 IBQ-B / IBQ-L，无分类器引导）：

模型配置	分词器基底	生成器架构	参数量	FID↓	Inception Score↑	Precision↑	Recall↑
Vanilla AR	IBQ	IBQ-B	342M	2.88	254.7	0.84	0.51
Vanilla AR	IBQ	IBQ-L	649M	2.45	267.5	0.83	0.52
Vanilla AR	StableVQ	IBQ-B	342M	2.35	256.0	0.82	0.58
Vanilla AR	StableVQ	IBQ-L	649M	2.18	250.4	0.82	0.59

消融实验¶

为了严格检验各项设计应对分布失配与严苛条件的鲁棒性，在极端码本扩张情境（初始码字尺度远小于特征尺度，极易诱发梯度爆炸）下对核心组件进行组合消融：

Region VQ	Dyn. STE	De. Sch.	峰值 Commit Loss	码本利用率↑	rFID↓	LPIPS↓	PSNR↑	SSIM↑
—	—	—	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.0%	1.72	0.2200	21.84	0.5854
✓	—	✓	<<0.1	100.0%	1.75	0.2206	21.85	0.5860
✓	✓	✓	<<0.1	100.0%	1.70	0.2208	21.81	0.5879

码本收缩情境消融（$\ell_2$ 归一化导致高维码字散布范围远大于真实特征分布，极易产生大量永久死码）：

基础架构	采用 Region VQ	码本初始化方式	码本利用率↑	rFID↓	PSNR↑	SSIM↑	LPIPS↓
FVQ (ViTBlock-2)	—	均匀分布 Uniform	18.75%	2.4355	20.9789	0.5677	0.2414
FVQ (ViTBlock-2)	—	高斯分布 Gaussian	62.50%	1.9530	21.3774	0.5836	0.2273
FVQ (ViTBlock-2)	✓	均匀分布 Uniform	100.0%	1.8176	21.5497	0.5896	0.2210
FVQ (ViTBlock-2)	✓	高斯分布 Gaussian	100.0%	1.8966	21.4603	0.5931	0.2255

关键发现¶

单独添加 Region VQ 而保留共享 Warmup 学习率会导致码本更新滞后而发生 NaN 崩溃；单独添加 Dynamic STE 虽能绝对压制 Commitment 梯度尖峰，但码本利用率暴跌至 1.27%，反向证明了传统训练高度依赖不稳定的编码器扰动来偶然激活冷码。
Dynamic STE、Region VQ 与解耦调度的协同运作，构成了闭环防御：Dynamic STE 保证编码器优化平稳，Region VQ 提供稠密且有依据的冷码拉引信号，解耦调度提供敏捷的分布追踪能力。三者合力使得最简单的单层线性投影器在 40 轮训练下即可达到 1.22 的 rFID，大幅超越需要 2 层 ViT 复杂投影器的 FVQ（1.70）。
在面对极端初始化失配时，StableVQ 的使用率恢复 AUC 达到 60.59，较 FVQ 的 8.08 提升逾 7 倍，彻底根除了对码本初始化技巧与特殊结构设计的依赖。

亮点与洞察¶

立论高度与本质解耦：跳出了以往仅在网络结构（如更深的非线性注意力映射）或经验性初始化 tricks 上打补丁的局限，从“职责分离”高度将不稳定根源锁定为编码器与码本的病态优化缠绕，给出了干净、自洽的系统级解释。
自适应退化的 Dynamic STE：通过批次内相对最小误差进行梯度截断加权，在出现严重未匹配时自适应压低伪影梯度，而在训练稳定期平滑退化为标准 STE，优雅规避了引入硬截断门限的工程调参难题。
非参数化的局部流形牵引：Region VQ Loss 创新性地在离散非活跃码字与高负载活跃码字间建立了以样本负荷为权重的定向目标投递，相比 Wasserstein VQ 或 MMD VQ 等全局分布匹配，既无需假设数据满足高斯混合先验，又免去了 $\mathcal{O}(N^2)$ 的核计算复杂度，兼具高效率与泛化力。

局限与展望¶

下游生成联合优化探索较浅：本文核心聚焦点在于保障分词器自身表征重构与码本利用率的稳定性，而离散视觉词元的语义聚集度与大语言模型/扩散模型下游生成友好的权衡机制仍需进一步联合探索。
跨模态与复杂维度延展：当前实验主要集中在二维图像的离散分词上。尽管文中所提原则具备通用性，但尚未在视频长时序建模、音频时频量化以及多模态统一表示学习中进行大规模实测。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从职责分离与优化解耦视角系统解决 VQ 训练失稳问题，提出的三项轻量组件逻辑闭环、优雅深入。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖多尺度码本、极端扩张/收缩失配压力测试、非高斯分布拟合测试及下游类条件生成，证据链极其扎实。
写作质量: ⭐⭐⭐⭐⭐ 结构清晰严谨，数学形式化精确，问题剖析入木三分，具有极佳的启发性与工程指导价值。
价值: ⭐⭐⭐⭐⭐ 彻底扫清了离散视觉分词器在超大码本缩放中的稳定性隐患，无需额外参数与计算开销，具备极强的工业界落地通用性。