Towards Efficient Exact Optimization of Language Model Alignment

TL;DR

This work reframes language-model alignment with human preferences as KL-regularized reward maximization and shows that the optimal policy under this objective is an energy-based distribution pi_beta^*. It introduces Efficient Exact Optimization (EXO), a probability-matching approach that minimizes the reverse KL between a data-driven proxy and the optimal policy, circumventing RL's high-variance training. The authors prove that EXO aligns in the same direction as traditional RL methods asymptotically and demonstrate, through extensive experiments on summarization, dialogue, and instruction-following tasks, that EXO outperforms Direct Preference Optimization (DPO) and PPO in both efficiency and alignment quality. They also reveal that DPO effectively optimizes a forward KL, which can miss critical modes of the target distribution under realistic model capacities. The work provides theoretical insight, empirical validation on realistic human-preference data, and publicly available code for reproducibility.

Abstract

The alignment of language models with human preferences is vital for their application in real-world tasks. The problem is formulated as optimizing the model's policy to maximize the expected reward that reflects human preferences with minimal deviation from the initial policy. While considered as a straightforward solution, reinforcement learning (RL) suffers from high variance in policy updates, which impedes efficient policy improvement. Recently, direct preference optimization (DPO) was proposed to directly optimize the policy from preference data. However, we show that DPO derived based on the optimal solution of the problem leads to a compromised mean-seeking approximation of the optimal solution in practice. In this paper, we propose efficient exact optimization (EXO) of the alignment objective. EXO is guaranteed to optimize in the same direction as RL algorithms asymptotically for arbitrary policy parametrization. This leads to the same mode-seeking solution, while enables efficient optimization by circumventing the complexities of RL. We also compare our method to DPO with both theoretical and empirical analyses, and further demonstrate the advantages of our method over existing approaches on realistic human preference data. Code is available at https://github.com/haozheji/exact-optimization.

Figures (9)

• Figure 1: Illustration of different characteristics of (a) $\pi_{\theta_{\textup{RKL}}}$ by minimizing the reverse KL (by EXO) and (b) $\pi_{\theta_{\textup{FKL}}}$ by minimizing the forward KL (by DPO).
• Figure 2: The frontier of oracle reward vs reverse KL to the SFT policy of different methods in the controlled experiment.
• Figure 3: Visualization of the estimated density ratio between the optimal and learned policy by EXO and DPO and the SFT policy on samples from the SFT policy sorted by their log probabilities.
• Figure 4: Win rates by comparing EXO to various baselines on the instruction-following task judged by GPT-4 and human labelers.
• Figure 5: Illustration of the relationship among the different objectives discussed in §\ref{['sec:method']}. : $\mathcal{J}_{\textup{lhf}}^{\beta_r}(\pi_\theta^{\beta_\pi})$ is a generalized version of $\mathcal{J}_{\textup{lhf}}^\beta(\pi_\theta)$ by distributing the KL regularization to both the learned policy $\pi_\theta$ and the reward model $r_\phi$ (§\ref{['sec:equivalence']}). : $\mathcal{L}_{\textup{dpo}}(\pi_\theta)$ is derived based on the optimal policy of $\mathcal{J}_{\textup{lhf}}^\beta(\pi_\theta)$ (§\ref{['sec:DPO']}). : $\mathcal{L}_{\textup{exo}}(\pi_\theta)$ is equivalent to $\mathcal{J}_{\textup{lhf}}^{\beta_r}(\pi_\theta^{\beta_\pi})$ in terms of their optimization directions (§\ref{['sec:exo']}). : $\mathcal{L}_{\textup{dpo-rw}}$ is the generalized version of $\mathcal{L}_{\textup{dpo}}$ by subsituting the pariwise loss with softmax loss over $K$ responses.(§\ref{['sec:compare-objs']}). The optimal policy, denoted by a dotted line, assumes unlimited model capacity. The solution, shown with a solid line, is the practically achievable policy within the realistic constraints of model capacity.

Theorems & Definitions (6)

• Theorem 3.1
• Theorem 3.2
• Theorem 3.3
• proof
• proof
• proof