由于 RLHF 的訓練過程中需要依賴大量的人類偏好數據進行學習，因此很難在訓練過程中要求人類標注者實時提供偏好反饋。為此，我們需要訓練一個模型來替代人類在 RLHF 訓練過程中實時提供反饋，這個模型被稱為獎勵模型

🔸一、 目標函數公式解釋

公式如下：

$$L=?{\mathbb{E}}_{(x,y^{+},y^{?})～D}\left[\log \sigma (r_{\theta }(x,y^{+})?r_{\theta }(x,y^{?}))\right]?\beta {\mathbb{E}}_{(x,y^{+})～D}\left[\sum _{t=1}^T\log p(y_t^{+}\mid x,y_{<t}^{+})\right]$$

含義拆解：

• x: 輸入（如問題或提示語）
• y+: 正例響應（由人類標注或偏好選擇的答案）
• y-: 負例響應（不好的答案）
• r_θ(x, y): 獎勵模型對 (x, y) 的打分（通常是最后一個 token 的輸出經過 reward head 得到）
• σ: Sigmoid 函數
• β: 權重超參，控制模仿學習（第二項）對總損失的影響程度

公式兩部分含義：

1. 對比損失（ranking loss）：

   $$

   • \log \sigma(r(x, y^+) - r(x, y^-))
     $$
   • 目標是使 正例得分 > 負例得分
   • 當 r(x, y+) ? r(x, y-) 時，sigmoid接近1，log接近0 → 損失小，說明模型學得好

2. 模仿學習損失（語言模型 loss）：

   $$

   • \sum_{t=1}^{T} \log p(y^+t \mid x, y^+{<t})
     $$
   • 即：語言模型在給定輸入 x 和前綴 y^+_{<t} 的條件下，預測下一個 token 的交叉熵損失
   • 起正則作用，防止獎勵模型過度擬合打分而喪失語言生成能力

---

🔸二、代碼結構分析

基于 LLaMA 的獎勵模型實現詳解（逐行解讀 + PyTorch 源碼分析）

📦 模塊導入

1  import torch
2  import torch.nn as nn
3  import torch.nn.functional as F
4
5  from transformers import LlamaForCausalLM

• torch：PyTorch 核心包
• nn：用于定義神經網絡模塊（如 Linear）
• F：包含函數式接口（如 loss 函數）
• LlamaForCausalLM：來自 Transformers 的 LLaMA 語言模型基類，支持自回歸文本生成

---

🧠 模型定義：獎勵模型類

7  class LlamaRewardModel(LlamaForCausalLM):
8      def __init__(self, config):
9          super().__init__(config)
10
11         # 初始化線性變換層，將隱狀態映射為標量，用于輸出最終獎勵
12         self.reward_head = nn.Linear(config.hidden_size, 1, bias=False)

• LlamaRewardModel 繼承自 HuggingFace 的 LlamaForCausalLM
• 增加了一個 reward_head 線性層，用于將模型輸出（hidden state）映射為 獎勵值（scalar）

---

🧾 正例/負例打分函數 _forward_rmloss

14 def _forward_rmloss(self, input_ids, attention_mask, **kargs):
18     output = self.model.forward(
19         input_ids=input_ids,
20         attention_mask=attention_mask,
21         return_dict=True,
22         use_cache=False
23     )
25     logits = self.reward_head(output.last_hidden_state).squeeze(-1)
26     return logits

• 輸入：拼接后的 [x, y] 序列
• self.model.forward(...)：獲得 LLaMA 模型輸出（hidden states）
• self.reward_head(...)：只對最后一層 hidden state 應用線性映射，輸出獎勵值
• squeeze(-1)：去除最后一維 [batch, 1] -> [batch]

squeeze(-1) 的作用是去掉張量的最后一個維度，前提是該維度的值是 1。
假設 logits 是一個 [batch_size, 1] 的張量：
logits = tensor([[0.73], [0.24], [0.91]]) # shape: [3, 1]
執行：
logits = logits.squeeze(-1)
結果為：
tensor([0.73, 0.24, 0.91]) # shape: [3]

?? 模仿學習損失函數 _forward_lmloss

29 def _forward_lmloss(self, prompt_ids, lm_attn_mask, response_ids):
35     outputs = self.model.forward(
36         input_ids=prompt_ids,
37         attention_mask=lm_attn_mask,
38         return_dict=True,
39         use_cache=False,
40     )
42     hidden_states = outputs.last_hidden_state
43     logits = self.lm_head(hidden_states)
44     loss_fct = nn.CrossEntropyLoss()
45     logits = logits.view(-1, self.config.vocab_size)
46     response_ids = response_ids.view(-1)
47     loss = loss_fct(logits, response_ids)
48     return loss

• prompt_ids 是 [x, y?] 拼接后的 token ID
• 輸出 logits：維度 [batch_size, seq_len, vocab_size]
• 計算交叉熵損失：對所有位置預測的 token 與 response_ids 進行對比

---

🚀 前向傳播函數：組合損失計算

50 def forward(self, sent1_idx, attention_mask_1, sent2_idx,attention_mask_2, labels, prompt_ids, lm_attn_mask, response_ids):

參數說明：

• sent1_idx: [x, y?] 拼接輸入（正例）
• sent2_idx: [x, y?] 拼接輸入（負例）
• labels: 全 0 標簽，用于對比損失
• prompt_ids: 與正例相關的 token（用于 LM Loss）
• response_ids: 正例的 target token（用于 LM Loss）

計算對比損失（Reward Loss）

61 reward0 = self._forward_rmloss(sent1_idx, attention_mask_1)
66 reward1 = self._forward_rmloss(sent2_idx, attention_mask_2)
71 logits = reward0 - reward1
72 rm_loss = F.binary_cross_entropy_with_logits(logits,labels.to(logits.dtype), reduction="mean")

• 分別計算 r(x, y?) 與 r(x, y?)

• 構造 logits = r? - r?

• 用 Binary Cross Entropy Loss 計算 reward loss

  公式對應：
  $$?\log (\sigma (r(x,y^{+})?r(x,y^{?})))$$

---

計算語言模型損失（Language Modeling Loss）

75 lm_loss = self._forward_lmloss(prompt_ids, lm_attn_mask, response_ids)

• 與傳統語言模型訓練一致，使用 CrossEntropyLoss

---

返回總損失

78 loss = rm_loss + lm_loss
79 return loss

• 二者直接加和（可選加權項 β，可自己加參數）
• 模型即同時優化打分能力 + 文本生成能力（聯合學習）

🔸四、總結

項目	描述
核心思想	同時學習獎勵模型 r_θ 和保持生成流暢性
優勢	1. 保留強化學習能力 2. 不失語義與流暢性
應用場景	RLHF 的 reward 模型訓練階段，如 OpenAI 的 GPT 訓練流程中 Step 2: Train Reward Model
可調參數	β 控制生成質量與偏好打分之間的權衡