Getting More Juice Out of the SFT Data: Reward Learning from Human Demonstration Improves SFT for LLM Alignment

Authors: jiaxiang  li
siliang  zeng
hoi-to  wai
chenliang  li
alfredo  garcia
mingyi  hong
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:20:00

Abstract: Aligning human preference and value is an important requirement for contemporary foundation models. State-of-the-art techniques such as Reinforcement Learning from Human Feedback (RLHF) often consist of two stages: 1) supervised fine-tuning (SFT), where the model is fine-tuned by learning from human demonstration data; 2) Preference learning, where preference data is used to learn a reward model, which is in turn used by a reinforcement learning (RL) step to fine-tune the model. Such reward model serves as a proxy to human preference, and it is critical to guide the RL step towards improving the model quality. In this work, we argue that the SFT stage significantly benefits from learning a reward model as well. Instead of using the human demonstration data directly via supervised learning, we propose to leverage an Inverse Reinforcement Learning (IRL) technique to {\it simultaneously} build an reward model and a policy model. This approach leads to new SFT algorithms that are not only efficient to implement, but are robust to the presence of low-quality supervised learning data. Moreover, we discover a connection between the proposed IRL based approach, and a recent line of works called Self-Play Fine-tune (SPIN, \cite{chen2024self}). Theoretically, we show that the proposed algorithms converge to the stationary solutions of the IRL problem. Empirically, we align 1B and 7B models using proposed methods and evaluate them on a reward benchmark model and the HuggingFace Open LLM Leaderboard. The proposed methods show significant performance improvement over existing SFT approaches. Our results indicate that it is beneficial to leverage reward learning throughout the entire alignment process. Our code is available at \url{https://github.com/JasonJiaxiangLi/Reward_learning_SFT}.

Main paper: 
 introduction: 
 Large Language Models (LLMs) have become the cornerstone of modern artificial intelligence applications. They are believed to lead the way towards artificial general intelligence [Bubeck et al., 2023] , also have shown great capabilities towards specialized domains such as math problem solving [Cobbe et al., 2021 , Trinh et al., 2024 , Wei et al., 2022 , Lewkowycz et al., 2022] , code generation [Chen et al., 2021 , Austin et al., 2021 , Li et al., 2022 ], text generation [Anil et al., 2023 , Touvron et al., 2023 , Thoppilan et al., 2022] , etc. Usually, researchers need to align the pre-trained LLMs with certain exquisitely prepared human-labeled data to achieve desired performance over certain tasks, a process which is thus known as alignment or fine-tuning. The alignment datasets can be categorized into two classes: the demonstration data, with the input prompt and a human response; and the preference data, with the input prompt and two responses, where human labeler will pick a chosen one and a rejected one. With the alignment datasets, one could employ methods like supervised fine-tune (SFT, Ouyang et al. [2022] , Tunstall et al. [2023] , Chung et al. [2024] ) for aligning demonstration datasets, and reinforcement learning from human feedback (RLHF, Christiano et al. [2017] , Ouyang et al. [2022] ) and direct preference optimization (DPO, Rafailov et al. [2024] ) for aligning preference datasets. More specifically, RLHF explicitly trains a reward model and uses reinforcement learning (in particular, policy optimization) methods to obtain a fine-tuned version of the LLM; on the other hand, DPO and many of its extensions simplifies the RLHF by training the LLM policy model directly, while implicitly learns the reward model via log of the ratio of likelihood between the learned model and a reference model. In practice, both types of methods exhibit better performance over SFT on the demonstration datasets, and they are adopted by state-of-the-art LLMs, for example ChatGPT benefited from RLHF (see Ouyang et al. [2022] ), zephyr benefited from DPO (see Tunstall et al. [2023] ).It is interesting to observe that, when dealing with preference data, state-of-the-art methods usually build an (explicit or implicit) reward model to evaluate the quality of responses for a given prompt. On the contrary, typically no reward modeling is done for demonstration datasets. Why this is the case? One may argue that, for a given set of prompts, preference datasets contain explicit preference information which is not found in the demonstration datasets; since this kind of information is harder to extract, it motivates the use of complicated methods such as reward modeling. However, since human preferences are also implicit in the demonstration data, one can argue that training a reward model that encodes human value distilled from these datasets may help to boost the alignment capability of the LLM. Indeed, in the RL literature, it is known that if the agents are given a set of demonstration data, then the so-called inverse RL methods (which learns the reward and policy simultaneously) can outperform the behavior cloning methods (which corresponds to supervised fine-tune in LLM alignments) by a large margin. In a Markov decision process (MDP), it is likely that supervised learning methods which naively fit the demonstration data will suffer from the distribution shift problem -the fine-tuned policy from supervised learning can produce unsatisfactory generations in certain states which were unseen in the training dataset [Ross et al., 2011] . Through formulating the learning from demonstration problem in a MDP setting, typical inverse reinforcement learning methods [Ziebart et al., 2008 , Ross et al., 2011 , Zeng et al., 2022] can alleviate such distribution shift issues. Witnessing the success in ChatGPT, where the alignment of LLMs is modelled in the MDP setting due to the auto-regressive process, one would expect that the LLM alignment with demonstration datasets can be improved as well through deploying imitation learning / inverse reinforcement learning methods.Inspired by the above observation, we pose the following question: Does building a reward model using the demonstration data benefit the alignment process?Contribution of this work. This paper answers the above question affirmatively. Specifically, by developing a framework based on certain IRL technique, we show that building a reward from demonstration datasets can significantly improve the quality of the resulting model, as compared to that obtained by standard reward-free SFT (1). Our main contributions are listed as below:• We develop a new reward-based SFT approach, which takes the form of a bilevel optimization, where in the lower-level, LLM policy is learned via policy optimization for a given reward, while in the upper-level, the reward model is optimized so to maximize the likelihood for observing the demonstration data. • Based on the above formulation, we propose two alignment algorithms, one learns the reward model explicitly, and the other implicitly. For the first algorithm, we show that the reward learned from only demonstration data already possesses strong capabilities in distinguishing between chosen and rejected responses; see Figure 1 and our experiment for details. For the second algorithm, we made an interesting observation that implicitly learning a reward is equivalent to improving the model by comparing the demonstration data with the synthetic data generated by the past models. Somewhat surprisingly, the resulting algorithm is closely related to the self-play fine-tune (SPIN, Chen et al. [2024] ) algorithm, recently proposed from a completely different viewpoint. It is worth pointing out that unlike SPIN, our proposed algorithms have finite-time convergence guarantees.• We demonstrate the power of the proposed approach theoretically and numerically. We prove that our implicit reward learning algorithm converges to some stationary point of our proposed formulation. We show that the proposed algorithms outperform vanilla SFT in almost all cases we have tested, for example the model performance on HuggingFace Open LLM Leaderboard increases from 59.47% to 61.03%. To our knowledge, this is the first work that formally demonstrate the power of reward learning when dealing with demonstration data for LLM alignment. Notations. We use π(y|x) to denote the LLM output probability for continuation y with input prompt x, and we refer to π as the policy. We use the notation π(y|x; θ) if the model π is directly parameterized by parameters θ. For the case when π is indirectly determined by parameter θ, we use notation π θ (y|x). We use D = {(x, y)} to denote the demonstration dataset and P = {(x, y w , y l )} for the preference dataset, where y w is preferred over y l . Since we assume that the demonstration continuations y are collected from a human expert distribution, we also denote (x, y) ∼ D as x ∼ ρ, y ∼ π E (•|x) when taking the expectations, where ρ is the distribution of the input prompts when collecting the data. We similarly have the notation x ∼ ρ, (y l ≺ y w ) ∼ π P (•|x) for the preference dataset. 
preliminaries: 
 Consider a Large Language Model (LLM) parameterized by θ and denote the output probability by π(y|x; θ) where x = [x 1 , ..., x n ] is the sequence of input prompts and y = [y 1 , ..., y m ] is the sequence of output continuation. Typical LLM is an auto-regressive model, meaning that it predicts the output probability of the y j given all tokens in x and y <j := [y 1 , ..., y j-1 ] (y <1 is null), i.e.π(y|x; θ) = m j=1 π(y j |x, y <j ; θ).In this paper, we do not focus on the architecture design of LLMs. We will fix the LLM architecture and always denote it as a probability model π(y|x; θ). The following discussions review two common procedures for fine-tuning θ: (1) supervised fine tuning (SFT) over demonstration dataset, (2) reinforcement learning with human feedback (RLHF) over preference dataset that consists of two steps: LLM alignment/fine-tuning based on a reward model using policy optimization; and reward learning process to learn the optimal reward for the preference dataset.SFT. Given a demonstration dataset D := {(x, y)}, the SFT optimizes the following problem:EQUATIONIt is easy to see that the above problem shares the same optimal solutions asmin θ E x∼ρ [D KL (π E (•|x) ∥π (•|x; θ))].The latter shows that SFT aims at imitating the demonstration dataset via minimizing the KL divergence. It is worth noting that the SFT stage described here is closely related to the imitation learning approach used in the RL literature for learning from demonstration [Osa et al., 2018] , whose goal is to mimic the policy of an expert.RLHF. Suppose that we have a reward model r(x, y; ϕ) (parameterized by ϕ and to be defined later) for any given input and output pair (x, y), the LLM can be fine tuned by the following RL problem:EQUATIONwhere π ref is a fixed reference model. Note that the KL regularization term in (2) is not computable given the sheer amount of possible output y (which could be corpus_size max_sequence_length in most language model tasks), therefore (2) is usually solved by standard policy optimization techniques such as REINFORCE [Ahmadian et al., 2024] or PPO [Schulman et al., 2017] .To find an appropriate reward model r(x, y; ϕ), RLHF (see e.g., Christiano et al. [2017] ) leverages a set of preference dataset P := {(x, y w , y l )}, where each data contains a pair of output y w , y l , and y w is preferred over y l by human labeler (denoted as y w ≻ y l ). The Bradley-Terry model [Bradley and Terry, 1952] assumes that the probability of choosing y w over y l isP (y w ≻ y l | x) = exp(r(y w ; x)) exp(r(y w ; x)) + exp (r (y l ; x)) = σ (r(y w ; x) -r (y l ; x)) .One could formulate the following problem to find the reward model:EQUATION)It is widely observed in the literature that, models trained via episodically learning the policy (2) and learning the reward (3) typically outperforms those that are only trained using SFT [Ouyang et al., 2022] . The reward model guides the performance of the LLM and allows a better generalization ability via the consistent input of the preference data from human labeler. Follow up works such as DPO proposes to incorporate reward learning implicitly by utilizing the structure of the optimal solution of the RL problem (2); for more details about the DPO, see Rafailov et al. [2024] .Discussion. At this point, let us take a step back and think about the above process. The LLM alignment problem takes human labeled demonstration and preference data to produce an aligned model. Clearly, both kinds of data encode information about how human would like the LLM output to be, but the processes of extracting such information is quite different (i.e., supervised learning vs RL). A series of questions naturally arises: Is supervised learning the best way to extract human inclination from the demonstration data? Can we also learn a reward model from the demonstration data to gauge human preference? Will policy model learned via such reward improve the supervised learning approach? In the next section, we will dive deep to carefully address these questions.3 Reward Learning and Policy Fine Tuning from Demonstration DataIn this section, we argue that reward learning from the demonstration dataset can benefit the LLM alignment problem. To do so, we develop a joint reward learning and policy fine tuning formulation and understand its capabilities in improving the LLM policy. The new formulation inspired us to develop two reward learning paradigms: i) Explicit reward learning, where a (parameterized) reward model is learned together with the language model policy, and ii) Implicit reward learning, where the reward model is learned implicitly through directly optimizing the policy, avoiding learning two models simultaneously.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We made it very clear that we focus on demonstration dataset. We also include a limitation section at the end.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
