Reward Modelling

Reinforcement learning needs a reward. For “write a helpful reply” there is no reward function, and asking annotators to produce one directly does not work — absolute scores drift between people, within a person, and across a session.

Comparisons are stable in a way scores are not. So collect comparisons and fit a scalar that explains them.

The model is the base transformer with the unembedding replaced by a scalar head reading the final token. Only the difference appears, so the absolute level is unidentifiable — a constant added to every reward changes nothing, which is worth remembering before quoting a reward number anywhere.

Overoptimisation

The reward model is a proxy, fitted to finitely many comparisons and wrong somewhere. Policy optimisation is a search for high reward, which is to say a search for exactly those places.

Gao et al. measured the shape. Plot true reward — from a much larger held-out “gold” model — against distance from the initial policy, and it rises, peaks, and falls.

Two things follow. The peak is at a finite KL, so more optimisation past it makes the model genuinely worse while the reported reward keeps climbing — the single most misleading number in an RLHF run. And $\alpha, \beta$ improve with reward model size, so a larger RM does not merely score better, it tolerates more optimisation before turning.

This is why KL regularisation is not an optional stabiliser. It is the term that keeps the policy inside the region where the proxy is still correlated with the thing it proxies for.

What goes wrong in practice

Length. Annotators prefer longer responses, so the RM learns that length is good, so the policy produces essays. It is the most reliably observed reward hack in the literature and it is a genuine property of the preference data, not a modelling error.

Sycophancy. Agreement is preferred to correction, and the policy learns to agree.

Format. Bullet points, headers and confident phrasing score well independently of content.

All three share a structure: a feature that correlates with quality in the training distribution and is cheap to produce in isolation. Any proxy has such features; the question is only how quickly optimisation finds them.

import torch
from torch import Tensor, nn
from torch.nn import functional as F


class RewardModel(nn.Module):
    def __init__(self, backbone: nn.Module, dim: int):
        super().__init__()
        self.backbone = backbone
        self.head = nn.Linear(dim, 1, bias=False)

    def forward(self, ids: Tensor, mask: Tensor) -> Tensor:   # [B, T]
        h = self.backbone(ids, attention_mask=mask)           # [B, T, D]
        # The last non-pad position: the reward is a property of the whole
        # response, and only that position has attended to all of it.
        last = mask.sum(-1) - 1
        return self.head(h[torch.arange(h.size(0)), last]).squeeze(-1)


def preference_loss(chosen: Tensor, rejected: Tensor, margin: float = 0.0) -> Tensor:
    return -F.logsigmoid(chosen - rejected - margin).mean()

Both responses in a pair must go through the model in the same batch and the same step. Splitting them across steps means the two rewards are computed under different parameters, and the difference the loss depends on is then partly a difference in the model rather than in the responses.

DPO shows that for the KL-constrained objective the optimal policy has a closed form in terms of the reward, which can be inverted — so the reward model can be substituted out and the policy trained on the preference pairs directly. No RM, no sampling loop, no PPO.

What that gives up is online data. A reward model can score responses the policy generates now; a fixed preference dataset cannot. Which is why RMs persist wherever the loop is online — GRPO and the reasoning-model pipelines built on it — and disappear wherever it is not.