Best-of-N Sampling with a Reward Model

easy
Checkpoint challenge
Solve this to show you know the material and jump straight to Unit 9. Everything before it counts as cleared.

The simplest alignment technique at inference time: sample N responses, score each with a reward model, and return the best. It's surprisingly strong, and its cost and effect are easy to analyse.

Implement:

  • best_of_n(candidates, reward_fn): the candidate with the highest reward (ties: the earliest)
  • kl_best_of_n(n): the KL divergence between the best-of-n distribution and the base policy, which has the closed form log⁡n−n−1n\log n - \frac{n-1}{n}
  • win_rate(n): the probability that the best of n samples beats one fresh sample, for continuous i.i.d. rewards. By symmetry, it's nn+1\frac{n}{n+1}.
  • simulate_win_rate(n, trials, rng): estimate the same thing by simulation, drawing rewards with rng.random(). For each trial, draw the n candidate rewards first, then one baseline reward.
#alignment#inference#sampling2 visible + 1 hidden tests
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
Best-of-N Sampling with a Reward Model · TechGradify