joint reward-learning and policy fine-tuning by inverse rl: 
 A challenge with learning only from the demonstration dataset is that the Bradley-Terry model (3) can no longer be used due to the lack of pairs of preference data. However, all is not lost as we recall that it is the value of the reward model that should be used in the fine-tuning process (2). Therefore, with only demonstration data D, a reasonable formulation is to combine the supervised learning problem (1) with the optimal policy generation problem (2), by requiring that the generated policy to 'match' with the demonstration dataset. With this intuition in mind, we consider the joint reward and policy learning problem via a maximum likelihood inverse reinforcement learning (ML-IRL) formulation [Ziebart et al., 2008 , 2013 , Zeng et al., 2022] :EQUATIONThe above problem has a bilevel structure which trains a reward model r (x, y; θ). At the upper level, its objective is similar to that of SFT (1), but is evaluated on the policy π θ induced by the reward model r (x, y; θ); meanwhile, this policy π θ is found in the lower level using the RL objective (2).There are several advantages of the bilevel formulation (4) over standard SFT (1). First, we notice formulating SFT as a RL / IRL problem can alleviate distribution shift and improve the generalization power [Ross et al., 2011] . In fact, we observe that (4) tends to give a less extreme policy even when the demonstration dataset is extreme. The latter is observed in the following stylized example.Example 1. Suppose we have only one state (input prompt) x and three actions (continuations) y 1 , y 2 , y 3 . Let the reference model π ref be a uniform distribution over all continuations, and the demonstration dataset is D = {y 3 }. One could easily compute the optimal solution for (1) and ( 4) by first-order optimality conditions. From Table 1 we can see that SFT (imitation learning) pushes all the likelihood toward the demonstration dataset, whereas ML-IRL (4) maintains non-zero weights for unseen data in the demonstration datasets. This is particular useful when we want to fine-tune from a pre-trained model, which is presumed to be powerful and have useful information already.Actiony 1 y 2 y 3 π ref 0.33 0.33 0.33 D {y 3 } π SFT 0.0 0.0 1.0 π IRL 2 2+e R/β 2 2+e R/β e R/β 2+e R/βTable 1: A state-less counter-example with three actions where IRL-based fine-tune (4) shows regularization effect over SFT (1) to maintain weights over unseen data in the demonstration dataset D.Here we assume r ∈ [0, R].Second, since the lower level problem in (4) encapsulates a generation process, it is anticipated that the proposed method can better distinguish between the preferred and non-preferred data than SFT, even if it is only trained on the demonstration dataset. The following numerical example highlights this point:Example 2. We compare the solution of SFT (1) and IRL (4) numerically, where the latter is solved using two algorithms RFT and IRFT (to be introduced shortly). We choose the preference based dataset Anthropic-HH and only keep the preferred continuation to form a demonstration dataset D = {(x, y w )} to implement SFT and IRL. We then compute the log probability gap log(π(y w |x)) -log(π(y l |x)) between the preferred y w and non-preferred y l on the test dataset; see Figure 1 right side. We observe that although all three methods are not exposed to the nonpreferred data y l during the training process, the IRL-based methods effectively distinguish the preferred continuation over the non-preferred one, while SFT assigns larger probability to the nonpreferred continuation (see Section 5 for the details of the implementation).Comparing to SFT (1), the bilevel problem (4) appears to be more complicated. In particular, solving standard bilevel optimization problem typically involves computation of Hessian matrices, which is too expensive for LLM related applications [Liu et al., 2023] . Fortunately, in our next result, we show that the bilevel problem can be significantly simplified (proof in Appendix B): Lemma 3.1. Problem (4) is equivalent to the following minimax optimization problem:max θ min π E x∼ρ,y∼π E (•|x),ỹ∼π(•|x) r(x, y; θ) -r(x, ỹ; θ) β + D KL π(•|x)∥π ref (•|x) . (5)The above reformulation is remarkable. First, minimax problem is much easier to solve as compared with bilevel problem, e.g., a simple alternating minimization can yield reasonably good solution; see Algorithm 1 for such an algorithm, and Sec. 3.3 for its theoretical analysis. More importantly, it shows that even only the demonstration data is available, the reward optimization problem takes a similar form as what has been used in RLHF (3), where not one but two reward functions are contrasted. The key difference here is that one reward is evaluated on the continuation y in D, the other is evaluated on ỹ, which is the continuation generated from the current policy π(•|x). We believe that such contrast is the key reason that enables the IRL based formulation to distinguish the preferred continuations over the non-preferred ones; see Example 2 and Figure 1 .Now that we have turned the original bilevel problem (4) into a minimax optimization problem (5), we can naturally develop a gradient-descent-ascent type algorithm for (5), which alternates between updating the policy according to the current reward, and updating the reward based on the current policy -an algorithm that we call Reward-learning Fine-tune (RFT), see Algorithm 1. Note that in Algorithm 1: Reward-learning Fine-Tune (RFT) Input: Initialize reward parameter θ 0 (θ -1,K = θ 0 ) and policy model π 0 , the stepsize of reward update η t , and T , K the outer and inner iterations.for t = 0, 1, . . . , T -1 do Take θ t,0 = θ t-1,K Data Sample: Sample state x t,k ∼ ρ, an expert response y t,k ∼ π E (•|x t,k ) and agent response ỹt,k ∼ π t (•|x t,k ), for k = 0, 1, ..., K -1 for k = 0, 1, ..., K -1 do Estimate Gradient: Calculate the stochastic gradient g t,k w.r.t. θ via g t,k = 1 β ∇ θ r(x t,k , y t,k ; θ t,k ) -1 β ∇ θ r(x t,k , ỹt,k ; θ t,k ) Reward Alignment: θ t,k+1 := θ t,k + η t g t,kend for Policy Alignment: Update the optimal π t (y|x) ∝ exp(r(x, y; θ t,K )) according to (9) end for the data sampling step, we sample the response from the current model for the next K iterations. If we take K = 1 and T = data size batch size * epoch, the sampling process would be done for every iteration. In practice however we take a relative small T and large K, because frequent on-line sampling is time consuming; see Section 4 for the implementation details.So far we have seen that (4) (equivalently ( 5)) can efficiently utilize the demonstration dataset for better alignment. However, the computation cost for training two models (reward and policy) is significantly higher than the standard SFT. It turns out that (4) can be simplified into a supervised learning problem. Observe the following property (see Appendix B for proof): Lemma 3.2. For the loss function ℓ in (4), we have:EQUATION)The proof of the above lemma uses the identity r(x, y; θ) = β log π θ (y|x) π ref (y|x) + β log Z θ (x) for some constants Z θ (x); see, e.g. Rafailov et al. [2024] . Again it is remarkable that, despite the fact that the IRL formulation only consumes the demonstration data D = x, y, the gradient of the IRL loss takes the form as the difference of two gradients, one related to the demonstration data, the other related to the data generated by the current policy. Lemma 3.2 leads to a simple scheme for implicit reward-based supervised fine-tune (IRFT) -for each training batch, it samples the response from the current model, and construct the gradient estimator (6) directly to update the parameters θ. This results in Algorithm 2, which is an SGD type algorithm for (4). In Algorithm 2, we use a double loop since generation at each step might again significantly take more time, similar to Algorithm 1. If K = 1 we get a single loop algorithm where we generate for every training step on the input batch.We conclude the section by theoretically inspecting the proposed algorithms. Note that details and proofs of convergence theorem are moved to Appendix B due to page limits. We observe:Theorem 3.1. Under Assumption B.1, for Algorithm 1 and 2 with η t = Θ(1/ √ T K) we have min t=1,...,T, k=1,...,K E[∥∇ℓ(θ t,k )∥ 2 ] ≤ O 1/ √ T K + 1/T . Theorem 3.1 indicates that the convergence dependency is O(1/ √ KT ) (assuming T > K), which indicates that the algorithm could converge to stationary point if we take both the inner loop and outer loop reasonably large. This is slightly contrary to the intuition since with larger inner loop number K, we are having more biased estimators. This theorem shows that this biasedness actually wouldn't harm the final convergence, thus validate our practice of having a relative large inner loop number K in practice (since generating at each training iteration is time-consuming).Algorithm 2: Implicit Reward-learning Fine-Tune (IRFT) 1: Input: Initialize model parameter θ 0 (θ -1,K = θ 0 ), the stepsize of reward update η t , and T , K the outer and inner iterations. 2: Output:θ 3: for t = 0, 1, ..., T -1 do 4: Take θ t,0 = θ t-1,K 5: Data Sample: Sample state x t,k ∼ ρ, an expert response y t,k ∼ π E (•|x t,k ) and agent response ỹt,k ∼ π θt,0 (•|x t,k ), for k = 0, 1, ..., K -1 6: for k = 0, 1, ..., K -1 do 7: Estimate Gradient: Calculate the stochastic estimator ∇ℓ(θ t,k ) via (6), i.e. ∇ℓ(θ t,k ) = ∇ θ t,k log π θ t,k (y t,k |x t,k ) π ref (y t,k |x t,k ) -∇ θ t,k log π θ t,k (ỹ t,k |x t,k ) π ref (ỹ t,k |x t,k ) . 8: Implicit Reward Alignment: Update θ t,k+1 = θ t,k + η t ∇ℓ(θ t,k ) 9:end for 10: end for 4 Discussions Implementation details of RFT. As mentioned, training a reward model and a policy at the same time is costly. In our experiments, we discovered that the reward alignment step can be completely separated from the policy alignment step. In particular, we take T = 1 and K = data size batch size * epoch so that we train the reward over the entire dataset and then switch to the policy alignment. In our experiments, we indeed observe that only one round of above procedure can readily show superior performance over SFT and implicit reward-learning methods for pythia-1.4b model.Implementation details of IRFT. It is worth noticing that in (6), the policy π is not parameterized by θ directly. In our numerical experiment, we directly parameterize the LLM π by θ, making (6) the gradient of an supervised optimization problem itself. Meanwhile, it is not straightforward to calculate the self-generation gradient (6) directly, thus we need to design a loss function for back-propagation in main-stream packages such as PyTorch and TensorFlow. In practice, at each training iteration we first sample ỹ ∼ π(•|x; θ) and pass the following loss functionEQUATIONinto the standard optimizers (such as SGD or Adam) for back-propagation. Here h is a nonlinear function. We take h = log σ where σ is the logistic loss function σ(t) := log(1 + exp(-t)) as in Rafailov et al. [2024] , Chen et al. [2024] for its non-negativity, smoothness and exponentially decaying tail to avoid excessive growth in the absolute value of the log-likelihood.Discussion on the computational costs. For Algorithm 1, we need to maintain a reward model and a policy model (which is the LLM), and this is doubling the standard LLM fine-tuning. Thus the memory consumption and computation time of Algorithm 1 is similar to the standard RLHF process (RLHF = reward learning + policy optimization); For Algorithm 2, we simply need to maintain the policy (LLM) model, and the memory consumption would be exactly the same as the standard SFT, whereas the computation time would involving generating for the entire training sample, which would be of similar level as the standard policy optimization process (same computational time as SPIN). Note that standard policy optimization process is equivalent to the time of standard SFT and a generation process toward all training input prompts.Algorithm 1 Forward+Backward 2SFT+GenerationAlgorithm 2 Backward SFT+GenerationTable 2 : Table summarizing the computational costs of proposed methods.We summarize the memory consumption and the computational time of the proposed methods in Table 2 , assuming that the reward and policy models are of same size. Here "Forward" means the memory required for storing a model in inference mode, and "Backward" is the memory required for storing a model in training mode, including weights, activations and gradients; also "SFT" means the computational time as standard SFT, and "Generation" means the time to generate continuations for each of the input training prompts. Therefore "2SFT+Generation" is roughly the same time as standard RLHF.Comparison to SPIN. We discuss here the connection between our proposed algorithms with the self-play fine-tune algorithm (SPIN in Chen et al. [2024] ), which also maximizes the gap between two rewards. First, SPIN is motivated by certain two-player games, while in our case, we show that the difference of two rewards in (5) naturally comes from a single, reward learning agent; see (4).Second, IRFT covers SPIN as a special case. In particular, if we take T = 1 and K as the total number of training iterations, the IRFT algorithm is equivalent to SPIN. In practice, we tested on different choices of T and show that a reasonable generation frequency can results in a strong model performance.Finally, since SPIN does not involve explicit reward learning, its connection to RFT is relatively remote. It is worth noting that the relation between the proposed Algorithm 1 and Algorithm 2 is similar to that of RLHF to DPO. There has been intensive discussions regarding whether reward-based or reward-free algorithm gives better model performances, but this topic is beyond the scope of the current paper. We refer to Xu et al. [2024] for a comprehensive study.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
