reward-model-from-scratch
A small reward model trained from scratch on pairwise preference comparisons: a
scalar reward head on top of distilbert-base-uncased, trained with the
Bradley-Terry loss (-log_sigmoid(r_chosen - r_rejected)), then evaluated by
held-out pairwise accuracy and a calibration analysis (reliability diagram + ECE).
Code: https://github.com/narinzar/reward-model-from-scratch
Task
Reward modeling / text classification. Given two responses to the same prompt, the model assigns each a scalar reward; the preferred response should score higher.
Architecture
- Base encoder:
distilbert-base-uncased - Masked-mean pooling over real tokens, then
Linear(hidden -> 1)producing one scalar reward per sequence - Trained with the Bradley-Terry loss on reward margins
Results (real, small-scale)
Measured on a single RTX 5090 (sm_120, CUDA 12.8 wheels). 1447 synthetic preference pairs (1158 train / 289 held-out eval), 3 epochs, batch size 16, lr 2e-5.
| Metric | Value |
|---|---|
| Held-out pairwise accuracy | 1.0000 |
| Expected Calibration Error (ECE, 10 bins) | 0.0078 |
The perfect held-out accuracy reflects that the synthetic label rule (longer-and-more-polite) is fully learnable at this scale. On a real, noisier preference corpus expect accuracy below 1.0 and a larger calibration gap.
Training data
Synthetic. Preference pairs are generated and labelled by a transparent rule (longer-and-more-polite is preferred), clearly marked as synthetic in the source. This is a from-scratch demonstration, not a production reward model. Swap in a real preference corpus by replacing one data loader.
Files
reward_model.pt— checkpoint dict withstate_dict,base_model_name, andmax_length. Load with theRewardModelclass from the GitHub repo.
License
MIT