Reward Models · L3 Code

Transformer backbone

usually initialized from SFT checkpoint

Pooling

take hidden state at last non-pad token

Linear head -> scalar

the reward

Anatomy of a reward model.

python

import torch
import torch.nn as nn

class RewardModel(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone                       # a pretrained causal LM (no LM head)
        hidden = backbone.config.hidden_size
        self.value_head = nn.Linear(hidden, 1, bias=False)

    def forward(self, input_ids, attention_mask):
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        hidden_states = out.last_hidden_state              # [B, T, H]
        # take the hidden state at the last real (non-padding) token per sequence
        last_idx = attention_mask.sum(dim=1) - 1
        pooled = hidden_states[torch.arange(hidden_states.size(0)), last_idx]
        return self.value_head(pooled).squeeze(-1)          # [B]

A reward model: backbone + scalar head.