基于人类反馈的强化学习 RLHF

1.强化学习和语言模型的联系

agent: 语言模型本身

state: prompt(input tokens)

action: 选择哪个token作为下一个token（贪婪，top k,top p）

reward model：当生成‘好的回复’语言模型应当被奖励，当生成‘差的回复’语言模型不会受到任何奖励

policy：语言模型本身，因为它对动作空间给出当前状态的概率进行建模：${a_t} \sim \pi(·|S_t)$

2. 语言模型的奖励模型

人类不擅长达成一致，但是我们很擅长比较

2.1 数据集

2.2 奖励模型架构

2.3 奖励模型损失

为了借助强化学习优化我们的语言模型表现，我们需要一个打分模型，对语言模型生成的每个回复给予一个数值。

现在我们有一个定义了基于一个query(prompt)下更喜欢的answer的数据集，我们可以搭建一个神经网络，使之可以对每一个response生成一个数字分数。

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

good answer : $r(x,y_w)$

bad answer : $r(x,y_l)$

如果 $r(x,y_w)$ > $r(x,y_l)$ -> sigmoid会返回一个value > 0.5 -> loss会很小

·如果顺序正确，loss会很小

如果 $r(x,y_w)$ < $r(x,y_l)$ -> sigmoid会返回一个value < 0.5 -> loss会很大

·如果顺序错误，loss会很大

Note:

sigmoid函数：

2.4 具体实现

在HuggingFace, 我们可以通过使用RewardTrainer和AutoModelForSequenceClassification训练一个custom reward模型，即拥有一个特殊先行层在顶部的语言模型。

我们只需要要求语言模型输出最后一个token的隐藏层，然后输入进线性层计算奖励，然后根据之前所说的Loss训练模型。

   def compute_loss(

        self,

        model: Union[PreTrainedModel, nn.Module],

        inputs: dict[str, Union[torch.Tensor, Any]],

        return_outputs=False,

        num_items_in_batch=None,

    ) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]:

        rewards_chosen = model(

            input_ids=inputs["input_ids_chosen"],

            attention_mask=inputs["attention_mask_chosen"],

            return_dict=True,

        )["logits"]

        rewards_rejected = model(

            input_ids=inputs["input_ids_rejected"],

            attention_mask=inputs["attention_mask_rejected"],

            return_dict=True,

        )["logits"]

        # calculate loss, optionally modulate with margin

        if "margin" in inputs:

            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean()

        else:

            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()

        if self.args.center_rewards_coefficient is not None:

            loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2)

        if return_outputs:

            return loss, {

                "rewards_chosen": rewards_chosen,

                "rewards_rejected": rewards_rejected,

            }

        return loss

3. 策略 Policy

3.1 强化学习中的轨迹

我们之前说过，强化学习的目标是选择一个policy,当agent根据这个策略行动可以最大化返回的期望值。

\[\pi^* = {arg\,max}_\pi J(\pi)
\]

策略的期望值，就是所有可能轨迹(trajectories)期望值之和：

\[J(\pi) = \int_\tau\ P(\tau | \pi) R(\tau) = E_{\tau \sim \pi} [R(\tau)]
\]

轨迹是从初始state开始的一系列（action, state)：

\[\tau = (s_0,a_0,s_1,a_1, ···)
\]

我们把下一个状态视作随机的并进行建模：

\[s_{t+1} \sim P(·|s_t,a_t)
\]

因此我们可以用下面的公式定义轨迹·：

\[P(\tau|\pi) = \rho_0(s_0) \prod_{t = 0}^{T - 1} P(s_{t+1}|s_t,a_t)\pi(a_t|s_t)
\]

我们通常倾向于使用折扣奖励（因为我们喜欢及时奖励而不是长远奖励）：

\[R(\tau) = \sum_{t = 0}^{\infty}\gamma^tr_t
\]

Note: $0<\gamma < 1$

如下如所示，计算的时候应该是 $R(\tau) = \gamma^{1} * (-1) + 0 +...+ \gamma^{8} * (100)
$

3.2 语言模型下的轨迹

当我们处理的是一个语言模型时，我们希望能微调语言模型使之以获取最大reward的路径选择下一个token：

\[\pi^* = {arg\,max}_\pi J(\pi)
\]

轨迹是从一系列prompts(state) 和他们下一个token(action)

\[\tau = (s_0,a_0,s_1,a_1, ···)
\]

目标：策略的期望最大

也就是：训练一个policy神经网络$\pi$,在所有状态S下，给出相应的Action，得到的Return期望最大

也就是：训练一个policy神经网络$\pi$,在所有的Trajectory中，得到的Return期望最大

3.3 使用策略梯度优化训练一个policy网络

\[J(\pi_\theta) = E_{\tau \sim \pi_\theta} [R(\tau)]
\]

