ShortOPD:用从短到长的在策略自蒸馏恢复被剪枝的大语言模型 ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation
用从短到长的在策略自蒸馏恢复结构化剪枝后崩溃的生成能力。
前置知识
结构化剪枝与 Block Influence 深度剪枝
结构化剪枝指直接移除整块耦合参数(典型是整个 Transformer 层),从而保留标准稠密架构,无需稀疏内核或特殊硬件支持,对部署最友好。本文用的是深度剪枝:用 Block Influence(BI)指标衡量第 $i$ 层对隐藏状态的改变程度,$\mathrm{BI}_i = 1 - \mathbb{E}[X_{i,t}^\top X_{i+1,t}/(\|X_{i,t}\|_2\|X_{i+1,t}\|_2)]$,BI 越低说明该层越冗余,删掉 BI 最低的若干层即可压缩模型。
本文所有方法都建立在'结构化剪枝后的学生模型'之上,理解 BI 才能看懂为什么删 4 层就让 Qwen3-4B-Instruct 的生成能力崩盘,以及教师为什么是剪枝前的同一个模型。
在策略蒸馏(On-Policy Distillation, OPD)
OPD 让被压缩的学生从自己的分布中采样 rollout $y\sim\pi_\theta(\cdot|x)$,再用其冻结的剪枝前自身 $\pi_T$ 作为教师在相同轨迹上打分,让学生在每个响应位置匹配教师的下一 token 分布,损失为 $\mathcal{L}_{OPD}=\mathbb{E}_x\mathbb{E}_{y\sim\pi_\theta}[|y|^{-1}\sum_t D_\alpha(\pi_T(\cdot|x,y_{<t}),\pi_\theta(\cdot|x,y_{<t}))]$。它在策略(轨迹来自学生)、稠密(每个 token 都有分布目标)、免标签、免验证器、免外部教师。
ShortOPD 就是 OPD 的改进版,核心论点是修复必须在学生实际访问的状态上做稠密 token 级监督,理解 OPD 才能看懂它相对 SFT/SeqKD/KD/RLVR 的本质区别。
广义 Jensen–Shannon 散度(Generalized JSD)
JSD 的推广形式 $D_\alpha(p,q)=\alpha\sum_i p_i\log(p_i/m_i)+(1-\alpha)\sum_i q_i\log(q_i/m_i)$,其中 $m=\alpha p+(1-\alpha)q$。本文取 $\alpha=0.5$,是有界折中:避免前向 KL 的模式覆盖在师生容量差距大时逼学生摊薄概率,也避免反向 KL 的模式寻求降低多样性。附录 C 还证明在师生分布一致时 $D_\alpha$ 的 logit 梯度严格为零并随偏差线性消失。
它是 OPD 损失的核心,也是论文定量刻画'重复后缀是低信息区'(JSD 仅 0.0014)的工具,理解它才能看懂为什么循环尾巴承载的蒸馏信号极弱。
best-of-k 采样与 pass@k
对同一 prompt 采样 $k$ 次响应并取最好结果,pass@$k$ 即 $k$ 次内至少一次通过的概率。它衡量'能力是否可达':若 pass@$k$ 随 $k$ 上升则说明正确轨迹存在于模型采样分布中,只是贪心解码没采到。表 1 显示剪枝模型 GSM8K 在 $k{=}64$ 时 pass@$k$ 达 91.2%,超过未剪枝教师的贪心 88.1%。
这是论文的核心诊断实验,支撑'能力被降级而非擦除'这一动机,直接决定了'在策略搜索 + 稠密修复'的正确性而非监督重学。
知识蒸馏(KD)与 SeqKD
经典知识蒸馏让学生匹配教师在固定序列上的下一 token 分布(teacher-forced forward KL)。SeqKD(序列级蒸馏)则先让教师生成固定响应再做 SFT。两者都是离策略的——训练序列不是学生自己生成的。
它们是论文的主要基线(KD 30.52、SeqKD 28.60 Avg),理解它们才能体会 ShortOPD(48.46 Avg)在策略 + 稠密两项同时满足带来的巨大差距。
研究动机
结构化剪枝(如删除整个 Transformer 层)是硬件友好的压缩手段,但绝大多数剪枝工作只在多选题识别类基准(MMLU、HellaSwag)上验证保留能力,而同样的剪枝检查点在部署真正需要的自由形式生成任务上会崩溃。论文给出一个触目惊心的例子:对 Qwen3-4B-Instruct 用 Block Influence 深度剪枝,仅删除 36 层中的 4 层(约 11% 参数),贪心解码生成能力就在数学、代码、开放生成任务上几乎全军覆没——表 1 显示 GSM8K 的 pass@1 从教师 88.1% 掉到 49.0%,HumanEval 从 75.0% 暴跌到 2.7%,MBPP 从 80.5% 跌到 6.6%。更严重的 25% 参数剪枝后,八任务归一化平均分从 75.17 直接塌缩到 5.71(仅教师的 7.6%)。这种'识别-生成鸿沟'——模型还能选出正确选择题答案却无法生成正确答案——正是阻碍结构化剪枝走向实际部署的根本障碍。
本文的目标是本文的目标是设计一种恢复训练配方,将被结构化剪枝损坏的学生模型的自由形式生成能力修复到接近原始未剪枝教师水平,且不依赖任何标注答案、奖励验证器或外部教师。具体地作者希望:(1) 验证'能力被降级而非擦除'假设——即正确轨迹仍存活在压缩模型自身采样分布中;(2) 找到能在学生实际访问的状态上做 token 级稠密修复的训练信号;(3) 在 25% 参数剪枝的 Qwen3-4B-Instruct 上把八任务平均分从 5.71 恢复到尽量接近教师 75.17 的水平;(4) 让恢复高效——既要质量高又要省 rollout token 和训练时间,避免在长 rollout 上浪费教师算力。最终 ShortOPD 在主表上把平均分恢复到 48.46(教师的 64.5%),约为未恢复值的 9 倍,并以 8.5 小时和 250M rollout token 匹配了需要 35.9 小时、869M token 的固定 8192 长 rollout。
与已有工作不同的是,本文独特切入有三点。第一,把 Wen 等人'短答案被降级而非擦除'的观点从短答案推广到整条生成轨迹——通过表 1 的 pass@k 实验证明剪枝模型的正确轨迹仍在其自身采样分布中(GSM8K 在 $k{=}64$ 时 pass@k 达 91.2%,甚至超过未剪枝教师贪心 88.1%),因此修复本质是'在学生自己的搜索空间内做分布性修复'而非监督式重学。第二,明确指出修复必须在学生实际访问的状态上做稠密 token 级监督:稀疏奖励(RLVR)和固定离策略标签(SeqKD)无法定位学生在解码路径上累积的局部错误。第三,提出并解决此前无人触及的问题——强剪枝下早期在策略 rollout 被重复后缀主导(25% 剪枝时后缀重复率达 84%),这些循环尾巴承载的蒸馏信号极弱(教师-学生 JSD 仅 0.0014,是普通 token 的 1/35;教师 NLL 仅 $3\times10^{-5}$ vs 普通 0.68),而 batch 均损失在重复最严重时反而最低,导致用损失调度会误判。ShortOPD 正是用独立的重复率门控来分配预算填补此空白。
核心方法
整体思路分两层。直觉上:既然剪枝后正确轨迹仍存活但被'降级',最合理的修复是让学生'沿自己的轨迹走一遍',让剪枝前的自己逐 token 纠正。这天然导出在策略蒸馏(OPD):学生从自身分布采样 rollout $y\sim\pi_\theta(\cdot|x)$,冻结教师 $\pi_T$ 在相同轨迹打分,让学生在每个响应位置匹配教师下一 token 分布,$D_\alpha$ 取 $\alpha{=}0.5$ 的广义 JSD,蒸馏 top-100 logits 加聚合尾质量并 clip 重要性权重。第二层创新是预算控制器:OPD 默认学生 rollout 可用,但强剪枝下早期长 rollout 会在恢复前先陷入重复后缀。ShortOPD 保留全局 response_length 作上下文上限,但为每步设动态预算 $H_s$;控制器观测末端循环率、干净截断率、有效长度三统计,以重复率作一级门控——高重复缩向有效长度 EMA,低重复且高截断才增长,第二个 EMA 平滑过渡,构成图 3 围绕 OPD 的第二闭环。
核心创新是 ShortOPD 的'重复门控、截断感知、从短到长'预算控制器(公式 4)。与已有方法的本质区别:几乎所有蒸馏/恢复工作(SeqKD、KD、Minitron)都是离策略的——在固定外部序列上训练;而 RLVR 虽在策略但奖励稀疏,当压缩模型很少采到正确答案时'几乎没什么可提升的'(PPO 在 GSM8K 仅 0.23%、GRPO 仅 1.59%,而 ShortOPD 同数据达 37.76%)。ShortOPD 同时拥有在策略状态(学生自己的 rollout)和稠密信号(教师 token 分布),是表 2 设计空间里唯一四项全中的方法。更关键的是它解决 OPD 自身的痛点:早期 rollout 被重复后缀主导,而 batch 均损失恰恰在重复最严重时最低(图 2b),导致用损失调度会误判。ShortOPD 的解法是用独立于损失的纯 token 末端周期循环探测器来决定是否进缩门,再用 EMA 平滑过渡。如此既保证早期教师算力集中在高信号前缀,又能在恢复后重新释放长生成能力。
方法步骤详情
ShortOPD 每步先以预算 $H_s$ 采样一批在策略 rollout;再用算法 2 在末端 512 token 上测周期 1–10,由循环标记 $a_i$ 与有效长度 $\ell_i$ 定位重复(OPD 损失与教师 NLL 仅精化边界、不参与判定);随后冻结教师 $\pi_T$ 对相同轨迹打分,对所有已生成 token 施加未修改的稠密 OPD 损失更新参数。控制器据此算出批次重复率、干净截断率与平均有效长度并更新 EMA:高重复触发缩门把目标压向有效长度的 $\lambda$ 倍,低重复且干净截断触发增长门乘以 $\gamma_\uparrow$,否则保持;最后用第二个 EMA 把预算平滑到目标、按 16 取整并 clip 到 $[H_{\min},H_{\max}]$。初始化 $H_1{=}H_{\max}$,超参 $\rho_{low},\rho_{high},\tau=0.20,0.45,0.10$、$\lambda{=}1.15$、$\gamma_\uparrow{=}1.25$,初始/边界 $2048/[1024,2048]$。
技术新颖性
新颖性体现在'把外环预算控制器嫁接到 OPD'及'用独立于损失的纯 token 探测器做门控'。与 Vision-OPD 把 OPD 引入多模态相比,本文是首个把 OPD 用于结构化剪枝后生成恢复、且教师是模型自身剪枝前快照的场景;此前恢复重训(Sheared LLaMA、Minitron)都离策略。其次,作者首次定量刻画'重复后缀是低信息区'——重复 token 的师生 JSD 仅 0.0014(普通的 1/35)、教师 NLL 仅 $3\times10^{-5}$ vs 普通 0.68,附录 C 还证明师生一致时广义 JSD 梯度严格为零并随偏差线性消失。第三,把恢复信号组织成'在策略 vs 离策略 × 稠密 vs 稀疏'二维设计空间(表 2),ShortOPD 是唯一在策略+稠密+免标签+免验证器的方法。第四,leave-one-domain-out 把'恢复语料只是 prompt 分布'提升为一等设计要素——去数学 prompt 使 GSM8K 从 62.70 暴跌到 8.04,去代码 prompt 使 HumanEval 从 43.90 崩到 0.61。
实验结果
核心发现按实验逐一分析。第一(主表 3,25% 剪枝 Qwen3-4B-Instruct):未恢复 Avg 仅 5.71,SFT 21.19,SeqKD 28.60,KD 30.52,ShortOPD 达 48.46(教师 64.5%),约为未恢复 9 倍、较 KD 提升 17.9 点,八任务全面领先,证明在策略恢复并非数学技巧。第二(图 4–5,控制器):相对固定 $H{=}2048$ 的 Vanilla OPD,ShortOPD 前百步每步生成时间 35.4→25.1s(-29%),端到端 11.1→8.5h;三调度 Avg 仅 48.5–50.2(< 2 点)但 token 从 250M 到 869M,比 8192 省 76% 时间。第三(表 4):训到 3 epoch 可扩到 Avg 55.41。第四(图 6):同 GSM8K-only 数据 RLVR-PPO 仅 0.23、GRPO 仅 1.59,而 ShortOPD 同数据达 37.76。第五(表 5):去数学 prompt 使 GSM8K 暴跌到 8.04、去代码 prompt 使 HumanEval 崩到 0.61。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 八任务平均生成恢复(25% 剪枝 Qwen3-4B-Instruct) | Avg(开放任务 judge 分×10,八任务平均) | ShortOPD 48.46(教师的 64.5%) | 未恢复 5.71;SFT w/o KD 21.19;SeqKD 28.60;KD 30.52 | 较最强离策略 KD +17.9 点,约为未恢复值的 9 倍 |
| 训练效率(rollout 调度对比) | 端到端训练时间 + 生成 rollout token | ShortOPD 8.5h / 250M token / Avg 48.5 | 固定 8192:35.9h / 869M token / Avg 50.2 | 时间 -76%、token -71%,平均分差仅 1.7 点 |
| GSM8K-only 稀疏 vs 稠密监督 | GSM8K pass@1 准确率(%) | ShortOPD 37.76(多域版 62.70) | RLVR-PPO 0.23;RLVR-GRPO 1.59 | 稠密教师信号在同数据同预算下碾压稀疏奖励 |
| 多 epoch 扩展性 | 八任务 Avg(教师占比 T%) | 1/2/3 epoch:48.46/53.45/55.41(T% 64.5→73.7) | 未恢复 5.71(T% 7.6) | 持续训练不断把 rollout 曝光转化为生成质量 |
| 恢复语料域依赖(leave-one-domain-out) | 各域 pass@1 / judge 分 | Full mixture:GSM8K 62.70、HumanEval 43.90 | w/o math:GSM8K 8.04;w/o code:HumanEval 0.61 | 证明语料覆盖是一等设计要素,缺域对应能力崩盘 |
局限与改进
作者明确承认:本文只在 Qwen3-4B 家族上、用 BI 深度剪枝、在约 25% 删除参数下做了实验,要声称普适性还需测试更多剪枝方法、模型家族、压缩比和更大模型;其次,论文未确定'轻量恢复训练在何种压缩程度下失效'的边界——附录 E 显示超过约 13 层剪枝后模型进入不连贯的 token 沙拉区,重复率反而崩塌、循环不再形成,此时 OPD 是否仍可恢复未知。我自己补充几点观察:第一,所有实验都是单一规模(4B)单一剪枝方法,缺少宽度剪枝、不同剪枝准则的对照,'教师是自身剪枝前快照'这一设定在宽度剪枝下是否同样有效未验证。第二,恢复质量仍远低于教师(生成 64.5%、多选题 67.6%),且固定 8192 调度还能再涨(50.2),说明动态预算控制器偏保守。第三,多选题表 9 揭示 ShortOPD(67.6 Avg)反而低于 SFT/KD(74.9/74.4),需靠 SFT-init 组合才能补齐到教师水平(76.8),说明该方法在'识别'维度并非优势。
独立分析的弱点
第一个弱点是'单一压缩方式与单一规模',改进方向是扩展到宽度剪枝、Sheared LLaMA 式学习式剪枝、稀疏量化组合,并在 7B–70B 上验证 $H_{\max}$、$\lambda$、$\rho_{high}$ 是否随规模漂移。第二个是'预算控制器偏保守'——固定 8192 还能涨分(50.2 vs 48.5),说明上界 $H_{\max}{=}2048$ 过低,改进方向是让增长门 $\gamma_\uparrow$ 自适应于干净截断历史或在后期抬高 $H_{\max}$。第三个是'识别维度劣化':多选题平均仅 67.6,落后 SFT/KD 约 7 点,需 SFT-init 组合才能补齐,改进方向是设计联合目标同时优化在策略分布与外部序列似然。第四个是'语料覆盖作隐性先验':要恢复某能力须预先有对应域 prompt,改进方向是自动 prompt 域扩展或课程式域注入。第五个是'循环探测器只看末端 $W{=}512$ token 且只测周期 1–10',对长周期或非周期同义复述会漏检,改进方向是引入基于 n-gram 覆盖率或 entropy 的更鲁棒重复度量。
未来方向
作者明确指出三方向:测试更多剪枝方法、模型家族、压缩比和更大模型以确立普适性;进一步研究恢复的扩展性(论文已展示 1→3 epoch 持续涨分,但受单 epoch 预算约束);确定模型被压得太狠、轻量恢复失效的边界。基于成果可延伸的方向:把 ShortOPD 的外环控制器思想迁移到其他在策略训练(RLHF、GRPO 的 rollout 调度、agentic 长轨迹 RL),解决'早期 rollout 质量差'的通病;把'教师是模型自身早期快照'这一自蒸馏框架与持续学习结合,做模型迭代中的能力保全;把重复门控推广为更通用的'信号密度门控'——用教师-学生 JSD/teacher NLL 直接估计每个 token 的信号强度并据此分配预算,而非仅依赖周期探测器;研究 OPD 在不同教师容量差距下的 KL 方向选择(forward KL 模式覆盖 vs reverse KL 模式寻求),并在恢复后期切换。
复现评估
复现评估偏中等。论文给出相当详尽细节:剪枝配置(删 36 层中索引 25–33 的 9 个最低 BI 层,约 25% 参数,PG-19 校准)、恢复语料(45,447 prompt:GSM8K 7,473+MATH 7,500+OpenCodeInstruct 15,000+MBPP 474+UltraChat/ShareGPT 15,000)、训练超参(rollout group 8、batch 64、长度各 4096、温度 0.8、学习率 $2\times10^{-6}$、top-100+tail JSD $\alpha{=}0.5$、clip 2.0、710 步/epoch、8×H20),Table 7 给出探测器与控制器全部超参。两个核心构建块都指向公开实现(ShortX 剪枝、Vision-OPD 蒸馏)。阻碍在于:恢复语料需自行构建且过滤规则繁琐(OpenCodeInstruct 要正单元测试、去重、移除与 HumanEval/MBPP 重名函数);开放任务 judge 用 GPT-5.5 引入未充分披露的随机性。综合看方法学层面可复现,端到端独立复现需中等工程量。
论文图表
两幅子图。(a) 横轴是移除层数(0–12),左纵轴是后缀重复率,右纵轴标注 code/math/open-ended 三条曲线,星标处为本文主设定 9 层(25% 参数),显示重复率随剪枝深度单调上升,到 25% 时高达 84%。(b) 横轴训练步 0–100,蓝线是蒸馏损失、红线是后缀重复率,阴影标出重复主导的 warm-up:前 80 步内 55–75% 的 rollout 陷入循环,而损失恰恰处于最低水平,直到约 80 步后重复消退损失才进入 plateau。
这张图揭示了 ShortOPD 要解决的核心矛盾:早期在策略 rollout 被重复主导(84%),而 batch 均损失在重复最严重时反而最低,证明不能靠损失来调度预算——这是 ShortOPD 引入独立重复探测器的根本理由。
横轴是 BI 顺序下移除层数 0–35,左纵轴后缀重复率、右纵轴 distinct-2 多样性。三阶段:$k\le12$ 时重复率随深度单调上升、distinct-2 下降(连贯生成,循环是主要退化模式,星标 9 层为本文设定);$k\approx13$ 之后输出崩解为不连贯 token 序列,distinct-2 回弹接近随机文本、循环不再形成、重复率反而崩塌;$k\ge35$ 时近空栈退化为单 token 循环,重复率饱和 100%。
这张图界定了方法的适用边界——重复是连贯生成区($k\le12$)的特征失效模式,更深剪枝进入不连贯区后 OPD 能否恢复未知,直接支撑论文承认的局限性。