StableVQ:把向量量化分词器的训练稳定性从“侥幸”变成“可保证”
核心概要
该工作提出 StableVQ,用三个不含可学习参数的组件——Dynamic STE、Region VQ Loss 与 Decoupled Schedule——把编码器-解码器与码本从相互纠缠的训练中解耦,使二者各自独立完成自身职责,从而在 ImageNet 上跨多种码本规模与初始化设置一致提升训练稳定性、码本利用率与重建质量。
深度剖析
论文把 VQ 分词器训练不稳定的根因定位为“编码器-解码器与码本的训练纠缠”:两个子系统都无法独立完成自身职责,系统只有在二者恰好配合时才能工作,而这种配合在训练压力最大时恰恰会失效。 此前共享投影方法(VQ-STE++、SimVQ、FVQ)主要从码本一侧解决梯度稀疏与利用率问题,本文把视角转向模块职责划分,将训练不稳定重新表述为模块职责失效,而非量化范式本身的固有限制。 属于概念性分析与失效模式刻画,论文以码尺度与 token 尺度关系的三类情形(码尺度小于 token 尺度、码尺度大于 token 尺度、尺度发散)组织论证,并配有示意图说明。
Dynamic STE 针对编码器的学习目标:当 token 被分配到较远的码时,直通估计器传来的重建梯度不再可靠,会与 commitment loss 冲突并可能升级为 NaN;该方法按 token 相对该码最佳匹配 token 的距离对梯度加权,越远衰减越多,最佳匹配 token 保留完整梯度。 与 VQ-STE++ 的交替优化、Rotation Trick 的旋转与重缩放不同,Dynamic STE 只做衰减、不做放大,且当批次内所有 token 都匹配良好时权重退化为 1,自动还原为标准 STE,无需阈值超参。 冻结码本、只训练编码器-解码器的对照实验显示,标准 STE 出现严重 commitment loss 尖峰与重建损失震荡并最终发散到 NaN,而 Dynamic STE 保持稳定;另有 15 轮受控实验与 Rotation Trick、VQ-STE++ 对比。
Region VQ Loss 针对码本的学习目标:标准 VQ 损失是不对称的——每个 token 都通过 commitment loss 获得明确目标,而只有被选中的码获得有效目标;该方法把活跃码收到的目标传播给邻近的非活跃码,用 FIFO 队列区分窗口活跃集与持续非活跃集,使码本无需依赖编码器震荡即可独立保证对编码器输出分布的完整跟踪。 共享投影虽让梯度经共享函数到达所有码,但信号间接、无方向且会衰减;Region VQ 改为从逐点对齐转向分布对齐,且不假设 token 分布形式,也避免了 MMD VQ 的成对核计算开销。 冻结编码器、只优化码本的对照实验中,标准 VQ 损失在 5000 步后仍停滞在约 12.5% 利用率,Region VQ 在 500 步即达到满利用率并保持;在非高斯混合基准上,当分布强烈非高斯时 Wasserstein VQ 与 MMD VQ 分别降至 34.8% 与 75.6%,Region VQ 保持 99.4%,且训练时间接近 Wasserstein VQ、快于 MMD VQ。
Decoupled Schedule 认为编码器-解码器与码本需要不同的优化动态:前者承担离散正则下的多目标重建任务,适合 warmup 加退火;后者需要持续跟踪不断演化的编码器分布,适合无 warmup 的恒定高学习率。 此前工作观察到 warmup 加退火有利于训练质量却损害码本利用率,FVQ 通过更复杂的共享投影器来补偿;本文把这一张力归因于两个模块目标不同,因而应各自独立调度。 两组受控实验显示:在 FVQ 架构上把码本改为恒定学习率不带来性能下降,而把编码器-解码器改为恒定学习率则明显掉点;在 SimVQ 架构上,码本利用率不因复杂调度而提升,却受益于稳定且足够高的恒定学习率。
启示与展望
该结果面向使用共享投影码本的离散视觉分词器训练,实验在 ImageNet 上以 VQGAN 风格编码器-解码器进行,覆盖不同码本规模与初始化设置,并额外在合成非高斯混合基准上检验 Region VQ;三个组件均为训练期干预,不改变分词器推理,因此可直接叠加到既有共享投影方法之上。论文把自身定位为提供稳定可靠的训练基础,使更强的训练目标与下游建模选择可以在其上探索,并指出同一视角可能延伸到视频、音频、多模态表示学习与压缩等依赖向量量化的领域。
论文的失效模式分析以码尺度与 token 尺度关系为组织线索,读者可留意这些情形在自身数据与分辨率下的对应关系;Dynamic STE 的权重依赖批次内相对量化误差,批次构成变化时其行为值得观察;Region VQ 的传播配额在合成基准上会因接收方重叠严重而自动放大,这一机制在真实数据上的触发频率尚待了解;此外,本次载入的文本为论文全文,但表格中的部分数值以占位形式呈现,具体数字需回到原文核对。
