A partition function normalizes nonnegative, reward-weighted reference probabilities into a valid distribution. For each fixed x, Z(x) sums the weights across all possible y. section 8.1.2
In KL-regularized reinforcement learning, Z(x) normalizes π_ref(y|x) exp(r(x,y)/β). The resulting reward-tilted distribution is the optimal policy, while Z(x) remains independent of π. section 8.1.2 section 15.2.2
The derivation uses the partition function to:
- Turn an unnormalized reward-tilted density into a probability distribution over responses. section 8.1.2
- Rewrite the RL objective as a KL divergence minus the constant log Z(x). section 15.2.2
- Show that minimizing the KL divergence produces the reward-tilted optimal policy. section 8.1.2
- Cancel from pairwise Bradley–Terry preference probabilities because the same Z(x) multiplies every candidate. section 8.1.2
Reward-tilted policy
Given Z(x) = Σ_y π_ref(y|x) exp(r(x,y)/β), the optimal policy is π*(y|x) = π_ref(y|x) exp(r(x,y)/β) / Z(x). section 8.1.2
Equation 8.5
\[\label{eq:rlhf_opt_eq_repeat} \max_{\pi} \mathbb{E}_{x \sim \mathcal{D}}\mathbb{E}_{y \sim \pi(y|x)} \left[r_\theta(x, y)\right] - \beta \mathcal{D}_{\text{KL}}\left(\pi(y|x) \| \pi_{\text{ref}}(y|x)\right).\]Trace the KL-regularized RL objective before normalization is introduced. section 8.1.2