当我们提到深度神经网络，我们的目标是迭代地改变神经网络的参数来最小化损失函数的值，这是典型随机梯度下降的用法，但是在我们的案例中，我们希望最大化这函数，因此我们使用随机梯度上升：

\[\theta_{k+1} = \theta_k + \alpha \nabla_\theta {J(\pi_\theta)|}_{\theta_k}
\]

策略的梯度叫做 policy gradient, 优化这种策略的算法叫做policy gradient algorithms.

现在的问题就是，如何计算梯度，我们需要考虑所有可能的轨迹，这在计算上是不可行的，因为这是一种computationally intractable问题，除非我们有一个很小的state space。

a. 梯度表达式：

\[\nabla_\theta J(\pi_\theta) = \nabla_\theta \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)]
\]

$R(\tau)$：一条轨迹的总回报。

$\tau \sim \pi_\theta$：轨迹是从策略$\pi_\theta$中采样得到的。

b. 期望扩展为积分（连续就是积分，离散就是求和）：

\[\nabla_\theta \mathbb{E}_{\tau \sim \pi_\theta} [R(\tau)]
= \nabla_\theta \int_\tau P(\tau|\theta) R(\tau) d \tau
\]

c. 将梯度移到积分内部（换序积分定理（也称为微分与积分的交换定理））：

\[\int_\tau \nabla_\theta P(\tau|\theta) R(\tau) d \tau
\]

d. 引入对数导数技巧（log-derivative trick）：

\[\nabla_\theta P(\tau|\theta) = P(\tau|\theta) \nabla_\theta \log P(\tau|\theta)
\]

代入后，得到：

\[\int_\tau P(\tau|\theta) \nabla_\theta \log P(\tau|\theta) R(\tau) d\tau
\]

e. 返回期望形式：

\[\mathbb{E}_{\tau \sim \pi_\theta} [ \nabla \theta \log P(\tau|\theta) R(\tau)]
\]

f. 分解轨迹的概率：

\[P(\tau|\theta) = \rho_0(s_0) \prod_{t=0}^T \pi_ \theta(a_t|s_t) P(s_{t+1}|s_t, a_t)
\]

Note：环境的动态模型（状态转移概率）与策略参数无关，所以对 $\theta$ 求导只作用在策略上。

如下所示，等式两边同时加上log:

g. 最终公式（策略梯度）：

\[\nabla_\theta J(\pi_\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) R(\tau) \right]
\]

$R(\tau)$：轨迹的总回报。

$\log \pi_\theta(a_t|s_t)$：每个时间步策略的对数概率。

对于这个期望值，我们可以通过对收集到的一组轨迹（也就是一系列行动的记录）进行采样平均来近似这个期望值。

换句话说，我们通过观察和记录多个回合的行动，然后计算这些记录的平均值，来估计这个期望值。

\[\hat{g} = \frac{1}{|D|} \sum_{\tau \in D} \sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) R(\tau)
\]

RECAP:截至目前，我们学会了什么？

目标：通过梯度方法优化策略的参数，以最大化预期奖励。

算法步骤：

a. 创建策略网络：

使用一个神经网络定义策略$\pi_\theta$，输入为当前状态$s_t$，输出为在动作空间上的概率分布$\pi_\theta(a_t|s_t)$。

b. 采样轨迹：

使用策略$\pi_\theta$来与环境交互，生成轨迹（包括状态$s_t$、动作$a_t$、奖励$R(\tau)$）。通常每条轨迹可以运行固定步数（如 100 步）或直到达到终止条件。

c. 计算梯度：

根据公式：

\[\hat{g} = \frac{1}{|D|} \sum_{\tau \in D} \sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) R(\tau)
\]

d. 更新策略参数：

使用随机梯度上升更新参数：

\[\theta_{k+1} =\theta_k + \alpha \nabla_\theta J(\pi_\theta)|_{\theta_k}
\]

e. 循环迭代：

返回第 b 步，继续采样新的轨迹并重复以上步骤，直到收敛。

4. 为语言模型生成轨迹

还记得我们为reward model建立的偏好数据集吗？我们使用这些问题询问模型以生成回答。然后我们计算生成回答的reward，并根据估计的策略梯度训练模型。

如何计算语言模型的策略log probabilities

计算每个轨迹的reward

4.1 策略梯度优化GPO存在的问题

问题：

梯度估计是无偏的(unbiased)（意思是他可以收敛到真实的梯度），但是它具有很高的方差（varience)

解决：

方法一 : 采用 "rewards to go"（剩余奖励）来替代整条轨迹上的累积奖励

$R(τ)$

核心思想：you can't alter the past

a. 原始梯度公式

最初的梯度公式：

