ARM : Alignment with Residual Energy-Based Model (2024)

Abstract: While large language models (LLMs) trained with large-scale unsupervised learning acquire a wide variety of world knowledge and skills, its behavior does not necessarily align with human preferences. RLHF methods achieve successes in aligning LLM responses with human preferences and improving the controllability of LLM behavior with human instruction. However, RLHF methods are considerably complicated to implement, computationally expensive to train, and notoriously tricky to tune. In this work, we propose Alignment with Residual Energy-Based Model (ARM), as a simple and flexible alternative to RLHF methods. Our method is driven by an observation that we can learn an aligned policy by minimizing a forward Kullback–Leibler (KL) divergence from a target policy (in the form of a residual energy-based model) to a parameteric policy (LLM), instead of a reverse KL as in RLHF methods. With samples from the energy-based target policy, we can leverage the power of DPO (or other offline methods) to learn an aligned policy efficiently. ARM is simple to implement and applicable in various data settings. Our extensive experiments demonstrate its strong performance across multiple datasets, compared to strong baselines like PPO, DPO.

## introduction large language models (llms) have become extremely powerful and demonstrated remarkable capacities in various domains @xcite . llms trained on very large unsupervised datasets acquire a wide range of capacities and skills, completing tasks zero-shot or few-shot @xcite . the large unsupervised corpus contains text with various goals and values, which are not necessarily aligned with human preferences. after unsupervised learning, instruction tuning @xcite @xcite figure 1: graphic illustration of arm. in arm, we sample from a target policy, as an exponential tilting of the sft policy, with self-normalizing importance sampling. we can then learn an aligned policy from these samples with dpo (or any other offline methods). @xcite is often applied to llms, which can significantly improve their capacities on instruction following and align their responses with human values or preferences. while instruct tuning is straightforward, the most successful class of methods for alignment is reinforcement learning from human feedback (rlhf) @xcite @xcite . to apply rlhf, human preferences data on model responses are first collected, and a reward function is learned with the preference data as a surrogate to human value. given the surrogate reward function, rl methods can be applied, where language models are optimized to produce responses that receive high rewards while not drifting too far away from a reference model @xcite . despite the success of rlhf, these methods are often complicated to implement, expensive to train, and tricky to tune. recently, there is a surge of interest in developing simpler alternatives to rlhf methods such as dpo @xcite , rrhf @xcite . these methods are straightforward to implement and easier to train, yet they maintain the performance of rlhf methods on human preference learning. the reliance on complicated online rl methods is because the reward maximization (with some conservative constraint) in preference learning amounts to minimizing a reverse kullback-leibler (kl) divergence d kl π θ || π * , where π * is the target response distribution or policy that aligns with human preference, and π θ is parameteric policy (e.g., llms) we aim to learn (see section 4 for more details). optimizing the reverse kl is not straightforward since sampling from π θ is not differentiable, and we have to resort to online rl methods to optimize this objective. in this work, we propose to optimize the forward kl, d kl π * || π θ . as we will show later (section 4), the target distribution π * is a residual energy-based model with a reference distribution (usually the sft distribution) as the base model and the surrogate reward function as the (negative) residual energy term. we can sample from π * given a learned reference distribution and reward function, and let's denote it as d π * . if we learn π θ from d π * with maximum likelihood estimation (mle), it is a variant of expert iteration @xcite . besides mle, we can learn it with any other offline methods such as dpo. this work focuses on dpo since it is simple and performs well @xcite . we call our method of learning policy from d π * as alignment with residual energy-based model (arm) due to the central role of ebm in our method. see figure 1 for a illustration of arm. we conduct extensive experiments and demonstrate that our method, arm, yields substantial improvements over sft policies and outperforms competitive baselines such as ppo and dpo. in addition to standard benchmarks, we also examine arm when only non-pairwise preference data are available and in low-resource settings. these experiments highlight the applicability of our method to diverse settings due to its simplicity and flexibility. as a preview, figure 2 displays the win-rate of sft and arm policy responses under various conditions, as compared to human preferred responses, on the anthropic helpful-harmless dataset @xcite . our contributions are summarized as follows: ## related works since unsupervised llms have demonstrated unprecedented potentials in a wide variety of tasks and domains @xcite , much research has dedicated to study how to improve the controllability of llm behavior, in order to align it with human value and ensure it to follow human instructions. one line of work focuses on instruction tuning. early works leverage academic datasets by transforming them into instructional formats with human-written prompt templates @xcite @xcite . this approach is scalable and exhibits potentials in making unsupervised llms follow instructions. however, its performance significantly lags behind proprietary models like gpt-4 (openai, 2023) which most likely collects a large scale of human-written high quality instruction data. recently, some researchers attempt to collect high-quality instruction data from strong proprietary models and use them to improve open-sourced models' instruction following capacities @xcite . besides instruction tuning, rlhf is another class of methods that demonstrate success in human preference learning @xcite @xcite . this approach learns a surrogate reward function from human preference data and considers llms as policy models, and then applies rl methods to maximize rewards assigned to the policy without excessively drifting from some reference model. one popular method, ppo @xcite , is often used in this setting. motivated by the complexity of rl methods, many recent works attempt to develop simpler alternatives. dpo @xcite and rrhf @xcite are developed along this line. our work shares the same motivation. since collecting human feedback is expensive, some works explore to use ai models to provide feedback, and they offer potential solutions to the scalability limitations of rlhf @xcite . our method, once an sft model and a reward function are trained with human (or other ai) data, our method characterizes sampling instruction data and preference data from ai models and learning from these samples. thus, our work is also in line with these rlaif works on this aspect. ## preliminaries large language model training starts with unsupervised learning where it is trained on very large datasets with next token prediction. with scaling, llms gain wide knowledge and capacities after unsupervised training @xcite . to improve llms' instruction following capacity, the next step is instruction tuning or supervised finetuning (sft), where models are finetuned on instructions and human-written completions @xcite @xcite . given a dataset, d @xmath0, y)} where x is an instruction or a prompt and y is a human-written completion, sft is often done by to align model behavior with human value, rlhf is applied after learning the sft policy. this framework assumes 1) a latent reward function, r : x × @xmath1, that captures human preference, and 2) preference models, bradley-terry model @xcite or more general plackett-luce model @xcite , that specify preference probability given the reward function. assume we have access to {(x, y 0 , y 1 )} where y 0 , y 1 ∼ π sft (y|x), the bradely-terry model assumes human preference is captured by the following distribution, given π sft (y|x) and r ϕ (x, y), we would like to learn a policy, π θ (y | x), with feedback from the reward model. the objective is often formulated as reward maximization with kl-constraint: ## experiments in this section, we empirically evaluate our proposed method, arm. we first examine its performance on three datasets. after validating its performance in standard settings, we next explore how arm works in two other interesting settings: 1) we only have access to non-pairwise human feedback 2) we only have access to a limited amount of pairwise human feedback. these experiments aim to demonstrate the flexibility of our method and applicability to realistic scenarios. in the aforementioned experiments, we focus on arm based on the bradley-terry model (see section 4.2.4). we then compare arm based on bradley-terry versus plackett-luce. in the end, we do an ablation on the number proposal samples used in the self-normalizing importance sampling, n (see algorithm 1). ## conclusion in this work, we propose alignment with residual energy-based model (arm) as a simple alternative to complicated rlhf methods. the core idea is to learn from samples drawn from a target policy in the form of a residual energy-based model, with powerful offline methods like dpo. our proposed method is characterized by its simplicity and high performance, as demonstrated in diverse tasks and various data settings.