Best-of-N Sampling with a Reward Model

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 - \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 $\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.