Tilted distributions

Reweight a distribution by exp(reward / β) and renormalise. That one operation is the exact solution of KL-regularised RLHF, the starting point of DPO, a Bayesian posterior, and the yardstick for best-of-n.

Key idea: Exponential tilting, the Boltzmann distribution, the closed-form RLHF optimum, RL as Bayesian inference, and best-of-n
Real world: In a gas, molecules occupy high-energy states less often, in proportion to exp(−energy / temperature). Turn the temperature down and they crowd into the lowest-energy states; turn it up and they spread out.
In an LLM: The best policy for "maximise reward, but stay close to the reference model" has a formula: take the reference model's probability of each answer and multiply by exp(reward / β). Low β: chase reward. High β: stay put. DPO, best-of-n and RL-as-inference all start from this formula.

Tilting a distribution

Take a distribution $p(y)$ and a score $r(y)$ . The tilted distribution multiplies each probability by $e^{r(y)/\beta}$ , then renormalises:

p_\beta(y) = \frac{p(y)\, e^{r(y)/\beta}}{Z}, \qquad Z = \sum_y p(y)\, e^{r(y)/\beta}

High-scoring outcomes gain probability, low-scoring ones lose it.
$\beta$ controls how hard. As $\beta \to \infty$ , nothing changes. As $\beta \to 0$ , all the probability goes to the highest score.
Outcomes with $p(y) = 0$ stay at zero. Tilting can only reweight what $p$ already does.

$Z$ , the partition function, is just the normaliser, but it's usually impossible to compute for text, because it's a sum over every possible answer. Much of the cleverness in this area is avoiding it.

example · molecules in a gas

Physics calls this the Boltzmann distribution: a state with energy $E$ is occupied with probability proportional to $e^{-E/kT}$ . It's the softmax with negative energy as the score, and temperature playing the role of $\beta$ . That's where the word "temperature" in sampling comes from.

example · a model's answers, tilted

The reference model writes a vague answer 35% of the time and the best answer 3% of the time, with rewards 0.1 and 1.0. With $\beta = 0.2$ , the vague answer is multiplied by $e^{0.5} \approx 1.6$ and the best by $e^{5} \approx 148$ . After renormalising, the best answer goes from 3% to about 31%, and the vague one falls to about 4%.

The RLHF optimum

RLHF maximises reward minus a KL penalty to the reference model:

\max_\pi\ \mathbb{E}_{y \sim \pi}[r(y)] - \beta\, \mathrm{KL}(\pi \,\|\, \pi_\text{ref})

The exact answer is the tilted reference: $\pi^*(y) \propto \pi_\text{ref}(y)\, e^{r(y)/\beta}$ . Nothing else achieves a higher value of that objective. Three consequences:

RL fine-tuning is sharpening the reference model towards high reward. It can't create answers the reference gives zero probability (Does RL teach new skills?).
DPO rearranges the formula to express the reward through the policy, $r(y) = \beta \log \tfrac{\pi^*(y)}{\pi_\text{ref}(y)} + \beta \log Z$ . In a comparison of two answers to the same prompt, $Z$ cancels, and the reward model disappears (DPO).
RL as Bayesian inference: read $\pi_\text{ref}$ as a prior and $e^{r/\beta}$ as a likelihood of "this answer is good". Then $\pi^*$ is the posterior. The KL penalty isn't an ad hoc regulariser; it's what makes the result a proper posterior instead of a collapse onto one answer.

Best-of-n

The simplest way to tilt at inference time: sample $n$ answers, score them, and keep the best. The result is a new distribution over answers, which you can compare to the tilted one. Two facts:

Its KL from the reference is at most $\log n - \frac{n-1}{n}$ , small for practical $n$ : best-of-16 is under 2 nats.
On a reward-versus-KL plot, best-of-n lands very close to the tilted optimum, which is why it's a strong baseline that RLHF methods have to beat.

Try it

β0.20best-of-n4

r = 0.1r = 0.3r = 0.5r = 0.6r = 0.8r = 1.0

reference π_reftilted π* ∝ π_ref · exp(r/β)best of 4

