LLM Engineering

A reward model that overrates the third answer (4 points against 1 and 2). A modest policy shift has a KL of 0.085 from the reference; collapsing onto the overrated answer has 1.48. With a weak penalty, β = 0.1, the collapse scores higher (3.80 against 2.19) — reward hacking. With β = 2 the modest policy wins (2.03 against 0.98). The KL term is what keeps the model honest to its starting point.

• The reward model is a learned proxy; optimising hard against it eventually exploits its errors (Gao et al. 2023).
• β is a dial between following the reward and staying close to the reference; there is no principled value.
• Human raters reward confident, longer answers; RLHF can amplify that.

ref = np.array([0.5, 0.3, 0.2])                 # reference policy over three answers
reward = np.array([1.0, 2.0, 4.0])              # the reward model's scores (the third one it overrates)
kl = lambda p: np.sum(p * np.log(p / ref))
objective = lambda p, beta: p @ reward - beta * kl(p)
modest, hacked = np.array([0.3, 0.45, 0.25]), np.array([0.01, 0.01, 0.98])
result = {beta: (round(objective(modest, beta), 3), round(objective(hacked, beta), 3)) for beta in (0.1, 2.0)}

Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.

import torch.nn.functional as F

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.model = base_model
        self.head = nn.Linear(d_model, 1, bias=False)

    def forward(self, input_ids):
        hidden = self.model(input_ids).last_hidden_state
        reward = self.head(hidden[:, -1, :])  # score from last token
        return reward.squeeze(-1)

# Bradley-Terry loss
def reward_loss(chosen_reward, rejected_reward):
    return -F.logsigmoid(chosen_reward - rejected_reward).mean()Copy