加速视觉策略内蒸馏：基于批处理推测性雅可比展开

# 通过批量雅可比推测展开加速视觉在线策略蒸馏  
来源：https://arxiv.org/html/2608.18183  

**作者**：Zhehao Yu、Kenhong Lin、Baoquan Zhang  
**单位**：互联网信息协作深圳重点实验室；哈尔滨工业大学（深圳）；中国深圳  
**邮箱**：[{23s051004,yuzhehao,linkenghong}@stu.hit.edu.cn](mailto:) [[email protected]](mailto:)  

#### 摘要  
视觉在线策略蒸馏（OPD）通过学习当前学生模型生成的轨迹，提升了紧凑型视觉自回归模型的训练效果。然而，这些在线展开生成仍依赖自回归解码逐词元生成，为每个在线策略训练步骤带来了巨大成本。雅可比推测解码（SJD）提供了一种替代方案，因为它无需辅助草稿模型即可并行处理多个词元，但原方法仅针对单序列推理设计。我们提出了HB-SJD，一种用于视觉OPD的批量SJD展开后端。HB-SJD允许每张图像根据自身解码进度独立推进，同时不同序列位置的图像仍可通过批量模型前向进行验证。当图像完成生成时，HB-SJD在完整执行与紧凑执行模式之间切换，以降低后期展开轮次的成本。HB-SJD仅替换学生模型的展开后端，而教师模型、蒸馏目标和优化流程保持不变。在LlamaGen上的实验表明，HB-SJD显著缩短了展开生成及端到端训练时间，同时保持了蒸馏学生模型的生成质量。  

## 1 引言  
视觉自回归（AR）模型在下一词元和下一尺度图像生成任务中展现出强大的生成质量与可扩展性。随着模型规模增大，知识蒸馏提供了一种将能力迁移至更小学生的实用途径。近期工作进一步将蒸馏扩展至生成模型，通过匹配序列级行为并在学生生成的轨迹上进行训练。然而，传统蒸馏主要在数据或教师轨迹上训练学生，而学生在推理时必须以其自身先前输出为条件。在线策略蒸馏通过让当前学生生成训练轨迹，并利用教师监督学生访问的状态，减少了这种不匹配。这虽更好地匹配了学生的推理行为，但也需要在线学生展开生成，为每次在线策略更新增加了显著成本。  

然而，训练与推理更紧密的匹配伴随着明确代价：OPD将自回归生成直接纳入训练循环。在教师-学生评分开始前，当前学生必须先生成展开后缀。即使使用KV缓存，生成LL词元后缀仍需LL次顺序解码步骤。如图1(a)所示，该展开生成位于真实前缀与完成的学生轨迹之间，属于训练关键路径的一部分。同样的瓶颈也出现在GKD及近期视觉自回归蒸馏方法（如VarKD）中。  

这产生了一个核心矛盾：OPD受益于当前学生生成的轨迹，但生成这些轨迹本身成本高昂。我们能否在不改变OPD学习过程的前提下加速学生展开生成？  

（a）视觉OPD中的学生展开替换。（b）成本与性能对比。  
图1：HB-SJD概述及其主要优势。（a）HB-SJD用并行提议、验证和多词元承诺替代了逐词元自回归学生展开。（b）HB-SJD在保持相当生成质量的同时，减少了展开生成及端到端训练时间。  

为解决这一问题，我们转向雅可比推测解码（SJD）作为视觉在线策略蒸馏的学生展开引擎。近期研究也探索了将推测解码用于在线强化学习展开，显示了其在训练中降低展开成本的潜力。与这些语言模型强化学习场景不同，视觉OPD需从当前学生重复生成大批量图像轨迹。SJD非常适合此场景，因为它重用同一学生的预测作为草稿词元，且无需单独训练草稿模型。然而，原始SJD专为单序列推理设计。为每张图像单独运行会丧失批量并行性，而同步处理所有图像则会导致较快图像等待较慢图像。因此，高效地将SJD适配至大批量视觉OPD是主要挑战。  

我们提出**混合批量雅可比推测解码（HB-SJD）**，一种用于视觉OPD的批量SJD展开引擎。HB-SJD允许每张图像根据自身解码进度推进，同时仍能通过批量学生前向验证不同序列位置的图像。它进一步采用混合完整与紧凑执行模式，以在图像完成时保持后期展开轮次的效率。如图1(a)所示，HB-SJD仅替换学生展开后端，而OPD训练流程的其余部分保持不变。图1(b)表明，该替换在保持相当生成质量的同时，显著减少了展开生成及端到端训练时间。  

