精读 · cs.LG cs.AI cs.CL · NeurIPS 2025

The Overthinker's DIET: Cutting Token Calories with DIfficulty-AwarE Training

Weize Chen†, Jiarui Yuan†, Tailin Jin, Ning Ding, Huimin Chen, Zhiyuan Liu, Maosong Sun — Tsinghua University THUNLP（†共同第一作者）

📄 摘要 📑 PDF 🌐 HTML 🔍 OpenReview

Chain-of-Thought Token Efficiency GRPO Difficulty-Aware RL Reasoning Inference Scaling

核心论点：对所有问题施加相同的 token 压力是错的——难题需要长推理，简单题不需要。DIET 在 GRPO 训练中实时估计任务难度，据此动态调整 token 惩罚强度。同时发现并修复了 GRPO 中朴素奖励合并的一个理论缺陷（方差交叉污染），提出 Advantage Weighting。最终效果：准确率 +3.3%（50.2% vs 48.6%）同时 token 用量 −40.7%——是唯一一种同时提升准确率并大幅降低 token 的方法。

🎯 问题：过度思考 + 均匀惩罚的根本缺陷

LLM 推理模型的"过度思考"问题

以 DeepSeek-R1 为代表的推理模型通过长思维链获得了出色的推理能力，但带来了一个系统性问题：对所有问题都生成过长的响应，包括简单问题。一道可以用 200 token 解决的数学题，模型可能生成 5000+ token 的自我检查、反复确认和冗余推理——这导致推理成本极高，同时对难题的实际提升有限。

已有方案的通病：均匀 token 惩罚

现有压缩方案的通用目标函数： Maximize E[P(π,x)] − α · E[f(L(π,x))] 其中 P 是任务性能，L 是 token 用量，α 是固定的惩罚强度。

核心问题：α 对所有问题相同。这意味着 AIME 竞赛题和简单算术题受到同等压缩压力。结果：

压缩力度轻（小 α）：简单题仍然过长，没有实质改善
压缩力度重（大 α）：难题推理被截断，准确率崩塌
中间点：性能和效率都次优

论文用先决观察证实：未经任何压缩训练的基础模型，天然存在"难题→长响应，简题→短响应"的正相关——均匀压缩会破坏这种已有的良好性质。

📐 难度估计：零额外开销的在线方案

在 RL 训练中实时估计难度

GRPO 每个 prompt 本来就要采样 N 条响应来计算 advantage。DIET 直接利用这 N 条响应估计难度： Ĉ(x, πθ) = (1/N) · Σ 𝟙(yᵢ 正确)，yᵢ ~ πθ(·|x) D̂(x, πθ) = 1 − Ĉ(x, πθ)

Ĉ ≈ 0 → 模型在这题几乎全错 → 高难度 → 少压缩
Ĉ ≈ 1 → 模型在这题几乎全对 → 低难度 → 多压缩

零额外计算开销：GRPO 采样的 N 条响应本来就要生成，难度估计只是在已有结果上做一次均值，不增加任何 forward pass。

关键假设：用当前策略 πθ 在当前 prompt 上的通过率来代理"这道题有多难"。这是动态估计（每次训练步都更新），而不是预定义的静态难度标签——随着模型能力提升，同一道题的估计难度会变化。

🔬 两种难度感知策略

策略一：自适应惩罚强度（αada）

根据难度调整惩罚系数

αada(x, πθ) = αbase · Ĉ(x, πθ) 其中 αbase 是全局超参数，乘以通过率 Ĉ（即难度的互补）。效果：

简单题（Ĉ → 1）：惩罚强度 → αbase（满力度压缩）
难题（Ĉ → 0）：惩罚强度 → 0（不压缩，保留推理空间）

底层 token 惩罚使用 Kimi 格式： γᵢ = 0.5 − (L(yᵢ) − min L) / (max L − min L + ε) fKimi(yᵢ) = γᵢ（答对时） | min(0, γᵢ)（答错时） 即：答对但太长 → 负奖励；答错但短 → 不惩罚（不想鼓励用牺牲正确率的方式缩短）。

策略二：动态长度目标（fdyn）

根据难度设置允许长度上限

另一个思路：不调整惩罚强度，而是为每道题设置一个难度相关的允许长度目标，超出目标才惩罚： t(x, πθ) ~ Uniform(max(0, Lmax · (D̂ − δ)), Lmax · D̂)

难度 D̂ 高 → t 大 → 允许更长的响应
难度 D̂ 低 → t 小 → 超出即惩罚

使用随机区间（而不是固定 t = Lmax · D̂）是为了增加训练时的探索多样性。

惩罚仅对超出目标的部分生效，并在 batch 内归一化： fdyn = (p(yᵢ, t) − μp) / (σp + ε)，其中 p = max(0, L(yᵢ) − t)

