RLHF

LLM 针对 question 给出答案，这样的 QA pair 要如何对齐到人类的需求？

用 RL 的方法, 我们需要一个奖励模型, 但衡量一个答案“好”的程度是非常难以统一的，因而我们要求 LLM 为一个问题产生多个答案，再由专家选择最好的答案，这样就形成了所谓“偏好”数据集。

LLM 如何生成多个答案？——Use High Temperature

预训练模型接受 Q+A 作为输入，输出一个 Reward, 表示这个 Q+A pair 的好坏。Loss 为

$Loss = -log \sigma(r(x,y_w)-r(x,y_l))$

其中 $r(x,y_w)$ 表示 Q+A pair 的 Reward， $r(x,y_l)$ 表示 Q+A pair 的负样本的 Reward（如果多个则取平均）。

这样的损失函数不仅强迫模型为好的答案打出高分，还强迫模型为不好的答案打出低分（以形成分差）。

对于 Answer 之中的每一个位置, 都可以用这样的方法得到一个 reward, 我们就得到了 r(t): llm 推理第 k 个 token 时的 reward。

状态价值函数 V(s): 同样和 reward 的 model 共享前置层（pretrained-model，实际上这个 model 就等效于 RL 之中的 policy $\pi$ ），只不过换了一个线性层，输出单个值作为 $V_{\pi}(s)$

这个 Linear Layer 的训练：在后文的 RLHF 之中，我们会用 PPO 算法来训练这个 Linear Layer，这个 Linear Layer 的目标是最小化 $||V_{\pi}(s)-\sum^{T}_{t=0}\gamma^{t}r_{t}||^2$ ，其中 $\gamma$ 是折扣因子， $r_t$ 是第 t 个 token 的 reward，这里相当于用 $Q$ 来近似 $V$ 。

GAE: Advantage function $A_{\pi}(s_t,a_t)=Q_{\pi}(s_t,a_t)-V_{\pi}(s_t)$

在不同时间步上做近似

1-step: $r_t+\gamma V_{\pi}(s_{t+1})-V_{\pi}(s_t)$

2-step: $r_t+\gamma r_{t+1}+\gamma^2 V_{\pi}(s_{t+2})-V_{\pi}(s_t)$

3-step: $r_t+\gamma r_{t+1}+\gamma^2 r_{t+2}+\gamma^3 V_{\pi}(s_{t+3})-V_{\pi}(s_t)$

...

采样越多步骤, bias 越小, 但 variance 越大。称为 bias-variance tradeoff，实际上我们可以用一个指数加权的方法来平衡 bias 和 variance。

$\delta_{t} = r_t+\gamma V_{\pi}(s_{t+1})-V_{\pi}(s_t)$

$A_{t} = \delta_t + \gamma \lambda A_{t+1}$

当 $\lambda$ 为 1 时，就是完全的 GAE，当 $\lambda$ 为 0 时，就是 1-step 的方法。

LLM 的“State”: 已有的 token 序列（包括输入+已经生成的输出）

LLM 的“Action”: 生成下一个 token

RLHF 之中, on-policy 不可接受, 因为模型的推理时间太长，也就是和环境交互的代价太高（而我们的梯度和学习率却不能太高！），我们需要 off-policy 带来的经验回放等机制来实现并行训练。

on-policy 转 off-policy 也是经典的重要性采样，再加上 PPO 经典的 clip 方法，就可以实现 off-policy 的训练。

总体上, RLHF 的训练过程如下：

offline-policy(high temperature model)生成多个答案, 即为多个 trajectory, 记录对应的 r, s, a, A 和梯度数据, 放入经验回放池
从池中取出 mini-batch
计算梯度，更新 online-policy
k 个 epoch 后，更新 offline-policy 为 online-policy，重复 1-3

$\nabla_{\theta_{online}}J(\theta_{online},\theta_{offline})=\frac{1}{N}\sum^{N}_{i=1}\sum^{T}_{t=0}\nabla_{\theta_{online}}\frac{log\pi_{\theta_{online}}(a_{i,t}|s_{i,t})}{log\pi_{\theta_{offline}}(a_{i,t}|s_{i,t})}A_{\pi}(s,a)$

PPO Loss 的设计:

$L_{policy}=-max(\frac{\pi_{\theta_{online}}(a|s)}{\pi_{\theta_{offline}}(a|s)}A_{\pi}(s,a), clip(\frac{\pi_{\theta_{online}}(a|s)}{\pi_{\theta_{offline}}(a|s)}, 1-\epsilon, 1+\epsilon)A_{\pi}(s,a))$

$L_{vf}=\frac{1}{2}||V_{\theta_{online}(s)}-(\sum^{T}_{t=0}\gamma^{t}r_{t}|s=0)||^2$

$L_{entropy}=\sum^{T}_{t=0}\pi_{\theta_{online}}(a|s)log\pi_{\theta_{online}}(a|s)$

$L_{ppo}=L_{policy}+c_1L_{vf}+c_2L_{entropy}$

这里的正负号是方便适配 torch 等框架, 只做梯度下降, 但 policy loss 本身是一个最大化问题，所以我们给 $L_{policy},L_{entropy}$ 取负号 min -> -max。

其中, $c_1$ 和 $c_2$ 是超参数，entropy loss 用于增加模型的探索性，vf loss 用于增加模型的稳定性（控制价值函数）。

另一个 Reward hack 是, 有时候模型会不断重复生成无意义的“cheat reward”的对齐答案, 例如让模型变得有礼貌的对齐之中一直回答“谢谢”, 因此我们需要约束模型的答案的语义不能偏离原始模型太多，使用的方法就是对原始的模型做一个拷贝，冻结其参数，再在每一个 token 的奖励处减去原始模型和训练模型的 KL 散度作为惩罚项。

# peusudo code
def off_policy_learning(frozen_model, model_totrain, reward_model, optimizer, num_epoches):
    for k in num_epoches:
        # this model act as offline-policy
        trajectories = sample_trajectories(model_totrain)
        # calculate the reward for each token, log probability, advantage, KL-divergence between frozen_model and model_totrain 
        trajectories = update_trajectories_with_more_info(trajectories, frozen_model, model_totrain, reward_model)
        for j in range(num_epoches):
            minibatch = get_random_minibatch(trajectories)
            loss = ppo_algorithm(minibatch, model_totrain, optimizer)
            loss.backward()
            optimizer.step()