在人工智能领域，始终存在着一种拉锯战：既要创造出足够聪明、能生成惊艳图像的模型，又要保持其体积足够小以便快速运行。大型系统通常被称为“教师”，它们通过一次预测图像的一个微小部分，从一个起点开始构建整个画面，从而生成细节极其丰富的图片。然而，这种循序渐进的过程既缓慢又昂贵。为了让这些系统变得更快，研究人员使用了一种名为“知识蒸馏”的技术，即让一个更小、更简单的“学生”模型去模仿“教师”。一种更新、更有效的过程被称为“策略内蒸馏”（on-policy distillation），它要求学生通过实时生成自己的训练样本来进行练习。这确保了学生学习的是它在实际创作艺术时的真实行为方式，而不仅仅是死记硬背旧的样本。问题在于，这种练习过程极其缓慢；学生仍然必须一次构建一个微小的视觉单元（token），这造成了一个瓶颈，拖慢了整个学习过程。
来自深圳的一个研究团队发现了一种方法，可以在不改变学生学习基本规则的情况下打破这一瓶颈。他们引入了一种名为 HB-SJD 的新方法，它充当了学生练习环节的快速引擎。该系统不再强迫学生一次只构建一个 token（或微小的视觉单元），而是允许模型同时提出多个未来的图像片段，然后检查这些猜测是否正确。把这想象成一个学生在考试：不再是写出一个答案，等待老师批改，然后再写下一个，而是学生一次性写出一整页答案，然后老师快速勾选出哪些是正确的。学生保留正确的答案，仅重写错误的答案，从而前进得更快。
研究人员发现，仅仅在同一时间对一组图像运行这种更快速的方法是不够的。在标准设置中，如果一张图像生成得很快，它必须等待组内最慢的一张图像完成后才能进入下一步。这种等待期浪费了宝贵的计算能力。为了解决这个问题，该团队设计了他们的系统，使每张图像都能以自己的速度推进。如果一张图像完成了，它就会退到一旁，系统会在不暂停的情况下继续处理剩余的图像。此外，他们发现随着活跃图像数量的减少，运行系统的最佳方式也会发生变化。当处理大量图像时，系统通过保持整个组处于活跃状态来达到最佳效果。但随着图像完成且组规模缩小，系统会切换到一种更精简的模式，仅处理剩余的活跃图像，从而避免管理空置槽位的开销。
在图像生成模型上的测试证明，这种新方法非常有效。研究人员发现，它能将学生生成练习图像所需的时间缩短近一半，加速倍数在 1.4 到 1.6 倍之间。至关重要的是，这种速度提升并没有以牺牲质量为代价。使用这种新方法训练的学生模型所生成的图像，与使用较慢旧方法训练的图像一样清晰、逼真。无论图像是长是短，无论处理图像的组规模是大是小，以及在不同的训练阶段，该系统都能稳定表现良好。通过仅替换生成练习数据的引擎，研究人员成功地显著加快了整个学习过程，证明了你可以在不损害最终结果质量的前提下，加速人工智能的训练。
技术摘要：通过批处理投机 Jacobi 展开加速视觉在策略蒸馏
1. 问题陈述
视觉自回归（AR）模型在图像生成方面展现了强大的可扩展性。知识蒸馏是将大型教师模型的能力迁移到小型学生模型的一种标准方法。最近的进展，如在策略蒸馏（On-Policy Distillation, OPD），通过让学生在训练期间生成自己的轨迹并由教师进行监督，提高了训练效果。这减少了训练行为与推理行为之间的不匹配。
然而，OPD 引入了一个显著的计算瓶颈：学生必须在训练过程中在线生成展开后缀（rollout suffixes），然后教师才能提供监督。传统的 AR 解码是逐 token 进行的，对于一个长度为 L 的 token 后缀，需要 L 个连续的解码步骤。这一过程直接处于训练的关键路径上，为每一次在策略更新增加了巨大的成本。虽然投机 Jacobi 解码（Speculative Jacobi Decoding, SJD）提供了一种在无需辅助草稿模型的情况下并行处理多个 token 的方法，但原始的 SJD 是为单序列推理设计的。将 SJD 直接应用于视觉 OPD 的大批量设置时面临两个挑战：
- 失去并行性： 为每张图像单独运行 SJD 会破坏 GPU 的批处理并行能力。
- 同步开销： 以步调一致的方式处理所有图像（批同步模式）会迫使较快的图像等待较慢的图像在共享窗口边界处完成，从而抵消了效率增益。
2. 方法论：HB-SJD
作者提出了混合批处理投机 Jacobi 解码（Hybrid Batched Speculative Jacobi Decoding, HB-SJD），这是一种专为视觉 OPD 设计的展开后端。HB-SJD 替换了标准的学生展开引擎，而保持教师、蒸馏目标和优化过程不变。
核心组件
- 独立的单图进度： 与批同步 SJD 不同，HB-SJD 允许批次中的每张图像根据其自身的解码进度进行推进。图像不会等待批次中最慢的成员完成验证窗口。
- 批处理验证： 尽管进度独立，HB-SJD 仍保持了 GPU 效率。它利用活跃图像的最后提交 token 和草稿窗口来构建验证器输入。通过使用单图位置索引和 KV 缓存位置，这些异构输入可以在单次批处理学生前向传播中共同处理。
- 混合执行策略： 随着图像完成其展开过程，活跃图像的数量会减少。HB-SJD 根据硬件校准的阈值 (γ) 在两种执行模式之间动态切换：
- 全量执行（Full Execution）： 保留原始物理批次形状。已完成的图像被视为逻辑上不活跃（不更新状态），但仍占用内存行。当活跃批次较大时，这种方式非常高效。
- 紧凑执行（Compact Execution）： 仅对剩余的活跃图像运行验证器，仅访问它们对应的 KV 缓存行。这避免了在批次较小时处理不活跃行的开销。
- 初始化与验证：
- 基于历史的初始化： HB-SJD 利用空间局部性，通过固定历史偏移量填充草稿窗口中的先前提交 token，而非随机初始化。
- 验证策略： 系统支持贪婪验证（接受与学生 argmax 预测匹配的草稿）和概率验证（基于概率比率接受草稿，允许非 argmax 提案）。
3. 主要贡献
- 用于视觉 OPD 的 SJD： 本文将投机 Jacobi 解码作为一种展开引擎引入视觉在策略蒸馏。通过重用当前学生在 Jacobi 迭代中的预测，它实现了多 token 展开，且无需额外训练辅助草稿模型。
- HB-SJD 架构： 作者提出 HB-SJD 将 SJD 从单序列推理扩展到大批量视觉 OPD。它独特地结合了独立的单图进度与高效的批处理执行，使得图像可以在保持 GPU 并行性的同时以不同速度推进。
- 即插即用集成： HB-SJD 被设计为学生展开后端的即插即用替代方案。它不会改变现有方法（如 GKD、VarKD）的蒸馏目标或训练程序。
- 混合执行： 该方法引入了混合的全量/紧凑执行策略，能够根据生成过程中不断缩小的活跃批次大小进行自适应，从而优化不同阶段的延迟。
4. 实验结果
研究使用 LlamaGen（B 和 L 版本）在 ImageNet 数据集上进行了实验，将 HB-SJD 与 GKD 和 VarKD 框架下的标准缓存自回归（AR）展开进行了对比。
- 效率： HB-SJD 大幅减少了展开和端到端训练时间。
- 加速比： 根据模型大小和蒸馏方法的不同，实现了 1.48× 至 1.65× 的展开加速（例如，在 LlamaGen-B 上使用贪婪验证的 VarKD 实现了 1.58× 加速）。
- 训练时间： 总训练时间减少了约 1.18×。
- 消融实验： 实验证明独立进度至关重要；批同步 SJD 由于同步等待，实际上比标准缓存 AR 还要慢（加速比为 0.939×）。混合执行策略相比仅使用全量执行模式提供了额外的 1.182× 加速。
- 质量： 生成质量与 AR 基线相当。
- 指标： FID、Inception Score (IS)、精度（Precision）和召回率（Recall）差异极小。例如，使用 GKD+HB-SJD 的 LlamaGen-B 取得的 FID 为 4.95，而 AR 基线为 4.96。
- 定性分析： 视觉样本显示出相当的语义内容、物体结构和局部细节。
- 鲁棒性： 加速效果在以下方面保持一致：
- 不同的展开长度（展开越长，加速效果越明显）。
- 不同的批次大小（8 到 64）。
- 不同的训练阶段（2K 到 22K 步）。
- 不同的 Jacobi 窗口大小（确定 W=16 为平衡设置）。
5. 重要性与主张
本文声称 HB-SJD 解决了视觉 OPD 中的核心矛盾：既需要学生生成的轨迹，又面临生成这些轨迹的高昂成本。通过将展开后端与学习过程解耦，HB-SJD 能够在不损害在策略蒸馏的理论优势或模型生成质量的前提下，实现大幅度的加速。
作者强调，由于 SJD 重用了相同的学生预测，因此其方法不同于语言模型强化学习中的投机解码（后者通常需要与变化的策略保持一致），因为它避免了对独立草稿模型的需求。此外，该方法被呈现为一种实用的工程解决方案，可以集成到现有的视觉自回归蒸馏流水线中，以降低训练成本并保持生成保真度。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题，再试一次？
无垃圾邮件，随时退订。