【论文阅读】Thinkless: LLM Learns When to Think

Thinkless: LLM Learns When to Think 这篇文章介绍了名为Thinkless的模型，该模型旨在通过自适应选择思考和非思考推理模式来提高推理语言模型（LLM）的效率和性能。Thinkless的核心方法是Decoupled Group Relative Policy Optimization (Decoupled GRPO)，它将混合推理目标分解为两个组件：模式选择（Mode Selection）和准确性提升（Accuracy Improvement）。模式选择决定了模型何时采用长链推理，而准确性提升则在选定的推理模式下优化响应内容以提高答案的正确性。

项目地址：https://github.com/VainF/Thinkless

本篇博客仅聚焦文章的方法部分

Method

文章分两个阶段实现：

• （1）Distillation for Warm-up，对预训练的推理模型进行微调以统一两种推理风格，
• （2）Reinforcement Learning with DeGRPO

Distillation for Warm-up

在作者提出的框架中，第一步是构建一个能够生成短回应和长回应的模型 ${\pi }_{\theta }$。作者利用两个预训练的专家模型进行蒸馏：一个推理模型 ${\pi }_{\text{think}}$，训练用于通过逐步推理生成详细的思维链；以及一个指令跟随模型 ${\pi }_{\text{short}}$，优化用于生成与用户意图对齐的简洁回答。给定一个提示语料库 X = { x i } i = 1 N X = \{x_i\}_{i=1}^N X={xi​}i=1N​，我们使用这些模型生成一个合成的配对数据集：

$$D_{\text{distill}}=\{(x_i,<\text{think}>a_{\text{think}}^i,<\text{short}>a_{\text{short}}^i){\}}_{i=1}^N,$$

其中 $a_{\text{think}}^i={\pi }_{\text{long}}(x_i)$ 和 $a_{\text{short}}^i={\pi }_{\text{short}}(x_i)$。每个response都以前缀控制标记 c ∈ C = { < short > , < think > } c \in C = \{<\text{short}>, <\text{think}>\} c∈C={<short>,<think>} 开头，以指示预期的推理风格。然后，我们通过监督细调（SFT）在该数据集上细调目标推理模型 ${\pi }_{\theta }$。目标是学习一个基于控制标记的多风格回应分布。这一蒸馏阶段确保模型能够以高保真度生成两种类型的回应，同时配对数据集的构建确保模型的回应分布是平衡的，这有助于后续的强化学习过程探索不同的解决方案。

Learning When to Think via Decoupled GRPO

在蒸馏阶段之后，模型可以生成长回应和短回应。然而，它仍然缺乏一种机制来决定哪种推理模式适合特定的输入 $x$。为了提供这种能力，作者将模式选择视为一个强化学习问题，并优化策略 ${\pi }_{\theta }(c,a|x)={\pi }_{\theta }(c|x){\pi }_{\theta }(a|x,c)$，其中第一个标记 c ∈ C = { < short > , < think > } c \in C = \{<\text{short}>, <\text{think}>\} c∈C={<short>,<think>} 作为控制标记，决定了推理模式，而后续的标记 $(a_{i,1},\dots ,a_{i,T_i})$ 构成了生成的回应。为了方便表示，我们将第 $i$ 个样本的整个长度为 $T_i+1$ 的序列表示为 $a_i=(a_{i,0},\dots ,a_{i,T_i})$，其中 $a_{i,0}\in C$ 是控制标记。

奖励设计

