深入解析《Generalizable Reward Model (GRM)》：提升奖励模型泛化能力的创新方法

深入解析《Generalizable Reward Model (GRM)》：提升奖励模型泛化能力的创新方法

近年来，大型语言模型（LLMs）的对齐研究成为自然语言处理领域的热点，尤其是在基于人类反馈的强化学习（RLHF）框架中，奖励模型的训练与优化显得尤为重要。然而，现有奖励模型在面对分布外（out-of-distribution, OOD）数据时，泛化能力不足，容易导致“奖励过度优化”（reward over-optimization），从而降低模型的实际性能。NeurIPS 2024 收录的论文《Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs》提出了一种名为 Generalizable Reward Model (GRM) 的新型方法，通过正则化隐藏状态显著提升奖励模型的泛化能力。本文将详细介绍该论文的贡献、技术细节和数学公式，面向熟悉 LLM 和 RLHF 的研究者，深入剖析其创新之处。

Paper：https://arxiv.org/pdf/2406.10216

---

1. 论文背景与问题

RLHF 是一种通过人类偏好数据训练奖励模型，并利用强化学习（如 PPO）优化语言模型以最大化奖励的对齐方法。尽管 RLHF 在对齐任务中取得了显著成功，但奖励模型的泛化能力仍然是一个关键挑战。现有奖励模型在面对未见过的提示-响应对时，往往表现出以下问题：

• 奖励过度优化：在优化过程中，代理奖励模型（proxy reward model）看似提升了奖励分数，但实际上偏离了真实的人类偏好，导致性能下降。这种现象被称为“奖励黑客”（reward hacking）。
• 泛化能力不足：奖励模型在训练数据分布外的表现不佳，尤其是在数据量有限的情况下，模型容易过拟合，导致 OOD 任务的准确率下降。

针对这些问题，论文提出了一种轻量且高效的解决方案，通过正则化奖励模型的隐藏状态，增强其对分布变化的鲁棒性和泛化能力。与以往通过约束策略优化或训练多个奖励模型的方法不同，GRM 专注于改进奖励模型本身，是一种更直接且成本效益更高的方案。

---

2. GRM 的核心贡献

GRM 的核心创新在于通过保留预训练语言模型的语言生成能力，对隐藏状态进行正则化，从而提升奖励模型的泛化能力。以下是论文的三大主要贡献：

1. 提出 GRM 框架：通过在隐藏状态上施加文本生成正则化，GRM 同时训练奖励头和语言模型头，增强奖励模型对未见数据的泛化能力。
2. 验证三种文本生成正则化：论文提出并验证了三种正则化方法（DPO、DPO 无参考模型、SFT），其中 SFT 正则化被证明是最有效且稳定的选择。
3. 显著提升泛化性能：实验结果表明，GRM 在多个 OOD 任务上显著提高了奖励模型的准确率，并有效缓解了 RLHF 中的过度优化问题。

---

3. GRM 的技术细节与数学公式

GRM 的核心思想是通过正则化隐藏状态，防止随机初始化的奖励头对预训练特征的扭曲，从而提升模型的 OOD 性能。以下是其技术细节和关键数学公式的详细解析。

3.1 GRM 模型结构

GRM 的结构如图 1 所示，包含一个共享的隐藏状态，以及两个输出头：

• 奖励头（$r_{\theta }$）：用于预测提示-响应对的奖励分数，优化奖励损失 ${\mathcal{L}}_{\text{reward}}$。
• 语言模型头（${\pi }_{{\theta }_{\mathrm{LM}}}$）：保留预训练模型的文本生成能力，优化一组文本生成损失 ${\mathcal{L}}_{\text{reg}}$。

总损失函数定义为：

$${\mathcal{L}}_{\text{total}}=(1-\alpha ){\mathcal{L}}_{\text{reward}}+\alpha {\mathcal{L}}_{\text{reg}}$$

其中，$\alpha$ 是平衡奖励损失和正则化损失的超参数，默认为 0.01。

3.2 奖励损失

奖励模型基于 Bradley-Terry 模型，通过比较偏好数据中的选择响应（$y_c$）和拒绝响应（$y_r$）来训练。奖励损失函数为：

$${\mathcal{L}}_{\text{reward}}(\theta )=-{\mathbb{E}}_{(x,y_c,y_r)\sim D}\left[\log \left(\sigma \left(r_{\theta }(x,y_c)-r_{\theta }(x,y_r)\right)\right)\right]$$

其中：

