Paper deep dive
Approximated Variational Bayesian Inverse Reinforcement Learning for Large Language Model Alignment
Yuang Cai, Yuyu Yuan, Jinsheng Shi, Qinhong Lin
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 94%
Last extracted: 3/12/2026, 6:26:53 PM
Summary
The paper introduces Approximated Variational Alignment (AVA), a novel training objective for Large Language Model (LLM) alignment based on Bayesian Inverse Reinforcement Learning (BIRL). By utilizing Approximated Variational Reward Imitation Learning (AVRIL), AVA enables intermediate reward modeling and direct reward estimation from single demonstrations, addressing limitations in existing methods like RLHF and DPO that rely on reward differences and end-of-sentence rewards. Experiments demonstrate that AVA improves reward modeling, RL fine-tuning, and direct optimization, effectively mitigating reward hacking.
Entities (5)
Relation Signals (3)
Approximated Variational Alignment → isbasedon → Bayesian Inverse Reinforcement Learning
confidence 95% · propose a novel LLM alignment training objective, Approximated Variational Alignment (AVA), based on Bayesian Inverse Reinforcement Learning (BIRL)
Approximated Variational Alignment → utilizes → Approximated Variational Reward Imitation Learning
confidence 95% · perform LLM alignment through Approximated Variational Reward Imitation Learning (AVRIL)
Approximated Variational Alignment → outperforms → Direct Preference Optimization
confidence 90% · AVA surpasses Bradley-Terry and Preference Transformer in reward modeling... and outperforms DPO and AfD in direct optimization
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:The alignment of large language models (LLMs) is crucial for generating helpful and harmless content. Existing approaches leverage preference-based human feedback data to learn the reward function and align the LLM with the feedback data. However, these approaches focus on modeling the reward difference between the chosen and rejected demonstrations, rather than directly modeling the true reward from each demonstration. Moreover, these approaches assume that the reward is only obtained at the end of the sentence, which overlooks the modeling of intermediate rewards. These issues lead to insufficient use of training signals in the feedback data, limiting the representation and generalization ability of the reward and potentially resulting in reward hacking. In this paper, we formulate LLM alignment as a Bayesian Inverse Reinforcement Learning (BIRL) problem and propose a novel training objective, Approximated Variational Alignment (AVA), to perform LLM alignment through Approximated Variational Reward Imitation Learning (AVRIL). The BIRL formulation facilitates intermediate reward modeling and direct reward modeling on each single demonstration, which enhances the utilization of training signals in the feedback data. Experiments show that AVA outperforms existing LLM alignment approaches in reward modeling, RL fine-tuning, and direct optimization.
Tags
Links
- Source: https://arxiv.org/abs/2411.09341
- Canonical: https://arxiv.org/abs/2411.09341
Trouble viewing inline? Open PDF directly →
Full Text
81,650 characters extracted from source content.
Expand or collapse full text
Approximated Variational Bayesian Inverse Reinforcement Learning for Large Language Model Alignment Yuang Cai, Yuyu Yuan, Jinsheng Shi, Qinhong Lin Abstract The alignment of large language models (LLMs) is crucial for generating helpful and harmless content. Existing approaches leverage preference-based human feedback data to learn the reward function and align the LLM with the feedback data. However, these approaches focus on modeling the reward difference between the chosen and rejected demonstrations, rather than directly modeling the true reward from each demonstration. Moreover, these approaches assume that the reward is only obtained at the end of the sentence, which overlooks the modeling of intermediate rewards. These issues lead to insufficient use of training signals in the feedback data, limiting the representation and generalization ability of the reward and potentially resulting in reward hacking. In this paper, we formulate LLM alignment as a Bayesian Inverse Reinforcement Learning (BIRL) problem and propose a novel training objective, Approximated Variational Alignment (AVA), to perform LLM alignment through Approximated Variational Reward Imitation Learning (AVRIL). The BIRL formulation facilitates intermediate reward modeling and direct reward modeling on each single demonstration, which enhances the utilization of training signals in the feedback data. Experiments show that AVA outperforms existing LLM alignment approaches in reward modeling, RL fine-tuning, and direct optimization. Introduction Large language models (LLMs) trained on massive corpus encode a large amount of knowledge and demonstrate powerful linguistic and reasoning capabilities in various domains (OpenAI 2022; Achiam et al. 2023). However, due to the inevitable harmful and useless information in the training data, LLMs can potentially generate content inconsistent with human values or requirements (Holtzman et al. 2019; Zhang et al. 2019; Weidinger et al. 2021). LLM alignment is a prevalent and effective approach for LLMs to generate harmless and helpful content. The alignment task typically relies on human feedback data in the form of preferences, where each preference data consists of a chosen sentence and a rejected sentence, labeled by human annotators (Zopf 2018; Tay et al. 2020). Reinforcement Learning from Human Feedback (RLHF) and Direct Preference Optimization (DPO) are two common approaches to align LLMs with human feedback data (Shen et al. 2023). RLHF first performs reward modeling to learn a reward function from the feedback data and then fine-tunes the LLM policy to maximize the expected reward achieved by its generated content using Reinforcement Learning (RL) (Ouyang et al. 2022; Bai et al. 2022; Touvron et al. 2023; Yang et al. 2023). DPO formulates the reward modeling objective as a ranking objective based on the LLM policy, which facilitates the joint performance of reward modeling and LLM policy fine-tuning through a unified training objective (Yuan et al. 2023; Rafailov et al. 2024; Song et al. 2024). The Inverse Reinforcement Learning (IRL) problem generally involves learning a reward model from observed demonstration data produced by a Markov Decision Process (MDP) (Ng, Russell et al. 2000). Conversely, the Natural Language Generation (NLG) process can be viewed as an MDP where the generated sentences are considered demonstration data (Ranzato et al. 2016). Therefore, the alignment task performed by RLHF and DPO can be seen as addressing an IRL problem that infers the implicit reward function hidden in the preference-based human feedback data and learns the LLM policy either separately or jointly. However, existing RLHF and DPO alignment approaches only model the reward difference between the chosen and rejected demonstrations without explicitly modeling the true reward of every single sentence. This limitation means that the demonstration data is not fully utilized, which restricts the representation ability of the reward model and can lead to reward hacking (Skalse et al. 2022; Gao, Schulman, and Hilton 2023; Coste et al. 2023; Zhang et al. 2024). Additionally, current approaches generally model the end-to-end sentence-level reward without considering the reward of intermediate states. Model generalization may be limited when confronted with data that have similar intermediate state distributions but different complete sentence distributions. It is more intuitive to model intermediate rewards since humans can not only provide overall feedback on the entire text but also explain which parts of the text influenced their feedback. In this paper, we propose a novel LLM alignment training objective, Approximated Variational Alignment (AVA), based on Bayesian Inverse Reinforcement Learning (BIRL) (Ramachandran and Amir 2007). Specifically, we formulate the reward distribution as a posterior distribution conditioned on the demonstration data and perform Approximated Variational Reward Imitation Learning (AVRIL) (Chan and van der Schaar 2021) to jointly approximate the reward distribution (i.e., the reward model) and the demonstration likelihood (i.e., the policy). Unlike most previous LLM alignment approaches, which only model the reward difference between chosen and rejected demonstrations, AVA directly models the reward of every single demonstration through the AVRIL training objective, thereby making better use of the training signals from feedback data. Additionally, we do not adhere to the assumption that the reward is only obtained at the end of the sentence. Instead, we leverage the AVRIL training objective to model the intermediate reward conditioned on the intermediate demonstration data. To demonstrate the flexibility of our approach, we use the AVA training objective on data in different formats through different pipelines. Our work makes the following main contributions: • We present a novel insight into LLM alignment by formulating the alignment task as a BIRL problem, which enhances the utilization of training signals and improves the representation and generalization ability of the LLM. • We demonstrate the flexibility of AVA by employing it for both reward modeling and direct optimization on either preference data or demonstration data. • We empirically show that AVA surpasses Bradley-Terry and Preference Transformer in reward modeling and downstream RL fine-tuning, and outperforms DPO and AfD in direct optimization, which indicates a reduction in the reward hacking issue and an improvement in representation and generalization ability. Related Work LLM Alignment The Bradley-Terry model (Bradley and Terry 1952) formulates preference likelihood using the reward model and is widely adopted by RLHF alignment approaches for reward modeling. After reward modeling, the LLM is fine-tuned to maximize the expected reward achieved by LLM-generated content through downstream RL training (Ouyang et al. 2022; Bai et al. 2022; Touvron et al. 2023; Yang et al. 2023). A more concise approach for preference alignment is Direct Preference Optimization (DPO) (Rafailov et al. 2024), which denotes preference as the relative log-likelihood difference between the chosen sentence and the rejected sentence. DPO unifies reward modeling and LLM fine-tuning into a single process, facilitating LLM alignment with a simple classification loss. In addition to aligning LLMs with pairwise human preference data, some recent works also align LLMs with non-pairwise demonstration data. Sun and van der Schaar (2024) propose Alignment from Demonstrations (AfD), which leverages high-quality demonstration data to overcome challenges such as noisy labels and privacy concerns in preference datasets. Intermediate Reward Modeling The above alignment approaches only model the end-to-end reward of a complete sentence, without considering the reward of intermediate states. This lack of intermediate reward modeling stems from the assumption that the reward is only achieved when the sentence is fully generated, regarding the Natural Language Generation (NLG) process as an MDP (Ranzato et al. 2016). To address this issue, we refer to related work on preference modeling in classic RL problems without the aforementioned assumption. Notably, the Preference Transformer (Kim et al. 2023) uses the attention weights computed by the Transformer architecture (Vaswani et al. 2017) to estimate the weighted non-Markovian reward of each intermediate state of the trajectory. The reward of a complete trajectory is then the weighted sum of all intermediate rewards. The preference between the chosen and rejected trajectories is formulated by their rewards and optimized through a contrastive training objective, similar to Bradley-Terry. Preliminaries MDP Formulation of NLG At time step t, the state is the previously generated tokens denoted as 1:t=(y1,y2,⋯,yt)subscript:1subscript1subscript2⋯subscripty_1:t=(y_1,y_2,·s,y_t)y1 : t = ( y1 , y2 , ⋯ , yitalic_t ), the action is the currently generated token yt+1subscript1y_t+1yitalic_t + 1. Note that in auto-regressive decoding, the output tokens are time-shifted. The action space is the vocabulary VV containing all possible tokens. In the text generation setting, the state transition is deterministic, so we do not consider the transition probability function. The reward of taking action yt+1subscript1y_t+1yitalic_t + 1 under state 1:tsubscript:1y_1:ty1 : t is denoted as R(1:t,yt+1)=R(1:t+1)subscript:1subscript1subscript:11R(y_1:t,y_t+1)=R(y_1:t+1)R ( y1 : t , yitalic_t + 1 ) = R ( y1 : t + 1 ), i.e., the reward can be the function of either the current state and the current action or merely the function of the next state due to the deterministic state transition. It is worth noting that for simplicity of denotation, we do not separately denote the prompt text and the response text but denote them as a whole sentence yy. The separation of prompt and response is trivial during implementation. The policy can be denoted as πw(yt+1|1:t)subscriptconditionalsubscript1subscript:1 _w(y_t+1|y_1:t)πitalic_w ( yitalic_t + 1 | y1 : t ), which is also the distribution of the language model parameterized by w. For simplicity, we sometimes denote the policy as πw()=∏t=1||−1πw(yt+1|1:t)subscriptsuperscriptsubscriptproduct11subscriptconditionalsubscript1subscript:1 _w(y)= _t=1^|y|-1 _w(y_t+1|y_1:% t)πitalic_w ( y ) = ∏t = 1| y | - 1 πitalic_w ( yitalic_t + 1 | y1 : t ), where |||y|| y | is the length of sequence |||y|| y |. Note that the accumulated product starts from πw(y2|1:1)subscriptconditionalsubscript2subscript:11 _w(y_2|y_1:1)πitalic_w ( y2 | y1 : 1 ) instead of πw(y1)subscriptsubscript1 _w(y_1)πitalic_w ( y1 ) since we assume that all sequences start with a special token denoting the start of the sequence. Bayesian Inverse Reinforcement Learning Inverse Reinforcement Learning (IRL) is the problem of extracting a reward function of a Markov Decision Process (MDP) given observed optimal behavior (Ng, Russell et al. 2000). Bayesian Inverse Reinforcement Learning (BIRL) regards the reward function R as the hidden variable affecting and motivating the behavioral data TT. The objective of BIRL is to learn the posterior distribution p(R|)conditionalp(R|T)p ( R | T ). Approximate Variational Reward Imitation Learning (AVRIL) (Chan and van der Schaar 2021) adopts variational inference to approximate the posterior distribution. Specifically, AVRIL employs a parameterized distribution qϕsubscriptitalic-ϕq_φqitalic_ϕ and minimizes the Kullback-Leibler (KL) divergence between qϕsubscriptitalic-ϕq_φqitalic_ϕ and the posterior distribution p(R|)conditionalp(R|T)p ( R | T ), as shown in Eq. 1. This KL divergence is hard to compute since the posterior distribution is intractable. A common solution is to maximize the Evidence Lower Bound (ELBO), as shown in Eq. 2, where the second term is to minimize the KL divergence between qϕsubscriptitalic-ϕq_φqitalic_ϕ and the tractable prior distribution. minϕDKL[qϕ(R)||p(R|)] _φD_KL[q_φ(R)||p(R|T)]minitalic_ϕ DKL [ qitalic_ϕ ( R ) | | p ( R | T ) ] (1) maxϕR∼qϕ(⋅)[logp(|R)]−DKL[qϕ(R)||p(R)] _φE_R q_φ(·)[ p(T|R)]-D_% KL[q_φ(R)||p(R)]maxitalic_ϕ blackboard_ER ∼ q start_POSTSUBSCRIPT ϕ ( ⋅ ) end_POSTSUBSCRIPT [ log p ( T | R ) ] - DKL [ qitalic_ϕ ( R ) | | p ( R ) ] (2) The first term of Eq. 2 is to maximize the log-likelihood of the observed optimal behaviors given any reward sampled from qϕsubscriptitalic-ϕq_φqitalic_ϕ. AVRIL denotes the action distribution as a Boltzmann policy, as shown in Eq. 3, where QRπsubscriptsuperscriptsubscriptQ _T_RQitalic_πcaligraphic_Titalic_R is the state-action value function following policy πsubscript _Tπcaligraphic_T under reward function R. Intuitively, we can approximate the state-action value using a Deep Q Network (DQN) (Mnih et al. 2013) QθsubscriptQ_θQitalic_θ parameterized by θ. An important problem is that in the RL setting, the reward function is fixed when optimizing QθsubscriptQ_θQitalic_θ. However, in the AVRIL setting, the reward function is also being optimized during the optimization of QθsubscriptQ_θQitalic_θ. The reward function and the state-action value function should satisfy R(s,a)=s′∼P(⋅|s,a),a′∼π(⋅|s′)[QRπ(s,a)−γQRπ(s′,a′)],∀s∈,a∈R(s,a)=E_s P(·|s,a),a π(·|s^% )[Q^π_R(s,a)-γ Q^π_R(s ,a )],~% ∀~s ,a ( s , a ) = blackboard_Es′ ∼ P ( ⋅ | s , a ) , a′ ∼ π ( ⋅ | s′ ) [ Qitalic_πitalic_R ( s , a ) - γ Qitalic_πitalic_R ( s′ , a′ ) ] , ∀ s ∈ S , a ∈ A, i.e., the reward should equal the expectation of the TD error. By adding a penalty term forcing the TD error to follow the reward distribution, the final objective to be maximized is shown in Eq. 4, where qϕ(R|s,a)subscriptitalic-ϕconditionalq_φ(R|s,a)qitalic_ϕ ( R | s , a ) denotes the distribution of reward values given the state s and the action a. In this way, the behavior is indirectly conditioned on the reward, which is consistent with the likelihood p(|R)conditionalp(T|R)p ( T | R ) in the ELBO (Eq. 2). Here, B(a|s;Qθ)conditionalsubscriptB(a|s;Q_θ)B ( a | s ; Qitalic_θ ) is the Boltzmann policy upon the state-action value function QθsubscriptQ_θQitalic_θ parameterized by θ. The third term in the square brackets is to restrict the TD error to satisfy the constraint R(s,a)=s′,a′[QRπ(s,a)−γQRπ(s′,a′)],∀s∈,a∈formulae-sequencesubscriptsuperscript′delimited-[]subscriptsuperscriptsubscriptsuperscriptsuperscript′formulae-sequencefor-allR(s,a)=E_s ,a [Q^π_R(s,a)-γ Q^π_R(% s ,a )],~∀~s ,a ( s , a ) = blackboard_Es′ , a′ [ Qitalic_πitalic_R ( s , a ) - γ Qitalic_πitalic_R ( s′ , a′ ) ] , ∀ s ∈ S , a ∈ A. qϕ(R|s,a)subscriptitalic-ϕconditionalq_φ(R|s,a)qitalic_ϕ ( R | s , a ) denotes the distribution of reward values given the state s and the action a. We refer to the training objective in Eq. 4 as the approximated variational training objective of the ELBO. B(a|s;QRπ)=exp(βQRπ(s,a))∑a′∈exp(βQRπ(s,a′))conditionalsubscriptsuperscriptsubscriptsubscriptsuperscriptsubscriptsubscriptsuperscript′subscriptsuperscriptsubscriptsuperscript′B(a|s;Q _T_R)= (β Q _T_R(s,% a)) _a (β Q _T_R(s,a^% ))B ( a | s ; Qitalic_πcaligraphic_Titalic_R ) = divide start_ARG exp ( β Qitalic_πcaligraphic_Titalic_R ( s , a ) ) end_ARG start_ARG ∑a′ ∈ A exp ( β Qitalic_πcaligraphic_Titalic_R ( s , a′ ) ) end_ARG (3) maxϕ,θ∑(s,a,s′,a′)∈[logB(a|s;Qθ)−DKL[qϕ(⋅|s,a)||p(⋅)]+λlogqϕ(Qθ(s,a)−γQθ(s′,a′)|s,a)] _φ,θ _(s,a,s ,a ) % [ aligned B(a|s;Q_θ)-D_KL [q_φ(% ·|s,a)||p(·) ]\\ +λ q_φ (Q_θ(s,a)-γ Q_θ(s ,a^% )|s,a ) aligned ]maxitalic_ϕ , θ ∑( s , a , s′ , a′ ) ∈ T [ start_ROW start_CELL log B ( a | s ; Qitalic_θ ) - DKL [ qitalic_ϕ ( ⋅ | s , a ) | | p ( ⋅ ) ] end_CELL end_ROW start_ROW start_CELL + λ log qitalic_ϕ ( Qitalic_θ ( s , a ) - γ Qitalic_θ ( s′ , a′ ) | s , a ) end_CELL end_ROW ] (4) Approximated Variational Alignment In this section, we formulate the LLM alignment tasks as the BIRL problems and perform alignment with the Approximated Variational Alignment (AVA) training objectives. The AVA training objectives involve AVA from Demonstration (AVA-d) and AVA from Preference (AVA-p), both of which are BIRL training objectives based on the Approximated Variational Reward Imitation Learning (AVRIL) training objective (Chan and van der Schaar 2021). AVA-d is the implementation of the AVRIL training objective under the NLG setting, which learns on non-pairwise demonstration datasets. AVA-p is a contrastive variant of AVA-d, which learns on pairwise preference datasets. Alignment from Demonstration We first consider the problem of aligning an LLM policy with the demonstration data DD, where each sentence ∈y ∈ D is the ground-truth sentence. The alignment objective is to encourage the LLM policy to generate sentences like the demonstration data. Instead of building a direct training objective (e.g., supervised fine-tuning) to optimize the LLM policy, we focus on performing BIRL to learn a reward function from the demonstration data DD, i.e., to learn the posterior p(R|)conditionalp(R|D)p ( R | D ) with a parameterized distribution qϕ(R)subscriptitalic-ϕq_φ(R)qitalic_ϕ ( R ). As illustrated in the preliminaries, the optimization of qϕsubscriptitalic-ϕq_φqitalic_ϕ can be achieved by maximizing the AVRIL training objective (Eq. 4), where each element (s,a,s′,a′)∈superscript′(s,a,s ,a ) ( s , a , s′ , a′ ) ∈ T is a state-action quadruplet consisting of the current state s, the current action a, the next state s′ and the next action a′. As for the Natural Language Generation (NLG) setting, at each time step t, the current state is the current sub-sentence 1:tsubscript:1y_1:ty1 : t, the current action is the token-to-be-generated yt+1subscript1y_t+1yitalic_t + 1, the next state is 1:t+1subscript:11y_1:t+1y1 : t + 1, the concatenation of 1:tsubscript:1y_1:ty1 : t and yt+1subscript1y_t+1yitalic_t + 1, and the next action is yt+2subscript2y_t+2yitalic_t + 2. By substituting the state-action quadruplet in Eq. 4 with the new quadruplet (1:t,yt+1,1:t+1,yt+2)subscript:1subscript1subscript:11subscript2(y_1:t,y_t+1,y_1:t+1,y_t+2)( y1 : t , yitalic_t + 1 , y1 : t + 1 , yitalic_t + 2 ) and rewrite the summation in timestep-wise form, we can obtain the AVRIL training objective applicable to the NLG setting, as shown in Eq. 5. We refer to this training objective as Approximated Variational Alignment from Demonstration (AVA-d), which is a variant of the AVRIL training objective in the natural language generation setting. ℱd()=∑∈∑t=1||−2[logB(yt+1|1:t;Qθ)−dt(ϕ)+λlogqϕ(δt(θ)|1:t+1)]subscriptℱsubscriptsuperscriptsubscript12delimited-[]conditionalsubscript1subscript:1subscriptsubscriptitalic-ϕsubscriptitalic-ϕconditionalsubscriptsubscript:11F_d(D)= _y _t=1^|% y|-2 [ aligned B(y_t+1|y_1:t;Q_θ% )-d_t(φ)\\ +λ q_φ ( _t(θ)|y_1:t+1 ) % aligned ]Fitalic_d ( D ) = ∑y ∈ D ∑t = 1| y | - 2 [ start_ROW start_CELL log B ( yitalic_t + 1 | y1 : t ; Qitalic_θ ) - ditalic_t ( ϕ ) end_CELL end_ROW start_ROW start_CELL + λ log qitalic_ϕ ( δitalic_t ( θ ) | y1 : t + 1 ) end_CELL end_ROW ] (5) dt(ϕ)=DKL[qϕ(⋅|1:t+1)||p(⋅)]d_t(φ)=D_KL [q_φ(·|y_1:t+1)||p(·) ]ditalic_t ( ϕ ) = DKL [ qitalic_ϕ ( ⋅ | y1 : t + 1 ) | | p ( ⋅ ) ] (6) δt(θ)=Qθ(1:t,yt+1)−γQθ(1:t+1,yt+2)subscriptsubscriptsubscript:1subscript1subscriptsubscript:11subscript2 _t(θ)=Q_θ(y_1:t,y_t+1)-γ Q_θ(% y_1:t+1,y_t+2)δitalic_t ( θ ) = Qitalic_θ ( y1 : t , yitalic_t + 1 ) - γ Qitalic_θ ( y1 : t + 1 , yitalic_t + 2 ) (7) Here, qϕ(R|1:t+1)subscriptitalic-ϕconditionalsubscript:11q_φ(R|y_1:t+1)qitalic_ϕ ( R | y1 : t + 1 ) is the reward distribution of the sub-sequence 1:t+1subscript:11y_1:t+1y1 : t + 1. The Boltzmann policy B(yt+1|1:t;Qθ)conditionalsubscript1subscript:1subscriptB(y_t+1|y_1:t;Q_θ)B ( yitalic_t + 1 | y1 : t ; Qitalic_θ ) built upon the Q-value model QθsubscriptQ_θQitalic_θ acts as the LLM policy for text generation. By maximizing ℱdsubscriptℱF_dFitalic_d, the Q-value model (i.e., the LLM policy) QθsubscriptQ_θQitalic_θ as well as the reward distribution qϕsubscriptitalic-ϕq_φqitalic_ϕ will be jointly optimized to be aligned with the demonstration dataset DD. Similar to the original AVRIL objective, the AVA-d objective consists of three sub-objectives: the log-likelihood maximization, the KL divergence minimization, and the TD-error constraint. The first objective trains the LLM policy to maximize the likelihood of the demonstration data, which is identical to supervised fine-tuning. The second objective is to ensure the reward distribution satisfies the prior distribution assumption. The third objective, TD-error constraint, distinguishes AVA-d from conventional supervised fine-tuning. With the constraint, the update of the Q-value model will not only increase the Q-value of the ground-truth token in demonstration data but also make the TD error of the Q-values 7 close to the reward obtained after generating the current token, which ensures the consistency between the reward and the policy. TQR Architecture Figure 1: Overview of the TQR architecture. The original AVRIL adopts the architecture with a reward encoder and a Q-value decoder. To compute the AVA-d training objective and leverage the pre-trained weights of the backbone transformer model, we add a reward head and a Q-value head at the top of the Transformer decoder, as shown in Fig. 1. We refer to this architecture as Transformer with Q-value and Reward Heads (TQR). The Q-value head takes the hidden states of the last decoder layer as input and outputs the Q-value of each action (i.e., token), as shown in Eq. 8. The reward is assumed to follow Gaussian distribution, and the reward head takes in the same hidden states and outputs the mean and standard deviation of the reward of each state, as shown In Eq. 9. Here, tsubscripth_thitalic_t is the hidden state vector of time step t, Qθ(1:t,⋅)∈ℝ||subscriptsubscript:1⋅superscriptℝQ_θ(y_1:t,·) ^|V|Qitalic_θ ( y1 : t , ⋅ ) ∈ blackboard_R| V | is a vector whose i-th element equals Qθ(1:t,v(i))subscriptsubscript:1superscriptQ_θ(y_1:t,v^(i))Qitalic_θ ( y1 : t , v( i ) ), where v(i)superscriptv^(i)v( i ) is the i-th token in the vocabulary, and μt,σt∈ℝsubscriptsubscriptℝ _t, _t μitalic_t , σitalic_t ∈ blackboard_R are mean and standard deviation of reward R(1:t+1)subscript:11R(y_1:t+1)R ( y1 : t + 1 ) at time step t. Now we can compute the training objective in Eq. 5 based on the above outputs of the Q-value head and the reward head. Qθ(1:t,⋅)subscriptsubscript:1⋅ Q_θ(y_1:t,·)Qitalic_θ ( y1 : t , ⋅ ) =QHead(t;θ),∀t∈1,⋯,||formulae-sequenceabsentQHeadsubscriptfor-all1⋯ =QHead(h_t;θ),∀~t∈\1,·s,|% y|\= QHead ( hitalic_t ; θ ) , ∀ t ∈ 1 , ⋯ , | y | (8) [μt;σt]subscriptsubscript [ _t; _t][ μitalic_t ; σitalic_t ] =RHead(t;ϕ),∀t∈1,⋯,||formulae-sequenceabsentRHeadsubscriptitalic-ϕfor-all1⋯ =RHead(h_t;φ),∀~t∈\1,·s,|% y|\= RHead ( hitalic_t ; ϕ ) , ∀ t ∈ 1 , ⋯ , | y | (9) R(1:t+1)subscript:11 R(y_1:t+1)R ( y1 : t + 1 ) ∼qϕ(R|1:t+1)=(R;μt,σt)similar-toabsentsubscriptitalic-ϕconditionalsubscript:11subscriptsubscript q_φ(R|y_1:t+1)=N(R; _t, _% t)∼ qitalic_ϕ ( R | y1 : t + 1 ) = N ( R ; μitalic_t , σitalic_t ) (10) Inspired by preference transformer (Kim et al. 2023), we further compute a reward weight for each time step of reward based on attention weights, as shown in Eq. 11, where isubscriptq_iqitalic_i is the i-th row of the query matrix of the attention mechanism, t′subscriptsuperscript′k_t kitalic_t′ is the t′-th row of the key matrix. We then apply reward weights to the outputs of the Q-value head (Eq. 8) and reward head (Eq. 9). Specifically, we simply multiply the output of the t-th position of the heads by the reward weight wtsubscriptw_twitalic_t, as shown by the red arrows in Fig. 1. wt=1||∑i=1||∑t=1||softmax(i⋅t′t′=1||)tsubscript1superscriptsubscript1superscriptsubscript1softmaxsubscriptsuperscriptsubscript⋅subscriptsubscriptsuperscript′1w_t= 1|y| _i=1^|y| _t=1^|y|% softmax ( \q_i·k_t \_% t =1^|y| )_twitalic_t = divide start_ARG 1 end_ARG start_ARG | y | end_ARG ∑i = 1| y | ∑t = 1| y | softmax ( qitalic_i ⋅ kitalic_t′ = 1| y | )t (11) Besides using a randomly initialized Q-value head, we can also construct a pre-trained Q-value model from the pre-trained LLM policy. The Boltzmann policy formulates the action probability as the softmax function of Q-values. Inversely, we can also formulate the Q-value as the log-softmax function of action probabilities, as shown in Eq. 12, where α is the temperature hyperparameter, πwsubscript _wπitalic_w is the LLM policy parameterized by w. Note that the log-softmax operation is a non-strict inversion of the softmax operation, which means we can tune α to find the best way to map token-level probabilities to token-level Q-values. Qw(1:t,yt+1)=logexp(απw(yt+1|1:t))∑y′∈exp(απw(y′|1:t))subscriptsubscript:1subscript1subscriptconditionalsubscript1subscript:1subscriptsuperscript′subscriptconditionalsuperscript′subscript:1Q_w(y_1:t,y_t+1)= (α _w(y_t+1|y% _1:t)) _y (α _w(y |% y_1:t))Qitalic_w ( y1 : t , yitalic_t + 1 ) = log divide start_ARG exp ( α πitalic_w ( yitalic_t + 1 | y1 : t ) ) end_ARG start_ARG ∑y′ ∈ V exp ( α πitalic_w ( y′ | y1 : t ) ) end_ARG (12) By substituting with the above Q-value model, the AVA-d training objective can be denoted as Eq. 13, and the TD error can be denoted as Eq. 14. This denotation facilitates us to initialize the Q-value model from a pre-trained LLM policy and adopt the AVA-d objective to fine-tune the LLM policy. ℱd()=∑∈∑t=1||−2[βlogsoftmax(απw(yt+1|1:t))−dt(ϕ)+λlogqϕ(δt(w)|1:t+1)]subscriptℱsubscriptsuperscriptsubscript12delimited-[]softmaxsubscriptconditionalsubscript1subscript:1subscriptitalic-ϕsubscriptitalic-ϕconditionalsubscriptsubscript:11F_d(D)= _y _t=1^|% y|-2 [ aligned β (α _w(y_t% +1|y_1:t))\\ -d_t(φ)+λ q_φ( _t(w)|y_1:t+1) % aligned ]Fitalic_d ( D ) = ∑y ∈ D ∑t = 1| y | - 2 [ start_ROW start_CELL β log softmax ( α πitalic_w ( yitalic_t + 1 | y1 : t ) ) end_CELL end_ROW start_ROW start_CELL - ditalic_t ( ϕ ) + λ log qitalic_ϕ ( δitalic_t ( w ) | y1 : t + 1 ) end_CELL end_ROW ] (13) δt(w)=logsoftmax(απw(yt+1|1:t))softmax(απw(yt+2|1:t+1))γsubscriptsoftmaxsubscriptconditionalsubscript1subscript:1softmaxsuperscriptsubscriptconditionalsubscript2subscript:11 _t(w)= softmax(α _w(y_t+1|y_1:t)% )softmax(α _w(y_t+2|y_1:t+1))^γδitalic_t ( w ) = log divide start_ARG softmax ( α πitalic_w ( yitalic_t + 1 | y1 : t ) ) end_ARG start_ARG softmax ( α πitalic_w ( yitalic_t + 2 | y1 : t + 1 ) )γ end_ARG (14) Alignment from Preference We then consider the problem of aligning an LLM policy πwsubscript _wπitalic_w with preference data PP, where each data item (+,−)∈superscriptsuperscript(y^+,y^-) ( y+ , y- ) ∈ P consists of the chosen sentence +superscripty^+y+ and the rejected sentence −superscripty^-y-. We denote the set of all chosen sentences as +=+|(+,−)∈superscriptconditional-setsuperscriptsuperscriptsuperscriptP^+=\y^+|(y^+,y^-) \P+ = y+ | ( y+ , y- ) ∈ P and the set of all rejected sentences as −=−|(+,−)∈superscriptconditional-setsuperscriptsuperscriptsuperscriptP^-=\y^-|(y^+,y^-) \P- = y- | ( y+ , y- ) ∈ P . The alignment objective is to encourage the LLM policy to generate sentences like the chosen demonstrations +superscriptP^+P+ while discouraging the LLM policy from generating sentences like the rejected demonstrations −superscriptP^-P-. Similar to the derivation of the AVA-d training objective, we first focus on performing BIRL to learn a reward function from the preference data PP. We need to consider not only the chosen sentences as positive demonstrations but also the rejected sentences as negative demonstrations. We consider two posterior distributions, which are the reward conditioned on the chosen demonstrations p(R|+)conditionalsuperscriptp(R|P^+)p ( R | P+ ) and the reward conditioned on demonstrations that differ from rejected demonstrations p(R|−¯)conditional¯superscriptp(R| P^-)p ( R | over¯ start_ARG P- end_ARG ). Here, −¯superscript P^-over¯ start_ARG P- end_ARG denotes demonstrations that differ from −superscriptP^-P-. Therefore, we define the training objective as Eq. 15, where the first term drives the reward distribution qϕsubscriptitalic-ϕq_φqitalic_ϕ close to rewards that motivate the positive behaviors +superscriptP^+P+, while the second term drives qϕsubscriptitalic-ϕq_φqitalic_ϕ close to rewards that motivate behaviors that differ from the negative demonstrations. We refer to the training objective as Contrastive Bayesian Inverse Reinforcement Learning (CBIRL). minϕDKL[qϕ(R)||p(R|+)]+DKL[qϕ(R)||p(R|−¯)] _φD_KL[q_φ(R)||p(R|P^+)]+D_KL[q_% φ(R)||p(R| P^-)]minitalic_ϕ DKL [ qitalic_ϕ ( R ) | | p ( R | P+ ) ] + DKL [ qitalic_ϕ ( R ) | | p ( R | over¯ start_ARG P- end_ARG ) ] (15) Unsurprisingly, the minimization of these two KL divergences is infeasible. We derive the equivalent ELBO objective, as shown in Eq. 16. The derivation is shown in the Technical Appendix. maxϕ[R∼qϕ(⋅)[logp(+|R)+log[1−p(−|R)]]−DKL[qϕ(R)||p(R)]] _φ [ aligned E_R q_φ(·) [ p% (P^+|R)+ [1-p(P^-|R)] ]\\ -D_KL[q_φ(R)||p(R)] aligned ]maxitalic_ϕ [ start_ROW start_CELL blackboard_ER ∼ q start_POSTSUBSCRIPT ϕ ( ⋅ ) end_POSTSUBSCRIPT [ log p ( P+ | R ) + log [ 1 - p ( P- | R ) ] ] end_CELL end_ROW start_ROW start_CELL - DKL [ qitalic_ϕ ( R ) | | p ( R ) ] end_CELL end_ROW ] (16) Towards implementation, we need to further derive the ELBO objective as an approximated variational objective. Note that the main difference between the ELBO of CBIRL and the ELBO of conventional BIRL is the second optimization term in Eq. 16, which minimizes the log-likelihood of the negative demonstrations −superscriptP^-P-. Therefore, the approximated variational objective also contains the minimization of the negative demonstrations, as shown in Eq. 17. We refer to this training objective as the Approximated Variational Alignment from Preference (AVA-p). By maximizing ℱp()subscriptℱF_p(P)Fitalic_p ( P ), on the one hand, the LLM policy πwsubscript _wπitalic_w will be encouraged to generate sentences like +superscriptP^+P+ and discouraged to generated sentences like −superscriptP^-P-; on the other hand, the policy and the reward will stay consistent under the TD-error constraint. ℱp()=∑+/−∈∑t[βlogsoftmax(απw(yt+1+|1:t+))−βlogsoftmax(απw(yt+1−|1:t−))−dt(ϕ)+λlogqϕ(δt(θ)|1:t+1)]subscriptℱsubscriptsuperscriptabsentsubscriptdelimited-[]softmaxsubscriptconditionalsuperscriptsubscript1superscriptsubscript:1softmaxsubscriptconditionalsuperscriptsubscript1superscriptsubscript:1subscriptitalic-ϕsubscriptitalic-ϕconditionalsubscriptsubscript:11F_p(P)= _y^+/- _t% [ aligned β (α _w(y_t+1^+|% y_1:t^+))\\ -β (α _w(y_t+1^-|y_1:t^-))\\ -d_t(φ)+λ q_φ ( _t(θ)|y_1:t+1% ) aligned ]Fitalic_p ( P ) = ∑y+ / - ∈ P ∑t [ start_ROW start_CELL β log softmax ( α πitalic_w ( yitalic_t + 1+ | y1 : t+ ) ) end_CELL end_ROW start_ROW start_CELL - β log softmax ( α πitalic_w ( yitalic_t + 1- | y1 : t- ) ) end_CELL end_ROW start_ROW start_CELL - ditalic_t ( ϕ ) + λ log qitalic_ϕ ( δitalic_t ( θ ) | y1 : t + 1 ) end_CELL end_ROW ] (17) To ensure the reward difference between the chosen and rejected demonstrations, we adopt a more intuitive auxiliary training objective, the Contrastive Expected Return (CER) training objective, as shown in Eq. 18, which encourages the reward of the positive demonstrations to be higher than the reward of the negative demonstrations. Note that we only consider the reward of the last timestep in the CER objective. Although we model the intermediate rewards, we still assume that the reward of the last timestep is decisive for the overall expected return, since empirical practice and research (Geva et al. 2023; Hanna, Liu, and Variengien 2024) show that the last position of the Transformer gathers most of the knowledge. ℱc()subscriptℱ _c(P)Fitalic_c ( P ) =∑+/−∈σ[qϕ(R|+)[R]−qϕ(R|−)[R]]absentsubscriptsuperscriptabsentdelimited-[]subscriptsubscriptitalic-ϕconditionalsuperscriptdelimited-[]subscriptsubscriptitalic-ϕconditionalsuperscriptdelimited-[] = _y^+/- σ [E_q_% φ(R|y^+)[R]-E_q_φ(R|y^-)[R] ]= ∑y+ / - ∈ P σ [ blackboard_Eq start_POSTSUBSCRIPT ϕ ( R | y+ ) end_POSTSUBSCRIPT [ R ] - blackboard_Eq start_POSTSUBSCRIPT ϕ ( R | y- ) end_POSTSUBSCRIPT [ R ] ] (18) AVA Pipelines The AVA training objectives facilitate the joint optimization of the reward function and the policy. Therefore, AVA can be leveraged for both reward modeling and direct optimization, which are two common pipelines in LLM alignment. Both pipelines have their advantages and disadvantages. The reward modeling pipeline can produce a lightweight and reusable reward function for downstream RL fine-tuning while it suffers from the high RL training cost. The direct optimization pipeline is more efficient than reward modeling with RL during training but cannot produce a lightweight reward function for other uses and may suffer from overfitting. AVA for Reward Modeling Data: Dataset DD, initial implicit policy πψ(1)subscriptsuperscript1 _ψ^(1)πitalic_ψ( 1 ), initial reward distribution qϕ(1)subscriptsuperscriptitalic-ϕ1q_φ^(1)qitalic_ϕ( 1 ), training epochs T Result: The trained reward distribution qϕ(T)subscriptsuperscriptitalic-ϕq_φ^(T)qitalic_ϕ( T ) 1 for i∈1,⋯,T1⋯i∈\1,·s,T\i ∈ 1 , ⋯ , T do 2 if DD is demonstration dataset then 3 ϕ(i+1)←ϕ(i)+∇ϕ(i)ℱd()←superscriptitalic-ϕ1superscriptitalic-ϕsubscript∇superscriptitalic-ϕsubscriptℱφ^(i+1)←φ^(i)+ _φ^(i)F_d(D)ϕ( i + 1 ) ← ϕ( i ) + ∇ϕ( i ) Fitalic_d ( D ); 4 ψ(i+1)←ψ(i)+∇ψ(i)ℱd()←superscript1superscriptsubscript∇superscriptsubscriptℱψ^(i+1)←ψ^(i)+ _ψ^(i)F_d(D)ψ( i + 1 ) ← ψ( i ) + ∇ψ( i ) Fitalic_d ( D ); 5 6 else 7 ϕ(i+1)←ϕ(i)+∇ϕ(i)ℱp()+∇ϕ(i)ℱc()←superscriptitalic-ϕ1superscriptitalic-ϕsubscript∇superscriptitalic-ϕsubscriptℱsubscript∇superscriptitalic-ϕsubscriptℱφ^(i+1)←φ^(i)+ _φ^(i)F_p(D% )+ _φ^(i)F_c(D)ϕ( i + 1 ) ← ϕ( i ) + ∇ϕ( i ) Fitalic_p ( D ) + ∇ϕ( i ) Fitalic_c ( D ); 8 ψ(i+1)←ψ(i)+∇ψ(i)ℱp()+∇ψ(i)ℱc()←superscript1superscriptsubscript∇superscriptsubscriptℱsubscript∇superscriptsubscriptℱψ^(i+1)←ψ^(i)+ _ψ^(i)F_p(D% )+ _ψ^(i)F_c(D)ψ( i + 1 ) ← ψ( i ) + ∇ψ( i ) Fitalic_p ( D ) + ∇ψ( i ) Fitalic_c ( D ); 9 10 end if 11 12 end for return qϕ(T)subscriptsuperscriptitalic-ϕq_φ^(T)qitalic_ϕ( T ) Algorithm 1 AVA for reward modeling. The pipeline of AVA for reward modeling is shown in Alg. 1. In the TQR architecture, the reward function shares the same backbone model with the policy. Our purpose of reward modeling is to obtain an accurate and lightweight reward model. Therefore, we initialize the TQR architecture with a lightweight backbone model. In other words, we initialize the policy with a lightweight pre-trained language model πψ(1)subscriptsuperscript1 _ψ^(1)πitalic_ψ( 1 ) instead of a large language model. Meanwhile, the reward distribution is also initialized and denoted by qϕ(1)subscriptsuperscriptitalic-ϕ1q_φ^(1)qitalic_ϕ( 1 ) After the initialization, we leverage either AVA-d or AVA-p training objectives to optimize the reward distribution according to the type of the dataset DD. Note that the AVA training objectives require us to jointly train the reward function with the policy, although finally we only need the reward function. After reward modeling, we can leverage RL algorithms to fine-tune the LLM policy πwsubscript _wπitalic_w to maximize the expected reward produced by the trained reward distribution qϕsubscriptitalic-ϕq_φqitalic_ϕ, as shown in Eq. 19. J(w)=∼πw(⋅)[∑t=1||−1R∼qϕ(⋅|1:t+1)[R]]J(w)=E_y _w(·) [ _t=1^|y|-1% E_R q_φ(·|y_1:t+1)[R] ]J ( w ) = blackboard_Ey ∼ π start_POSTSUBSCRIPT w ( ⋅ ) end_POSTSUBSCRIPT [ ∑t = 1| y | - 1 blackboard_ER ∼ q start_POSTSUBSCRIPT ϕ ( ⋅ | y1 : t + 1 ) end_POSTSUBSCRIPT [ R ] ] (19) AVA for Direct Optimization The pipeline of AVA for direct optimization is shown in Alg. 2. Here, we directly initialize the policy with the pre-trained LLM πw(1)subscriptsuperscript1 _w^(1)πitalic_w( 1 ) and leverage the AVA training objectives to jointly optimize the policy and the reward distribution qϕ(1)subscriptsuperscriptitalic-ϕ1q_φ^(1)qitalic_ϕ( 1 ). After training, the LLM policy and the reward distribution are both aligned with the demonstration or preference dataset DD. Data: Dataset DD, initial LLM policy πw(1)subscriptsuperscript1 _w^(1)πitalic_w( 1 ), initial reward distribution qϕ(1)subscriptsuperscriptitalic-ϕ1q_φ^(1)qitalic_ϕ( 1 ), training epochs T Result: The finally trained LLM policy πw(T)subscriptsuperscript _w^(T)πitalic_w( T ) 1 for i∈1,⋯,T1⋯i∈\1,·s,T\i ∈ 1 , ⋯ , T do 2 if DD is demonstration dataset then 3 ϕ(i+1)←ϕ(i)+∇ϕ(i)ℱd()←superscriptitalic-ϕ1superscriptitalic-ϕsubscript∇superscriptitalic-ϕsubscriptℱφ^(i+1)←φ^(i)+ _φ^(i)F_d(D)ϕ( i + 1 ) ← ϕ( i ) + ∇ϕ( i ) Fitalic_d ( D ); 4 w(i+1)←w(i)+∇w(i)ℱd()←superscript1superscriptsubscript∇superscriptsubscriptℱw^(i+1)← w^(i)+ _w^(i)F_d(D)w( i + 1 ) ← w( i ) + ∇w( i ) Fitalic_d ( D ); 5 6 else 7 ϕ(i+1)←ϕ(i)+∇ϕ(i)ℱp()+∇ϕ(i)ℱc()←superscriptitalic-ϕ1superscriptitalic-ϕsubscript∇superscriptitalic-ϕsubscriptℱsubscript∇superscriptitalic-ϕsubscriptℱφ^(i+1)←φ^(i)+ _φ^(i)F_p(D% )+ _φ^(i)F_c(D)ϕ( i + 1 ) ← ϕ( i ) + ∇ϕ( i ) Fitalic_p ( D ) + ∇ϕ( i ) Fitalic_c ( D ); 8 w(i+1)←w(i)+∇w(i)ℱp()+∇w(i)ℱc()←superscript1superscriptsubscript∇superscriptsubscriptℱsubscript∇superscriptsubscriptℱw^(i+1)← w^(i)+ _w^(i)F_p(D)+% _w^(i)F_c(D)w( i + 1 ) ← w( i ) + ∇w( i ) Fitalic_p ( D ) + ∇w( i ) Fitalic_c ( D ); 9 10 end if 11 12 end for 13return πw(T)subscriptsuperscript _w^(T)πitalic_w( T ); Algorithm 2 AVA for Direct Optimization. Experiment Experiment Setup Datasets For preference datasets, we consider Anthropic-Harmless, Anthropic-Helpful, and OpenAI-Summary and perform reward modeling, RL fine-tuning, and direct optimization on these datasets. For demonstration datasets, we consider Alpaca-GPT-4 and Math-GPT-4o and only perform direct optimization on these datasets. Metrics For reward modeling, we evaluate the accuracy at which the reward of the chosen sentence is greater than that of the rejected sentence, as well as the win rates of the Best-of-N sampling (Stiennon et al. 2020; Nakano et al. 2021) results. For RL fine-tuning, we evaluate the win rates of the LLMs fine-tuned with different reward models (i.e., AVA-p/d and baselines). For direct optimization, we evaluate the win rates of LLMs fine-tuned with AVA-p/d against LLMs fine-tuned with baseline approaches. Pre-trained Models For reward modeling, we initialize the implicit policy with GPT-2 (117M) and BART-base (140M) to see the reward modeling performance with different initializations. For RL fine-tuning and direct optimization, we initialize the LLM policy with Llama-2-7b-chat-hf. The reward models adopted in RL fine-tuning only involve those initialized with GPT-2. Baselines For reward modeling, we adopt Bradley-Terry (Bradley and Terry 1952) and Preference Transformer (Pref-Trans) (Kim et al. 2023) as baselines. For direct optimization from preference, we adopt DPO (Rafailov et al. 2024) as the baseline. For direct optimization from demonstration, we adopt AfD (Sun and van der Schaar 2024) as the baseline. Since AfD constructs preference data from demonstration data and relies on preference-based training objectives, we combine AfD with different preference-based training objectives. Specifically, for reward modeling, we construct AfD w/ Bradley-Terry, Afd w/ Pref-Trans, and AfD w/ AVA-p. For direct optimization, we construct AfD w/ DPO. For win rate evaluations of aligned LLMs, we also adopt supervised fine-tuning (SFT) as the baseline. Ablation Variants We construct the following variants of the AVA-p and AVA-d training objectives for ablation studies: • AVA-p/d w/o rwt: AVA-p/d without reward weighting, which removes the computation of reward weights and the weighted rewards from the TQR architecture. • AVA-p w/o neg: AVA-p without the negative demonstration, which removes the minimization of the likelihood of the negative demonstrations. Note that the objective does not completely degenerate into the AVA-d training objective since we still keep the CER auxiliary objective. • AVA-p w/o irl: AVA-p without inverse reinforcement learning, which removes the TD-error constraint and the reward prior assumption and only keeps the likelihood optimization, which can be regarded as contrastive supervised fine-tuning. • AVA-p w/o cer: AVA-p without CER auxiliary objective. • AVA-p/d w/o ptq: AVA-p/d without pre-trained Q-value head, which does not reuse the LM head of the pre-trained policy as the Q-value head but initializes the Q-value head from scratch. For detailed experiment setup, please refer to our code and the Experiment Details section of the Technical Appendix. Reward Modeling Harmless Helpful Summary gpt2 bart gpt2 bart gpt2 bart Baselines Bradley-Terry 70.02 68.96 69.39 67.56 59.27 59.27 Pref-Trans 70.26 71.32 71.37 72.37 59.31 56.91 Ours AVA-p 70.27 72.30 72.37 74.84 61.79 64.31 AVA-p w/o rwt 70.06 70.73 69.81 69.32 60.55 58.89 AVA-p w/o neg 70.54 70.36 69.75 69.15 62.06 58.65 AVA-p w/o irl 69.48 67.46 68.87 65.38 58.96 58.46 AVA-p w/o cer 70.06 70.73 69.81 69.32 60.55 58.58 AVA-p w/o ptq 68.67 68.69 68.51 67.60 61.25 57.76 AVA-d 70.54 70.36 69.75 69.15 62.06 59.00 Table 1: Reward accuracy of AVA and baseline training objectives. Table 1 reports the reward accuracy of baseline and AVA training objectives. The results show that AVA-p surpasses Bradley-Terry and Pref-Trans in reward accuracy on all reported reward modeling tasks with different initial models and datasets. The ablation results further reveal that AVA-p achieves the highest reward accuracy on the greatest number of tasks compared to ablated training objectives, which suggests that removing any module from AVA-p diminishes the reward accuracy on most tasks. Furthermore, we consider the chosen half of the preference data as demonstration data and train the reward model on it using the AVA-d training objective. Surprisingly, AVA-d achieves the best performance on 2 out of 6 tasks, despite learning solely from the chosen demonstrations. To further evaluate reward modeling performance, we employ Best-of-N (BoN) sampling. We evaluate the win rates of BoN w/ AVA-p against BoN w/ Bradley-Terry and BoN w/ Pref-Trans, where “BoN w/ x” means that the reward model used for BoN is trained with the “x” training objective. Additionally, we evaluate the win rate of BoN w/ AVA-p against the stochastic sampling results without BoN. Table 2 reports the win rates of the reward model trained with AVA-p against reward models trained with baseline objectives in BoN sampling. The results further demonstrate that AVA-p surpasses Bradley-Terry and Pref-Trans in reward modeling. Task Opponent Win↑ Tie Lose↓ Harmless Stochastic 43.0 17.4 39.6 BoN w/ Bradley-Terry 28.8 42.6 28.6 BoN w/ Pref-Trans 35.6 31.9 32.5 Helpful Stochastic 26.1 50.4 23.5 BoN w/ Bradley-Terry 13.2 76.1 10.7 BoN w/ Pref-Trans 19.3 62.3 18.4 Summary Stochastic 60.2 0.8 39.0 BoN w/ Bradley-Terry 34.6 34.5 30.9 BoN w/ Pref-Trans 43.6 25.8 30.6 Table 2: Win rates of BoN with AVA-p reward model. RL Fine-tuning Task Opponent Win↑ Tie Lose↓ Harmless SFT 42.5 23.4 34.1 PPO w/ Bradley-Terry 9.2 81.7 9.1 PPO w/ Pref-Trans 9.0 83.2 7.8 Helpful SFT 23.3 58.8 18.0 PPO w/ Bradley-Terry 1.8 97.2 1.0 PPO w/ Pref-Trans 2.6 95.8 1.6 Summary SFT 73.8 1.4 24.7 PPO w/ Bradley-Terry 18.5 66.3 15.2 PPO w/ Pref-Trans 33.9 34.6 31.5 Table 3: Win rates of PPO with AVA-p reward model. We adopt the PPO algorithm (Schulman et al. 2017) to fine-tune LLMs to maximize the reward produced by different reward models. We evaluate the win rates of PPO w/ AVA-p against PPO w/ Bradley-Terry and PPO w/ Pref-Trans, where “PPO w/ x” means that the reward model used for PPO fine-tuning is trained with the “x” training objective. We also evaluate the win rate of PPO w/ AVA-p against supervised fine-tuning (SFT), where the LLM is fine-tuned on the chosen half of the preference data with supervised learning. The results in Table 3 show that AVA-p outperforms the baseline reward modeling objectives on all reported tasks in downstream RL fine-tuning of the LLM. Direct Optimization Task Opponent Win↑ Tie Lose↓ Harmless SFT 37.1 28.9 34.0 DPO 13.7 73.8 12.5 Helpful SFT 22.5 59.6 17.9 DPO 14.4 72.4 13.2 Summary SFT 59.0 7.3 33.7 DPO 44.9 11.0 44.1 Table 4: Win rates of direct optimization with AVA-p. From Preference We adopt AVA-p and DPO (Rafailov et al. 2024) to directly optimize the LLM from preference data and evaluate the win rates of AVA-p against DPO and SFT. The results in Table 4 show that AVA-p outperforms DPO in direct optimization from preference data. From Demonstration We adopt AVA-d and AfD (Sun and van der Schaar 2024) to directly optimize the LLM from demonstration data. We evaluate the win rates of AVA-d against SFT, DPO w/ AfD, and AVA-p w/ AfD, where “x w/ AfD” means applying the “x” training objective on AfD-format data. The results in Table 5 show that AVA-d outperforms the AfD approaches in direct optimization from demonstration data. Moreover, AVA-d is more training-efficient since AfD requires supervised fine-tuning and sampling from LLM policies. Task Opponent Win↑ Tie Lose↓ Alpaca SFT 58.1 7.2 34.7 DPO w/ AfD 57.2 6.9 35.9 AVA-p w/ AfD 56.5 7.1 36.4 Math SFT 47.0 9.7 43.3 DPO w/ AfD 44.3 11.4 44.3 AVA-p w/ AfD 45.4 11.4 43.1 Table 5: Win rates of direct optimization with AVA-d. Conclusion We present AVA, a flexible novel LLM alignment objective with enhanced capabilities. The flexibility of AVA is evident in two aspects. Firstly, AVA can utilize either preference data or demonstration data for alignment purposes. Secondly, AVA can be integrated into the reward modeling and RL fine-tuning pipeline or used to directly optimize the LLM. The representation and generalization capabilities of AVA are also evident in two aspects. Theoretically, AVA formulates reward modeling as a BIRL problem, facilitating both intermediate reward modeling and direct reward modeling on demonstration. Experimentally, AVA achieves superior reward accuracy in reward modeling tasks and higher win rates in RL fine-tuning and direct optimization of LLMs, which demonstrates the alleviation of the reward hacking issue and improved alignment performance. References Achiam et al. (2023) Achiam, J.; Adler, S.; Agarwal, S.; Ahmad, L.; Akkaya, I.; Aleman, F. L.; Almeida, D.; Altenschmidt, J.; Altman, S.; Anadkat, S.; et al. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774. Bai et al. (2022) Bai, Y.; Jones, A.; Ndousse, K.; Askell, A.; Chen, A.; DasSarma, N.; Drain, D.; Fort, S.; Ganguli, D.; Henighan, T.; et al. 2022. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862. Bradley and Terry (1952) Bradley, R. A.; and Terry, M. E. 1952. Rank analysis of incomplete block designs: I. The method of paired comparisons. Biometrika, 39(3/4): 324–345. Chan and van der Schaar (2021) Chan, A. J.; and van der Schaar, M. 2021. Scalable Bayesian Inverse Reinforcement Learning. In International Conference on Learning Representations. Coste et al. (2023) Coste, T.; Anwar, U.; Kirk, R.; and Krueger, D. 2023. Reward model ensembles help mitigate overoptimization. arXiv preprint arXiv:2310.02743. Gao, Schulman, and Hilton (2023) Gao, L.; Schulman, J.; and Hilton, J. 2023. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, 10835–10866. PMLR. Geva et al. (2023) Geva, M.; Bastings, J.; Filippova, K.; and Globerson, A. 2023. Dissecting recall of factual associations in auto-regressive language models. arXiv preprint arXiv:2304.14767. Hanna, Liu, and Variengien (2024) Hanna, M.; Liu, O.; and Variengien, A. 2024. How does GPT-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model. Advances in Neural Information Processing Systems, 36. Holtzman et al. (2019) Holtzman, A.; Buys, J.; Du, L.; Forbes, M.; and Choi, Y. 2019. The curious case of neural text degeneration. arXiv preprint arXiv:1904.09751. Kim et al. (2023) Kim, C.; Park, J.; Shin, J.; Lee, H.; Abbeel, P.; and Lee, K. 2023. Preference transformer: Modeling human preferences using transformers for RL. arXiv preprint arXiv:2303.00957. Mnih et al. (2013) Mnih, V.; Kavukcuoglu, K.; Silver, D.; Graves, A.; Antonoglou, I.; Wierstra, D.; and Riedmiller, M. 2013. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602. Nakano et al. (2021) Nakano, R.; Hilton, J.; Balaji, S.; Wu, J.; Ouyang, L.; Kim, C.; Hesse, C.; Jain, S.; Kosaraju, V.; Saunders, W.; et al. 2021. Webgpt: Browser-assisted question-answering with human feedback. arXiv preprint arXiv:2112.09332. Ng, Russell et al. (2000) Ng, A. Y.; Russell, S.; et al. 2000. Algorithms for inverse reinforcement learning. In Icml, volume 1, 2. OpenAI (2022) OpenAI. 2022. Introducing ChatGPT. https://openai.com/index/chatgpt/. Ouyang et al. (2022) Ouyang, L.; Wu, J.; Jiang, X.; Almeida, D.; Wainwright, C.; Mishkin, P.; Zhang, C.; Agarwal, S.; Slama, K.; Ray, A.; et al. 2022. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730–27744. Rafailov et al. (2024) Rafailov, R.; Sharma, A.; Mitchell, E.; Manning, C. D.; Ermon, S.; and Finn, C. 2024. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36. Ramachandran and Amir (2007) Ramachandran, D.; and Amir, E. 2007. Bayesian Inverse Reinforcement Learning. In IJCAI, volume 7, 2586–2591. Ranzato et al. (2016) Ranzato, M.; Chopra, S.; Auli, M.; and Zaremba, W. 2016. Sequence level training with recurrent neural networks. In 4th International Conference on Learning Representations, ICLR 2016. Schulman et al. (2017) Schulman, J.; Wolski, F.; Dhariwal, P.; Radford, A.; and Klimov, O. 2017. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. Shen et al. (2023) Shen, T.; Jin, R.; Huang, Y.; Liu, C.; Dong, W.; Guo, Z.; Wu, X.; Liu, Y.; and Xiong, D. 2023. Large language model alignment: A survey. arXiv preprint arXiv:2309.15025. Skalse et al. (2022) Skalse, J.; Howe, N.; Krasheninnikov, D.; and Krueger, D. 2022. Defining and characterizing reward gaming. Advances in Neural Information Processing Systems, 35: 9460–9471. Song et al. (2024) Song, F.; Yu, B.; Li, M.; Yu, H.; Huang, F.; Li, Y.; and Wang, H. 2024. Preference ranking optimization for human alignment. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, 18990–18998. Stiennon et al. (2020) Stiennon, N.; Ouyang, L.; Wu, J.; Ziegler, D.; Lowe, R.; Voss, C.; Radford, A.; Amodei, D.; and Christiano, P. F. 2020. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33: 3008–3021. Sun and van der Schaar (2024) Sun, H.; and van der Schaar, M. 2024. Inverse-RLignment: Inverse Reinforcement Learning from Demonstrations for LLM Alignment. arXiv preprint arXiv:2405.15624. Tay et al. (2020) Tay, Y.; Ong, D.; Fu, J.; Chan, A.; Chen, N.; Luu, A. T.; and Pal, C. 2020. Would you rather? a new benchmark for learning machine alignment with cultural values and social preferences. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, 5369–5373. Touvron et al. (2023) Touvron, H.; Martin, L.; Stone, K.; Albert, P.; Almahairi, A.; Babaei, Y.; Bashlykov, N.; Batra, S.; Bhargava, P.; Bhosale, S.; et al. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017. Attention is all you need. Advances in neural information processing systems, 30. Weidinger et al. (2021) Weidinger, L.; Mellor, J.; Rauh, M.; Griffin, C.; Uesato, J.; Huang, P.-S.; Cheng, M.; Glaese, M.; Balle, B.; Kasirzadeh, A.; et al. 2021. Ethical and social risks of harm from language models. arXiv preprint arXiv:2112.04359. Yang et al. (2023) Yang, A.; Xiao, B.; Wang, B.; Zhang, B.; Bian, C.; Yin, C.; Lv, C.; Pan, D.; Wang, D.; Yan, D.; et al. 2023. Baichuan 2: Open large-scale language models. arXiv preprint arXiv:2309.10305. Yuan et al. (2023) Yuan, Z.; Yuan, H.; Tan, C.; Wang, W.; Huang, S.; and Huang, F. 2023. Rrhf: Rank responses to align language models with human feedback without tears. arXiv preprint arXiv:2304.05302. Zhang et al. (2024) Zhang, X.; Ton, J.-F.; Shen, W.; Wang, H.; and Liu, Y. 2024. Overcoming reward overoptimization via adversarial policy optimization with lightweight uncertainty estimation. arXiv preprint arXiv:2403.05171. Zhang et al. (2019) Zhang, Y.; Sun, S.; Galley, M.; Chen, Y.-C.; Brockett, C.; Gao, X.; Gao, J.; Liu, J.; and Dolan, B. 2019. Dialogpt: Large-scale generative pre-training for conversational response generation. arXiv preprint arXiv:1911.00536. Zopf (2018) Zopf, M. 2018. Estimating summary quality with pairwise preferences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), 1687–1696. Reproducibility Checklist This paper: • Includes a conceptual outline and/or pseudocode description of AI methods introduced (yes) • Clearly delineates statements that are opinions, hypothesis, and speculation from objective facts and results (yes) • Provides well marked pedagogical references for less-familiare readers to gain background necessary to replicate the paper (yes) Does this paper make theoretical contributions? (yes) If yes, please complete the list below. • All assumptions and restrictions are stated clearly and formally. (yes) • All novel claims are stated formally (e.g., in theorem statements). (yes) • Proofs of all novel claims are included. (yes) • Proof sketches or intuitions are given for complex and/or novel results. (yes) • Appropriate citations to theoretical tools used are given. (yes) • All theoretical claims are demonstrated empirically to hold. (yes) • All experimental code used to eliminate or disprove claims is included. (yes) Does this paper rely on one or more datasets? (yes) If yes, please complete the list below. • A motivation is given for why the experiments are conducted on the selected datasets (NA) • All novel datasets introduced in this paper are included in a data appendix. (NA) • All novel datasets introduced in this paper will be made publicly available upon publication of the paper with a license that allows free usage for research purposes. (NA) • All datasets drawn from the existing literature (potentially including authors’ own previously published work) are accompanied by appropriate citations. (yes) • All datasets drawn from the existing literature (potentially including authors’ own previously published work) are publicly available. (yes) • All datasets that are not publicly available are described in detail, with explanation why publicly available alternatives are not scientifically satisficing. (NA) Does this paper include computational experiments? (yes) If yes, please complete the list below. • Any code required for pre-processing data is included in the appendix. (yes). • All source code required for conducting and analyzing the experiments is included in a code appendix. (yes) • All source code required for conducting and analyzing the experiments will be made publicly available upon publication of the paper with a license that allows free usage for research purposes. (yes) • All source code implementing new methods have comments detailing the implementation, with references to the paper where each step comes from (yes) • If an algorithm depends on randomness, then the method used for setting seeds is described in a way sufficient to allow replication of results. (yes) • This paper specifies the computing infrastructure used for running experiments (hardware and software), including GPU/CPU models; amount of memory; operating system; names and versions of relevant software libraries and frameworks. (yes) • This paper formally describes evaluation metrics used and explains the motivation for choosing these metrics. (yes) • This paper states the number of algorithm runs used to compute each reported result. (yes) • Analysis of experiments goes beyond single-dimensional summaries of performance (e.g., average; median) to include measures of variation, confidence, or other distributional information. (no) • The significance of any improvement or decrease in performance is judged using appropriate statistical tests (e.g., Wilcoxon signed-rank). (no) • This paper lists all final (hyper-)parameters used for each model/algorithm in the paper’s experiments. (yes) • This paper states the number and range of values tried per (hyper-) parameter during development of the paper, along with the criterion used for selecting the final parameter setting. (NA)