tilted: reward, KL0.717, 0.91

best of 4: reward, KL0.626, 0.61

log n − (n − 1)/n0.64

The purple curve is the best any policy can do: for each KL budget, tilting gets the highest expected reward, because it is the exact solution of “maximise reward minus β × KL”. Best-of-n (green, dashed) lands just below it, and its KL never exceeds log n − (n − 1)/n, which is why a best-of-16 sample counts as a fairly small move away from the base model: at most 1.84 nats.

The reference model’s answers, each with a reward. Tilting multiplies each probability by exp(r/β) and renormalises; best-of-n samples n answers and keeps the highest-scoring one.

example · hiring from a shortlist

Interview 5 candidates and hire the best. You haven't changed who applies; you've reshaped the distribution of who gets hired towards the top of the applicant pool. Interview 50 and you'd get even better hires, but each extra interview buys less.

example · reward overoptimisation

Gao et al. (2022) measured progress against KL from the reference. With a learned reward model, true quality first rises, then falls as KL grows, for both best-of-n and RL. The proxy reward keeps rising. KL is the natural x-axis because it's the budget in the tilting objective, and it's how far you've moved from the distribution the reward model was trained on.

Where this shows up

The KL-regularised objective and its optimum: KL and reward hacking.
DPO and variants derive their losses from $\pi^*$ : DPO.
Best-of-n as a baseline and as training data: distillation.
Power distributions $p^\alpha$ are tilting by the model's own log-probability: MCMC.

going deeper

Deriving the optimum

\mathbb{E}_\pi[r] - \beta\, \mathrm{KL}(\pi \| \pi_\text{ref}) = -\beta \sum_y \pi(y) \log \frac{\pi(y)}{\pi_\text{ref}(y)\, e^{r(y)/\beta}} = -\beta\, \mathrm{KL}(\pi \,\|\, \pi^*) + \beta \log Z

with $\pi^*(y) = \pi_\text{ref}(y) e^{r(y)/\beta} / Z$ . The KL is minimised, at zero, by $\pi = \pi^*$ . The maximum value is $\beta \log Z$ .

DPO from the optimum

Solving for $r$ and plugging into the Bradley–Terry model, $Z(x)$ cancels between the chosen answer $y_w$ and the rejected answer $y_l$ :

\mathcal L_\text{DPO} = -\log \sigma\Big(\beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_\text{ref}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_\text{ref}(y_l \mid x)}\Big)

RL as Bayesian inference

Korbak et al. write the target as a posterior with a binary "optimality" variable $\mathcal O$ whose likelihood is $p(\mathcal O = 1 \mid y) \propto e^{r(y)/\beta}$ :

\pi^*(y) = p(y \mid \mathcal O = 1) = \frac{\pi_\text{ref}(y)\, p(\mathcal O = 1 \mid y)}{p(\mathcal O = 1)}

KL-regularised RL is then variational inference: minimising $\mathrm{KL}(\pi_\theta \| \pi^*)$ , a reverse KL, which explains its mode-seeking behaviour.

Best-of-n

If rewards are distinct and $F$ is the reference's CDF of reward, the best of $n$ samples has $P(\text{best} \le y) = F(y)^n$ . Beirami et al. show

\mathrm{KL}(\pi_\text{BoN} \,\|\, \pi_\text{ref}) \le \log n - \frac{n - 1}{n}

with equality for continuous rewards. The widget computes the exact discrete version.

Papers

RL with KL penalties is better viewed as Bayesian inference (Korbak et al., 2022)
Direct Preference Optimization (Rafailov et al., 2023)
Theoretical guarantees on the best-of-n alignment policy (Beirami et al., 2024)
BoNBoN Alignment for Large Language Models and the Sweetness of Best-of-n Sampling (Gui et al., 2024): best-of-n is nearly optimal on the reward–KL frontier, and how to train towards it
Scaling Laws for Reward Model Overoptimization (Gao et al., 2022)

Tilting a distribution

The RLHF optimum

Best-of-n

Try it

Where this shows up

Deriving the optimum

DPO from the optimum

RL as Bayesian inference

Best-of-n

Papers

Related