• $r_{\theta }(x,y)$ 是奖励模型对提示 $x$ 和响应 $y$ 的打分。
• $\sigma (\cdot )$ 是 sigmoid 函数。
• $D$ 是偏好数据集，包含 $(x,y_c,y_r)$ 三元组。

通过最小化该损失，奖励模型学习为人类偏好的响应分配更高分数。

3.3 文本生成正则化

为了防止奖励头扭曲预训练特征，GRM 引入了三种文本生成正则化方法，分别基于 DPO 和 SFT 损失。这些正则化方法通过语言模型头优化隐藏状态的文本生成能力。以下是三种正则化的定义：

3.3.1 DPO 正则化

DPO（Direct Preference Optimization）是一种无需显式奖励模型的偏好优化方法。GRM 借鉴其损失函数作为正则化项：

$${\mathcal{L}}_{\text{DPO}}({\theta }_{\mathrm{LM}})=-{\mathbb{E}}_{(x,y_c,y_r)\sim D}\left[\log \sigma \left(\beta \log \left(\frac{{\pi }_{{\theta }_{\mathrm{LM}}}(y_c\mid x)}{{\pi }_{\text{ref}}(y_c\mid x)}\right)-\beta \log \left(\frac{{\pi }_{{\theta }_{\mathrm{LM}}}(y_r\mid x)}{{\pi }_{\text{ref}}(y_r\mid x)}\right)\right)\right]$$

其中：

• ${\pi }_{{\theta }_{\mathrm{LM}}}$ 是语言模型头的策略。
• ${\pi }_{\text{ref}}$ 是参考模型（通常为预训练模型）。
• $\beta$ 是控制 KL 散度的超参数，默认为 0.1。

该正则化通过比较选择和拒绝响应的相对概率，鼓励隐藏状态保留文本生成能力。

3.3.2 DPO 无参考模型正则化

为了降低内存消耗，GRM 提出了一种无需参考模型的 DPO 正则化：

$${\mathcal{L}}_{\text{DPO-noref}}({\theta }_{\mathrm{LM}})=-{\mathbb{E}}_{(x,y_c,y_r)\sim D}\left[\log \sigma \left(\beta \log \left(\frac{{\pi }_{{\theta }_{\mathrm{LM}}}(y_c\mid x)}{{\pi }_{{\theta }_{\mathrm{LM}}}(y_r\mid x)}\right)\right)\right]$$

该方法直接比较选择和拒绝响应的概率，减少了对参考模型的依赖，适合大型模型的训练。

3.3.3 SFT 正则化

SFT（Supervised Fine-Tuning）正则化是最简单的形式，仅优化选择响应的概率：

$${\mathcal{L}}_{\text{SFT}}({\theta }_{\mathrm{LM}})=-{\mathbb{E}}_{(x,y_c)\sim D}\left[\log \sigma \left(\beta \log \left({\pi }_{{\theta }_{\mathrm{LM}}}(y_c\mid x)\right)\right)\right]$$

尽管形式简单，SFT 正则化在实验中表现出最稳定和高效的性能，成为 GRM 的默认选择。

3.4 理论动机

GRM 的正则化方法源于一个对抗优化问题，目标是学习一个对对抗策略鲁棒的奖励模型。论文考虑以下优化目标：

$$\theta =\arg \underset{\theta }{\min }\left\{{\mathcal{L}}_{\text{reward}}(\theta )+\gamma \underset{\pi }{\max }J(\theta ,\pi )\right\}$$

其中，$J(\theta ,\pi )$ 是策略优化的目标函数，定义为：

$$J(\theta ,\pi )={\mathbb{E}}_{x\sim D,y\sim \pi (\cdot \mid x)}\left[r_{\theta }(x,y)\right]-\beta {\mathbb{E}}_{x\sim D}\left[\mathrm{KL}\left(\pi (\cdot \mid x)\parallel {\pi }_{\text{ref}}(\cdot \mid x)\right)\right]$$

通过求解内层最大化问题，得到最优策略 ${\pi }_{\theta }^{\ast }$：

$${\pi }_{\theta }^{\ast }=\frac{1}{Z_{\theta }(x)}{\pi }_{\text{ref}}(y\mid x)\exp \left(r_{\theta }(x,y)/\beta \right)$$

其中，$Z_{\theta }(x)$ 是归一化常数。将 ${\pi }_{\theta }^{\ast }$ 代入优化目标，并引入校准策略 ${\pi }_{\text{cal}}$，最终推导出正则化损失形式：

