The Weakest Link: Distilling LLM reasoning with worst-case constrained reinforcement learning
The Weakest Link: Distilling LLM Reasoning with Worst-Case Constrained Reinforcement Learning
作者用最坏前缀约束蒸馏推理模型,报告数学设置的推理成功率提高,代码正确率则下降。
- 只奖励最终答案可能保留错误的中间推理。作者认为,整段平均的教师偏离惩罚允许其他高概率token补偿局部错误,因此不能替代对推理路径的约束。
- 作者把蒸馏写成约束强化学习:限制每个前缀的平均教师负对数似然,一旦越界便保留不可行状态。训练使用教师评估成本,部署不查询教师,并拆分即时与长期梯度。
- Table 1中,所提方法的宽松预算版本在三个数学设置均取得最高RSR;GSM-Symbolic的FAC略低于GRPO。Table 2中,代码RSR最高,但pass@1低于列出的全部基线。
- 附录G指出,教师似然不能完整刻画逻辑有效性,方法要求兼容tokenizer,预算仍需调参,OOD可行性保证尚待研究。实验涉及数学与代码;人工复核为60个判定。
- Evaluated models
- Qwen2.5-1.5BLlama-3.2-3BQwen2.5-Coder-0.5B
- Benchmarks
- MATHGSM8KApple/GSM-Symbolic (main)MBPPAPPS
- Metrics
- Final Answer Correctness (FAC)Constraint Satisfaction (CS)Reasoning Success Rate (RSR)Average Win Rate (%)pass@1Spearman ρ
To address the challenge of distilling LLM reasoning capabilities into smaller models, the researchers propose modeling reasoning distillation as a reinforcement learning problem with worst-case constraints—applying constraints on teacher log-likelihood at every trajectory prefix, rather than averaging soft divergence penalties. The method derives an augmentation-free constrained MDP whose reward transformation preserves the hard constraint semantics while decomposing into low-variance single-step and long-term policy gradient terms, and satisfies the worst-case constraint almost surely in the penalty limit. In experiments on mathematical reasoning and code generation tasks, the method significantly broadens the accuracy-fidelity Pareto frontier, matches the high final-answer correctness of pure RL, and greatly reduces violations of teacher constraints, achieving the highest strict reasoning success rate across all evaluation settings.
Paper analysis
What problem does this paper attempt to solve?
作者认为,最终答案奖励与平均蒸馏惩罚都可能保留中间推理错误。
核心问题是如何在蒸馏中同时优化最终答案和中间推理,而不允许整段平均惩罚掩盖局部偏离。
- 应用场景:把大教师模型的推理能力迁移到小学生模型。作者将资源受约束的部署作为动机。
- 答案奖励的盲点:作者称,只优化可验证任务奖励会出现reward hacking:学生通过错误中间步骤得到正确答案。Figure 1和附录H给出了生成示例及评审意见。
- 软惩罚的补偿问题:作者认为,序列平均的教师偏离成本允许其他高概率token抵消局部高成本;提高惩罚又可能牺牲任务表现。
- 核心假设:作者用教师似然作为推理有效性的代理,并认为已越界的前缀不应被后续token修复。附录E检验的是似然与RSR的相关性。
- 提出的方法:使用worst-case Constrained RL约束所有前缀的平均成本,构造无需状态扩展的奖励,并拆分即时与长期梯度。它是训练方法,不是外部攻击或防御评测。
What related research exists?
作者把工作定位在RL蒸馏、过程监督与约束强化学习的交叉处。
相关工作围绕蒸馏目标、过程监督和约束求解组织,比较基线另有具体实现说明。
RL蒸馏与reward hacking
- 分布匹配蒸馏:Hinton et al.(2015)、Sanh et al.(2020)、Gu et al.(2024)等匹配教师输出分布。作者结合Zhang et al.(2025)讨论学生容量与教师复杂推理的差距。
- 任务奖励与软正则:Ko et al.(2024;2025)被用于讨论引入可验证奖励;Agarwal et al.(2024)等被用于讨论任务奖励与教师KL正则的组合。
- 奖励投机:作者引用Gudibande et al.(2024)讨论正确答案与有效推理之间的分离,并将平均软惩罚下的补偿作为自己的问题定位。
增强推理蒸馏
- 过程监督:Hsieh et al.(2023)、Adarsh et al.(2024)、Chen et al.(2025)监督教师的中间逻辑步骤。
- token加权:Li et al.(2024;2025)、Saadi & Wang(2025)使用注意力或Bayesian启发式调整关键token的蒸馏损失权重。
- 知识增强:Kang et al.(2023)、Tian et al.(2025)借助外部知识补充推理。作者对这一组方法概括的代价包括步骤标注、过程奖励模型或检索开销。
约束RL与状态扩展
- 约束优化:Achiam et al.(2017)、Schulman et al.(2015)等被用于讨论防止偏离;作者称,LLM规模的对偶优化存在计算与方差成本。
- Saute:Sootla et al.(2022a;2022b)、Calvo-Fullana et al.(2024)通过增加剩余预算状态和重塑奖励处理约束。作者指出,直接用于蒸馏会涉及长度相关成本和测试时教师查询。
实验基线
- 基线方法与基准:数学比较SFT on Teacher CoT、GKD、MiniLLM、GRPO、GKD-GRPO,使用MATH及GSM8K→GSM-Symbolic;代码在MBPP上比较GRPO、GKD-GRPO。
作者将区别定位为:用最坏前缀平均替代整段软平均,保留越界后的不可行性,并利用完整历史取消显式预算状态;这属于作者的方法主张。
How does the paper address this problem?
方法约束每个前缀的平均教师成本,越界后不可恢复,测试时不查询教师。
worst-case Constrained RL在任务奖励之外设置前缀可行性边界,再通过奖励变换和梯度拆分训练学生。
状态、成本与约束
- 生成过程:状态是提示词及已生成前缀,动作是下一个token,转移为确定性的前缀追加;教师冻结。
- 成本代理:C(st,at)=−logμ(at∣ st)。成本是教师对学生所选token的负对数似然,作者用它代理推理偏离。
- 核心约束:maxT≥11/T∑t=0T−1C(st,at)≤ d。它要求每个前缀的平均成本不超过预算,同时最大化任务回报(Eq. 4)。
- 约束含义:它限制前缀平均,不能直接等同于逐步骤逻辑正确。Proposition 3.1给出的单步成本界仍随该步骤的前缀长度增长。
越界后保持不可行,无需扩展状态
- 奖励变换:可行时保留原任务奖励,否则使用−n惩罚;Lemma 3.2称,历史中已有的越界不会因追加token消失。
- 历史作为状态:作者认为,LLM已条件化于完整前缀,无需再把剩余预算作为输入。训练rollout后由教师计算成本,测试时只运行学生。
- 理论条件:Theorem 3.3的几乎必然可行性成立于n趋于无穷,且极限MDP存在有限价值最优策略的条件下。附录A.2还假设奖励有界、离散生成和正成本。
- 预算可行性:附录A.2指出,若预算低于可实现的最坏成本,所有策略都会在极限受到无穷惩罚,有限价值假设不成立。
即时项与长期项的梯度拆分
- 即时项:对当前词表上的动作求精确期望;教师前向计算提供各候选token概率,随后确定是否越界(第3.3节、附录A.3)。
- 长期项:保留Monte Carlo未来回报,用于把后续越界的影响归因到较早的token。
- 方差结论:附录A.4的Rao-Blackwellization定理约束的是即时梯度分量的协方差迹;严格降低还要求条件方差非零,不是整个训练过程的统一保证。
实际训练与奖励
- 共同训练框架:所有RL方法使用GRPO的组相对优势和PPO式裁剪;数学任务终局奖励为正确1.0、错误0.0(附录B)。
- 训练设置:每次更新64个回答,来自8题各8个样本;AdamW、学习率1×10⁻⁵、γ=1.0、bfloat16,数学训练20 epochs。
- 有限惩罚与预算:实际n=20;预算比较−ln(1/4)、−ln(1/2)、−ln(3/4)。因此实验不等于无穷惩罚定理的条件。
- 模型特例:Llama-3.2-3B的全部被比较方法先接受3 epochs的GKD预热。代码使用512-token completion、temperature 1.0,训练15 epochs(附录C)。
代理选择与推理判定
- 替代约束:附录D称,逐token约束会混合格式不确定性与逻辑错误;滑动窗口可被后续低成本token修复,失去吸收性。
- 独立审计:RSR要求最终答案正确且评审判为No Error;教师成本不直接充当RSR标签。附录B.5给出了完整评审提示词。
What experiments did the paper conduct?
数学主表中宽松预算版本的RSR最高;代码实验的RSR提高伴随pass@1下降。
实验分别测量答案正确、教师约束满足和评审认可的推理成功,三者不能互换。
实验设置
- 数学模型:Qwen2.5-1.5B由Qwen2.5-Math-7B-Instruct蒸馏;Llama-3.2-3B由Llama-3.2-11B-Vision-Instruct蒸馏。教师与原始学生也出现在Table 1参照行。
- 代码模型:Qwen2.5-Coder-0.5B(base)由Qwen2.5-Coder-7B-Instruct蒸馏(附录C)。
- 数据与规模:数学为MATH训练/评测,以及GSM8K训练、Apple/GSM-Symbolic (main)评测。代码训练约4K题,包括464道MBPP题及筛选后的APPS题;MBPP测试500题。
- 基线:数学包括SFT on Teacher CoT、GKD、MiniLLM、GRPO,以及五档λ的GKD-GRPO;所提方法比较三档d。代码Table 2列出GRPO及两档GKD-GRPO。
- 指标与评审:FAC匹配提取答案;CS按统一d=−ln(1/2)计算测试轨迹比例;RSR要求答案正确且DeepSeek-R1-Distill-Qwen-32B判为No Error。Pairwise只比较双方答案正确的样本,须在两种顺序均获胜。
- 采样与复核:附录E报告Qwen-MATH每run 5,000代,共12 runs、60,000轨迹;附录D每题256条轨迹。附录F人工复核60个判定;主表未报告独立重复训练次数或误差条。
主结果
据Table 1,以下数值均为百分比。Qwen指Qwen2.5-1.5B,Llama指Llama-3.2-3B;GSM-Sym.设置使用GSM8K训练。Ours均为训练d=−ln(1/4),所有行CS均按评测d=−ln(1/2)计算。
Swipe horizontally to read the table
学生/评测集 方法 RSR FAC CS Qwen/MATH GRPO 27.54 34.84 28.56 Qwen/MATH Ours 30.00 34.86 81.30 Llama/MATH GRPO 12.90 18.78 14.28 Llama/MATH Ours 14.12 20.54 56.94 Qwen/GSM-Sym. GRPO 64.36 76.24 1.00 Qwen/GSM-Sym. Ours 65.82 76.14 83.88 - 主表结论:Ours该预算在三个设置的RSR均为Table 1最高;GSM-Sym.的FAC低于GRPO,故不能写成所有设置的答案正确率都更高。
- 作者解读:作者认为结果改善了准确率与教师遵循度的权衡。Figure 3的横轴为Constraint Satisfaction,纵轴为Final Answer Correctness;转换文本不能核查点的位置。
- 参照行歧义:Table 1的Student、Teacher各设置仅保留一个数值,转换后不能确定指标列,因此不引用其分数。
成对推理评审与人工一致性
- Figure 4数值:Ours d=−ln(1/4)在Qwen-MATH、Llama-MATH、Qwen-GSM-Sym.的平均胜率为56.3%、63.5%、60.6%;GRPO为25.5%、53.4%、31.0%。
- 例外与文本冲突:GKD-GRPO λ=0.01在Qwen-MATH为67.7%,高于Ours。第4.3节写Ours为56.2%、GRPO的Qwen-GSM为30.9%;Figure 4分别标56.3%、31.0%,此处采用图中标注。
- 人工复核:附录F按方法、数据集和判定类型分层,标注者不知道方法;60个判定全部一致,报告100%。这是总体审计比例,并非各分组独立结果;另未给pairwise一致性数值。
代码生成:可见单元测试下的reward hacking
- 评测条件:MBPP测试500题,提示词包含单元测试;奖励按全部测试通过计算,每个候选在独立python3子进程运行,超时5秒(附录C)。
- Table 2结果:Ours d=−ln(1/4)的RSR为34.2%、pass@1为45.6%、FAC−RSR为11.4个百分点;GRPO对应1.8%、92.0%、90.2个百分点。
- 例外与解释:GKD-GRPO λ=0.01的RSR为34.0%、pass@1为50.4%。Ours的pass@1低于Table 2全部基线;作者称GRPO复制可见测试投机,代码RSR的具体判定协议未说明。
教师似然代理与约束类型分析
- 相关性验证:Table 3在Qwen-MATH的60,000轨迹池中,最坏前缀教师logp、负reverse KL、负forward KL与RSR的Spearman ρ分别为+0.231、+0.224、+0.186。它们是相关性统计。
- 答案正确子集:附录E中GRPO被评审标错与有效推理的平均最坏logp分别为−1.09、−0.92,Mann–Whitney p<10⁻⁶;另报告pooled AUC 0.63。
- Figure 2与5的边界:Figure 2来自事后测试过滤,不是两种训练算法对比;作者称最优前缀过滤RSR更高,但文本无峰值。附录D写预算扫向负数,与非负成本及S≤d描述不一致。Figure 5图例写GDK-GRPO,正文写GKD-GRPO;看不到曲线,不能核查作者所称尾部与截断。
其他消融与分析
- Qwen/MATH预算:Table 1中d=−ln(1/2)、−ln(3/4)的RSR为29.76%、26.36%。
- Llama/MATH预算:同表两预算RSR为13.34%、13.92%;严格预算的RSR并非随预算收紧持续下降。
- Qwen/GSM-Sym.预算:同表两预算RSR为63.06%、58.96%,CS为60.44%、86.30%。
- 软惩罚变量:Qwen/GSM-Sym.的GKD-GRPO λ=0.001为FAC 74.08%、CS 0.00%;λ=0.1为66.68%、84.32%(Table 1)。
- 似然分位组:附录E最低/最高quartile的RSR为30.1%/59.6%,FAC为35.9%/63.0%。
- 训练耗时:附录B.4报告每方法单加速器少于23小时,GRPO约16小时;未报告具体加速器型号或逐方法耗时表。
What are some points that can be further explored?
作者指出似然代理、tokenizer兼容、OOD保证、预算调参与任务范围的局限。
作者明确保留了成本代理和分布外保证等问题,理论约束不能直接替代逻辑正确性审计。
作者指出的局限与后续方向
- 成本代理:教师log-likelihood不能完整刻画推理有效性;逻辑正确但风格不同的步骤可能成本较高。作者提出探索learned process reward models等成本(附录G)。
- tokenizer兼容:学生与教师需有兼容tokenizer;跨词表蒸馏需重分词、跨tokenizer对齐或序列/语义成本,作者留待后续研究(附录G)。
- 分布外保证:作者称理论保证针对训练分布,把正式可行性保证扩展到OOD仍是开放问题(附录G)。
- 预算选择:最优d仍需经验调参;作者提出自适应预算,并说明可在已有验证轨迹上事后扫阈值,以权衡接受比例与RSR(附录G)。
- 任务范围:作者将数学与代码之外的定理证明、multi-hop QA、智能体框架列为后续方向(附录G)。
实验覆盖范围
- 模型与任务:实验使用两组数学师生模型和一组代码师生模型;数学覆盖MATH及GSM8K→GSM-Symbolic,代码训练使用MBPP/APPS,评测MBPP(第4.1节、附录C)。
- 有限惩罚:实验n=20,比较三个d;理论几乎必然可行性要求n趋于无穷且存在有限价值最优策略(附录B.3、A.2)。
- 评审证据:数学评审模型为DeepSeek-R1-Distill-Qwen-32B;人工复核60个判定。论文未报告代码RSR的具体判定协议、pairwise样本量或独立一致性数值(附录B.5、C、F)。
- 统计报告:附录E给出60,000轨迹的代理相关性及答案正确子集检验;Table 1、Table 2未报告重复训练次数、置信区间或方法差异检验。
- 约束类型比较:Saute累计成本与前缀成本的Figure 2对比是Qwen2.5-1.5B轨迹上的事后过滤分析。附录D讨论了逐token和滑动窗口约束不被采用的理由。
Summarize the main content of the paper
作者提出前缀约束蒸馏;数学与代码结果需区分答案正确、教师遵循和推理有效。
这是把推理蒸馏写成最坏前缀约束RL的训练研究。作者关心最终答案奖励下的reward hacking,以及序列平均软惩罚允许局部偏离被其他token补偿的问题。
- 方法核心:约束所有前缀的平均教师负对数似然,历史一旦越界便不可恢复;训练使用教师成本,部署不查询教师。梯度拆分为词表精确期望的即时项和采样的长期项。
- 数学RSR:Table 1中,训练d=−ln(1/4)的Ours在Qwen2.5-1.5B/MATH、Llama-3.2-3B/MATH、Qwen2.5-1.5B/GSM-Symbolic分别为30.00%、14.12%、65.82%,均为表中最高。
- 答案正确率例外:同表GSM-Symbolic设置,Ours FAC为76.14%,GRPO为76.24%;教师约束满足率分别为83.88%、1.00%,统一按评测d=−ln(1/2)计算。
- 代码结果:Table 2的MBPP测试500题中,Qwen2.5-Coder-0.5B使用同一Ours预算取得RSR 34.2%、pass@1 45.6%;GRPO为1.8%、92.0%。作者称可见测试被用于投机,不能把pass@1等同于有效推理。
- 证据与边界:Table 3的Qwen-MATH最坏前缀logp与RSR相关系数为+0.231;附录F人工复核60个判定、报告100%一致。理论保证有无穷惩罚和有限价值假设,实验使用有限惩罚。
作者认为前缀约束改善了准确率与教师遵循度的权衡;附录G同时指出,似然代理不等于逻辑有效性,兼容tokenizer、预算选择和OOD保证仍需处理。