# Reward Models ## Reward Models A standalone guide to the scalar Bradley–Terry reward model: score a full prompt–response sequence, then learn which answer should rank higher. Based on [RLHF Book — Reward Modeling](https://rlhfbook.com/c/05-reward-models) and the existing [PPO reward-model notes](https://www.qiyuzhi.com/study/rlhf-rlvr/ppo/#2-reward-model). ### 1. Bradley–Terry derivation For the same prompt \(x\), let \(y_1\) be the preferred response and \(y_2\) the rejected response. The Bradley–Terry model assumes: $$ \boxed{ P(y_1\succ y_2\mid x) =\frac{e^{r(x,y_1)}}{e^{r(x,y_1)}+e^{r(x,y_2)}} } $$ This may look complicated, but it is simply a **binary softmax**: exponentiate the two scalar reward scores, then normalize them into probabilities that sum to one. #### Rewrite as a sigmoid Write \(r_1=r(x,y_1)\) and \(r_2=r(x,y_2)\). Dividing the numerator and denominator by \(e^{r_1}\) gives: $$ P(y_1\succ y_2) =\frac{1}{1+e^{r_2-r_1}}. $$ The sigmoid function is: $$ \sigma(z)=\frac{1}{1+e^{-z}}. $$ Therefore: $$ \boxed{P(y_1\succ y_2)=\sigma(r_1-r_2)} $$ **This is the most important form:** the model only needs the reward-score difference \(r_1-r_2\) to predict which response is preferred. ### 2. Loss formula **The formula worth remembering for interviews — one preference pair:** $$ \boxed{L=-\log\sigma(r_c-r_r)} $$ - \(r_c=r_\theta(x,y_c)\): reward for the **chosen** answer. - \(r_r=r_\theta(x,y_r)\): reward for the **rejected** answer, for the same prompt \(x\). - \(\sigma\): sigmoid, converting the score difference into a preference probability. $$ P(y_c>y_r\mid x)=\sigma(r_c-r_r),\qquad L_{\mathrm{batch}}=-\frac{1}{B}\sum_{i=1}^{B}\log\sigma(r_{c,i}-r_{r,i}). $$ A larger chosen-minus-rejected margin means higher preference probability and lower loss. For equal scores, the probability is \(0.5\) and the loss is about \(0.693\). Raw reward scores are **not probabilities**: only their difference determines the preference likelihood. Adding the same constant to both scores leaves the loss unchanged. [Source: RLHF Book](https://rlhfbook.com/c/05-reward-models). **Example from the PPO notes:** \(r_c=2.4\), \(r_r=0.9\), so the margin is \(1.5\), preference probability is approximately \(0.818\), and loss is approximately \(0.201\). During reward-model training, gradients update the head and any unfrozen backbone parameters; during PPO, the trained reward model is normally frozen to score policy responses. [PPO reward-model section](https://www.qiyuzhi.com/study/rlhf-rlvr/ppo/#2-reward-model). ### 3. Process map \(B\) = sequences per forward pass (or preference pairs during training), \(T\) = padded sequence length, \(H\) = hidden size. Chosen and rejected answers use the **same model**, in two forward passes. [[RL_REWARD_MODEL_SHAPES]] ### 4. Implementation The supplied implementation and all its comments are preserved below. **Assumptions:** sequences are right-padded and nonempty, so \(\text{attention\_mask.sum}(1)-1\) is the last valid index. This indexing is not valid for left-padded batches. Supply a backbone that returns `outputs.hidden_states`, an initialized `model` and `optimizer`, and the chosen/rejected input tensors; this is a teaching snippet, not a standalone training script. ```python import torch import torch.nn as nn import torch.nn.functional as F class BradleyTerryRewardModel(nn.Module): """ Reward Model flow: prompt + response ↓ Transformer ↓ hidden states [B, T, H] ↓ select ONE hidden vector per sequence (usually last real token / EOS) ↓ seq_repr [B, H] ↓ Linear(H -> 1) ↓ one scalar reward per sequence [B] B = batch size T = sequence length H = hidden size """ def __init__(self, base_model, hidden_size): super().__init__() # Example: Qwen / Llama transformer. # # It converts token sequences into hidden states. self.base_model = base_model # Reward head: # # [H] -> [1] # # Example: # [2048] -> [1] # # Mathematically: # reward = w^T h + b # # This is different from an LM head: # # LM head: # [H] -> [vocab_size] # # Reward head: # [H] -> [1] self.head = nn.Linear(hidden_size, 1) def _sequence_rep(self, hidden, attention_mask): """ Select ONE hidden-state vector to represent each sequence. hidden: [B, T, H] Example: B = 3 T = 7 H = 2048 hidden.shape = [3, 7, 2048] We want: one [2048] vector per sequence output: [B, H] = [3, 2048] """ # ------------------------------------------------------------ # 1. Find the last valid token position for every sequence # ------------------------------------------------------------ # Example attention_mask: # # [ # [1, 1, 1, 1, 1, 0, 0], # 5 real tokens # [1, 1, 1, 1, 1, 1, 1], # 7 real tokens # [1, 1, 1, 1, 1, 1, 0], # 6 real tokens # ] # # shape = [3, 7] # Sum each row: # # attention_mask.sum(dim=1) # -> [5, 7, 6] # # But indexing starts at 0: # # 5 tokens -> positions 0,1,2,3,4 # last token position = 4 # # Therefore subtract 1: lengths = attention_mask.sum(dim=1) - 1 # lengths: # tensor([4, 6, 5]) # ------------------------------------------------------------ # 2. Create batch row indices # ------------------------------------------------------------ # hidden.size(0) == B == batch size # # If B = 3: # # torch.arange(3) # -> tensor([0, 1, 2]) # # This simply identifies: # # batch row 0 # batch row 1 # batch row 2 batch_idx = torch.arange( hidden.size(0), device=hidden.device, ) # device=hidden.device ensures these indices are on # the same device as hidden, e.g. GPU. # ------------------------------------------------------------ # 3. Select one hidden vector from each sequence # ------------------------------------------------------------ # Suppose: # # batch_idx = [0, 1, 2] # lengths = [4, 6, 5] # # Then: # # hidden[batch_idx, lengths] # # is advanced indexing / SELECTION. # # PyTorch pairs the two index arrays: # # (0, 4) # (1, 6) # (2, 5) # # So this is conceptually equivalent to: # # [ # hidden[0, 4], # hidden[1, 6], # hidden[2, 5], # ] # # hidden[0, 4] means: # batch sequence 0, # token position 4, # keep the entire hidden dimension # # Therefore each selected item has shape: # # [H] = [2048] # # Selecting one for each of 3 sequences gives: # # [3, 2048] # # IMPORTANT: # this operation is only SELECTING hidden vectors. # It does NOT average or combine tokens. # # The last token hidden state already contains information # about previous tokens through causal self-attention. seq_repr = hidden[batch_idx, lengths] return seq_repr def forward(self, input_ids, attention_mask): """ input_ids: [B, T] attention_mask: [B, T] returns: rewards [B] One scalar reward for each entire sequence. """ # ------------------------------------------------------------ # 1. Run sequences through Transformer # ------------------------------------------------------------ outputs = self.base_model( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, ) # Take hidden states from final Transformer layer. # # Example: # # hidden.shape = [B, T, H] # = [3, 7, 2048] # # There is one H-dimensional hidden vector PER TOKEN. hidden = outputs.hidden_states[-1] # ------------------------------------------------------------ # 2. Convert [B, T, H] -> [B, H] # ------------------------------------------------------------ # Select last valid-token representation from each sequence. # # Before: # [3, 7, 2048] # # After: # [3, 2048] # # Conceptually: # # sequence 0 -> one [2048] vector # sequence 1 -> one [2048] vector # sequence 2 -> one [2048] vector seq_repr = self._sequence_rep( hidden, attention_mask, ) # ------------------------------------------------------------ # 3. Convert each [H] sequence representation to one scalar # ------------------------------------------------------------ # self.head = nn.Linear(H, 1) # # It independently processes every batch row: # # [B, H] # ↓ # [B, 1] # # Example: # # seq_repr.shape: # [3, 2048] # # output: # # [ # [0.64], # reward for sequence 0 # [0.52], # reward for sequence 1 # [0.82], # reward for sequence 2 # ] # # shape = [3, 1] rewards = self.head(seq_repr) # ------------------------------------------------------------ # 4. Remove unnecessary last dimension # ------------------------------------------------------------ # Before: # # shape = [B, 1] # # [ # [0.64], # [0.52], # [0.82], # ] # # squeeze(-1): # # [B, 1] -> [B] # # After: # # [0.64, 0.52, 0.82] # # One scalar reward per sequence. rewards = rewards.squeeze(-1) return rewards # ================================================================== # Bradley-Terry Reward Model Training # ================================================================== # Suppose batch size B = 3. # # Each training example is a preference pair: # # pair 0: # prompt + chosen_0 # prompt + rejected_0 # # pair 1: # prompt + chosen_1 # prompt + rejected_1 # # pair 2: # prompt + chosen_2 # prompt + rejected_2 optimizer.zero_grad() # ------------------------------------------------------------ # 1. Score chosen responses # ------------------------------------------------------------ rewards_chosen = model( input_ids=chosen_input_ids, attention_mask=chosen_attention_mask, ) # Shape: # [B] # # Example: # # rewards_chosen = # [1.2, -0.3, 2.1] # # Meaning: # # chosen sequence 0 -> reward 1.2 # chosen sequence 1 -> reward -0.3 # chosen sequence 2 -> reward 2.1 # ------------------------------------------------------------ # 2. Score rejected responses using SAME reward model # ------------------------------------------------------------ rewards_rejected = model( input_ids=rejected_input_ids, attention_mask=rejected_attention_mask, ) # Shape: # [B] # # Example: # # rewards_rejected = # [0.2, 0.7, 0.1] # ------------------------------------------------------------ # 3. Compare each chosen/rejected pair # ------------------------------------------------------------ # Element-wise subtraction: # # pair 0: # 1.2 - 0.2 = 1.0 # # pair 1: # -0.3 - 0.7 = -1.0 # # pair 2: # 2.1 - 0.1 = 2.0 # # Therefore: # # reward_diff = # [1.0, -1.0, 2.0] # # Positive: # RM thinks chosen > rejected # # Negative: # RM currently thinks rejected > chosen # -> this pair should receive a larger loss reward_diff = rewards_chosen - rewards_rejected # ------------------------------------------------------------ # 4. Bradley-Terry probability # ------------------------------------------------------------ # Bradley-Terry assumes: # # P(chosen > rejected) # = sigmoid(r_chosen - r_rejected) # # Example: # # diff = 1 # sigmoid(1) ≈ 0.73 # # -> RM assigns 73% probability that chosen is preferred. # # # diff = -1 # sigmoid(-1) ≈ 0.27 # # -> RM only assigns 27% probability that chosen is preferred. # # # diff = 2 # sigmoid(2) ≈ 0.88 # ------------------------------------------------------------ # 5. Bradley-Terry loss # ------------------------------------------------------------ # We know from the dataset that: # # chosen SHOULD beat rejected. # # Therefore target probability is effectively 1. # # Loss: # # L = -log P(chosen > rejected) # # = -log sigmoid(r_chosen - r_rejected) # # # Intuition: # # r_chosen >> r_rejected # ↓ # difference large positive # ↓ # sigmoid ≈ 1 # ↓ # loss ≈ 0 # # # r_chosen < r_rejected # ↓ # difference negative # ↓ # sigmoid close to 0 # ↓ # large loss # # # Goal of training: # # r_chosen - r_rejected # # should become increasingly positive. # # logsigmoid(x) is mathematically: # # log(sigmoid(x)) # # but is more numerically stable than manually doing: # # torch.log(torch.sigmoid(x)) # # .mean() averages the loss across preference pairs in batch. loss = -F.logsigmoid(reward_diff).mean() # ------------------------------------------------------------ # 6. Backprop + parameter update # ------------------------------------------------------------ loss.backward() # Gradients flow: # # Bradley-Terry loss # ↓ # reward scores # ↓ # Linear(H -> 1) reward head # ↓ # selected final-token hidden states # ↓ # Transformer # # So unless the base model is explicitly frozen, # both the Transformer and reward head are trained. optimizer.step() ``` ## Outcome Reward Model (ORM) ### 1. Differences - **Training target:** from “which answer is preferred?” → “is this answer correct?” Each solution has its own binary outcome: `1` for correct, `0` for incorrect. - **Where the head is applied:** from one selected EOS / last-valid-token hidden state → every completion-token hidden state. In the implementation, the same scalar head actually runs at **all token positions**; the loss mask keeps only completion positions. This is the **per-token ORM variant** described in the [RLHF Book](https://rlhfbook.com/c/05-reward-models#outcome-reward-models), not a claim that every ORM uses this architecture. Each completion prefix predicts the final outcome; it is not a label for whether that individual token or reasoning step is correct. ### 2. Process map No chosen/rejected pairing is required. The dataset can contain independently labeled solutions such as `[(prompt + solution_A, 1), (prompt + solution_B, 0)]`. If a solution is correct, repeat its outcome `1` at **every completion-token label**. If it is incorrect, repeat `0`. Prompt and padding positions get `-100`; `attention_mask` alone does not exclude them from the BCE loss. [[RL_ORM_SHAPES]] ### 3. Implementation The class below keeps your model logic and comments, with the pasted constructor/indentation repaired and the numerical loss walkthrough added as comments. Supply token-aligned `labels [B, T]` and at least one unmasked completion token per training batch. ```python import torch.nn as nn import torch.nn.functional as F class OutcomeRewardModel(nn.Module): def __init__(self, base_lm): super().__init__() self.lm = base_lm # Same kind of scalar head as before: # # H -> 1 self.head = nn.Linear( self.lm.config.hidden_size, 1 ) def forward( self, input_ids, attention_mask=None, labels=None, ): outputs = self.lm( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, return_dict=True, ) # [B, T, H] hidden = outputs.hidden_states[-1] # IMPORTANT DIFFERENCE FROM BRADLEY-TERRY: # # Bradley-Terry: # first select EOS: # [B,T,H] -> [B,H] # # ORM: # DO NOT select EOS. # # Apply H->1 head to EVERY token. # # [B,T,H] # ↓ Linear(H->1) # [B,T,1] # ↓ squeeze # [B,T] logits = self.head(hidden).squeeze(-1) if labels is None: return None, logits # Example labels: # # [-100,-100,-100, 1,1,1,1,1] # # Ignore prompt tokens. # Padding labels must also be -100. mask = labels != -100 # Binary classification: # # each completion token predicts: # # probability final outcome is correct # # label = 1: final answer correct # label = 0: final answer incorrect # ------------------------------------------------------------ # Loss example: a separate four-token completion, outcome = 1 # ------------------------------------------------------------ # # logits = [-0.5, 0.0, 0.8, 1.5] # # Convert each token logit w_t to a probability: # # p_t = sigmoid(w_t) # # token logit p(correct eventually) -log(p_t) # 3 -0.5 0.38 0.974 # + 0.0 0.50 0.693 # 2 0.8 0.69 0.371 # =5 1.5 0.82 0.201 # # The actual outcome is 1, so every token's target is 1: # # labels = [1, 1, 1, 1] # # BCE wants ALL FOUR probabilities to move toward 1. # For an incorrect solution, all four targets would be 0. # # Binary cross-entropy for one token: # # L_t = -[y * log(p_t) + (1-y) * log(1-p_t)] # # If y = 0: L_t = -log(1-p_t) # If y = 1: L_t = -log(p_t) # # Average this correct solution's completion-token losses: # # L = (1 / T_completion) * sum_t -log(p_t) # = (0.974 + 0.693 + 0.371 + 0.201) / 4 # ≈ 0.560 # # T_completion = 4 here, not the padded input length T. # Use logits directly below: BCEWithLogits combines sigmoid # and BCE numerically stably. Do NOT pass sigmoid(logits). # # Boolean masking selects all valid completion tokens in # the batch. The default reduction averages across those # tokens, not prompt/padding positions or sequence pairs. loss = F.binary_cross_entropy_with_logits( logits[mask], labels[mask].float(), ) return loss, logits ``` ## Process Reward Model (PRM) ### 1. Definition - **ORM asks:** “Will the whole solution eventually be correct?” - **PRM asks:** “Was this particular reasoning step correct?” The distinction is the **training target**, not just the number of scores. In the per-token ORM above, every completion token inherits the final outcome. In the binary PRM below, only annotated step-ending positions receive a correctness label for that individual step. This is the binary version; some PRMs use three classes instead. [Source: RLHF Book — Process Reward Models](https://rlhfbook.com/c/05-reward-models#process-reward-models). ### 2. Data and formula **Question:** John has 12 apples. He gives away 4, then buys 3. How many does he have? The correct answer is **11**, but suppose the model produces this solution: | Reasoning step | ORM: final-outcome label | PRM: individual-step label | | --- | --- | --- | | Step 1: `12 - 4 = 8` | `0` — whole solution is wrong | `1` — correct | | Step 2: `8 + 3 = 12` | `0` — whole solution is wrong | `0` — incorrect | | Step 3: `Therefore, the final answer = 12` | `0` — whole solution is wrong | `0` — incorrect | **ORM:** the whole-solution outcome is `WRONG`, so all completion tokens are labeled `0`—even the tokens in the perfectly correct Step 1. That label means “the final answer is wrong,” not “this step is wrong.” **PRM:** this example has step labels `[1, 0, 0]`. It can reward the valid first step and flag the arithmetic error in the second. Here, the final-answer step is also annotated incorrect because 12 is not the answer. Each step is scored using the question and the solution prefix **through that step**, without looking at future steps. For example, illustrative predictions could be: ```text Question + Step 1 ↓ P(Step 1 correct) = 0.96 Question + Step 1 + Step 2 ↓ P(Step 2 correct) = 0.08 ``` With a causal transformer, one forward pass over the full sequence can produce both scores: each step-ending hidden state can only attend to the prefix up to its own position. **Binary PRM loss for one solution with K annotated steps:** $$ \boxed{L_{\mathrm{PRM}}=-\sum_{i=1}^{K}\left[y_i\log p_i+(1-y_i)\log(1-p_i)\right]} $$ - \(y_i\in\{0,1\}\): whether reasoning step \(i\) is correct. - \(p_i=\sigma(z_i)\): predicted probability that step \(i\) is correct, using its boundary logit \(z_i\), conditioned on the question and the prefix through that step. - \(K\): number of annotated steps, **not** the token count \(T\). For a correct step (\(y_i=1\)), the loss is \(-\log p_i\). For an incorrect step (\(y_i=0\)), it is \(-\log(1-p_i)\). This is BCE summed over annotated steps; using a mean instead divides by the number of included steps. The book additionally averages the summed loss over sampled solutions. [Source: RLHF Book, binary PRM loss](https://rlhfbook.com/c/05-reward-models#process-reward-models). ### 3. Process map **Only annotated step-ending positions have labels `0` or `1`. Everything else is `-100`:** prompt tokens, intermediate tokens within a step, padding, and any unannotated positions. The scalar head can still run at **every token**. The key change from ORM is which labels are present and which positions enter the loss. [[RL_PRM_SHAPES]] **An illustrative slice of row 0** (ellipses omit other positions; these logits are a separate toy example from the probabilities above): ```text logits[0] = [-0.2, 0.3, 0.1, 2.5, 0.7, 0.4, ... -2.1, ...] labels[0] = [-100, -100, -100, 1, -100, -100, ... 0, ...] step 1 boundary → label 1 step 2 boundary → label 0 step 3 boundary → label 0 (not shown in this slice) ``` `squeeze(-1)` removes only the size-1 dimension: `[B, T, 1] → [B, T]`. **It does not throw away scores.** Boolean masking selects the annotated boundary scores afterward: ```python hidden = outputs.hidden_states[-1] # [B, T, H], e.g. [2, 30, 2048] logits = self.head(hidden) # [B, T, 1] logits = logits.squeeze(-1) # [B, T]: all token scores remain mask = labels != -100 # True only at annotated step ends # Assume the batch contains at least one annotated step. loss = F.binary_cross_entropy_with_logits( logits[mask], labels[mask].float(), # [N] each: N labeled steps in batch reduction="sum", # Sum step losses; "mean" averages them ) ``` Pass **raw logits**, not sigmoid probabilities, to `binary_cross_entropy_with_logits`. `-100` is a sentinel that the explicit mask removes; BCE does not automatically ignore it. Intermediate positions have no direct classification-loss term, but their representations can still influence later boundary predictions through attention. ## Comparing Reward Model Types and Value Functions Summary adapted from [RLHF Book — Comparing Reward Model Types](https://rlhfbook.com/c/05-reward-models#comparing-reward-model-types-and-value-functions). | Model Class | What They Predict | How They Are Trained | LM Structure | | --- | --- | --- | --- | | **Reward Models** | Sequence-level quality score | Contrastive loss between pairwise (or N-wise) comparisons between completions to the same prompt | Linear head on EOS/last-token hidden state | | **Outcome Reward Models** | Per-token probability that the final answer is correct | Labeled outcomes (e.g., success/failure on verifiable domains); each sample is labeled independently, with no need for paired comparisons on the same prompt | Per-token scalar head trained with binary cross-entropy; labels repeat the outcome label | | **Process Reward Models** | A reward or correctness score at the end of each reasoning step | Intermediate feedback or stepwise annotations; loss is applied at annotated step-ending positions, with other positions masked | Per-token head predicting step correctness: binary, or three-class (−1 = incorrect, 0 = neutral, 1 = correct) | | **Value Functions** | Expected remaining return given the current state under a policy | Regression to return targets at sequence positions, typically computed from rollouts | A scalar regression head with per-token outputs | **PRM label convention:** the binary example above uses `0 = incorrect, 1 = correct`. In the three-class variant, `0 = neutral` instead; these are different label schemes.