Ster/Docs

Reward model

A model that judges text instead of writing it: one scalar per sequence, fitted to a preference set under the Bradley-Terry objective, and carried in an artifact of a different kind that generation refuses.

Overview#

A reward model scores a sequence. It is the piece policy optimization needs and cannot supply for itself: a reward that is a language model's own likelihood is circular, since raising it is what the policy was already doing.

Ster's is a linear head over the last position's residual state after the final norm — one row of hidden_size weights producing one number — trained together with the LoRA adapters beneath it. The last position is the only one that has attended to the whole sequence, so it is the only one that can score it. `ster tune reward` trains it.

The Bradley-Terry objective#

Bradley-Terry models a comparison: the probability that the chosen response beats the rejected one is the sigmoid of the difference of their scores.

text
loss = -log sigmoid(score_chosen - score_rejected)

Two properties fall straight out of that difference rather than being chosen. The head has no bias, because a bias is added to both scores and cancels — it would be a parameter with an identically zero gradient, dead weight in every sense. And only differences are identified: the absolute scale and offset of the scores are not determined by the objective, so the numbers a run reports are meaningful against each other and against no external unit.

The head is initialised to zeros rather than drawn. A single output row has no symmetry for a random draw to break, and the gradient at zero is the difference of the two residual states, so it learns on the first step. Zero also pins the starting point at “no opinion”: every sequence scores exactly zero, so the first loss is exactly ln 2, the same identity check the preference losses give.

That starting point makes one reported number look alarming when it is not. A run short enough that the head never moved reports accuracy 0.0 and mean_score_margin 0.0000, which reads as a head that ranked every pair backwards and means a head that has not moved: it starts at zeros, so every pair scores an exact tie, and a tie is not a strict win. tied_pairs is what tells the two apart from the document alone. A head that never moved ties every pair it trained on, so tied_pairs equals trained_pairs; a head that learned the order backwards ties none of them and reports a negative mean_score_margin beside the same accuracy. It counts exact equality rather than closeness, because the tie it names is two sides of one pair run through identical weights, not two scores a trained head happens to find similar. The accuracy itself was left alone rather than redefined to count ties as half-wins: it is a number people already consume, and the fix belonged beside it. Read the tie count first and the accuracy second, and give the run more than one epoch before either number means anything.

Why the artifact is a different kind#

A reward model is an adapter and a head, and the two are useless apart: the head reads a residual stream the adapters shaped, and the adapters were shaped to make that head separate. Pairing one with adapters it never saw would produce scores that mean nothing. So they travel in one safetensors file — the adapter factors plus one extra tensor named reward.head, shaped [1, hidden_size] — and the artifact does not offer the alternative.

The sidecar carries a kind that says which of the two things a reader is holding:

kindContainsAccepted by
adapterAdapter factors only.generate --adapter, tune merge, tune evaluate --adapter
rewardAdapter factors plus reward.head.tune grpo --reward

The field is #[serde(default)] and defaults to adapter, so every sidecar written before it existed still loads and still means what it meant. It is additive, which is why the artifact schema version does not move.

Applying an artifact where it does not belong is a refusal rather than a wrong answer, in both directions. generate, tune merge and tune evaluate refuse a reward model with adapter artifact is a reward model, not a generation adapter, because attaching its adapters and dropping its head would produce a decode nobody trained. tune grpo --reward refuses a generation adapter with adapter artifact is a generation adapter, not a reward model, because it has no head to score with. Validation also refuses both malformed halves before a file reaches disk: reward artifact is missing tensor reward.head, which is the head it exists to carry, and adapter artifact carries a reward.head tensor but declares kind adapter.

What it costs#

One extra tensor and one extra optimizer variable. On the four-layer toy checkpoint a supervised run trains 16 tensors and 7168 parameters; a reward run trains 17 and 7232, the difference being the 64-wide head. The head is in the same VarMap as the adapters, so one optimizer steps both, which is why the progress line counts tensors rather than adapter tensors.

A reward pass computes no vocabulary projection at all. The head reads the residual stream and never looks at a token distribution, and on a real checkpoint that projection is the widest matmul in the pass — computing and backpropagating it only to drop the result would be the single most expensive thing the run did.