\[\nabla_{\theta} J(\pi_\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T} \nabla a_{\theta} \log \pi_\theta(a_t|s_t) R(\tau) \right]
\]

问题：公式将整个轨迹的累积奖励 $ R(\tau)$与所有时间步的梯度相乘。
结果：由于 $R(\tau)$包含了从 $t=0$到$T$的所有奖励，梯度的方差非常高，尤其是在奖励稀疏或延迟时。

b. 改进：按时间步分解奖励

将公式重写为：

\[\nabla_{\theta} J(\pi_\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_\theta(a_{i,t} | s_{i,t}) \left( \sum_{t=0}^{T} r(s_{i,t}, a_{i,t}) \right)
\]

每个动作梯度乘以整条轨迹的奖励总和。
发现的问题：对于一个具体时间步 ( t )，轨迹中之前的奖励（即 ( t'<t ) 的奖励）对当前动作无关紧要，却仍然被考虑进了公式（因为过去已经无法改变）。

c. 引入 “rewards to go”

为了减少方差，使用 "rewards to go" $\sum_{t'=t}^{T} r(s_{i,t'}, a_{i,t'})$ 来替换原始的轨迹总奖励$R(\tau)$。改写后的公式为：

\[\nabla_{\theta} J(\pi_\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_\theta(a_{i,t} | s_{i,t}) \left( \sum_{t'=t}^{T} r(s_{i,t'}, a_{i,t'}) \right)
\]

解释：
- 对于每个时间步 $ t$，只计算从当前时间步到轨迹结束的累积奖励。
- 忽略过去的奖励，因为它们在期望中会相互抵消，不影响最终梯度的无偏性。

d. 减少方差的原因

原始公式中，过去的奖励$t' < t$对当前动作梯度没有直接影响，但由于被包含进公式中，会引入额外的噪声（方差）。
改用 "rewards to go" 后，公式更准确地反映了当前动作对未来奖励的直接影响，从而降低了梯度估计的方差。

方法二：引入基准函数（Baseline）来进一步减少梯度估计方差。

目标：通过引入基准函数$b$，减少梯度估计的方差，改善优化过程的稳定性。
常用基准：状态价值函数 $V^\pi(s)$，也可以使用其他函数，具体取决于任务和模型的设计。
优势：减小了无关噪声的影响，让策略更专注于改进未来表现。

a. 梯度公式中的基准函数

基础的梯度公式为：

\[\nabla_{\theta} J(\pi_\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_\theta(a_{i,t} | s_{i,t}) \left( \sum_{t'=t}^{T} r(s_{i,t'}, a_{i,t'}) \right)
\]

通过引入一个基准函数 (b)（Baseline），公式变为：

\[\nabla_{\theta} J(\pi_\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_{\theta} \log \pi_\theta(a_{i,t} | s_{i,t}) \left( \sum_{t'=t}^{T} r(s_{i,t'}, a_{i,t'}) - b \right)
\]

b. 为什么引入基准函数？

无偏性：可以证明，减去一个基准函数$b$并不会改变梯度估计的无偏性。
- 也就是说，基准函数的引入不会对梯度的方向或期望值产生影响。
减少方差：引入基准函数后，公式中引入了一种对奖励的“校正”，从而有效减少了梯度估计中的方差。

c. 选择基准函数的方式

常用的基准函数是状态价值函数$V^\pi(s)$：

\[V^\pi(s) = \mathbb{E}_\pi \left[ \sum_{t'=t}^{T} r(s_{t'}, a_{t'}) \, | \, s_t = s \right]
\]

解释：
- $V^\pi(s)$表示从当前状态$s_t$出发，按照策略 $\pi$所能获得的期望未来奖励。
- 使用 $V^\pi(s)$作为基准，可以有效去掉与当前状态 $s_t$无关的波动。

例子：高值状态与低值状态

对于某个状态 $s_t$，如果期望的未来奖励（价值）较低，基准值会相应较低。
如果某个状态的未来奖励较高（即更接近目标或有利），基准值会较高。
减去基准值后，梯度更新会更聚焦于当前动作的实际贡献，而不是受环境噪声或总体奖励波动的影响。

估计$V^\pi (s)$:

注意：语言模型已经有一个线性层（把隐藏层转化为logits值）这里添加的是另一个线性层。

d. 具体优化效果

引入基准函数后：
- 当 $\sum_{t'=t}^{T} r(s_{i,t'}, a_{i,t'})$ 大于基准值 $b$时，表示当前动作比预期更好，梯度更新会倾向于强化该动作。
- 当 $\sum_{t'=t}^{T} r(s_{i,t'}, a_{i,t'})$小于基准值 $b$时，表示当前动作比预期更差，梯度更新会倾向于削弱该动作。

e. 引入 Q 和 V：

(未完待续)