reward-model-from-scratch

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.

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 with state_dict, base_model_name, and max_length. Load with the RewardModel class from the GitHub repo.

License

MIT