θ = arg ⁡ min ⁡ θ { ( 1 − α ) L reward ( θ ) + α DPO L DPO ( π θ ∗ ) + α SFT L SFT ( π θ ∗ ) } \theta = \arg \min_\theta \left\{ (1 - \alpha) \mathcal{L}_{\text{reward}}(\theta) + \alpha_{\text{DPO}} \mathcal{L}_{\text{DPO}}(\pi_\theta^*) + \alpha_{\text{SFT}} \mathcal{L}_{\text{SFT}}(\pi_\theta^*) \right\} θ=argθmin​{(1−α)Lreward​(θ)+αDPO​LDPO​(πθ∗​)+αSFT​LSFT​(πθ∗​)}

论文进一步放宽了 $r_{\theta }$ 和 ${\pi }_{\theta }^{\ast }$ 的关系，提出共享隐藏状态的奖励模型和语言模型头，从而实现高效的正则化。
(具体请看论文附录)

---

4. 实验结果与分析

论文在多个数据集和基准测试上验证了 GRM 的性能，包括 Unified-Feedback（ID 数据）、HHH-Alignment、MT-Bench 和 RewardBench（OOD 数据）。以下是关键实验结果：

4.1 泛化性能

GRM 在 ID 和 OOD 任务上均显著优于基线模型。例如，在 400K Unified-Feedback 数据上，GRM w/ SFT 将 HHH-Alignment 分数从 73.4 提升至 79.8，MT-Bench 分数从 71.2 提升至 73.4。尤其在数据量较小时（40K），GRM 的优势更加明显，表明其对有限数据的鲁棒性。

4.2 正则化方法比较

三种正则化方法（DPO、DPO 无参考、SFT）在大数据量下性能相当，但在小数据量（40K）时，SFT 正则化表现最佳，计算成本也最低。这使得 SFT 成为 GRM 的首选正则化方法。

4.3 缓解过度优化

在 BoN（Best-of-n Sampling）和 PPO 实验中，GRM 显著缓解了过度优化问题。图 2 和图 3 显示，GRM 的代理分数和金标准分数随 KL 散度增加而持续上升，而基线模型的金标准分数在高 KL 时下降，表明存在过度优化。GRM 在 7B 模型上的金标准分数达到 1.5，远超基线的 0.5。

4.4 噪声鲁棒性

GRM 在 25% 标签噪声下仍表现出色，金标准分数保持在 1.0 以上，未出现显著下降，证明了其对偏好数据噪声的鲁棒性。

4.5 大规模模型扩展

在 llama3-8b-instruct 模型上，GRM 在 RewardBench 上达到 87.0 的平均分数，超越了 Starling-RM-34B（82.7）和 GPT-4（85.9），尤其在推理任务中表现突出（92.3）。

---

5. GRM 的优势与局限性

5.1 优势

1. 缓解特征扭曲：文本生成正则化保留了预训练模型的生成能力，防止隐藏状态被奖励头扭曲。
2. 防止过拟合：对抗性正则化避免模型过拟合到错误特征，尤其在数据量有限或存在噪声时效果显著。
3. 高效性：GRM 无需训练多个奖励模型或额外数据，SFT 正则化进一步降低了计算成本。

5.2 局限性

1. 合成噪声实验：论文使用 25% 合成噪声评估鲁棒性，可能无法完全反映真实人类标注的复杂性。
2. 计算限制：受限于计算资源，GRM 未在超过 10B 参数的模型上测试，未来扩展到更大模型可能带来进一步提升。

---

6. 对未来研究的启发

GRM 的提出为奖励模型的泛化研究开辟了新方向。未来研究可以从以下方面进一步探索：

• 真实人类标注数据：使用真实偏好数据验证 GRM 的鲁棒性，解决合成数据可能引入的偏差。
• 更大模型扩展：将 GRM 应用于更大规模的模型（如 70B 或更高），验证其在复杂任务中的性能。
• 奖励头结构优化：探索奖励头的不同结构（如多层非线性头）对泛化能力的影响。
• 结合其他正则化技术：将 GRM 与标签平滑、集成方法等结合，进一步提升奖励模型的鲁棒性。

---

7. 结论

《Generalizable Reward Model (GRM)》通过正则化隐藏状态，提出了一种高效且鲁棒的奖励模型训练方法，显著提升了 LLMs 在 RLHF 框架下的泛化能力。其核心贡献在于通过共享隐藏状态的奖励头和语言模型头，结合 DPO 和 SFT 正则化，缓解了奖励过度优化问题，并在 OOD 任务中取得了优异表现。论文的数学推导清晰，实验设计严谨，为研究者提供了宝贵的参考。GRM 的轻量设计和高效性使其在实际应用中具有广泛潜力，同时也为未来的对齐研究提供了新的思路。

