用批量并行解码加速视觉自回归模型的在线蒸馏训练。
Accelerating Visual On-Policy Distillation with Batched Speculative Jacobi Rollouts
- 提出批量并行的推测雅可比解码，支持多图像独立推进。
- 相比原方法，训练时间减少40%以上，生成质量不变。
- 适合追求高效训练的视觉生成模型研究者使用。
视觉在线蒸馏（OPD）通过当前学生模型生成的轨迹来提升紧凑型视觉自回归模型的训练效果。然而，这些在线回溯仍采用逐标记的自回归解码，显著增加了每个训练步骤的开销。推测雅可比解码（SJD）提供了替代方案，可无需辅助草稿模型并行处理多个标记，但原始方法仅适用于单序列推理。本文提出HB-SJD，一种用于视觉OPD的批量SJD回溯后端。HB-SJD允许每张图像根据自身解码进度独立推进，同时不同序列位置的图像仍通过批量模型前向传播进行验证。当图像完成时，HB-SJD在完整与紧凑执行模式间切换，降低后续回溯轮次的成本。该方法仅替换学生模型的回溯后端，不改变教师模型、蒸馏目标和优化流程。在LlamaGen上的实验表明，HB-SJD显著降低了回溯和端到端训练时间，同时保持了蒸馏学生模型的生成质量。
原文摘要 · Abstract (English)
Visual on-policy distillation (OPD) improves the training of compact visual autoregressive models by learning from trajectories generated by the current student. However, these online rollouts are still produced token by token with autoregressive decoding, which adds substantial cost to every on-policy training step. Speculative Jacobi Decoding (SJD) provides an alternative because it can process multiple tokens in parallel without an auxiliary draft model, but the original method is designed for single-sequence inference. We introduce HB-SJD, a batched SJD rollout backend for visual OPD. HB-SJD allows each image to advance independently according to its own decoding progress, while images at different sequence positions are still verified in batched model forwards. As images finish, HB-SJD switches between Full and Compact execution to reduce the cost of later rollout rounds. HB-SJD only replaces the student rollout backend and leaves the teacher, distillation objective, and optimization procedure unchanged. Experiments with LlamaGen show that HB-SJD substantially reduces rollout and end-to-end training time while preserving the generation quality of the distilled student.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。