我们的贡献总结如下：  
- • 我们引入雅可比推测解码作为视觉在线策略蒸馏的展开引擎。通过重用当前学生的预测，SJD无需辅助草稿模型即可实现多词元展开。  
- • 我们提出HB-SJD，将SJD从单序列推理扩展至大批量视觉OPD。HB-SJD将独立的每图像进度与高效的批量执行相结合，允许不同图像独立推进，同时保留GPU并行性。  
- • 我们将HB-SJD集成为现有视觉OPD方法的即插即用展开后端，不改变其蒸馏目标或训练流程。实验表明，HB-SJD在保持相当生成质量的同时，显著缩短了展开生成及端到端训练时间。  

## 2 相关工作  
### 2.1 视觉自回归蒸馏  
知识蒸馏通过分布级或序列级监督，将知识从大型教师迁移至小型学生。对于生成模型，近期方法进一步研究序列级匹配及在学生生成轨迹上的训练。MiniLLM改进了语言模型的序列级蒸馏，而GKD则引入了使用当前学生生成样本的在线策略训练。同时，VAR和LlamaGen等视觉自回归模型通过下一尺度和下一词元图像生成展现了强大的可扩展性。VarKD进一步研究了视觉自回归模型的监督蒸馏、序列级蒸馏及在线策略蒸馏。在视觉OPD中，学生生成的轨迹成为训练循环的一部分，使得展开效率成为蒸馏目标之外的重要实际因素。这些工作主要关注使用何种监督以及如何优化学生模型。我们的工作与其互补：我们保持蒸馏框架不变，专注于减少视觉在线策略蒸馏中生成学生轨迹的成本。  

### 2.2 自回归生成加速  
推测解码通过提议多个未来词元并并行验证，降低了自回归生成的顺序成本。类似思想已扩展至视觉自回归生成。ZipAR利用空间结构，连续推测解码研究连续视觉分布，而LANTERN、LANTERN++和GSD则采用松弛验证以提升图像生成加速。雅可比推测解码（SJD）采用不同方法：它使用先前雅可比迭代的预测作为提议，因此无需辅助草稿模型。后续变体进一步改进了提议重用、稳定性、验证和延续性。其他基于雅可比的方法探索了额外的空间或路径级并行，如并行雅可比解码和PathRelax。这些方法大多关注推理时的生成加速。批量推测解码也在推理服务中被研究，其中不同接受长度导致不同请求的解码进度差异。近期，推测解码被用于降低语言模型强化学习训练中的在线展开成本。我们的场景不同：视觉OPD需要当前学生在训练中重复生成大批量图像轨迹。因此，我们通过独立的每图像进度和高效的批量执行将SJD适配至该场景，同时保持教师模型和蒸馏目标不变。  

## 3 基础知识  
### 3.1 视觉在线策略蒸馏  
令$x_{1:T} = (x_1, \ldots, x_T)$表示词元化的图像。视觉自回归模型将其分布分解为 $p_\theta(x_{1:T}) = \prod_{t=1}^T p_\theta(x_t \mid x_{<t})$。学生模型 $p_\theta$ 从教师模型 $p_\phi$ 处学习，最小化蒸馏损失 $L_{KD}$。在线策略蒸馏中，学生首先自回归生成轨迹 $x_{1:T}^{(s)}$，随后教师对该轨迹的前缀 $x_{<t}^{(s)}$ 提供监督信号，引导学生生成每个词元 $x_t^{(s)}$。  

### 3.2 雅可比推测解码  
雅可比推测解码（SJD）将自回归生成视为求解方程组。令 $x_1, \ldots, x_T$ 为待生成词元序列。在每次雅可比迭代 $k$ 中，模型预测所有位置的词元 $\hat{x}_t^{(k)} = \arg\max_{x_t} p_\theta(x_t \mid \hat{x}_1^{(k-1)}, \ldots, \hat{x}_{t-1}^{(k-1)})$。当序列收敛（即 $\hat{x}_t^{(k)} = \hat{x}_t^{(k-1)}$ 对所有 $t$ 成立）时终止。为加速收敛，SJD将雅可比迭代解释为推测解码的提议步骤：当前迭代的预测作为草稿，下一次迭代的验证决定接受或拒绝草稿词元。  

## 4 方法  
### 4.1 批量雅可比推测展开（HB-SJD）  
我们提出HB-SJD，一种针对视觉OPD的批量SJD展开引擎。HB-SJD在保留GPU并行性的同时，允许每张图像独立推进解码。  

