Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Contrastive Policy Gradient: Aligning LLMs on sequence-level scores in a supervised-friendly fashion

Yannis Flet-Berliac, Nathan Grinsztajn, Florian Strub, Bill Wu, Eugene Choi, Chris Cremer, Arash Ahmadian, Yash Chandak, Mohammad Gheshlaghi Azar, Olivier Pietquin, Matthieu Geist

TL;DR

The paper tackles the challenge of aligning LLMs while enabling optimization of arbitrary rewards in an offline setting, addressing limitations of traditional RLHF and direct alignment methods. It introduces Contrastive Policy Gradient (CoPG), an off-policy, supervised-friendly RL algorithm that uses a contrastive baseline to avoid importance sampling and to unify IPO with policy-gradient-like updates. A key theoretical result shows that the unique KL-regularized optimum is achieved by an exponential tilt of the reference policy with respect to the reward, and CoPG reduces to known methods in special cases. Empirically, CoPG improves reward optimization on both a toy bandit and an offline LLM summarization task, outperforming IPO and DPO under realistic offline constraints and demonstrating favorable KL-reward trade-offs. Overall, CoPG offers a scalable, stable approach for optimizing arbitrary rewards in LLM fine-tuning, with potential online extensions and multi-reward settings.

Abstract

Reinforcement Learning (RL) has been used to finetune Large Language Models (LLMs) using a reward model trained from preference data, to better align with human judgment. The recently introduced direct alignment methods, which are often simpler, more stable, and computationally lighter, can more directly achieve this. However, these approaches cannot optimize arbitrary rewards, and the preference-based ones are not the only rewards of interest for LLMs (eg., unit tests for code generation or textual entailment for summarization, among others). RL-finetuning is usually done with a variation of policy gradient, which calls for on-policy or near-on-policy samples, requiring costly generations. We introduce Contrastive Policy Gradient, or CoPG, a simple and mathematically principled new RL algorithm that can estimate the optimal policy even from off-policy data. It can be seen as an off-policy policy gradient approach that does not rely on important sampling techniques and highlights the importance of using (the right) state baseline. We show this approach to generalize the direct alignment method IPO (identity preference optimization) and classic policy gradient. We experiment with the proposed CoPG on a toy bandit problem to illustrate its properties, as well as for finetuning LLMs on a summarization task, using a learned reward function considered as ground truth for the purpose of the experiments.

Contrastive Policy Gradient: Aligning LLMs on sequence-level scores in a supervised-friendly fashion

TL;DR

The paper tackles the challenge of aligning LLMs while enabling optimization of arbitrary rewards in an offline setting, addressing limitations of traditional RLHF and direct alignment methods. It introduces Contrastive Policy Gradient (CoPG), an off-policy, supervised-friendly RL algorithm that uses a contrastive baseline to avoid importance sampling and to unify IPO with policy-gradient-like updates. A key theoretical result shows that the unique KL-regularized optimum is achieved by an exponential tilt of the reference policy with respect to the reward, and CoPG reduces to known methods in special cases. Empirically, CoPG improves reward optimization on both a toy bandit and an offline LLM summarization task, outperforming IPO and DPO under realistic offline constraints and demonstrating favorable KL-reward trade-offs. Overall, CoPG offers a scalable, stable approach for optimizing arbitrary rewards in LLM fine-tuning, with potential online extensions and multi-reward settings.

Abstract

Reinforcement Learning (RL) has been used to finetune Large Language Models (LLMs) using a reward model trained from preference data, to better align with human judgment. The recently introduced direct alignment methods, which are often simpler, more stable, and computationally lighter, can more directly achieve this. However, these approaches cannot optimize arbitrary rewards, and the preference-based ones are not the only rewards of interest for LLMs (eg., unit tests for code generation or textual entailment for summarization, among others). RL-finetuning is usually done with a variation of policy gradient, which calls for on-policy or near-on-policy samples, requiring costly generations. We introduce Contrastive Policy Gradient, or CoPG, a simple and mathematically principled new RL algorithm that can estimate the optimal policy even from off-policy data. It can be seen as an off-policy policy gradient approach that does not rely on important sampling techniques and highlights the importance of using (the right) state baseline. We show this approach to generalize the direct alignment method IPO (identity preference optimization) and classic policy gradient. We experiment with the proposed CoPG on a toy bandit problem to illustrate its properties, as well as for finetuning LLMs on a summarization task, using a learned reward function considered as ground truth for the purpose of the experiments.
Paper Structure (14 sections, 2 theorems, 47 equations, 7 figures, 1 algorithm)

This paper contains 14 sections, 2 theorems, 47 equations, 7 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Contrastive Policy Gradient
  4. General objective
  5. A simple sample-based objective
  6. Related works
  7. Toy experiment
  8. LLM experiments
  9. Can CoPG optimize a reward?
  10. How does CoPG compare to direct alignment?
  11. KL-Reward trade-off
  12. Discussion and perspectives
  13. Proofs of theoretical results
  14. Qualitative results

Key Result

Theorem 1

Assume that $\pi_\text{ref}$, $\mu_1$ and $\mu_2$ all have the same support. Then, the unique maximizer of $L(\pi)$, defined Eq. eq:copg_objective_function, is $\pi_*(y)\propto \pi_\text{ref}(y)\exp\frac{R(y)}{\beta}$, which is also the unique maximizer of $J(\pi)$.

Figures (7)

  • Figure 1: Bandit experiment. CoPG achieves zero regret, converging to the optimal solution. IPO converges to a biased solution, as it optimizes for the expected preference. PG without a baseline has increasing regret, and PG with a value baseline converges to a biased solution.
  • Figure 2: CoPG: Rewards of generations along training.
  • Figure 3: DPO: Rewards of generations along training.
  • Figure 4: IPO: Rewards of generations along training.
  • Figure 5: Final reward as a function of $\beta$.
  • ...and 2 more figures

Theorems & Definitions (6)

  • Theorem 1: CoPG solves the right problem
  • Theorem 1: CoPG solves the right problem.
  • proof
  • proof
  • proof
  • proof