设 $y^{\ast }$ 表示输入 $x$ 对应的正确答案。我们考虑一个最小设计的奖励函数 $r(a,y^{\ast },c)$，其奖励值如下：
$$r(a,y^{\ast },c)=\{\begin{array}{ll}1.0,& \text{if }c=<\text{short}>\text{ and }\text{Extract-Answer}(a)=y^{\ast },\\ 1.0-\gamma ,& \text{if }c=<\text{think}>\text{ and }\text{Extract-Answer}(a)=y^{\ast },\\ -1.0,& \text{if }\text{Extract-Answer}(a)\ne y^{\ast },\end{array}$$
其中 $1>\gamma >0$ 引入了对短正确答案的偏好，相对于长回应。

解耦策略优化

基于简单的奖励函数，作者采用基于 GRPO 的框架进行训练。设 { a i } i = 1 G \{a_i\}_{i=1}^G {ai​}i=1G​ 表示从当前策略 ${\pi }_{{\theta }_{\text{old}}}$ 中采样的一个 mini-batch 轨迹。目标函数定义为：

$$J_{\text{GRPO}}(\theta )={\mathbb{E}}_{x,a_i}\left[\frac{1}{G}\sum _{i=1}^G\left(\frac{1}{T_i+1}\sum _{t=0}^{T_i}{L}_{i,t}(\theta )-\beta D_{\text{KL}}\left[{\pi }_{\theta }(\cdot |x)\parallel {\pi }_{\text{ref}}(\cdot |x)\right]\right)\right]\text{\hspace{1em}(1)}$$

其中 $L_{i,t}(\theta )$ 表示token-level surrogate loss，形式上定义为：

$$L_{i,t}(\theta )=\min \left(\frac{{\pi }_{\theta }(a_{i,t}|x,a_{i,<t})}{{\pi }_{{\theta }_{\text{old}}}(a_{i,t}|x,a_{i,<t})}{\hat{A}}_{i,t},\text{clip}\left(\frac{{\pi }_{\theta }(a_{i,t}|x,a_{i,<t})}{{\pi }_{{\theta }_{\text{old}}}(a_{i,t}|x,a_{i,<t})},1-ϵ,1+ϵ\right){\hat{A}}_{i,t}\right)\text{\hspace{1em}(2)}$$

在本工作中，作者使用 ${\hat{A}}_{i,t}=r-\text{mean}(r)$ 计算相对优势，这一选择是基于观察到训练数据包含不同难度的问题，使用标准差归一化可能会引入偏差。

当应用于训练时，方程$(1)$的目标函数有两个目的：

• 学习适当的控制标记以进行模式选择，
• 提高回应标记的准确性：
  $$\frac{1}{T_i+1}\sum _{t=0}^{T_i}{L}_{i,t}(\theta )=\frac{1}{T_i+1}{L}_{i,0}(\theta )\text{(Control Tokens)}+\frac{1}{T_i+1}\sum _{t=1}^{T_i}{L}_{i,t}(\theta )\text{(Response Tokens)}\text{\hspace{1em}(3)}$$

对于模式选择，回应风格基于第一个标记 $a_{i,0}$，该标记在之前的蒸馏阶段进行训练。因此，调整这个单一控制标记的概率就足以在不同的推理模式之间切换。因此，这个标记控制推理模式的学习。对于回应准确性，优化目标是提高剩余标记 $a_{i,1:T_i}$ 的生成。然而，上述方程 $(3)$在优化过程中引入了两种不平衡：

1. 模式-准确性不平衡 - 每个轨迹只包含一个控制标记，但有 $T_i$ 个回应标记，不成比例地减少了模式选择的影响；

2. 长短不平衡 - 更长的序列 $T_{\text{think}}^i\gg T_{\text{short}}^i$ 进一步抑制了控制标记的梯度贡献，因为归一化因子 $1/(T_i+1)$ 导致 <think> 标记相对于 <short> 标记的优化不足。

这种不平衡可能导致训练初期严重的模式崩溃。为了解决这些不平衡，作者提出了一种解耦的 GRPO 变体，记为 $J_{\text{DeGRPO}}$，分别归一化控制标记和回应标记的贡献：
$$J_{\text{DeGRPO}}(\theta )={\mathbb{E}}_{x,a_i}\left[\frac{1}{G}\sum _{i=1}^G\left(\alpha L_{i,0}(\theta )\text{(Control Tokens)}+\frac{1}{T_i}\sum _{t=1}^{T_i}{L}_{i,t}(\theta )\text{(Response Tokens)}-\beta D_{\text{KL}}\left[{\pi }_{\theta }(\cdot |x)\parallel {\pi }_{\text{ref}}(\cdot |x)\right]\right)\right]$$
在 DeGRPO 中，模式选择$L_{i,0}(\theta )$ 和回应准确性改进 ${\sum }_{t=1}^{T_i}{L}_{i,t}(\theta )$ 独立归一化。引入了一个与长度无关的权重系数 $\alpha$ 以平衡模式选择和回应生成的优化。这种公式确保控制标记在短序列和长序列中都获得一致的梯度规模，从而解决模式-模式和长短不平衡，实现更稳定的推理模式选择优化。如实验所示，适当大的 $\alpha$ 可以使模式更新更高效。在实验中，设置 $\alpha =1/1000$ 以实现稳定的训练。

方法总结

总之，保留了标准 GRPO 框架的整体结构。对于每个查询，从当前策略中抽取一个 mini-batch 样本以估计标记级别的优势。为了解决模式选择和回应生成之间的不平衡，我们独立归一化与控制标记和回应标记相关的优势。这种分离允许在优化过程中显式平衡它们的贡献，从而实现更稳定和有效的训练。