Darkweb007/reward-model-from-scratch

🧠 Reward Model from Scratch

Bradley-Terry preference loss implemented from scratch. DistilBERT + scalar head trained on Anthropic HH-RLHF. Includes reward hacking simulation showing Goodhart's Law in action.

What This Is

Most RLHF tutorials use pre-built RM libraries. This project implements the full reward model training pipeline from first principles:

• —Bradley-Terry loss: derived mathematically, implemented in pure PyTorch
• —Architecture: DistilBERT backbone + scalar head (66M params)
• —Dataset: Anthropic HH-RLHF (human-helpfulness preference pairs)
• —Reward hacking: empirically demonstrated with KL divergence tracking

The Math

The Bradley-Terry model assigns each response a latent quality score r. The probability that response A is preferred over B:

P(A ≻ B) = σ(r_A - r_B)

Loss = negative log-likelihood of observed preferences:

L = -log σ(r_chosen - r_rejected)
  = softplus(-(r_chosen - r_rejected))

Why not MSE? Humans assign relative preferences ("A is better than B"), not absolute quality scores ("A is 7.3/10"). Bradley-Terry models exactly ordinal comparisons, not cardinal values.

Training Results

Reward Hacking — Goodhart's Law

"When a measure becomes a target, it ceases to be a good measure."

When PPO optimizes against the RM, the policy eventually learns to exploit proxy reward signals. After ~400 steps, RM score continues rising while true quality degrades.

Mitigation: KL penalty in the RLHF objective:

reward = RM(prompt, response) - β * KL(π_θ || π_ref)

Architecture

reward_model/
├── model.py   # RewardModel: backbone + scalar head
├── loss.py    # BradleyTerryLoss, RewardHackingDetector, KL penalty
└── train.py   # Full training loop, PreferenceDataset, training curves

Running Locally

git clone https://github.com/data-geek-astronomy/reward-model-from-scratch
cd reward-model-from-scratch
pip install -r requirements.txt

# Train from scratch on HH-RLHF (requires ~45min on GPU)
python -c "
from reward_model.train import RewardModelTrainer, TrainingConfig, load_hh_rlhf_data
train_data, eval_data = load_hh_rlhf_data(max_train=10000, max_eval=1000)
trainer = RewardModelTrainer(TrainingConfig())
trainer.train(train_data, eval_data)
"

# Run the Gradio demo
python app.py

References

• —Training language models to follow instructions with human feedback (InstructGPT / RLHF paper)
• —Scaling Laws for Reward Model Overoptimization (reward hacking paper)
• —Anthropic HH-RLHF Dataset

License

MIT