numerical experiments: 
 In this section we study the proposed Algorithm 1 and 2 numerically. Our experiment mainly show the advantages of the proposed methods in the following aspects: (1) Reward learning is key to improve over standard SFT, even if we do not have preference dataset;(2) The double loop design in both Algorithm 1 and 2 enable us to explore appropriate parameter settings that could break the performance limits of the state-of-the-art methods, including SFT and SPIN.Model and Datasets. Since reward-based methods can be costly by training two models at the same time, we mainly test Algorithm 1 on pythia-1b reward model and pythia-1.4b policy model [Biderman et al., 2023] . We tested pythia on Anthropic-HH dataset [Bai et al., 2022] . Anthropic-HH is a preference dataset that provide two continuations based on helpfulness and harmlessness, and we only pick 10k chosen/preferred continuation data to form the demonstration dataset, which enable us to check the log likelihood of the non-preferred continuation without feeding the model with such data. At each iteration, we train our model for 2 epochs (seeing each data for two times).Algorithm 2 is tested on two models: pythia-1.4b and zephyr-7b-sft-full [Tunstall et al., 2023] . We tested on Ultrachat200k dataset by HuggingFace, which is a subset of the high quality demonstration UltraChat dataset [Ding et al., 2023] for text generation and dialogue. For Ultra-chat200k, we adopt the same strategy as Chen et al. [2024] to pick up 50k data for training. At each iteration, we again train our model for 2 epochs.For the Anthropic-HH dataset, we show the reward evaluated by the PKU-Alignment/beaver-7b-v3.0-reward [Dai et al., 2024 , Ji et al., 2023] model which is a popular 7b model fine-tuned from meta-llama/Llama-2-7b tailored for evaluating human preferences regarding helpfulness and harmlessness. We also record win rate of the two proposed methods over base model and SFT model. For the Ultrachat200k dataset, we follow the widely used Hugging-Face Open LLM Leaderboard [Beeching et al., 2023] . This evaluation package assess an LLM based on six tasks: LLMs on commonsense reasoning (Arc Clark et al. [2018] , HellaSwag Zellers et al.[2019], Winogrande Sakaguchi et al. [2021] ), multi-task language understanding (MMLU Hendrycks et al. [2020] ), human falsehood mimic (TruthfulQA Lin et al. [2021] ) and math problem solving (GSM8K, Cobbe et al. [2021] ). See the appendix for more implementation details.We present the result of Algorithm 1 over Anthropic-HH dataset. We first fine-tuned pythia-1.4b using supervised fine-tune over the entire dataset (160k training data in total) using only the preferred/chosen data for 10 epochs and pick up the checkpoint with the best testing accuracy as our base model. We then use PKU-Alignment/beaver-7b-v3.0-reward model as our ground truth reward model. We use this model to pick 10k data from Anthropic-HH dataset with the highest reward scores. Next, we fine-tune the base model using SFT and Algorithm 1. Figure 2 shows the experiment results on averaged reward and win rate, where we record the average score (by PKU-Alignment/beaver-7b-v3.0-reward) of the continuation generated for test datasets, also the win rate (ratio of samples where the reward of our model's generation is higher than the model compared) of the proposed Algorithm 1 over the full SFT base model and the top 10k SFT model. The figures show that the proposed algorithm improves over the SFT models in terms of effectively improve the helpfulness and harmlessness of the model continuation. We remind the readers that the advantage of Algorithm 1 over SFT in Figure 2 can be partially explained by Figure 1 right side: despite the fact that SFT, Algorithm 1 and 2 are only observing chosen/preferred data, the latter two still outperforms SFT since they discourage the likelihood of the synthetic non-preferred data, thus bringing better performance and robustness for the model.Different from the time consuming Algorithm 1, Algorithm 2 is more capable of handling large data and models. We first present the result for pythia-1.4b models over Ultrachat200k data. We remind the reader again that T = 1 in Algorithm 2 is equivalent to SPIN [Chen et al., 2024] 1 . We tested on different choices of T and identify that T = 5 to 8 gives the best performance in the Open LLM Leaderboard evaluations.The Open LLM Leaderboard result is presented in Table 3 . We have the following main observations based on the results in Table 3: 1. SFT is not efficient in terms of boosting the pre-trained model performance on downstream tasks comparing to methods which promote the decreasing of the likelihood of synthetic data, namely SPIN and IRFT; 2. SPIN and IRFT (Algorithm 2) are both capable of further improving the performance of pythia model over downstream tasks, whereas IRFT shows better results due to more frequent generation comparing to SPIN. IRFT with T > 1 outperforms both SFT and SPIN on most of the tasks as well as the average score; 3. More frequent generation might also result in more variances, therefore a reasonable T (around 5) results in the best evaluation performance. Careful hyperparameter tuning might be needed for different models and datasets when applying our method, while we recommend using T = 5 as the default setting.Apparently 1b model is not strong enough to handle hard tasks, e.g. GSM8k and all model performances are not desirable. Now we present the result for zephyr-7b-sft-full. We remind the reader that this is a fully SFT-ed model and further SFT would only detriment the model performance (see Chen et al. [2024] ). The results are presented in Table 4 where we can see that similar to the 1b case, both SPIN and IRFT could effectively improve the performance of SFT-ed model and the average performance of IRFT with T = 5 stands out. The success of IRFT and SPIN further suggest that reward learning is indeed beneficial for aligning with demonstration data.  
conclusions and limitations: 
 In this paper we proposed reward-learning approaches for aligning LLMs with demonstration datasets. We show both theoretically and numerically the great potential of reward-learning for alignment even without preference dataset. Our theory only indicate the convergence of the proposed algorithm to stationary point, and it is not clear what the policy converges to. The additional computation resources required for tuning two models or generate synthetic data in our algorithms are not negligible. Future works include exploring reward-learning for larger models and more complicated demonstration tasks, boosting the algorithm efficiency, and understanding how synthetic negative sample helps the LLMs to distinguish the preference dataset, etc. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: See the appendix for the implementation details.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: Yes Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: We don't think there are direct societal impacts from this work.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pre-trained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper doesn't contain there high risk models or data.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The experiment part acknowledges all models and data we use.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Fine-tuning language models is prevailing to improve LLMs performance on various instructional tasks, and has shown great success in enabling LLMs to generalize to efficiently respond out-ofsample instructions [Chung et al., 2024] . Despite many successful applications of SFT, people soon realized the great potential of reward learning and reinforcement learning based fine-tuning over preference datasets for different tasks, including text-summarizing [Liu et al., 2020 , Ziegler et al., 2019] , story-telling [Ziegler et al., 2019] , instruction-following [Ouyang et al., 2022 , Ramamurthy et al., 2022] , etc. Equipped with the popular Bradley-Terry model [Bradley and Terry, 1952] , RLHF fine-tune a language model using policy optimization methods, such as REINFORCE [Williams, 1992] , proximal policy optimization (PPO, Schulman et al. [2017] ) and a lot more. On major obstacle for preference dataset fine-tuning is the costly and time-consuming process of human labeling, and methods such as self-play fine-tune (SPIN, Chen et al. [2024] ), synthetic data with binary feedback in self-training [Singh et al., 2023] , weak-to-strong generalization [Burns et al., 2023] and self-rewarding fine-tuning [Yuan et al., 2024] seek for improvement over SFT under weaker data supervisions comparing to preference datasets. In particular, SPIN generates synthetic samples for input prompts in the demonstration dataset and use them as the rejected data to for a 'pseudo' preference data. As we will see, SPIN actually coincides with our implicit reward learning approach where we motivate the synthetic data in a more natural way.In the reinforcement learning (RL) literature, inverse reinforcement learning (IRL) proposes to jointly learn the reward r which best explains an expert policy π E and the policy π which in turn mimics this expert policy π E from demonstration data. The most popular framework is the maximum entropy IRL (MaxEnt-IRL) framework Ziebart et al. [2008] , Levine et al. [2011] , Ziebart et al. [2013] , Bloem and Bambos [2014] , Zeng et al. [2022] , which seeks for a policy maximizing the entropic-regularized reward that matches the empirical averages in expert's demonstrations data. MaxEnt-IRL utilizes only the demonstration dataset for reward learning and already yields superior performance over the plain behavior cloning [Pomerleau, 1988 , Osa et al., 2018] approach on various RL tasks.We restate and prove Lemma 3.1: Lemma B.1. Problem (4) is equivalent to the following minimax optimization problem:max θ min π E x∼ρ,y∼π E (•|x),ỹ∼π(•|x) r(x, y; θ) -r(x, ỹ; θ) β + D KL π(•|x)∥π ref (•|x) . (8)Proof. It is straightforward to see that the lower-level problem in (4) enjoys a closed-form solution:π θ (y|x) = π ref (y|x) exp 1 β r(x, y; θ) ỹ∈A π ref (ỹ|x) exp 1 β r(x, ỹ; θ) (9)where A is the set of all possible responses. Plugging ( 9) into (4), we obtain:EQUATIONUtilizing the following identity:log   ỹ∈A π ref (ỹ|x) exp 1 β r(x, ỹ; θ)   = max π E ỹ∼π(•|x) [ 1 β r(x, ỹ; θ)] -D KL π(•|x)∥π ref (•|x)we obtain the following max-min problem (omitting some constant terms):EQUATIONThe proof is completed.Next, we restate and prove Lemma 3.2: Lemma B.2. For the loss function ℓ in (4), we have:EQUATIONwhich we refer to as the self-generation gradient, since at each iteration one need to generate one sample output ỹ from the current policy π θ and calculate the difference of the two rewards.Proof. Omitting the constant terms not related to θ in (10), we havemax θ ℓ(θ) = E x∼ρ,y∼π E (•|x)   1 β r(x, y; θ) -log   ỹ∈A π ref (ỹ|x) exp 1 β r(x, ỹ; θ)     (13)Calculating the derivative we get∇ θ ℓ(θ) = 1 β E x∼ρ,y∼π E (•|s) [∇ θ r(x, y; θ)] -E x∼ρ   ∇ θ log   ỹ∈A π ref (ỹ|x) exp 1 β r(x, ỹ; θ)     = 1 β E x∼ρ,y∼π E (•|x) [∇ θ r(x, y; θ)] - 1 β E x∼ρ   y∈A π ref (y|x) exp 1 β r(x, y; θ) ỹ∈A π ref (ỹ|x) exp 1 β r(x, ỹ; θ) ∇ θ r(x, y; θ)   = 1 β E x∼ρ,y∼π E (•|x) [∇ θ r(x, y; θ)] - 1 β E x∼ρ,y∼π θ (•|s) [∇ θ r(x, y; θ)] = 1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ (•|x) [∇ θ r(x, y; θ) -∇ θ r(x, ỹ; θ)]which implies that to minimize ℓ(θ), one should always generate samples based on the current estimation of the policy ỹ ∼ π θ (•|x) and then update. Now from (9) we get:EQUATIONwhere Z θ (x) is the denominator of (9). In the view of ( 14), we can actually directly estimate:EQUATIONThe proof is completed. Now we move to the proof for Section 3.3. We state the assumption needed for proving the final result: Assumption B.1. For Algorithm 1 and 2, we assume that 1. The policy distribution π θ is uniformly lower and upper bounded, i.e.π min ≤ ∥π θ (•|x)∥ ∞ ≤ π maxwhere 0 < π min < π max , for all x;2. ∇π θ is bounded, i.e. ∥∇π θ (•|x)∥ ≤ L 0 for all x;3. ∇π θ is Lipschitz, i.e. ∥∇π θ1 (y|x) -∇π θ2 (y|x)∥ ≤ L 1 ∥θ 1 -θ 2 ∥, for all x and y;where π θ is as defined in (9).The above assumption can readily establish the assumption below, which is needed for our final convergence result.Assumption B.2. For Algorithm 2, we assume that1. ℓ is L-Lipschitz smooth w.r.t. θ, i.e. ∥∇ℓ(θ 1 ) -∇ℓ(θ 2 )∥ ≤ L∥θ 1 -θ 2 ∥ 2.The stochastic estimator ∇ℓ is bounded, i.e.∥ ∇ℓ(θ)∥ ≤ GThese are all standard assumptions in nonconvex smooth stochastic optimization. We have the following lemma:Lemma B.3. If Assumption B.1 holds, Assumption B.2 also holds with the following parameters:L = L 0 (3L 0 + L 1 ) π 2 min , G = 2L 0 π minProof. We just show the value for L since G can be similarly computed. Since∇ℓ(θ) = 1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ (•|x) ∇ θ log π θ (y|x) π ref (y|x) -∇ θ log π θ (ỹ|x) π ref (ỹ|x)we haveEQUATIONFor the first part, since∇ log π θ (y|x) = ∇π θ (y|x) π θ (y|x)we have1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ 1 (•|x) ∇ θ log π θ1 (y|x) π ref (y|x) -∇ θ log π θ1 (ỹ|x) π ref (ỹ|x) -∇ θ log π θ2 (y|x) π ref (y|x) + ∇ θ log π θ2 (ỹ|x) π ref (ỹ|x) ≤ 1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ 1 (•|x) ∇π θ1 (y|x) π θ1 (y|x) - ∇π θ2 (y|x) π θ2 (y|x) - ∇π θ1 (ỹ|x) π θ1 (ỹ|x) + ∇π θ2 (ỹ|x) π θ2 (ỹ|x) ≤ 1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ 1 (•|x) ∇π θ1 (y|x) π θ1 (y|x) - ∇π θ2 (y|x) π θ2 (y|x) + 1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ 1 (•|x) ∇π θ1 (ỹ|x) π θ1 (ỹ|x) - ∇π θ2 (ỹ|x) π θ2 (ỹ|x) ≤ 1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ 1 (•|x) ∥π θ2 (y|x)∇π θ1 (y|x) -π θ1 (y|x)∇π θ2 (y|x)∥ π θ1 (y|x)π θ2 (y|x) + (same term for ỹ) ≤ 2 β π max L 1 + L 2 0 π 2 min ∥θ 1 -θ 2 ∥For the second term in the last line of ( 16), we have1 β E x∼ρ,y∼π E (•|x),ỹ∼π θ 1 (•|x) ∇ θ log π θ2 (y|x) π ref (y|x) -∇ θ log π θ2 (ỹ|x) π ref (ỹ|x) -E x∼ρ,y∼π E (•|x),ỹ∼π θ 2 (•|x) ∇ θ log π θ2 (y|x) π ref (y|x) -∇ θ log π θ2 (ỹ|x) π ref (ỹ|x) ≤ 1 β E x∼ρ,y∼π E (•|x) ỹ∈A ∇ θ log π θ2 (y|x) π ref (y|x) -∇ θ log π θ2 (ỹ|x) π ref (ỹ|x) (π θ1 (ỹ|x) -π θ2 (ỹ|x)) ≤ 1 β 2L 0 π min E x∼ρ,y∼π E (•|x) ỹ∈A (π θ1 (ỹ|x) -π θ2 (ỹ|x)) = 1 β 2L 0 π min E x∼ρ,y∼π E (•|x) ỹ∈A π θ1 (ỹ|x) π θ1 (ỹ|x) -π θ2 (ỹ|x) π θ1 (ỹ|x) ≤ 1 β 2L 2 0 π 2 min ∥θ 1 -θ 2 ∥Plugging these back to ( 16) we get∥∇ℓ(θ 1 ) -∇ℓ(θ 2 )∥ ≤ 2 β π max L 1 + 2L 2 0 π 2 min ∥θ 1 -θ 2 ∥Now since we generate at the beginning of the inner loop, the estimator ∇ℓ(θ t,k ) is not an unbiased estimator of ∇ℓ(θ t,k ) for any k > 0, i.e.EQUATIONEQUATIONWe thus need to carefully analyze this biasedness so that the convergence can be boosted by a large K, since a large K will result in a very large bias. E[∥∇ℓ(θ t,k )∥ 2 ] ≤ O ∆ 0 + LG 2 √ T K + L2 G 2T where ∆ 0 = ℓ * -ℓ(θ 0 ) and we omit constant factors in Õ.Proof. We prove directly for Algorithm 2 since the gradient estimator ( 12) and the estimator g t,kAlgorithm 1 (we do solve the π subproblem to its optimum) are both for the original bilevel problem (4).From the Lipschitz gradient of ℓ we haveℓ(θ t,k+1 ) ≥ ℓ(θ t,k ) + η t ⟨ ∇ℓ(θ t,k ), ∇ℓ(θ t,k )⟩ - η 2 t L 2 ∥ ∇ℓ(θ t,k )∥ 2 i.e. η t ∥∇ℓ(θ t,k )∥ 2 ≤ (ℓ(θ t,k+1 ) -ℓ(θ t,k )) + η t ⟨∇ℓ(θ t,k ) -∇ℓ(θ t,k ), ∇ℓ(θ t,k )⟩ + η 2 t L 2 ∥ ∇ℓ(θ t,k )∥ 2Taking expectation to θ t,k and by Assumption B.2, we haveη t E∥∇ℓ(θ t,k )∥ 2 ≤ (Eℓ(θ t,k+1 ) -ℓ(θ t,k )) + η t ⟨∇ℓ(θ t,k ) -E ∇ℓ(θ t,k ), ∇ℓ(θ t,k )⟩ + η 2 t LG 2 2where the expectation is taken w.r.t. the sample ỹt,k to generate the estimator of current iteration.Sum up from k = 0 to k = K we getK-1 k=0 η t E∥∇ℓ(θ t,k )∥ 2 ≤ (Eℓ(θ t,K-1 )-ℓ(θ t,0 ))+η t K-1 k=0 ⟨∇ℓ(θ t,k )-E ∇ℓ(θ t,k ), ∇ℓ(θ t,k )⟩+ η 2 t LG 2 K 2 (19)Since the expectation is taken only on the random sample at current iteration, and we know that the true gradient and the approximated gradient are ( 17) and ( 18), we have the following estimate:∥∇ℓ(θ t,k ) -E ∇ℓ(θ t,k )∥ = 1 β E x t,k ∼ρ,y t,k ∼π E (•|x t,k ),ỹ t,k ∼π θ t,k (•|x t,k ) ∇ θ log π θ t,k (y t,k |x t,k ) π ref (y t,k |x t,k ) -∇ θ log π θ t,k (ỹ t,k |x t,k ) π ref (ỹ t,k |x t,k ) -E x t,k ∼ρ,y t,k ∼π E (•|x t,k ),ỹ t,k ∼π θ t,k (•|x t,k ) ∇ θ log π θ t,k (y t,k |x t,k ) π ref (y t,k |x t,k ) -∇ θ log π θ t,k (ỹ t,k |x t,k ) π ref (ỹ t,k |x t,k ) = 1 β E x t,k ,y t,k ∇ θ log π θ t,k t,k |x t,k ) π ref (y t,k |x t,k ) -∇ θ log π θ t,k (ỹ|x t,k ) π ref (ỹ|x t,k ) π θ t,k (ỹ|x t,k ) -π θt,0 (ỹ|x t,k ) dỹ ≤ 1 β 2L 2 0 π min ∥θ t,k -θ t,0 ∥ Denote L = 1 β 2L 2 0πmin , we thus have:K-1 k=0 ⟨∇ℓ(θ t,k ) -E ∇ℓ(θ t,k ), ∇ℓ(θ t,k )⟩ ≤ K-1 k=0 ∥∇ℓ(θ t,k ) -E ∇ℓ(θ t,k )∥∥∇ℓ(θ t,k )∥ ≤ 1 2 K-1 k=0 ∥∇ℓ(θ t,k ) -E ∇ℓ(θ t,k )∥ 2 + 1 2 K-1 k=0 ∥∇ℓ(θ t,k )∥ 2 ≤ L2 2 K-1 k=0 ∥θ t,k -θ t,0 ∥ 2 + 1 2 K-1 k=0 ∥∇ℓ(θ t,k )∥ 2 = η 2 t L2 2 K-1 k=0 k-1 i=0 ∇ℓ(θ t,i ) 2 + 1 2 K-1 k=0 ∥∇ℓ(θ t,k )∥ 2 Therefore K-1 k=0 ⟨∇ℓ(θ t,k ) -E ∇ℓ(θ t,k ), ∇ℓ(θ t,k )⟩ ≤ η 2 t L2 G 2 2 K(K -1) 2 + 1 2 K-1 k=0 ∥∇ℓ(θ t,k )∥ 2Substituting back into (19) leads to1 2 K-1 k=0 η t E∥∇ℓ(θ t,k )∥ 2 ≤ (Eℓ(θ t,K-1 ) -ℓ(θ t,0 )) + η 3 t L2 G 2 K(K -1) 2 + η 2 t LG 2 K 2 Summing up from t = 0 to T -1 gives 1 2 T -1 t=0 K-1 k=0 η t E∥∇ℓ(θ t,k )∥ 2 ≤ Eℓ(θ T -1,K-1 )-ℓ(θ -1,0 )+ T -1 t=0 η 3 t L2 G 2 K(K -1) 2 + T -1 t=0 η 2 t LG 2 K 2With a constant step size η t = η > 0, we have1 2T K T -1 t=0 K-1 k=0 E∥∇ℓ(θ t,k )∥ 2 ≤ Eℓ(θ T -1,K-1 ) -ℓ(θ -1,0 ) ηT K + η 2 L2 G 2 K -1 2 + η LG 2 2 Taking η = Θ(1/ √ T K), we get 1 T K T -1 t=0 K-1 k=0 E∥∇ℓ(θ t,k )∥ 2 = O Eℓ(θ T -1,K-1 ) -ℓ(θ -1,0 ) + LG 2 √ T K + L2 G 2 T Hence as T → ∞, the rate is O(1/ √ T K).We follow the code as in SPIN [Chen et al., 2024] , where we utilize DeepSpeed ZeRO-3 [Rajbhandari et al., 2020] and FlashAttention-2 [Dao, 2023] to reduce the memory cost. We use RMSProp [Hinton et al., 2012] optimizer with no weight decay. For 1b models, we use two NVIDIA A100-40G to do the training with per device batch size of 4 for Algorithm 1 and per device batch size of 8 for Algorithm 2. For 7b models we use eight NVIDIA A100-40G to do the training with per device batch size of 2. We train all models with bfloat16 prevision. We set the peak learning rate to be 5e-7 for first two epochs and 1e-7 for the next two epochs. We fix β = 0.1 and consider the max sequence length to be 1024 for 1b models and 2048 for 7b models. We use the same prompt template "### Instruction: prompt\n\n### Response: " as in Chen et al. [2024] . For the policy optimization step in Algorithm 1, we use the PPO trainer in the TRL package [von Werra et al., 2020] . For the HuggingFace Open LLM Leaderboard evaluation, we use the Language Model Evaluation Harness library (v0.4.2) [Gao et al., 2023] , and we also use the same number of few-shots as in Chen et al. [2024] .Finally, in Table 5 , we further provide the generation examples of our fine-tuned model in Table 4 .During the reviewing and discussion periods, we conducted extra experiments with LoRA to provide further evidences for this work.We first provide the result of 7b experiments with LoRA in Table 6 . In this setting we see a significant improvement over the pretrained model (zephyr-7b-sft-full), where we observe 2.3% lift from the baseline and a 1% lift from SPIN. SFT in contrast can only achieve less than 1% lift from the base line. The reason here might be due to the limited model size when using LoRA, a contrastive training better helps the model distinguishing the referred and non-preferred continuations, yielding better performance over the standard SFT.As a side note, we do not anticipate to significantly outperform SPIN since algorithmically our proposed IRFT method includes SPIN as a special case. Rather, one of our main objective is to provide a theoretical foundation for the contrastive-type training algorithm, such as SPIN, which can all be studied under the bilevel inverse RL framework. The comparison with SPIN largely indicates that SPIN is still a RL-based fine-tuning method, suggesting an alternative interpretation that leads to provable convergence in lieu of the two-player game interpretation in Chen et al. [2024] .Reconciling our result with SPIN [Chen et al., 2024] Readers may realize that the result in Section 5.3 is different from SPIN [Chen et al., 2024] . We believe that this is due to two reasons:1. First, the baseline in our paper is different from SPIN paper: We believe a different baseline model is used in SPIN as evidenced in the Github discussions and the SPIN paper was released before the baseline is fully trained (Jan 2 vs Jan 10)2 . In particular, in Table 3 of the SPIN paper, the base model yields a "26.76" accuracy for GSM8k dataset, but we observe "31.92" which is significantly higher. We notice that a newer version of both the model zephyr-7b-sft-full3 and the lm_eval evaluation package4 which both our paper and SPIN use for evaluation are used in our paper. We run test on different versions of base model and eval codebase and obtain Table 7 , where we indeed see that the new version of the base model has a significant lift in the performance on GSM8k comparing to the old version. We remark that SPIN paper observes the most significant increase of SPIN algorithm on GSM8k task (from 26.76 to 35.10). Since we use the newest model in all our experiments, it leaves much less space for us to improve from.2. Second, we should not compare the iter3's 8.63% increase in Table 3 of SPIN with our paper's 2.66% increase directly. When we say we take T = 5 and epoch=2 as in Table 4 , we essential split the data into 5 chunks and generate more frequently than SPIN, but still consume and generate for all the training data for 2 epochs in total (SPIN iter0 also trained for 2 epochs). So what we need to compare is the first iteration of SPIN (which is SPIN Table 5 : Generation example of fine-tuned models in Table 4 .Craft an intriguing opening paragraph for a fictional short story. The story should involve a character who wakes up one morning to find that they can time travel.zephyr-7b-sft-full It was just another ordinary morning for Emily, until she opened her eyes and found herself in a different time and place. She was lying on her bed, surrounded by strange decor and people dressed in clothes from a bygone era. Emily rubbed her eyes, thinking she must be dreaming. But as she got up, she realized that she could feel the weight of the day ahead of her, much like she always did. It was only then that she realized that she had the power to time travel.IRFT T = 1 (SPIN iter 0) As soon as her eyes fluttered open, something was different. The room looked the same, the sun streaming in through the same window, but something felt off. It wasn't until she tried to sit up that she realized what it was -her body felt lighter, as if she were floating on air. Confused, she looked around, searching for an explanation, and that's when she saw it -the clock on her bedside table was ticking backwards. Panic set in as she tried to sit up, but her body refused to obey her commands. And just like that, she was hurled through time, landing with a thud in a completely different era.IRFT T = 1 (SPIN iter 1) As soon as the alarm clock shrilled, something felt different to James, but he brushed it aside and reached out to turn it off, just like any other morning. Yet, as soon as he sat up, he felt a strange pull as if he was being dragged backwards, and his room began to spin in a dizzying whirlwind of colors. In the blink of an eye, he found himself standing in a different place altogether, and he knew that he had just woken up in a different time.IRFT T = 5 (2 epochs) I opened my eyes groggily, ready to start another mundane day. But as I rubbed my eyes and stretched my arms, something strange happened. I felt a sudden jolt of energy course through my veins, and before I knew it, the world around me had transformed. The room that I was in was no longer my own, and I was standing in the middle of a bustling medieval town square. I blinked in disbelief, trying to make sense of what had just happened. Had I gone mad, or was I dreaming? But as I looked around me, I realized with awe that I had somehow managed to time travel.IRFT T = 5 (4 epochs) It was just another regular morning for Emily, until as soon as she opened her eyes the room started to blur and spin. Panicking, she tried to stand up and grab the nightstand but her hand went right through it. Confused and terrified, she tried to scream, but as soon as the sound started to come out of her throat, she was engulfed by a bright light that covered her body. When Emily opened her eyes again, she realized that she had traveled back in time, to a moment when she was 9 years old and standing at the foot of her parents' bed, ready to tell them the good news about acing her history test.iter0 in SPIN's original paper). Table 4 essentially indicates that, under our fair comparison setting, every iteration of our proposed algorithm outperforms every iteration of SPIN.Last, we provide a simple experiment to show how accurate the reward learned by our model is. This experiment also addresses the generalization ability of the reward. Since we train our 7b model with a high-quality dataset (ultrachat) and we believe that the corresponding implicit reward r = log(π θ /π ref ) should already be pretty accurate in terms of distinguish the good over the rejected continuations. Therefore we did a simple test: we construct the implicit reward by Justification: We show both theoretical and numerical results that support our claims in Section 3, 3.3 and 5.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.IRFT with T = 1 andepochs is equivalent to SPIN iteration 0, and T = 2 with 2 epochs for each T is equivalent to SPIN iteration 1, etc.see for example the discussion in this link.See their model commit history.See their codebase. 


Decision: This work proposes a novel large language model (LLM) alignment approach using Reward Learning from Human Demonstration. All reviewers consistently recommended accepting this work. AC agrees that this work is interesting and deserves to be published on NeurIPS 2024. The reviewers did raise some valuable concerns that should be addressed in the final camera-ready version of the paper. The authors are encouraged to make the necessary changes in the final version.

Reviews: review: paper_summary: paper_summary: This paper proposes two methods focusing on RLHF, namely RFT and IRFT. The takeaway message is that the SFT stage also significantly benefits from learning a reward model instead of using the human demonstration data directly via supervised learning.
strength_weakness: strength_weakness: 1. The training procedure is complicated, meaning there are many issues in tunning. Is there any computation cost analysis for better understanding the limitations of this method?
2. Though the author discusses SPIN, the truth is that equation 4.7 in SPIN is very similar to equation 6 in this paper. 
3. The experimental results are not strong. It did not surpass the SPIN by a large margin. Are there any other baselines that can be incorporated such as DPO?
questions: questions: How can the reward model obtain the generalization power? Can the author say more about it, in my understanding, the main issue that affects the generalization is the scale of the demonstration dataset.
limitations: limitations: See weaknesses. I have to say I am not an expert in this field, therefore, I am willing to change my attitude if I find I am wrong.
paper_summary: paper_summary: The authors propose using inverse reinforcement learning (IRL) on demonstration data in place of supervised learning, as is typically done for LLMs. The intuition is that human preferences are also encoded in demonstrations collected for SFT. Concretely, the authors propose a bilevel optimization approach with policy learning at the lower level and reward learning at the upper level. Two methods are proposed for alignment: with either implicit or explicit reward learning. The authors demonstrate improved performance when used to finetune smaller language models, compared to existing SFT approaches.
strength_weakness: strength_weakness: - The experimental results seem to support the efficacy of the proposed algorithms, however the gains across various tasks with IRFT seem fairly marginal (it would help if Table 2 and 3 contained error bars, perhaps with policies finetuned with different seeds). As it stands, it seems like different benchmarks results are sensitive to different hyperparameters (choice of T), and the baselines outperform IRFT on 3/6 of the benchmarks in Table 3, suggesting that the additional complexity of the method does not necessarily lead to improved performance across all tasks. For example, the best average performance of IRFT is 61.03, whereas the best baseline achieves 61.02.
- Furthermore, the results seem more mixed in Table 3 compared to Table 2, suggesting that it's not clear the method is still effective as model sizes increase.
- It would help to present the SFT results in Table 3. Even though the authors note that further SFT could degrade performance, it would help to see how effectively using the demonstration data with IRFT or SPIN compares against naive SFT. 
- Minor note: line 157 seems incomplete, "even when the demonstration policy is"...(extreme)?
- Minor note: the notation is a bit confusing, with theta referring to both the model parameters and the reward model parameters. For clarity, it would help if separate variables were used.
- Minor note: missing reference to appendix in line 226.
- Minor note: grammar error in line 320 -- "SPIN and IRFT are both capable of further improv(ing) the performance.."
questions: questions: See Weaknesses section for suggestions.
limitations: limitations: The paper is missing more thorough discussion on the method’s limitations and weaknesses (e.g. sensitivity to T, more complicated training process compared to SFT, etc.).
paper_summary: paper_summary: This paper proposes to study if IRL techniques can be used for aligning large language models. They propose two different algorithms: one that explicitly learns a reward model and one that implicitly learns the reward in the policy. These reward models are learned by contrasting expert generations and the policy generations.
strength_weakness: strength_weakness: - The experiments are not very convincing. For the Open LLM leaderboard experiments, the gain is really small (around 1 or 2%). It seems likely that this gain is due to variance in the training process, especially since IRFT has ~1% variance in different hyperparameter settings.
- There isn’t much discussion of the computational cost of the algorithms.
questions: questions: - How accurate is the learned reward model?
- The experimental setting is not very clear. Is Algorithm 1 used on the top 10K data selected by the reward model? If so this is really problematic: shouldn’t you fine tune on the whole dataset (or a randomly downsampled subset)? Can you assume that you will have access to 10K samples with high reward scores? And wouldn't just be better to use the reward model used for data selection for RL?
- How robust is IRFT to hyperparameter setting compared to SFT?
- Is 1-2% increase on the OpenLLM benchmark meaningful?
limitations: limitations: Yes, the authors discuss limitations.
paper_summary: paper_summary: This paper proposes an approach to aligning Large Language Models (LLMs) with human preferences by integrating Inverse Reinforcement Learning (IRL) into the supervised fine-tuning (SFT) stage. The proposed IRL-based method simultaneously builds a reward model and a policy model during the SFT stage, enhancing efficiency and robustness against low-quality data.
strength_weakness: strength_weakness: 1. The paper does not address the scalability of the proposed methods to even larger models or different types of LLMs beyond the ones tested.

2. The evaluation results only present the empirical performance but lack experiments on computation costs introduced by this method.
questions: questions: 1. It seems that RFT and IRFT perform equally. What's the takeaway from these two methods? Moreover, how to choose them wisely.
2. Typo in L226, there's an error in referring to the Appendix.
3. In L204-206 "In 204 practice however we take a relatively small T and large K, because frequent on-line sampling is time 205 consuming" How does this parameter choice affect the final performance, especially compared to large T and Large K?
limitations: limitations: The limitations are included in the conclusion section.

Comments: title: A gentle reminder on the discussion period
Thank you for the discussion
Concern about OpenLLM benchmark and Comparison to SPIN
comment: Thank you for your clear response, I will keep my positive score.
Thanks to the author's rebuttal, most of my doubts have been resolved, but I still don't think the performance gap is significant. In the discussion phase, I would discuss this point. No matter what the discussion result is, my final score will align with other reviewers.
I thank the authors for responding to my questions. The additional experiments on the accuracy of the implicit reward model are interesting to see, and overall the proposed method presents an interesting perspective on introducing IRL in the SFT stage. However, the improvements in the LoRA setting do not fully address my concerns around the lack of clear improvement coming from the method. As such, I have adjusted my score accordingly.
Dear reviewers,

We thank all of your constructive feedbacks toward this work, and we tried our best to address your concerns and questions. As the author-reviewer discussion period is ending in three days, we gently remind you to check out our rebuttal and let us know if you still have any comments. Thank you very much!

Best,
Authors of submission 14483
Thank you very much for re-evaluating our work and we will incorporate the discussion to the revised paper!
Thank you for the comprehensive experiments and discussion of the SPIN codebase. All of my concerns are cleared up and I will raise my score.
(Continuing above)

We also tried to revert back to v0.4.0 of the evaluation codebase (which we believe should be the version of evaluation codebase used by SPIN) and strictly following the metrics SPIN paper mentioned. We have the following result for the same setting as the Table in our general rebuttal:

*(base model: newest version; evaluation codebase: v0.4.0; test with LoRA r=64)*
| Task     | Arc_challenge | TruthfulQA_mc2 | Winogrande | GSM8k | Hellaswag | MMLU | Average |
| -------- | ------- | ------- | ------- | ------- | ------- | ------- | ------- |
| zephyr-7b-sft-full  |  58.02  |  40.4   |  76.16   |  34.19   |  80.89   |   57.46  | 57.85   |
| SFT |  58.87   |  40.34   |  76.56   | 34.5   |  81.42   | 57.93    |   58.27  |
| IRFT (T=1, SPIN iter0)    |  59.64   |   40.4  |  75.85   | 34.5    |  81.98   |  57.46   |  58.31   |
| IRFT (T=5, iter0)    | 61.09    |  42.92   | 76.95   |  35.25   |  82.48   |  57.66   |  59.39   |

In the above table, the baseline performancec in GSM8k is "34.19", which is still significantly higher than "26.76" in that of SPIN paper. Therefore the baseline is still significantly different from SPIN paper, also our algorithm (last row) is still outperforming SPIN (second last row) by more than 1 point of average accuracy in iter0. We hope this clarification makes the comparison clear. We will point this out and revise our experiment setting section accordingly.

> Furthermore, is it possible to re-run some of your experiments with different seeds and conduct significance tests on the results?

**Response**: Due to limited computational resources, we did not conduct extensive experiment on multiple seeds (note that SPIN also did not do significance test). Here we provide 3 trials on IRFT iter0 (with T=5) with three different seeds:

*(base model: newest version; evaluation codebase: v0.4.3; test without LoRA)*
| Task     | Arc | TruthfulQA | Winogrande | GSM8k | Hellaswag | MMLU | Average |
| -------- | ------- | ------- | ------- | ------- | ------- | ------- | ------- |
| zephyr-7b-sft-full  | 74.83 | 34.07 | 76.09 | 31.92 | 81.09 | 58.86 | 59.48 |
| IRFT (T=5, iter0) trial1   |  75.82 | 39.99 | 77.19 | 31.24 | 82.07 | 57.93 | 60.71
| IRFT (T=5, iter0) trial2   |  76.66 |	36.25 |	76.48 |	32.37 |	82.13 |	58.30|	60.37  |
| IRFT (T=5, iter0) trial3   |  75.42 |	35.85 |	76.80 |	32.60 |	82.27 |	58.92 |	60.31  | 

Above table indicates a variance of 0.05 (a standard deviation of 0.22) for the average performance. We wish this gives you a flavor of the level of variance for our experiments.

It might be very hard for us to systematically test the random seed effect starting from now. For each 7b experiment, we need to generate continuation for each prompt, which will take roughly 8 hours for the entire training dataset and another 8 hours to do the full-fine-tune on 8 A100-40G; For LoRA (r=64) we need 4 A100-40G for a similar time. We would need to wait in the queue for resources, which makes it impossible to estimate the time for us to finish. The same situation applies for running the SPIN algorithm.
> **Concern about OpenLLM benchmark and Comparison to SPIN**: 

Thanks for checking the rebuttal. TL;DR: 1. we suspect that a different version of baseline "zephyr-7b-sft-full" has been used in the SPIN paper which led to a significantly higher lift of 8-9% improvement, compared to ours of 2.6%. 2. when evaluated on the same baseline and codebase, every iteration our proposed algorithm outperforms every iteration of SPIN.

The SPIN paper reported the following improvements:

*(base model: older version; evaluation codebase: v0.4.0; test without LoRA)*
| Iterations  |  zephyr-7b-sft-full   | iter0 | iter1 | iter2 | iter3 |
| -------- | ------- | ------- | ------- | ------- | ------- |
| Average Performance | 58.14 | 60.80 | 62.12 | 62.97 | 63.16 |
| Absolute increase |  | 2.66 | 1.32 | 0.85 | 0.89 |
| Relative increase |  | 4.56% | 6.85% | 8.31% | 8.63% |

We believe the result should be interpreted in the following way:
1. First, the baseline in our paper is different from SPIN paper: We believe a different baseline model is used in SPIN as evidenced in the Github discussions & the SPIN paper was released before the baseline is fully trained (Jan 2 vs Jan 10), see for example [this discussion](https://github.com/uclaml/SPIN/issues/12). In particular, in Table 3 of the SPIN paper, the base model yields a "26.76" accuracy for GSM8k dataset, but we observe "31.92" which is significantly higher. We notice that a newer version of both the model zephyr-7b-sft-full (see [their model commit history](https://huggingface.co/alignment-handbook/zephyr-7b-sft-full/commits/main)) and the lm_eval evaluation package which both our paper and SPIN use for evaluation (see [their codebase](https://github.com/EleutherAI/lm-evaluation-harness)) are used in our paper. We run test on different versions of base model and eval codebase and obtain a table as follows:
*(Performance of different version of zephyr-7b-sft-full, note we also change the first two tasks to make it fully aligned with SPIN paper)*
| Base model version| Eval package version     | Arc_challenge | TruthfulQA_mc2 | Winogrande | GSM8k | Hellaswag | MMLU | Average |
| -------- | -------| ------- | ------- | ------- | ------- | ------- | ------- | ------- |
| [Newest](https://huggingface.co/alignment-handbook/zephyr-7b-sft-full)  | [v0.4.0](https://github.com/EleutherAI/lm-evaluation-harness/tree/v0.4.0) |  58.02  |  40.40   |  76.16   |  34.19   |  80.89   |   57.46  | 57.85   |
| [Version in SPIN](https://huggingface.co/alignment-handbook/zephyr-7b-sft-full/commit/90e0792328bc522e1662a3a7c611b030d563bf5b)  | [v0.4.0](https://github.com/EleutherAI/lm-evaluation-harness/tree/v0.4.0) |  60.84 | 43.74  |  78.69  |  26.23   |  82.79   |  58.97   | 58.54   |

where we indeed see that the new version of the base model has a significant lift in the performance on GSM8k comparing to the old version. We remark that SPIN paper observes the most significant increase of SPIN algorithm on GSM8k task (from 26.76 to 35.10). **Since we use the newest model in all our experiments, it leaves much less space for us to improve from.**

2. Second, we should not compare the iter3's 8.63% increase with our paper's 2.66% increase directly. When we say we take $T=5$ and epoch=2 as in Table 3 of our paper, we essential split the data into 5 chunks and generate more frequently than SPIN, but still consume and generate for all the training data for **2 epochs in total** (SPIN iter0 also trained for 2 epochs). So what we need to compare is the first iteration of SPIN (which is SPIN iter0 in SPIN's original paper). In view of Table 3 in our paper, a fair comparison is the following table, note that the baseline is newest version of the model on the newest evaluation codebase (we follow the "iteration" notion in SPIN paper):

*(base model: newest version; evaluation codebase: v0.4.3; test without LoRA)*
| Iterations  |  zephyr-7b-sft-full   | iter0 | iter1 |
| -------- | ------- | ------- | ------- |
| SPIN | 59.48 | 60.32 | 61.02 |
| IRFT | 59.48 | 60.71 | 61.03 |

Essentially **under our fair comparison setting**, every iteration of our proposed algorithm outperforms every iteration of SPIN. We didn't further do experiment for iter2 and iter3 because under our experiment setting, both SPIN and IRFT seem to have less significant improvement from iter0 to iter1 (less than 1 point in average accuracy), and also due to our limited computing resources.

(To be continued)
Thank you for providing more details on the computational cost of the algorithm and showing the accuracy of the obtained reward model. I am still concerned over the weak performance on the OpenLLM leaderboard. I checked the SPIN paper, and they reported a 8-9% improvement on the OpenLLM leaderboard. Do you know why there is a large discrepancy between the results reported in SPIN and those in this paper? Furthermore, is it possible to re-run some of your experiments with different seeds and conduct significance tests on the results?