Contrastive Policy Gradient: Aligning LLM s on sequence-level scores in a supervised-friendly fashion (2024)

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 CoPGon 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.

## introduction reinforcement learning from human feedback (rlhf) @xcite is a classic finetuning step intended at aligning a large language model (llm) with human judgment @xcite . the underlying principle is to learn a reward model from a preference dataset, and to optimize this reward with a regularized † equal contribution. ⋆ corresponding author: matthieu@cohere.com. reinforcement learning (rl) approach @xcite @xcite , usually a policy gradient (pg) approach @xcite or a variation like proximal policy optimization (ppo) @xcite . these methods require on-policy or near-on-policy samples, and thus require costly generations from the llm. they can also be hard to tune and computationally heavy, for example through the use of an additional value network. more recently, the field of direct alignment methods has surged with the introduction of direct preference optimization (dpo) @xcite , sequence likelihood calibration (slic-hf) @xcite or identity preference optimization (ipo) @xcite . these approaches allow directly learning a policy optimizing for preferences, from a given preference dataset, in an offline manner and without using a proxy reward function. they are usually considered simpler, more stable, and computationally more lightweight than classic rlhf. however, by design, they cannot optimize for arbitrary reward functions. we posit that preference-based rewards are not the only rewards worth considering when finetuning an llm. not everything can be measured through preferences, which are also costly to label. such examples are using unit tests as a reward for code generation @xcite or a reward measuring textual-entailment for summarization @xcite . the aim of this paper is to propose an rl approach able to optimize an arbitrary reward while being as simple as direct alignment. it is important to note that we do not introduce any specific reward here. our intent is to provide a convenient and efficient tool for optimizing an arbitrary reward function. in particular, we take inspiration from the contrastive learning objective, which had tremendous success in the self-supervised learning technique @xcite , and extend it to rl techniques. to this end, we introduce contrastive policy gradient, or copg. it minimizes a supervised-friendly loss, of which we show the optimal policy of interest (optimizing the initial rl problem) to be the unique minimizer. it can be interpreted as a form of off-policy policy gradient, not relying on importance sampling (an approach that can easily lead the gradient variance to explode), but exploiting a specific state baseline that can be seen as a contrastive term to the reward being optimized. our proposed approach is versatile as it regroups ipo and policy gradient as special cases. notably, we obtain as a special case an offline off-policy generalization of reinforce leave-one-out (rloo) @xcite . to illustrate its properties, we experiment with the proposed copg on a toy bandit problem. we also test it for finetuning an llm on a summarization task. for this case, we train a reward model from the preference dataset and consider it the ground truth to be optimized. ## background we denote a prompt x and a generation y, and we call the llm to be trained a policy π(y|x). we assume the prompts to be sampled according to some unknown distribution ρ. we also assume to have access to some reference model π ref , typically the llm pretrained and supervisedfinetuned (sft model), used both for initializing π and regularizing the rl problem. we consider having access to a reward model r(x, y) to be maximized under π, with some regularization toward the reference model through a kl-divergence kl(π(•|x)||π ref (•|x)). to lighten notations, we drop the prompt x in the main text (it appears explicitly again in the proofs in appx. a). a classic approach is policy gradient, which maximizes j by gradient ascent. this is not a supervised-friendly loss because the expectation depends on the optimized policy, not on some fixed dataset of generations. the gradient is given by in practice, an empirical estimate of this gradient requires fresh generations from π, making it a costly method. it is common to subtract a baseline b: this does not bias the gradient as long as it does not depend on y (because e y∼π [∇ ln π(y)] = 0), and this is generally introduced to reduce the variance of the empirical gradient @xcite . a classic baseline is an estimate of the expected reward (that is the value, b ≈ e y∼π [r π β (y)]). objective (2) can be made supervised-friendly by relying on some fixed sampling distribution µ (e.g., underlying a dataset), by using importance sampling. indeed, we have the corresponding empirical gradient suffers from a larger variance: when the policy π becomes different from the sampling model µ (which happens, as π is trained), the probability ratio can explode. training is thus not stable or efficient. moreover, this requires having access to the probabilities µ(y), which may not be possible, for example when the underlying generations have been made by humans rather than by an llm. an approach for alleviating this stability issue is to clip the probability ratio in objective (4), to kill the gradient whenever π becomes too different from µ. this is the core idea behind ppo, which considers for µ older copies of the π network and uses a value estimate as the baseline. however, this still requires fresh generations, even if they can be used for more than one gradient step. ## related works contrastive policy gradient is related to policy gradient @xcite . it can be seen as a sound off-policy policy gradient, but crucially not relying on importance sampling, and thus not requiring clipping techniques such as proximal policy optimization @xcite or variations @xcite , allowing for broader applicability (notably, ppo cannot be applied offline to a dataset of unknown density). this link is even stronger: property 1 (copg and policy gradient). copg generalizes policy gradient in the sense that the expectation of the gradient of ℓ copg according to the current policy is exactly (up to the scaling) the policy gradient of eq. ( 3 ). so, we can retrieve policy gradient as a special case of the proposed approach, which is much less restrictive as it does not require the generations to be sampled according to the current policy. this result is asymptotic, as the expectation requires an infinite number of generations, but we have a similar connection to a more practical policy gradient approach. reinforce leave-one out is a sample-based policy gradient, using a monte carlo estimate of the expected reward from k generations as a baseline @xcite . it is remarkably effective for finetuning llms, simpler than ppo while providing better results, but still relying on fresh generations for each mini-batch @xcite . when using only two generations, the gradient is naturally symmetrized not to waste information and matches exactly eq. ( 9 ). property 2 (copg and rloo). the sample-based gradient ∇ℓ copg (y, y ′ ; π) is exactly the gradient of rloo for @xmath0, when both y and y ′ are sampled from the current policy π. a core difference is that copg is valid for any sampling distribution, while rloo critically relies on using on-policy generations when derived from objective (2). this result shows that it is valid and principled to use the rloo gradient in an offpolicy manner. this is highly non-trivial, new to the community, and made possible thanks to the proposed principled approach. contrastive policy gradient can also be related to direct alignment methods @xcite @xcite , and more especially to identity policy optimization @xcite , in the following sense. property 3 (copg and ipo). for a pair of generations (y, y ′ ), assume without loss of generality that y is preferred to y ′ according to the reward model, and redefine r(y) = -r(y ′ ) = 1 4 , then we have ∇ℓ copg (y, these results show that if we replace the reward in our objective with a binary signal depending on which generation is preferred, we follow the same gradient as ipo. in that sense, our approach also subsumes direct alignment approaches, allowing us to optimize for an arbitrary reward. ## toy experiment for an illustrative purpose, we consider a simple bandit problem, with 3 arms rewarded by @xmath1.5, 2, 1). we choose the data distributions to be µ 1 = (0.1, 0.2, 0.7) and µ 2 = (0.05, 0.05, 0.9). using these distributions, we sample a dataset of 10 4 pairs of rewarded arms. we set β = 0.5 and π ref (y) = 1 3 for @xmath2, 2, 3}. the analytical solution to this bandit problem is π * (y) ∝ exp r(y) β . we consider the practical copg objective of eq. ( 8 ), for which we recall the gradient eq. ( 9 ). given that copg generalizes policy gradient (pg, prop. 1), we also experiment with it to illustrate the importance of choosing the right baseline. for pg, we consider the following sample-based gradient: 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. as explained before, this gradient is valid whenever both y and y ′ are sampled according to the current policy π. however, here we use pairs of arms sampled from the dataset. in other words, this can be seen as a naive off-policy policy gradient. we consider two kind of baselines, @xmath3. the first case corresponds to vanilla policy gradient, and the second case corresponds to the baseline most often used in the literature. notice that in practice this should be estimated (typically with a value network), but we compute it exactly in this experiment. given the link between copg and ipo (prop. 3), we also experiment with ipo. for each approach, we train the policy π with stochastic gradient descent. we use adam (kingma and ba, 2014) with learning rate 10 -3 , batches of size 512, and train for 100 epochs. we measure the performance of the trained policy with the regret: with j being defined eq. ( 2 ). if copg and pg both rely on a reward function, ipo can only use preferences. we simply set them to be sampled according to a bradley-terry model, p (@xmath4, with σ the logistic function. results are presented in fig. 1 . we can observe that copg converges to the right solution, as predicted theoretically (thm. a). ipo converges to a biased solution. this was to be expected, as it can be shown to optimize the reward e y ′ ∼µ 2 [σ(r(y) -r(y ′ )] (azar et al., 2024), which is different from the reward of interest r. regarding policy gradient, we can observe that without baseline, naively applying the policy gradient on off-policy data leads to an increase in the regret: learning deteriorates the initial policy. adding the value baseline helps, but it still converges to a biased solution (and it is an ideal algorithm, as here the baseline is analytically computed, while it has to be estimated in practice). sample-based copg converges to the right solution, showing the importance of choosing the right baseline in an off-policy context. ## llm experiments in this section, we demonstrate the ability of copg to optimize a reward function for finetuning an llm. as depicted in sec. 3.2, we consider a pure offline objective, where one has to learn from a fixed dataset of pairs of generations. classic rlhf approaches, such as policy gradient or ppo, do not work in such a pure offline setting. ppo could possibly do a single-step policy improvement, but it would require access to the underlying probabilities µ(y) of elements in the dataset, which are not available. moreover, it would require an additional costly value network. therefore, as baselines, we consider direct alignment methods, specifically dpo and ipo, for which the preferred completion is chosen according to the same reward model being optimized by copg. dataset. we consider the reddit tl;dr dataset foot_0 of @xcite . it is a summarization dataset with an sft split, consisting of human-written summaries, and a preference split, made of human-annotated preference pairs. we will rerank the preferences according to the reward model, so that copg and direct alignment methods optimize for consistent objectives. policy model. we use llama2-7b as the base model foot_1 @xcite . we supervise finetune it on the sft split of the tl;dr dataset, giving π ref as a result. this is both the initial policy and the reference model for copg as for direct alignment baselines. this model is trained for 2 epochs with adam, with a cosine decay scheduler (2.10 -5 to 0), warmup of 10%, using a batch of size 128. reward model. our objective is to provide an approach to optimize arbitrary reward functions. as a proof of concept, for this empirical study, we train a reward model using the preference split of the tl;dr dataset and will consider it as the ground truth reward function to be optimized. we insist right away that we do not claim such a reward model to be the best thing to optimize for improving the llm, we use it as a proxy for assessing if the proposed approach can indeed optimize a reward at scale. for training the reward function, we use a classic bradley-terry model @xcite , optimizing for the loss the reward model is trained for two epochs on the train split of the preference dataset, with adam, the learning rate of 10 -6 , a batch of size 128, and a warm-up of 10% of the total number of training steps. the trained reward model achieves an accuracy of 89.1% on the train set and of 72.8% on the validation set. ## discussion and perspectives we have introduced contrastive policy gradient, a new rl approach for finetuning llms. it is a form of policy gradient that contrasts the reward with a specific baseline. the corresponding objective function is supervised-friendly, in the sense that it does not (necessarily) rely on fresh generations from the model. this allows to learn a policy in a pure offline setting, without relying on importance sampling or clipping of log-probability ratios, and does not require the introduction of an additional value network. we have demonstrated that copg indeed optimizes for the optimal kl-regularized policy (thm. 1), and we have shown that it generalizes policy gradient (prop. 1), rloo (prop. 2), or ipo (prop. 3). on a controlled but simple bandit experiment (sec. 5), we have illustrated empirically the convergence properties of the proposed approach, the importance of choosing the right baseline (the one coming from our derivations rather than the classic value baseline, and even less no baseline), and the advantage of optimizing a reward function rather than preferences derived from these rewards, which leads to a biased solution. on a larger scale llm experiment (sec. 6), we have shown that copg is able to optimize a reward function, in a fully offline and off-policy manner, conversely to other rl-finetuning approaches, and that it can achieve a higher reward than by using a direct alignment approach on preferences ranked by the reward model. our experimental study also suggests that this higher reward is achieved at a lower kl cost, compared to ipo or dpo. a core perspective is to study copg in an online setting. indeed, if it works in the offline setting, being the first such rl approach in the context of llms (to the best of our knowledge), it is not restricted to this setting. recent works have highlighted the benefits of using online (or fresh) data for direct alignment @xcite , and we hypothesize that these findings can also benefit to copg. typically, one can consider using a replay buffer, as classically done in off-policy rl but less common for llm finetuning. as copg uses pairs of generations, we think that this can also open new perspectives by using heterogeneous distributions for the completions, as briefly discussed in sec. 3.2. for example, one distribution could correspond to the replay buffer, while the other could correspond to exploratory generations (for addressing the exploration/exploitation dilemma), or by using a dataset of good but suboptimal generations, in the spirit of rl from demonstrations. another important perspective is to experiment copg on more tasks and rewards, and we also plan to study its possible extension to the multi-objective rl setting.