两种策略的区别：αada 改变"惩什么力度"；fdyn 改变"惩什么范围"。两者可以单独使用，论文最终的 DIET 是两者合并。

⚠️ GRPO 陷阱：方差交叉污染 + Advantage Weighting 修复

朴素方案的理论缺陷

在 GRPO 中，最直接的做法是把任务奖励和 token 惩罚合并后再做 advantage 归一化： r'ᵢ = routcome,ᵢ − αada · pᵢ Â'ᵢ = (r'ᵢ − μr') / (σr' + ε)

问题：方差交叉污染。 合并奖励的方差是：σ²r' ≈ σ²outcome + α²ada · σ²p

因此实际有效惩罚缩放因子变成了： τ̂p = αada / √(σ²outcome + α²ada · σ²p + ε') 其中 σ²outcome = Ĉ(1−Ĉ)（二项分布方差），这在 Ĉ = 0.5（中等难度）时最大，在 Ĉ ≈ 0 或 1（极难/极易）时趋近零。

这意味着：

中等难度题：outcome 方差大 → 分母大 → 惩罚被意外削弱
极简单题（Ĉ ≈ 1）：outcome 方差趋零 → 分母小 → 惩罚被意外放大

这与我们设计 αada 的意图完全相反：我们想在简单题上强惩罚（αada 大），但 GRPO 归一化把这个效果"对冲"掉了。

Advantage Weighting：分开归一化再合并。 Âoutcome,ᵢ = (routcome,ᵢ − μoutcome) / (σoutcome + ε) Âp,ᵢ = (pᵢ − μp) / (σp + ε) Â'ᵢ = Âoutcome,ᵢ − α' · Âp,ᵢ 每个分量用自己的均值和方差归一化，消除方差交叉污染。α' 此时能准确控制两个分量的相对权重，不受对方方差影响。

当 α 固定且不依赖 Ĉ 时，τ̂p 对所有题的偏差方向相同，效果是"惩罚整体被稀释但仍然存在"——虽然次优但不会导致训练崩溃。但当 αada 变成 Ĉ 的函数后，τ̂p 的实际效果对不同难度的偏差方向开始反转，最终使 advantage weighting 中的分量比例完全失控，训练无法收敛（论文实验证实：不加 Advantage Weighting 的 DIET 训练失败，从结果表中剔除）。

🔄 循环压缩压力（Cyclical Compression Pressure）

用余弦调度避免过早收敛到"极度简洁"

如果全程保持最大压缩压力，模型会在训练早期就收敛到某种"简洁但不够灵活"的模式，失去探索更好推理路径的能力。DIET 用余弦调度振荡压缩压力： c(t) = 0.5 · (1 + cos(2πt/T)) Â'ᵢ(t) = Âoutcome,ᵢ − c(t) · αada · Âp,ᵢ

c(t) = 1（余弦峰值）：最大压缩压力
c(t) = 0（余弦谷值）：零压缩压力，模型自由推理

这造成"压缩 → 自由探索 → 压缩 → ..."的周期，防止模型在某一个压缩状态下过早固化。

类似于退火调度，但作用在奖励维度而非学习率。在"无压缩"阶段，模型可以发现更好的推理结构，然后在"压缩"阶段被强化为更简洁的形式。

📊 实验结果

全量对比（5个 benchmark 宏平均）

方法	宏平均 P@1	vs 基础	宏平均 Token	Token 变化
基础模型	48.6%	—	10,280	—
DIET（αada + fdyn）	50.2%	+3.3%	6,097	−40.7%
自适应惩罚强度（仅 αada）	49.9%	+2.8%	5,723	−44.3%
动态长度目标（仅 fdyn）	48.6%	±0%	5,657	−45.0%
Kimi 1.5 DPO	49.9%	+2.7%	8,441	−17.9%
O1-Pruner	45.9%	−5.4%	4,829	−53.0%
CosFn	46.8%	−3.5%	6,798	−33.9%
Kimi 1.5 RL	35.0%	−27.9%	4,050	−60.6%
TokenSkip（SFT）	30.9%	−36.4%	1,697	−83.5%
Kimi 1.5 SFT	42.7%	−12.1%	9,865	−4.0%

关键观察：只有 DIET 家族同时提升了两个指标

表格分为两类：

激进压缩类（TokenSkip、Kimi RL、O1-Pruner）：token 大幅降低，但准确率崩塌或显著下降。这些方法在"省钱"和"答题"之间强制取舍。
DIET 类：token 降低同时准确率提升——这是"帕累托改进"，不是权衡。原因：简单题上省下来的 token 相当于把计算资源重新分配给难题。

Kimi 1.5 DPO 接近 DIET 的准确率，但 token 降幅只有 −17.9%（vs DIET 的 −40.7%）。

方法	MATH500 P@1	MATH500 Token	AIME P@1	AIME Token
基础模型	82.1%	5,534	28.5%	16,590
DIET	83.0%	3,061 (−44.7%)	31.8%	10,578 (−36.2%)
TokenSkip	64.1%	1,120	6.8%	2,231
O1-Pruner	79.1%	2,531	25.0%	8,961

注意 DIET 在 AIME（难题）上的 token 降幅（−36%）比 MATH500（简单题，−45%）小——这正是难度感知压缩应该有的行为：简单题多压，难题少压。

推理时 Scaling 优势

固定 token 预算下的 Majority Voting 性能

在固定总 token 预算下，DIET 可以采样更多样本进行 majority voting（因为每个样本更短），同时每个样本的质量没有下降——两个效果叠加。

对比：TokenSkip/Kimi SFT 的单样本质量太差，即使采样再多次，准确率也很快停止提升。DIET 在低 token 预算下就能达到比基础模型高得多的 majority voting 准确率，scaling 曲线的斜率显著更陡。

→ 实际部署含义：如果系统有固定推理算力（比如 TPU 小时），DIET 是目前最高效的利用方式。

长度-难度相关性

DIET 是唯一增强了模型本有的"难题→长响应"规律的方法

用 Pearson 相关系数衡量"估计难度 vs 平均响应长度"的相关性（所有 p < 0.01）：

DIET：高于基础模型——压缩之后，难度和长度的相关性反而更强
CosFn：接近 DIET，但准确率次之
基础模型：正相关（自然有但未被训练加强）
Kimi 1.5 DPO / SFT：相关性比基础模型更弱——均匀压缩破坏了难度感知

这是 DIET 的质性证明：它确实让模型"学会了"按难度分配思考资源，而不只是整体压缩。

💡 核心洞察

为什么"难度感知"这个想法正确但之前没人做

"按难度压缩"在直觉上几乎是显然的——没有人会说"解方程和证明黎曼猜想应该用同样的字数"。但在 RL 实现层面有两个技术障碍：

难度怎么估计：如果要在训练时查询外部难度数据库，需要对所有训练题预标注难度，且标注会过时（随模型能力变化）。DIET 的洞察是"用当前策略的通过率作为当前难度的动态代理"——这既准确又零成本。
GRPO 中怎么实现：直接把难度感知惩罚加到奖励里会被 GRPO 的方差归一化破坏（方差交叉污染问题）。Advantage Weighting 是让这件事真正可行的工程关键。

所以 DIET 的贡献是双层的：①正确的问题定义（难度自适应压缩），②让这个想法在 GRPO 中真正能训练的技术实现。两者缺一不可——论文实验证明，没有 Advantage Weighting 的 DIET "训练失败"。

GRPO Pitfall 的普适性警告

Advantage Weighting 解决的方差交叉污染问题，并不是 DIET 特有的——它存在于任何在 GRPO 中合并多个奖励维度的场景：

任务准确率 + 格式奖励（如 JSON 格式正确）
任务准确率 + safety 惩罚
任务准确率 + 多样性奖励

所有这些场景都面临"合并方差影响每个分量的实际权重"的问题。DIET 的 Advantage Weighting 是一个通用修复，应该成为 multi-objective GRPO 的标准实践。

与 SVR-R1 对照（今日日报 2607.10966）：SVR-R1 用"自验证"构造了另一种内在奖励（不依赖外部 verifier）。两篇都在探索"如何设计更丰富的 GRPO 奖励信号"——DIET 的方向是效率维度，SVR-R1 的方向是一致性维度。两篇都需要面对"多维度奖励在 GRPO 中的实现稳定性"问题，Advantage Weighting 对 SVR-R1 类工作同样有参考价值。

🧬 与 Optima 的谱系关系

从 MAS 通信 → 单 agent 推理：继承与超越

维度	Optima（ACL 2025）	DIET（NeurIPS 2025）
Token 惩罚对象	multi-agent 之间的通信消息	单 agent 内部的思维链
Token 惩罚方式	复合奖励中均匀惩罚 token 用量	难度感知的自适应惩罚
训练算法	SFT + DPO 迭代	GRPO + Advantage Weighting
GRPO 陷阱	未讨论（DPO 不存在此问题）	发现并修复
理论贡献	MCTS-DPO 数据生成框架	Advantage Weighting 的方差分析

DIET 补上了 Optima 没有涉及的两件事：①把场景从 MAS 通信扩展到单 agent 内部推理（解决了更广泛的问题）；②针对 GRPO（比 Optima 的 DPO 更主流的当下 RL 路线）提供了理论上正确的多目标实现方案。