根据论文《Regularizing Hidden States Enables Learning Generalizable Reward Model for LLMs》（GRM），训练奖励模型的实现需要基于预训练的大型语言模型（LLM），并在其隐藏状态之上同时训练一个奖励头（reward head）和一个语言模型头（LM head），以通过文本生成正则化增强奖励模型的泛化能力。论文明确提到，GRM 的结构是将预训练模型的隐藏状态共享给两个头：奖励头用于优化奖励损失，语言模型头用于优化文本生成正则化损失（如 SFT 或 DPO）。这意味着 LLM 的最后一层隐藏状态（而不是直接使用最后一层的输出）被用作两个头的输入。

代码

请注意，下面是伪代码，仅用作抛砖引玉！！！！具体请看原作者开源的实现。

以下是基于论文描述，使用 Python 和 PyTorch 实现的 GRM 奖励模型训练代码示例。代码假设使用 Hugging Face 的 Transformers 库，并以 gemma-2b-it 或 Mistral-7B-Instruct 作为基础模型（论文中使用的模型）。我们将重点展示如何实现共享隐藏状态、构造奖励头和语言模型头，以及如何结合奖励损失和正则化损失进行训练。由于论文推荐 SFT 正则化为最稳定和高效的选择，以下代码以 SFT 正则化为例。

---

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
from typing import Tuple, Dict
import uuid

# 自定义数据集类，处理偏好数据 (x, y_c, y_r)
class PreferenceDataset(Dataset):
    def __init__(self, data: list, tokenizer, max_length: int = 1024):
        self.data = data
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        prompt, chosen, rejected = self.data[idx]
        # 编码提示、选择响应和拒绝响应
        prompt_enc = self.tokenizer(prompt, max_length=self.max_length, truncation=True, padding=False, return_tensors="pt")
        chosen_enc = self.tokenizer(chosen, max_length=self.max_length, truncation=True, padding=False, return_tensors="pt")
        rejected_enc = self.tokenizer(rejected, max_length=self.max_length, truncation=True, padding=False, return_tensors="pt")
        return {
            "prompt_ids": prompt_enc["input_ids"].squeeze(),
            "prompt_mask": prompt_enc["attention_mask"].squeeze(),
            "chosen_ids": chosen_enc["input_ids"].squeeze(),
            "chosen_mask": chosen_enc["attention_mask"].squeeze(),
            "rejected_ids": rejected_enc["input_ids"].squeeze(),
            "rejected_mask": rejected_enc["attention_mask"].squeeze(),
        }

# GRM 模型定义
class GRMModel(nn.Module):
    def __init__(self, base_model_name: str, hidden_size: int, reward_head_nonlinear: bool = True):
        super(GRMModel, self).__init__()
        # 加载预训练模型（去掉原始 LM 头）
        self.base_model = AutoModel.from_pretrained(base_model_name)
        self.hidden_size = hidden_size

        # 奖励头：可配置为线性或非线性
        if reward_head_nonlinear:
            self.reward_head = nn.Sequential(
                nn.Linear(hidden_size, 1024),
                nn.ReLU(),
                nn.Linear(1024, 1)
            )
        else:
            self.reward_head = nn.Linear(hidden_size, 1)

        # 语言模型头：与预训练模型的词汇表大小一致
        self.lm_head = nn.Linear(hidden_size, self.base_model.config.vocab_size)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        compute_lm: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # 获取基础模型的隐藏状态
        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
        last_hidden_state = outputs.hidden_states[-1]  # 最后一层隐藏状态 [batch_size, seq_len, hidden_size]

        # 提取序列最后一个 token 的隐藏状态（通常用于分类任务）
        batch_size = input_ids.size(0)
        seq_lengths = attention_mask.sum(dim=1) - 1  # 每个序列的最后一个有效 token 位置
        last_hidden = last_hidden_state[torch.arange(batch_size), seq_lengths]  # [batch_size, hidden_size]

        # 奖励头输出
        reward = self.reward_head(last_hidden)  # [batch_size, 1]

        # 语言模型头输出（仅在需要正则化时计算）
        lm_logits = None
        if compute_lm:
            lm_logits = self.lm_head(last_hidden_state)  # [batch_size, seq_len, vocab_size]

        return reward, lm_logits

