大模型强化学习：从稀疏奖励到稠密奖励
Published:
在大模型后训练的RL过程中，通常由奖励模型或基于规则的奖励机制来给出奖励信号，但这种机制通常只能作用于模型生成的一个完整的序列上，给整个响应打一个单一的标量分数，这就导致无法给每个生成的token提供具体的奖励值，模型只能知道这个生成结果总体上怎么样，而难以知道生成的每一步的好坏，提高了其学习长推理链的难度。为了解决奖励稀疏化的问题，近两年出现了如PRM（过程奖励模型）、OPD（在线蒸馏）等相关优化方法来实现稠密奖励，它们致力于实现给生成的每个step甚至每个token都打出具体的分数，从而让模型更细粒度地学习复杂的推理流程。
目录
PRM
https://zhuanlan.zhihu.com/p/15540962086
https://zhuanlan.zhihu.com/p/16027048017
PRM（Process Reward Model，过程奖励模型）是OpenAI o1最早提出的，与之相比的是ORM（Outcome-supervised Reward Model）：
PRM：在生成步骤中分为若干步骤，对每一步进行打分，最终再将一个响应所有步骤的分数汇总（如取平均值或最低值等）得到该响应的总得分，是一种细粒度的奖励模型（dense reward）
ORM：只根据完整的响应结果给出一个分数，是一种系数奖励模型（sparse reward）
假设已经拥有了一个可以对模型输出按步打分的PRM，则其既可以作为RL训练中的奖励模型使用，也可以在推理阶段作为test time scaling（tts）的手段使用。
tts本质上是通过增大LLM计算量来得到更好的reasoning输出，例如CoT提升了推理深度，而这里则是通过生成多个候选输出并选择最优输出来提升推理广度。
其在LLM test time scaling生成中的用法为：
首先，根据一个按步骤回答的数据集，训练LLM使得其具有按步输出的能力。此时模型的输出正确率可能不高，但其确保按照step1, step2, …的格式进行输出，使得其可以被PRM进行打分。
然后在面对输入时，LLM进行N次采样输出，并通过PRM对每个采样的每步推理进行打分
最终，对于每个采样的所有步骤分数进行汇总（如取平均或最低值等），使得每个采样都获得一个整体分数，然后选取整体分数最高的响应作为模型的最终输出
例如，上图中有3个采样回答，如果按照所有步骤分乘积作为样本整体分数的话则answer2分数最高被选取，如果按照取所有步骤中得分最小值（反映该样本的下限）作为整体分数则answer3分数最高被选取。
PRM的训练：
考虑如何训练一个PRM（也即能够给LLM的step by step输出打分的奖励模型），首先需要考虑如何收集带有按步骤打分的标签的数据集。可以使用人工给大量解题步骤标注分数（如openai o1），也可以使用自动化样本标注方法（如Math-Sheperd）。这里主要介绍基于MCTS的自动化标注方法。
MCTS（Monte Carlo Tree Search）是RL领域一种通过树状采样来估计当前动作或状态的value的方法。在每个状态（节点）下，其会进一步探索多个action，得到多个后继节点，依此类推即可形成一个树状搜索空间。当到达叶子节点（结束状态）后，对其最终收益进行量化（例如输出答案正确则+1，错误则-1），然后再进行反向传播，从而反推所有中间节点的价值（直观来看，引向更多正确答案路径的中间节点价值更高），最终得到整棵树所有节点的价值。
具体而言，分为如下4个步骤：
采样（sample）
选择一个未被探索的节点（初始时只有一个表示原始问题的根节点），如下图中的$s_1$节点
扩展（expand）
从采样到的节点出发，展开其所有可能的子节点，如下图中的$s_{2,1},s_{2,2},s_{2,3}$。当然在文本生成中不可能穷举所有子节点，因此一般设置一个最大生成次数，在有限生成次数内所有不同的输出作为子节点的集合
模拟（simulate）
从展开的子节点里，再随机选择一个节点，再展开它的子节点，不断重复expand过程，直到最终到达叶子节点（最终答案）。
回传（backprop）
通过多次模拟得到了一个从根节点（原始问题$s_1$）到若干个最终生成答案的树，然后即可反向传播来推出各个节点的价值。一般来说有hard estimation (HE)和soft estimation (SE)两种估计中间节点的奖励值的方法。以SE为例，其计算“从当前节点出发到正确答案的路径数”除以“从当前节点出发的总路径数”，用这个比值来代表从当前节点出发能得到正确答案的概率。例如下图中$s_1$出发总共3条路径，有2条最终得到了正确答案，因此其价值$y_{s_1}^{SE}=\frac{2}{3}$。当然也可以根据规则或阈值做hard处理，每个节点的价值归为1或0。
可见，使用按步骤生成的数据集调好一个能按步骤生成（虽然正确率不一定高）的LLM generator后，即可使用它产生大量推理样本，并使用MCTS进行标注，从而得到带有“ground truth”分数的按步骤推理样本，可以用于拿来训练PRM了。下面考虑hard estimation的情况，也即一个步骤得分可以是+或-，表示其好或不好。
对于一条训练样本而言，其每个step后都附带一个+或-token，将它们换成占位符ки后即可作为训练输入数据。
PRM模型本身也是一个Causal LM，给定上述挖掉label后的输入数据，其预测各个步骤后的占位符处是+还是-对应的token（loss只统计这些占位符位置）：
本质上，PRM模型训练是在做分类任务，也即根据前序token序列，预测占位符位置输出+,-的概率，再与真实label做交叉熵。
扩展：ORM的训练
相比之下，普通的ORM reward model训练时做的是回归任务，其模型本身也是一个Causal LM，模型末尾处有一个regression head，来将最后一个预测的token（也即eos的位置）转换为一个分数值。
具体而言，每条训练数据为一个<prompt, chosen_response, rejected_response>三元组，prompt+chosen表示一个好的问答对，prompt+rejected表示一个坏的问答对。模型训练的loss是pairwise loss，希望尽可能拉大两个问答对的奖励值差异：
On-Policy Distillation
https://thinkingmachines.ai/blog/on-policy-distillation/
OPD（On-Policy Distillation，在线蒸馏），由Thinking Machine Lab提出，其结合了RL中“让模型自己探索”和KD中“教师模型密集指导”，使得学生模型在自己推理的同时让教师模型实时指导其每一步思维，而不是直接背诵教师模型生成的标准思路，从而能给出token级的稠密奖励信号，通过引入KD的指导思路来增强了RL的细粒度学习效果。
在传统KD中，teacher先生成一个标准答案轨迹，然后让student去学习拟合其token概率分布。然而在这种情况下，student在训练时只能看到teacher走出来的正确轨迹，也即永远在完全正确的上下文里进行训练，而其自己推理时会生成自己的轨迹，和教师生成的轨迹会出现分布偏差。例如，教师生成的正确轨迹为step1,step2,step3，推理时student的step1就错了，于是其就进入了teacher从未出现过的状态，使得接下来的生成崩掉。
而在OPD中，则让student自己生成轨迹（这也是On-Policy的来源），然后由teacher对这个student生成轨迹的每一个token进行打分。这样student就能在自己真实会进入的状态空间里学习：
从RL的角度来看，RL中奖励模型给的reward通常存在稀疏性的问题，其会给模型生成的整个结果打一个标量的reward分数，而难以告诉模型具体每一步的好坏。而OPD则可以由teacher给出稠密的、token级的奖励信号。换句话说，OPD可以看成用teacher logits替代sparse reward，由teacher模型提供的稠密奖励信号替代传统奖励模型的稀疏奖励信号，teacher模型本质上就变成了“奖励模型”。
设输入prompt为$x$，student policy为$\pi_\theta(y_t|x,y_{<t})$，teacher policy为$\pi_T(y_t|x,y_{<t})$。KD和OPD的实现区别如下：
在KD中，先由teacher生成轨迹：
\[y^{(T)}\sim \pi_T\]然后让student学习来最大化其生成这个教师轨迹的概率，也等价于最小化二者输出分布的前向KL散度：
\[\begin{aligned} \theta^*&=\max_\theta\log\pi_\theta(y^{(T)})\\ &=\max_\theta\sum_t\log\pi_\theta(y^{(T)}_t\|x,y^{(T)}_{<t})\\ &=\min_\theta D_{KL}(\pi_T\|\|\pi_\theta)\\ &=\min_\theta \mathbb{E}_{y\sim\pi_T}[\log\pi_T(y^{(T)})-\log\pi_\theta(y^{(T)})] \end{aligned}\]本质上就是把teacher logits当成静态的label，然后让student学习拟合它，是普通的有监督学习。
而OPD中，轨迹是student生成的：
\[y^{(\theta)}\sim\pi_\theta\]此时目标变为了最小化Reverse KL散度，student在它自己生成的（on-policy的）轨迹上训练：
\[\begin{aligned} \theta^*&=\min_\theta D_{KL}(\pi_\theta\|\|\pi_T)\\ &=\min_\theta\mathbb{E}_{y\sim\textcolor{red}{\pi_\theta}}[\log\pi_\theta(y^{(\theta)})-\log\pi_T(y^{(\theta)})]\\ &=\max_\theta\mathbb{E}_{y\sim\textcolor{red}{\pi_\theta}}[\log\pi_T(y^{(\theta)})-\log\pi_\theta(y^{(\theta)})]\\ &=\max_\theta\mathbb E_{y\sim\textcolor{red}{\pi_\theta}}[\sum_t\log\pi_T(y^{(\theta)}_t\|x,y^{(\theta)}_{<t})-\log\pi_\theta(y^{(\theta)}_t\|x,y^{(\theta)}_{<t})]\\ &=\max_\theta J(\theta) \end{aligned}\]对该目标求梯度可见，其形式本质上就是policy gradient：
\[\nabla_\theta J(\theta)=\mathbb{E}_{y\sim\pi_\theta}[(\log\pi_T(y^{(\theta)})-\log\pi_\theta(y^{(\theta)}))\nabla_\theta\log\pi_\theta(y^{(\theta)})]\]可见，此时teacher不再是静态的label提供者，而是奖励提供者，其提供的稠密奖励使得每个token都具有自己的reward值：
\[R(y^{(\theta)})=\sum_t\log\pi_T(y^{(\theta)}_t\|x,y^{(\theta)}_{<t})-\log\pi_\theta(y^{(\theta)}_t\|x,y^{(\theta)}_{<t})\]对于每个由student生成的token来说，teacher生成它的概率减去student生成它的概率，就是它获得的reward值，也即teacher喜欢的token能获得更高奖励。
从KL散度的角度考虑，普通KD最小化的是前向KL散度$D_{KL}(\pi_T||\pi_\theta)$，其倾向于“寻找均值”，迫使student学会覆盖teacher的所有模式。而OPD最小化的是反向KL散度$D_{KL}(\pi_\theta||\pi_T)$，倾向于“模态搜寻”，student会学习teacher最有自信的那些模式，而这恰恰对应了reasoning中少数的正确推理路径，因此OPD很适合做reasoning模型的蒸馏。