★ FEATURED ARTICLE 第一次把 RQ-VAE 真正从零开始完整跑起来的时候我的第一反应是这玩意儿和普通 VAE 的训练过程完全是两个画风。不是说 loss 数值长得不一样而是 rq-vae 训练过程里那种“阶段感”特别明显——前一段是编码器和码本在互相试探中间一段重建效果开始肉眼可见地变清晰后一段几乎是跟残差量化细节死磕。网上搜 rq-vae 训练过程出来的大多是结构图和公式真正讲“代码里每一步在干什么、曲线长什么样、遇到问题往哪查”的内容反而很少。这篇就把我自己从复现到调参跑通的完整链路拆开讲一遍给正准备上手的同学一份能直接照着排错的地图。RQ-VAE 全称是 Residual Quantized Variational AutoEncoder简单说就是在 VQ-VAE 的思路上加了一手“残差量化”把编码器出来的连续特征先量一次得到误差再量一次如此反复多轮。这样做的直接好处是每个码本可以做得不大但多级组合出来的表示空间却大得离谱。训练过程的复杂点也集中在这个“多级残差量化”上而不是编码器或解码器本身。这篇文章适合这几类人读看过 VQ-VAE 觉得原理不难但想迁移到 RQ-VAE 的算法工程师、正在复现图像离散表示模型但训练曲线一直乱七八糟的研究生、以及准备训完量化器之后接自回归先验模型的朋友。我会把损失函数、前向路径、反向梯度、码本更新、训练阶段、踩坑排查全部过一遍最后给出一份可抄作业的配置。1. 训练目标拆解重建、量化误差、码本使用到底在平衡什么1.1 从“一次量化”到“残差量化”为什么要执着于误差的误差想理解 rq-vae 训练过程得先搞清楚它和 VQ-VAE 的分岔点在哪。VQ-VAE 的做法是编码器把图像压缩成一批连续向量然后每个向量从码本里找一个最近的码字用这个码字的索引去替代原向量后续的解码器、自回归模型都是在这个离散索引上干活。这个流程的问题也很直观——一个码字的表达能力有限想把重建做到足够好要么把码本做得特别大要么把向量切得特别碎。RQ-VAE 走的是另一条路第一轮量化获得一个码字然后把原始向量和这个码字的差值当作新的“待量化向量”继续在第二个码本里找最近码字反复多次。每一级码本只需要覆盖上一级没表达好的残差范围所以码本容量可以显著缩小。K 级量化组合出来的表达能力是相乘的比如 4 级、每级 1024 个码字表示空间就有 1024 的四次方这个数字大到基本不用愁容量问题。训练过程里真正要盯住的是这个残差链条上的每一级都不能掉链子。如果第二级码本根本没怎么被用到那 RQ-VAE 本质上就退化成了加大版 VQ-VAE失去了多级量化的意义。我在实际训练中就遇到过这个情况后面排查部分会展开讲。1.2 三项损失的分工L_recon、L_quant、L_commitRQ-VAE 的总损失通常写成这样L L_recon Σ_i ||sg(z_i) - q_i||² β * Σ_i ||z_i - sg(q_i)||²这里 i 表示第几级残差量化层。第一项 L_recon 是重建损失衡量解码器输出的图片和原始输入差多少具体实现可以是 MSE也可以是加了感知损失后的组合。第二项是量化损失让码本里的码字去靠近编码器给出的特征注意前面有 sg也就是 stop-gradient表示这部分的梯度不会回传到编码器。第三项是承诺损失恰恰相反它让编码器的输出不要离选中的码字太远beta 用来控制这个约束的强弱。看到这个公式很多第一次上手的人会误以为把 beta 调大就万事大吉实际上这会严重拖慢编码器的学习速度。我自己常用的范围是 0.25 到 2.0 之间取决于重建损失和量化损失的数值量级。如果两者量级差太多要先做归一化或者加权重而不是贸然去调 beta。1.3 为什么要单独监控每一级量化损失单独看总 loss 曲线很容易骗自己。总 loss 在下降不代表每一级码本都在正常工作。比如第一级码本学得特别好几乎把特征表达完了后面几级的残差量级变得很小量化损失看起来也低但第二级、第三级码字可能根本没被选过几次——这种状态短期内 loss 很健康长远看却浪费了大量参数也削弱了后续离散 token 序列的表达多样性。我在训练时会在日志里把每一级的量化损失、每一级码本的使用频次单独拉出来看而不是只盯 tensorboard 里的总 loss 标量。这个习惯帮我抓出过不止一次“假收敛”的问题。2. 训练循环中的残差量化模块前向选码、反向直通、码本更新2.1 前向计算先选码再减残差继续选码要理解 rq-vae 训练过程必须先清楚它的前向路径不是一次性完成的。具体到一批特征张量上流程大致是这样的def residual_quantize(z_e, codebooks, K): z_e: 编码器输出的连续特征, 形状 [B, D] 或 [B, L, D] codebooks: 列表, 每个元素形状 [M, D]M 是码本大小 K: 量化级数, 一般是 2~8 residual z_e indices [] quantized_sum 0 for k in range(K): # 1. 计算当前残差和该级码本中每个码字的距离 dist (residual.unsqueeze(1) - codebooks[k].unsqueeze(0)).pow(2).sum(-1) # [B, M] # 2. 找出距离最近的码字索引 idx dist.argmin(dim-1) q_k codebooks[k][idx] # 3. 收集索引更新残差 indices.append(idx) quantized_sum quantized_sum q_k residual z_e - quantized_sum.detach() # 残差用停止梯度的量化结果计算 return quantized_sum, indices, residual每一级量化器的输入都是上一级的“减速剩余值”也就是原始特征扣除前面所有已选码字后的结果。这个设计使得第一级码字负责大致的轮廓第二级负责轮廓误差的修正第三级再修正第二级的误差逐级逼近原始特征。有一点容易被忽略残差的回传计算里quantized_sum通常要做detach()。原因很简单如果残差路径把梯度回传到前面级联的量化输出整个多级量化器的梯度会纠缠在一起训练会明显更不稳定。做 detach 之后每一级码本的更新更像独立拟合自己的残差子空间原本的思路反而更干净。2.2 反向传播直通估计器如何绕过不可导的 argmax所有向量量化方法都要面对同一个问题argmax没有梯度。你选出码字索引这个动作本身是离散的可编码器和解码器又都需要连续梯度来更新。解决办法是直通估计器也叫 STE。它的思路简单粗暴前向时正常做量化反向时假装量化操作不存在把解码器传来的梯度直接复制到编码器的输出上。落实到代码上就是z_q z_e (q - z_e).detach()。这样反向时z_q对z_e的梯度是 1梯度就能从解码器的重建损失一路传回编码器。多级量化的反向路径比单级 VQ 复杂在每一级量化操作都会产生一次“直通”解码器拿到的输入是多级量化结果的和梯度传回时实际上分摊给了每一个参与级。但在 RQ-VAE 的通常实现里因为前文提到残差计算时对量化结果做了 detach每一级的码本更新和编码器梯度其实是相对独立的这种独立性反而让训练更稳定。2.3 码本更新方式梯度式更新还是 EMA 式更新码本不是随机初始化就放着不管的它也需要学习。常见的有两条路第一条是梯度更新。量化损失||sg(z_i) - q_i||²对码字求导然后把梯度回传到码本参数上用 Adam 之类的优化器更新。这种方式的好处是能跟编码器共用一套优化器实现简单坏处是码本容易被大梯度带着跑飞特别是码本数量多但单级码本容量小的时候。第二条是 EMA 更新指数滑动平均。思路是累积统计哪些编码器输出被划分到了某个码字名下然后用这些输出的平均值不断平滑逼近码字本身。EMA 方式明显更稳几乎不会出现码本振荡。这也是不少公开实现默认走的路子。我自己的习惯是小型实验用梯度更新省事正式训练或码本超过 1024 后用 EMA 更新稳定性差异非常明显。EMA 的 decay 通常取 0.99 左右训到后期可以适当调成 0.995让码本更新更慢但更细。2.4 多级量化特有的“利益分配”问题RQ-VAE 多出来的训练难点是每级码本之间会争抢表达能力。如果第一级码本过于强大残差很小后面各级码本就学不到东西反过来如果第一级太弱残差始终很大后面各级会越来越难收敛。实践中可以通过控制每一级码本的使用熵来观察。理想情况下每一级码本的选择分布应该相对均匀不要出现某一个码字被狂用、其余码字闲置的情况。我在训练中一旦发现某级码本的有效使用数低于总数的 60%基本可以断定这一级正在“躺平”。修复的手段可以是重置不活跃码字也可以增大该级码本的学习率但最有效的往往还是先检查前面几级码本是不是学得太狠了。3. 三个训练阶段把模型一段段喂熟3.1 阶段一结构收敛期前几十个 epoch 的粗磨合rq-vae 训练过程很少是一条平稳下降的曲线开头一大段更像是“磨合期”。这个阶段编码器的输出分布还在大范围漂移码本也在快速调整重建损失偶尔会不降反升甚至出现脉冲式的抖动都不用太慌。我在这个阶段主要观察三件事一是重建图片的轮廓有没有从“纯噪声”慢慢变成“模糊的物体形状”二是量化损失的中位数有没有在降低三是每一级码本的使用率有没有开始分化。这里尤其不要看单次 step 的 loss要看平滑后的曲线通常用 200-500 步的滑动平均。如果磨合期重建 loss 一开始就在快速下降反而不一定是好事有可能是编码器直接绕过了量化模块把信息以某种隐式方式塞给了解码器也就是所谓的“端到端作弊”。需要定期检查解码器输入是不是真的来自量化后的离散索引序列。3.2 阶段二码本收敛与率失真调优期过了磨合期重建 loss 的下降速度会明显放缓进入码本精细收敛阶段。这个阶段最典型的现象是第一级码本的量化损失基本不动了后面几级的量化损失开始成为主导。换句话说模型正在通过多级残差来抠细节。到了这个时期我会开始逐步下调学习率比如从 3e-4 降到 1e-4 再降到 3e-5每次调整间隔大概 10-20 个 epoch。同时 beta 承诺损失可以适度调小一点给编码器更多空间去表达更高频的细节。不少复现代码在这里训练直接跑飞就是因为学习率没降或者 beta 一直顶在太高的值上编码器在细节层怎么都学不进去。有条件的可以引入率失真指标的监控也就是重建失真和离散 token 数之间的比值。RQ-VAE 的价值本来就是做高压缩比下的离散表示如果重建失真下不去压缩比再高也没意义。3.3 阶段三冻结量化器为自回归先验铺路RQ-VAE 训练通常不是终点后面往往还要接一个自回归 Transformer 来建模离散 token 序列。关键点在于训练先验模型之前要把编码器和码本全部冻结住只用已训练好的离散编码器把训练集转成 token 序列。不要小看这件事。如果先验模型训练过程中码本还在继续漂移离散 token 的语义可能整个变了之前学到的分布直接作废。我在实际项目里吃过一次亏量化器没冻结先验模型训到一半发现 token 分布跟之前采样时的统计对不上最后只能从头来白白浪费了算力。所以 rq-vae 训练过程的最后一步通常不是模型本身而是“编码器锁定 token 预计算 校验码本稳定性”。我自己会做一个简单校验固定住量化器对同一批图片分别今天和明天各生成一次 token 序列比对结果是否完全一致。如果一致才认为量化器进入了可交付状态。4. 训练中最容易踩的坑从症状反推根因4.1 码本坍缩loss 没降但重建全是灰块的典型原因码本坍缩是各种向量量化训练的通病RQ-VAE 只是表现得隐蔽一点。症状是训练中某一步开始重建图片全部变成模糊的色块或者平均颜色量化 loss 逐渐归零但重建 loss 下不去。根因通常是大量码字在一段时间内完全没有被选中导致这些码字的梯度消失慢慢失去意义活跃码字集中在极少数几个上。处理路径我一般按三步走先统计每个码字在一批数据中的被选中次数输出直方图。如果 90% 的选中都落在不到 20% 的码字头上基本实锤。然后重置那些不活跃的码字用当前批次的随机编码器输出或相近活跃码字加噪声来重新初始化。最后在训练配置上补丁要么改用 EMA 码本更新要么把承诺损失 beta 调高一点要么加入码本使用率的均匀性惩罚。4.2 残差量化退化成普通 VQ后级码本利用率极低这类问题在长训练中经常出现。第一级码本学得越来越强残余量越来越小后级码本可能一整批数据都用不上几次。从指标上看最后一级码本的熵接近 0用它解码的请求数少得可怜。修复思路是让前面的码本“让出一部分表达空间”或者反过来说提高后级码本的存在感。一个简单做法是把第一级码本的大小刻意调小一点比如从 1024 降到 512强迫一部分表达能力外溢到后续级。另一个方案是给每一级的残差加一个可学习的缩放权重让模型自己学到残差在不同级之间的能量分配。这种退化问题在论文里不太会被强调但它直接决定多级量化是否真的有效。建议每训练 20 个 epoch 输出一张“各级码本使用热度表”尽早发现退化的苗头。4.3 重建 loss 卡住不动固化的排查顺序重建 loss 迟迟不降这个问题大多数时候不在量化器身上而在前期最容易被忽视的细节里。我自己总结了一套固定的排查顺序照着查一遍基本能定位输入图片归一化是否和预训练骨干一致。数值范围不对重建目标本身就是乱的。解码器是否太浅。RQ-VAE 的量化结构给解码器施加了很大压力解码器如果容量不够训练后期 loss 会顶在某个固定值。量化维度和编码器输出维度是否对齐。这里最容易出 bug某个线性层输出维度和码本维度不一致损失算出来都是错的。重建损失是否被感知损失主导。感知损失尺度可能远大于 MSE导致量化器完全不甘于学习细节。症状可能原因排查手段重建全灰块码本坍缩token 直方图分布检查重建模糊但曲线平滑解码器容量不足增大解码器通道数重建清晰但量化失真大码本维度太低提高码字特征维度后级码本利用率低残差能量分配不均统计各级码本使用熵5. 一份能直接抄作业的参考训练配置5.1 网络与量化器参数我常用的 RQ-VAE 骨干是类似 VQGAN 的卷积编码器解码器结构输入图像 256x256编码器输出特征图 16x16特征维度设为 256。量化部分用 4 级残差量化每级码本大小 1024码字维度也是 256。这样一个位置最后对应 4 个 token ID序列长度比单级 VQ-VAE 稍长但表示精度高得多。注意一个细节这里说的码本大小 1024指的是每一级分别有 1024 个码字而不是 4 级共 1024。之前有人按“总码本”理解实际并没有发挥出残差量化的容量优势。5.2 训练超参数参考表超参推荐值备注优化器AdamWbetas 用 0.9 / 0.99 更稳学习率3e-4预热后逐步余弦衰减到 1e-5总步数 warmup5000 步线性从 1e-5 涨到 3e-4batch_size64根据显存调整最好不要低于 32beta_commit1.0观察量化 loss 量级可微调EMA decay0.99EMA 更新码本时使用梯度裁剪norm 1.0有效防止重建梯度爆炸训练分辨率256如果显存不够先训 128关于学习率有一个自己的经验RQ-VAE 对学习率的敏感度比普通 VAE 高很多。太高会让码本和编码器互相追逐产生持续的振荡太低则会影响多级残差结构的收敛速度。先跑 5 个 epoch 尝试 1e-3、3e-4、1e-4 三条曲线通常一眼就能看出哪个区间最合适。5.3 用这几个指标判断训练是否健康光看 loss 曲线完全不够我在实验里会同时盯五个指标重建图像样例、各级量化损失、码字使用熵、码字使用覆盖率和 token 序列的重建稳定性。前四个都好理解最后一个提一下固定模型把同一张图反复编解码十次看结果是否一致。如果出现同图不同 token 的现象大概率是推理阶段 drop 了一些细节或者量化器还有随机行为没有去掉。下面这一组是我判定“可以开始训先验模型”的门槛重建 PSNR 达到 26dB 以上不同数据集会有差异每一级码本的使用覆盖率达到 80% 以上且后三级码本的熵不低于第一级的 50%。达不到这个门槛就继续训练或调结构不建议硬着头皮接下游任务。6. 一些实操感悟论文不会告诉你的训练经验6.1 调参顺序比单点搜索更重要我见过很多人一上来就把 beta、码本大小、学习率全部拧一遍最后根本分不清是哪个参数起作用。自己走过的弯路多了以后建议调参顺序固定为先定网络结构再定码本容量和级数然后调学习率最后微调 beta。没跑通基线之前不要加感知损失、对抗损失这些花活否则问题定位会难上加难。有了稳定的 RQ-VAE 基线之后你会发现它的训练过程其实比普通 VAE 更好控制。普通 VAE 的后验坍塌需要各种花式技巧去调节RQ-VAE 只要盯着码本使用率和重建曲线大多数情况都能收敛到可用状态。6.2 先跑小规模验证再上全量RQ-VAE 的多级量化机制在 128x128、4 级、每级 512 码字的配置下一台消费级显卡就能轻松跑通。我强烈建议任何新数据集都先用这个规模验证花半天时间看重建效果和码本利用情况而不是直接冲到 256 分辨率、8 级量化的大配置。小规模验证和大规模训练之间通常存在一个缩放规律码本规模不变的情况下分辨率提升让每一级码本的角色分工更明确第一级偏向整体结构后级偏向纹理细节。如果小规模下后级码本的利用率就不健康放大到全量只会更严重务必先解决再说。6.3 不要迷信“更大码本”码本越大越好这个直觉在 RQ-VAE 这里不完全成立。码本太大会让每一级的选择分布变稀疏反而加速码本坍缩。正确思路是保持每一级码本在 512-2048 之间然后通过增加级数 K 来提升表达容量。级数 K 对训练时间的影响很小主要影响 token 序列长度的下游建模成本所以一般取 4-6 足够8 级以上的收益开始明显递减。最后分享一个实用小技巧训练快结束时把 EMA decay 从 0.99 调到 0.999多训几个 epoch码本和编码器之间的匹配度会明显变得更好下游先验模型的训练也会更省心。这是我在对比多个实验后收获的稳妥手段几乎不用额外付出训练成本。
阅读完成 · 觉得有帮助？