# 奖励损失函数（Bradley-Terry 模型）
def compute_reward_loss(
    model: GRMModel,
    prompt_ids: torch.Tensor,
    prompt_mask: torch.Tensor,
    chosen_ids: torch.Tensor,
    chosen_mask: torch.Tensor,
    rejected_ids: torch.Tensor,
    rejected_mask: torch.Tensor
) -> torch.Tensor:
    # 计算选择响应的奖励
    chosen_reward, _ = model(chosen_ids, chosen_mask, compute_lm=False)
    # 计算拒绝响应的奖励
    rejected_reward, _ = model(rejected_ids, rejected_mask, compute_lm=False)
    # Bradley-Terry 损失
    reward_diff = chosen_reward - rejected_reward
    reward_loss = -F.logsigmoid(reward_diff).mean()
    return reward_loss

# SFT 正则化损失
def compute_sft_loss(
    model: GRMModel,
    chosen_ids: torch.Tensor,
    chosen_mask: torch.Tensor,
    beta: float = 0.1
) -> torch.Tensor:
    # 获取语言模型头的 logits
    _, lm_logits = model(chosen_ids, chosen_mask, compute_lm=True)
    # 移位以预测下一个 token
    shift_logits = lm_logits[..., :-1, :].contiguous()
    shift_labels = chosen_ids[..., 1:].contiguous()
    # 计算交叉熵损失
    loss_fct = nn.CrossEntropyLoss(reduction="mean", ignore_index=self.tokenizer.pad_token_id)
    sft_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
    # 应用 beta 缩放并通过 sigmoid
    sft_loss = -torch.logsigmoid(beta * sft_loss)
    return sft_loss

# 训练函数
def train_grm(
    model: GRMModel,
    tokenizer,
    dataset: PreferenceDataset,
    epochs: int = 2,
    batch_size: int = 16,
    learning_rate: float = 1e-5,
    alpha: float = 0.01,
    beta: float = 0.1,
    device: str = "cuda"
):
    model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for batch in dataloader:
            # 移动数据到设备
            batch = {k: v.to(device) for k, v in batch.items()}

            # 计算奖励损失
            reward_loss = compute_reward_loss(
                model,
                batch["prompt_ids"],
                batch["prompt_mask"],
                batch["chosen_ids"],
                batch["chosen_mask"],
                batch["rejected_ids"],
                batch["rejected_mask"]
            )

            # 计算 SFT 正则化损失
            sft_loss = compute_sft_loss(model, batch["chosen_ids"], batch["chosen_mask"], beta)

            # 总损失
            total_loss = (1 - alpha) * reward_loss + alpha * sft_loss

            # 反向传播
            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

            total_loss += total_loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1}/{epochs}, Average Loss: {avg_loss:.4f}")

# 主函数
def main():
    # 模型和分词器配置
    base_model_name = "google/gemma-2b-it"
    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    hidden_size = 2048  # gemma-2b-it 的隐藏状态大小

    # 加载模型
    model = GRMModel(base_model_name, hidden_size, reward_head_nonlinear=True)

    # 模拟偏好数据 [(prompt, chosen, rejected), ...]
    data = [
        ("What's the capital of France?", "The capital of France is Paris.", "Florida."),
        # 添加更多数据...
    ]
    dataset = PreferenceDataset(data, tokenizer)

    # 训练模型
    train_grm(
        model,
        tokenizer,
        dataset,
        epochs=2,
        batch_size=16,
        learning_rate=1e-5,
        alpha=0.01,
        beta=0.1,
        device="cuda" if torch.cuda.is_available() else "cpu"
    )

if __name__ == "__main__":
    main()

---

代码说明与论文细节对应

1. 模型结构：

   • 共享隐藏状态：GRMModel 类中，self.base_model（预训练模型）输出最后一层隐藏状态（last_hidden_state），该状态被共享给奖励头（self.reward_head）和语言模型头（self.lm_head）。这与论文中描述的结构一致（见论文图 1 和第 3 节）。
   • 奖励头：默认使用非线性头（Linear -> ReLU -> Linear），如论文中所述（附录 B）。也可以切换为线性头（reward_head_nonlinear=False）。
   • 语言模型头：与预训练模型的 vocabulary 大小一致，用于计算文本生成损失。