**独立进度管理**。每张图像维护自身进度指针 $pos^{(i)}$，记录已接受的词元数。在每轮验证中，仅当图像 $i$ 的 $pos^{(i)} < T$ 时才参与。这避免了快速图像等待慢速图像。  

**批量验证**。我们将不同进度位置的图像组织成批量，执行一次学生模型前向传播。对于位置 $t$ 的验证，输入为所有图像在 $pos^{(i)} = t$ 处的上下文 $x_{<t}^{(i)}$，输出为预测分布 $p_\theta(x_t \mid x_{<t}^{(i)})$。  

**混合执行模式**。设 $A^{(r)}$ 为第 $r$ 轮验证中活跃图像数。若 $A^{(r)} > \gamma$（$\gamma$ 为阈值），采用**完整执行**：模型计算所有活跃图像在对应位置的预测。若 $A^{(r)} \leq \gamma$，则切换至**紧凑执行**：仅对剩余图像运行模型，避免对已完成图像进行不必要计算。形式上：  

$$
\mathcal{E}^{(r)} =
\begin{cases}
\mathrm{Full}, & A^{(r)} > \gamma, \\
\mathrm{Compact}, & A^{(r)} \leq \gamma.
\end{cases}
$$  

该混合策略在活跃图像较多时保持完整批量路径，而在仅剩少量图像时避免不必要的计算。  

## 5 实验  
### 5.1 实验设置  
#### 训练与评估。  
我们遵循VarKD的LlamaGen实验设置，包括ImageNet训练与评估协议。所有比较的展开方法使用相同的教师-学生配置、优化设置及训练步数。生成质量在50K样本上通过FID、Inception Score（IS）、精度和召回率评估。报告的主要延迟结果取三次独立运行的平均值。  

#### 展开设置。  
我们将HB-SJD与相同训练设置下带KV缓存的自回归展开进行比较。我们将这种带KV缓存的逐词元自回归展开称为***Cached AR***。除非另有说明，HB-SJD使用雅可比窗口大小 $W=16$、历史偏移 $H=24$ 及无分类器指导尺度2.0。我们评估两种验证策略：***贪心验证***，根据当前学生的贪心预测验证草稿；***概率验证***，遵循SJD的概率接受与重采样规则。我们报告展开延迟和端到端训练时间以评估效率。对于延迟分布，P95表示第95百分位延迟。  

#### 基线与评估。  
对于LlamaGen、KD和SeqKD，我们报告VarKD在其ImageNet评估协议下的生成质量结果。对于GKD、VarKD及其HB-SJD变体，我们在相同训练与评估设置下复现结果。所有展开效率结果在相同的NVIDIA H200 GPU上以相同设置测量。  

### 5.2 主要结果  
#### 定量结果。  
表1比较了HB-SJD与视觉自回归蒸馏基线在LlamaGen-B和LlamaGen-L上的表现。HB-SJD仅替换学生展开后端，蒸馏目标和训练流程保持不变。使用贪心验证时，HB-SJD在LlamaGen-B和LlamaGen-L上的GKD实现1.48×和1.56×的展开加速，VarKD的加速比达1.58×和1.65×。概率验证也提供一致加速，在不同设置下实现1.46–1.54×的加速。同时，FID、IS、精度和召回率与AR展开基线相近。这些结果表明，HB-SJD在不同模型规模和蒸馏方法下一致地降低了展开成本，同时在所有评估训练设置中保持了相当的生成质量。  

#### 定性结果。  
图4展示了以LlamaGen-XL为教师、LlamaGen-L为学生的定性结果。与相应的GKD和VarKD基线相比，使用HB-SJD训练的模型在语义内容、物体结构和局部细节上生成了视觉上可比的样本。这些结果与定量评估一致，表明用HB-SJD替代自回归展开并未引入明显的生成质量下降。  

### 5.3 消融研究  
#### 独立进度对批量SJD是否必要？  
为回答此问题，我们在相同展开工作量下比较了Cached AR、**批量同步SJD**和**独立批量SJD**。其中独立模式使用完整执行。如表3所示，批量同步SJD需要241轮验证，耗时4239毫秒，略慢于Cached AR的3982毫秒。相比之下，允许每张图像独立推进将验证轮数减少至176轮，展开延迟降至3275毫秒，相比Cached AR实现1.216×的加速。该结果表明，仅保持批量模型执行并不足够；独立的每图像进度对于高效将SJD适配至大批量视觉OPD至关重要。