2. 损失函数：

   • 奖励损失（compute_reward_loss）：基于 Bradley-Terry 模型，计算选择响应和拒绝响应的奖励差，优化公式 (1)：
     $${\mathcal{L}}_{\text{reward}}(\theta )=-{\mathbb{E}}_{(x,y_c,y_r)\sim D}\left[\log \left(\sigma \left(r_{\theta }(x,y_c)-r_{\theta }(x,y_r)\right)\right)\right]$$
   • SFT 正则化损失（compute_sft_loss）：实现公式 (10)，优化选择响应的生成概率：
     $${\mathcal{L}}_{\text{SFT}}({\theta }_{\mathrm{LM}})=-{\mathbb{E}}_{(x,y_c)\sim D}\left[\log \sigma \left(\beta \log \left({\pi }_{{\theta }_{\mathrm{LM}}}(y_c\mid x)\right)\right)\right]$$
     代码中通过交叉熵损失近似，并应用 beta 缩放和 sigmoid 函数。

3. 总损失：

   • 总损失结合奖励损失和正则化损失，公式 (4)：
     $${\mathcal{L}}_{\text{total}}=(1-\alpha ){\mathcal{L}}_{\text{reward}}+\alpha {\mathcal{L}}_{\text{reg}}$$
     其中，$\alpha =0.01$，$\beta =0.1$（论文默认值，附录 B）。

4. 训练细节：

   • 使用 LoRA（Low-Rank Adaptation）进行微调（论文附录 B），但为简单起见，代码中未实现 LoRA，可通过 Hugging Face 的 peft 库添加。
   • 训练 2 个 epoch，学习率 $1\times 10^{-5}$，batch size 16，符合论文设置（附录 B，表 6）。
   • 输入截断至 1024 个 token，模型使用 bf16 精度。

5. 数据集：

   • 代码中的 PreferenceDataset 假设偏好数据为 (prompt, chosen, rejected) 三元组，实际应用中应使用 Unified-Feedback 数据集（论文第 4 节）。

---

是否将 LLM 的最后一层给到两个头？

是的，GRM 的设计明确是将 LLM 的最后一层隐藏状态（last_hidden_state）作为两个头的输入，而不是直接使用 LLM 的最后一层输出（例如，原始语言模型头的 logits）。具体来说：

• 隐藏状态提取：代码中通过 outputs.hidden_states[-1] 获取最后一层隐藏状态，并从中提取最后一个有效 token 的状态（last_hidden）用于奖励头，完整序列的隐藏状态用于语言模型头。
• 论文依据：论文第 3 节提到，奖励头和语言模型头共享相同的隐藏状态（“sharing the same hidden states”），并且奖励头基于最后一个 token 的隐藏状态预测奖励分数（常见于序列分类任务）。语言模型头则基于整个序列的隐藏状态生成 token 概率。

因此，LLM 的最后一层隐藏状态被直接传递给两个头，而不是将原始语言模型头的输出（logits）拆分或重用。

---

扩展与注意事项

1. 添加 LoRA：

   • 论文中使用 LoRA 进行高效微调（LoRA rank=32, alpha=64，附录 B）。实际实现中，可使用 peft 库的 LoraConfig 和 get_peft_model 包装 base_model。

2. DPO 正则化：

   • 如果需要实现 DPO 或 DPO 无参考模型正则化，可参考公式 (8) 和 (9)，修改 compute_sft_loss 为对应的 DPO 损失函数。需要额外加载参考模型（${\pi }_{\text{ref}}$）或直接比较选择和拒绝响应的概率。

3. 实际数据集：

   • 代码中的 data 为模拟数据，实际应使用 Unified-Feedback 数据集（https://huggingface.co/datasets/llm-blender/Unified-Feedback）。数据集需预处理为 (prompt, chosen, rejected) 格式。

4. 计算资源：

   • 论文提到使用 NVIDIA RTX A6000（49GB）训练，2B 模型需要约 30.4 GPU 小时（40K 数据，2 epoch）。确保有足够 GPU 内存，或使用模型并行/量化技术。

5. 超参数调优：

   • 论文强调正则化权重 $\alpha$ 是关键超参数（附录 C.4）。建议在 [0.001, 0.1] 范围内进行网格搜索，尤其对于线性奖励头（$\alpha =0.001$ 更优）。

---

总结

以上代码实现了 GRM 奖励模型的训练，核心是共享 LLM 最后一层隐藏状态给奖励头和语言模型头，通过 SFT 正则化增强泛化能力。代码与论文的数学公式和结构设计一一对应，适合研究者进一步扩展（例如添加 DPO 正则化或 LoRA）。

后记

2025年5月16日于上海，在grok 3大模型辅助下完成。