Would I Lie To You? Inference Time Alignment of Language Models using Direct Preference Heads

Authors: avelina  asada
hadji  -kyriacou
ognjen  arandjelović
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 15:48:44

Abstract: Pre-trained Language Models (LMs) exhibit strong zero-shot and in-context learning capabilities; however, their behaviors are often difficult to control. By utilizing Reinforcement Learning from Human Feedback (RLHF), it is possible to fine-tune unsupervised LMs to follow instructions and produce outputs that reflect human preferences. Despite its benefits, RLHF has been shown to potentially harm a language model's reasoning capabilities and introduce artifacts such as hallucinations where the model may fabricate facts. To address this issue we introduce Direct Preference Heads (DPH), a fine-tuning framework that enables LMs to learn human preference signals through an auxiliary reward head without directly affecting the output distribution of the language modeling head. We perform a theoretical analysis of our objective function and find strong ties to Conservative Direct Preference Optimization (cDPO). Finally we evaluate our models on GLUE, RACE, and the GPT4All evaluation suite and demonstrate that our method produces models which achieve higher scores than those fine-tuned with Supervised Fine-Tuning (SFT) or Direct Preference Optimization (DPO) alone.

Main paper: 
 introduction: 
 Reinforcement Learning from Human Feedback (RLHF) is a technique that can be used to align an agent -such as a Large Language Model (LLM) -to human preferences and lead to more truthful, more helpful, less harmful and more preferred outputs [29] . Proximal Policy Optimization (PPO) [36] and Direct Preference Optimization (DPO) [31] are two such alignment techniques which have been extensively used to improve the quality of LLM outputs, leading to instruction following agents or chat assistants which are quickly approaching human-baselines in a variety of knowledge and reasoning tasks [4, 10, 42, 18, 24, 35, 11] .However, recent research has shown that RLHF may actually hurt an LLM's reasoning abilities rather than improve it. One study [5] discovered that performing alignment during the Supervised Fine-Tuning (SFT) stage of training may lead to worse performance on reasoning benchmarks, and another [3] that SFT alone outperforms RLHF for smaller models with the benefits of RLHF only emerging for models with more than 1 billion parameters. Ouyang et al. [29] also report an increased tendency for RLHF models to make up information in closed domain tasks ("hallucination") compared to models trained with SFT alone.To combat the the risk of RLHF compromising the abilities of an LLM in favour of producing preferable outputs we introduce Direct Preference Heads (DPH), a novel feature based approach that optimises a reward score produced by the LLM rather than optimising the logits produced by language modelling head. DPH can be used in combination with (or without) existing alignment techniques to allow language models to self-evaluate outputs sampled at inference time and select the highest scoring candidate.We evaluate the performance of DPH using an efficient 551M parameter LM on a variety of commonsense reasoning and Natural Language Understanding (NLU) tasks. All code used to train our models is available on GitHub 1 and we release our model weights on Hugging Face 2 . 
prior approaches: 
 Prior approaches to language model alignment involve directly optimizing the logits produced by the language modelling head to increase the likelihood of producing preferable responses while decreasing the likelihood of undesirable responses.Reinforcement Learning from Human Feedback seeks to learn a reward model from human feedback on completions generated by a language model which can be used to align an LM with human preferences. A typical RLHF pipeline consists of 3 steps: (1) supervised fine-tuning, (2) preference sampling and reward modelling, and (3) RL fine-tuning.Supervised Fine-Tuning The first step of a standard RLHF pipeline is fine-tuning a pre-trained LM on high quality data for downstream tasks to obtain a model π SFT .Reward Modelling Next, the SFT model is prompted with input tokens x to produce completions y. These answers are then rated by human labellers which rate the answers based on one or more criteria. A reward model r ϕ (x, y) is then trained to estimate the scores assigned by human labellers using maximum likelihood estimation.RL Fine-Tuning During the RL phase the learned reward function is used to provide feedback to the language model using the following optimization problemEQUATIONwhere β controls the deviation from the base reference policy π ref , which is typically initialized from π SFT . Due to the non-differentiable nature of language generation this objective must be optimized using a reinforcement learning algorithm such as PPO [36] .Direct Preference Optimization was introduced as a reparametrization of RLHF which eliminates both the sampling stage and the reward modelling stages and reformulates alignment procedure as a loss function which can be optimized directly on a dataset of pairs of preferred and dispreferred completions to given prompts. This allows DPO to stably and efficiently converge on an optimal policy using what is effectively a classification loss over positive and negative pairs.Given a dataset {(x, y w , y l )} where x is the prompt and y w , y l are the preferred and dispreferred completions, we introduce the following loss function:EQUATIONwhere π θ (y * |x) and π ref (y * |x) are the probabilities of completions y * for prompt x given by the policy model and reference models respectively, and the β parameter controls the deviation from the reference policy.There also exists an augmentation of DPO namely Conservative DPO (cDPO) [26] which is designed to be more robust to noisy labels through the introduction of label smoothing parameter ϵ. The objective function for cDPO is given by:Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We include a limitations sub-section in our discussion, and discuss the fact that -as with all probabilistic models -there is no guarantee the outputs will be correct. We also discuss the fact that models aligned using our method are still capable of producing unsafe outputs if the end user ignores the reward scores. 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. 3. Theory Assumptions and Proofs some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
direct preference heads: 
 The hypothesis underlying the Direct Preference Optimization framework of Rafailov et al. [31] is that a "language model is secretly a reward model" thereby making the purpose of Direct Preference Heads to exploit this and extract explicit reward signals without the need of an additional reward model.To obtain the rewards from a sequence x; y three components are required: an aggregated hidden state h which is conditioned on the intermediate representations of the language model, a pooling function f which transforms the hidden state, and a learnable vector w dph with the same dimension as the output of f . We then compute the reward r as follows:EQUATIONTo obtain the hidden state we take the output of the last transformer layer for the final token of the sequence, and we experiment with three choices of f : (1) the identity mapping following the convention established by OpenAI's GPT for sequence classification [30] , (2) a learnable affine projection with tanh nonlinearity following BERT's pooling function [21] , and (3) an inverted bottleneck FFN with SwiGLU activation mirroring the FFN blocks used within the transformer backbone followed by tanh nonlinearity [37] .We formulate two novel objective functions for our method: a separable objective which maximises positive rewards and minimises negative rewards, and a contrastive objective which maximises the margin between positive and negative rewards. The loss landscapes are illustrated by Figure 1 in the appendix.The Separable DPH loss function given by ( 5) is a function of the preferred and dispreferred rewards r w , r l , and the label smoothing parameter 0 ≤ ϵ ≤ 0.5 which controls the reward margin:EQUATION)Theorem 1. For all ϵ ∈ (0, 0.5] the objective function L SepDPH is convex and will optimize the policy π θ such that the preferred rewards r w produced by the preference head converge towards log 1-ϵ ϵ and the dispreferred rewards r l converge to log ϵ 1-ϵ .This can be proven by observing the first and second partial derivatives of the loss function with respect to the rewards. The first partial derivative is equal to zero at the points r w = log 1-ϵ ϵ and r l = log ϵ 1-ϵ respectively, and the second partial derivative is strictly positive for all values of r w , r l . A full proof is included in Appendix A.1.Like Separable DPH, the loss function for Contrastive DPH given by ( 6) is function of the preferred and dispreferred rewards r w , r l and the label smoothing parameter 0 ≤ ϵ ≤ 0.5. This version of the loss function optimizes the relative margin between the rewards rather than optimizing the absolute positive and negative rewards as in Separable DPH.EQUATIONTheorem 2. For all ϵ ∈ (0, 0.5] the objective function L ConDPH is convex and will optimize the policy π θ such that the difference between preferred rewards r w and dispreferred rewards r l produced by the preference head will converge to a fixed margin, given by r ∆ = r w -r l = log 1-ϵ ϵ . This can be proven by reparametrising the loss function such that r ∆ = r w -r l and then by considering the first and second partial derivatives with respect to this reward margin. It can be observed that the first partial derivative is equal to zero when r ∆ = log 1-ϵ ϵ , and the second partial derivative is strictly positive for all values of r ∆ . A full proof is included in Appendix A.2.The properties of both Contrastive DPH and Separable DPH show a strong relationship with Conservative DPO: SepDPH will converge to optimal fixed reward margins above zero for r w and below zero for r l ; ConDPH will converge to optimal fixed reward margins between r w and r l , and cDPO will converge to a fixed delta from the reference model [26] . Like Conservative DPO, this makes both Separable DPH and Contrastive DPH robust to preference label noise and makes training more stable than naive maximum likelihood estimation without label-smoothing.Although similar to the reward modelling phase of an RLHF pipeline, DPH has some distinct differences which set it apart. DPH does not require an SFT sampling and human labelling stage meaning it can take advantage of pre-constructed preference datasets such as those used for DPO. Typical RLHF also requires multiple models -a reward model, a reference model and a policy model -while DPH requires only a single model to produce both responses and rewards.Unlike other RLHF pipelines such as PPO [36] , the rewards produced by DPH are not used for RL fine-tuning; instead, the DPH rewards are to be used to prune candidate generations sampled from the LM at inference time to select the candidate which aligns most with human preferences. This makes DPH an excellent choice for small language models which are (1) more lightweight -and therefore can be efficiently used to generate multiple samples -and, (2) are more prone to degradation when aligned using typical RL techniques [5, 3] . 
datasets: 
 We make use of a variety of datasets for fine-tuning and evaluation which are outlined below. The specific prompt templates used for fine-tuning and evaluation are described in Appendix C.Natural Language Understanding (NLU) For general NLU we make use of the standard GLUE benchmark [40] . The overall score for GLUE is computed by the macro-average of unweighted metric averages for all 9 tasks, however we also include a secondary score which does not included the 'problematic' WNLI task following the evaluation used for BERT [21] . We opted to omit WNLI during fine-tuning due to the low sample size.Commonsense Reasoning In accordance with the GPT4All [1] evaluation suite, we use the following datasets to evaluate commonsense reasoning abilities: HellaSwag [42] , OpenBookQA [25] , WinoGrande [35] , ARC [10] , BoolQ [9] , and PIQA [7] .Reading Comprehension To evaluate reading comprehension abilities we use the RACE dataset [22] , a multiple-choice task which requires reasoning over provided passages.Instruction Following We include the Alpaca [38] , OpenOrca [23] , and UltraFeedback [12] datasets to train our models for instruction following. We make use of OpenOrca and a cleaned version of Alpaca for SFT, and binarized versions of OpenOrca and UltraFeedback for alignment.Auxiliary Datasets To provide additional training data for SFT we include the MMLU [18] , SQuAD V2 [33, 32] , Tiny Stories [15] , CNN-Dailymail [27] and CoQA [34] training splits. For alignment we only include MMLU and SQuAD V2.Prompts We make use of the ChatML prompt templating scheme [28] with handcrafted system, user and assistant prompts specific to each task. During fine-tuning we mask out the loss for all tokens of the prompt and condition the model on the content of assistant messages including the final 
 token. During evaluation we select the highest scoring answer using the average log-probabilities of the tokens in the final assistant message, or compute the reward scores on the final 
 token when evaluating with DPH.SFT Sampling When sampling from the datasets for SFT we randomly shuffle each dataset and uniformly interleave samples from all tasks in the mix. To control the weighting of samples from each task we fill the context window with n consecutive samples from the same task before sampling from a different task, where n is chosen to be 5 in our experiments. To maximise compute utilisation and minimize unused portions of the context window we make us of Transformer-XL [13] style training with a context window size of 2048 tokens and a recurrent memory size of 2048 tokens.DPH Sampling When sampling from datasets for DPH alignment we switch from the Transformer-XL style pipeline to typical SFT training, opting to only include single samples in the context window padded to a fixed maximum length. As some of the datasets we use for DPH are intended for SFT rather than alignment (namely GLUE, GPT4All, RACE, MMLU and SQuAD) we synthesise preference pairs where the 'correct' answer is used as the preferred completion and we uniformly sample an 'incorrect' answer from the available choices for the dispreferred completion. This is trivial for most datasets, however we use a special process for the SQuAD V2 dataset; for answerable questions we use "unanswerable" as the dispreferred completion, and for unanswerable questions we use SpaCy to randomly sample a noun span from the context to use as the dispreferred completion.The hidden states h used to compute the reward scores are likely sub-optimal for computing rewards when initialising π θ from π SFT . As such, it may be desirable to fine-tune some or all parameters in the language model to learn better reward signals. This necessitates the use of regularization to prevent degradation of the models generative capabilities while learning to predict rewards.Prior Regularization Typical parameter regularization strategies such as weight decay make the assumption that parameters θ follow a zero-mean Normal distribution p(θ) ∼ N (0, 1 β I) leading to an auxiliary loss term β 2 ||θ|| 2 2 . However, when performing transfer-learning or fine-tuning on a pre-trained model this assumption can be harmful and aid in catastrophic forgetting of the model's previously learnt abilities.An alternative regularization scheme is Prior Regularization [8, 14, 17] which instead makes the assumption that the fine-tuned parameters are normally distributed around the original parameters θ ref , that is θ ∼ N (θ ref , 1 β I), leading to the auxiliary loss termβ 2 ||θ -θ ref || 2 2 .We employ Prior Regularization to limit the divergence of π θ from π SFT while still facilitating the learning of improved hidden state representations for the Direct Preference Head. Pseudocode for optimizer based decoupled prior regularization is included in Appendix B.1.cDPO Regularization Rather than directly employing a KL divergence penalty similar to that used in (1) we find that it is possible -and even beneficial -to use Conservative DPO as a means of limiting the divergence of the policy model to a fixed delta from the reference model, and 'nudging' the model towards generating more preferable outputs which increases the chance of generating a better candidate completion at inference time with fewer sampling steps.We progressively fine-tune the models in 3 stages: vocab extension, supervised fine-tuning, and DPH alignment. The details of the pre-trained model are included in Appendix D.1.Vocab Extension Since our model was pre-trained without a chat structure it is necessary to train the embeddings for additional 
 and 
 tokens: we freeze all non-embedding parameters and use the same datasets as SFT. We fine-tune the embeddings for 4096 steps with a batch size of 128, a max LR of 6e-5 which warms up over 200 steps followed by cosine decay down to zero, and clip the global gradient norm to 1.Supervised Fine-Tuning After vocab extension we move onto the SFT step which conditions the model for NLU tasks and instruction following using the sampling and loss masking method described in Section 4.2. We fine-tune the model for 6144 steps with a batch size of 128, a max LR of 3e-5 which warms up over 200 steps followed by cosine decay down to zero, prior-regularization applied to all non-embedding parameters with coefficient 0.5, and clip the global gradient norm to 1.DPH Alignment Using the sampling method described in section 4.2 we jointly learn DPH rewards and perform cDPO alignment. The goal here is to gently push the model towards producing preferable outputs without compromising the model's reasoning abilities, and the priority is to attain the highest validation metrics from the DPH rewards. This requires balancing the two objectives, and as such we introduce weighting parameters α 1 , α 2 to our final joint objective in (7) where L DPH is either L sepDPH or L conDPH . We find α 1 , α 2 = 1 to be a good balance between DPO and DPH in our experiments.EQUATIONWe align the model for 23,040 steps with a batch size of 64 pairs, a max LR of 3e-6 which warms up over 200 steps followed by cosine decay down to 3e-7, prior-regularization applied to all parameters with coefficient 0.5, and clip the global gradient norm to 1. Following the optimal DPO parameters for OpenHermes-7b-2.5 [20] we use β = 0.6 and chose cDPO ϵ = 0.25 and DPH ϵ = 0.1 for regularisation. Additionally, we apply dropout with p = 0.1 to the outputs of the pooler.All fine-tuning was performed using an NVIDIA A100 SXM4 80GB GPU on a compute cluster, with jobs allocated 24 cores and 160GB of memory. Each checkpoint is saved in FP16 format which consumes about 1.1GB of storage; and the datasets require minimal storage space.For vocab extension we train for 4096 steps with an average of 7.99 seconds of compute per step which translates to about 9 hours. For supervised fine-tuning we train for 6144 steps with an average of 9.26 seconds of compute per step which translates to about 16 hours. For DPH alignment we train for 23040 steps with an average of 7.21 seconds of compute per step which translates to about 46 hours. The DPH ablations with our models use about 140 hours of compute, and the Qwen ablations use about 60 hours of compute. In total, we used approximately 270 hours of A100 compute to train our models and collect the results included in our paper. We used additional compute for preliminary tests and fixing bugs for silently failing experiments although this wasn't tracked. 
evaluation methodology: 
 As described in Section 4 we use NLU, commonsense reasoning and reading comprehension tasks to measure model capabilities, while the instruction following and auxiliary tasks are used to provide additional training signals. For the NLU tasks we evaluate on the test set of GLUE, providing average scores both with and without WNLI. For reading comprehension we evaluate on the RACE test set.For commonsense reasoning we follow the LM Evaluation Harness [16] implementations of these tasks, evaluating on the test sets of ARC and OpenBookQA and the validation sets of HellaSwag, WinoGrande, BoolQ and PIQA, which brings our evaluations in line with other models.For vocab extension and SFT checkpoints we obtain model predictions from the completions with the highest scoring log-probabilities. For the DPH checkpoints we report metrics for both logprobability predictions (Ours DPO ) and predictions chosen from the DPH rewards (Ours DPH ). We use the SwiGLU-based pooler with the separable objective function for all our experiments as we found this combination to perform best overall as shown in Section 5.2.1.Our results for NLU performance are included in Table 1 . Note that the results for GPT-1 [30] and BERT [21] are from sub-task specific fine-tunes.It is unsurprising that our model does not outperform BERT Large even though it has more parameters; this is likely due to BERT's task specific fine-tunes in comparison to our model which was jointly trained on several tasks. Despite this our instruction following DPH model achieves a 2.2% higher average GLUE score compared to task-specific GPT-1 fine-tunes and manages to attain the highest overall accuracy and macro-average on RTE and STS-B respectively. Our results for commonsense reasoning are summarized in Table 2 . Note the Pythia [6] and TinyLlama [43] models were not fine-tuned for any specific task but received significantly more pre-training and have much higher parameter counts. With SFT alone we are able to attain comparable performance to TinyLlama using half as many parameters, and when applying DPH alignment we achieve a 7.2% increase over the TinyLlama average score and the highest accuracy in 5 of the 7 tasks.Our results for reading comprehension are included in Table 3 . The results for GPT-1 were taken from a RACE specific fine-tune, and the results for LLaMA [39] were zero-shot without fine-tuning. Our SFT baseline achieves a higher average accuracy on RACE compared with the non fine-tuned LLaMa models but cannot match the accuracy of the RACE specific GPT-1 fine-tune; however after alignment our model attains a 3.5% higher average over GPT-1 while still maintaining excellent scores on other tasks using the same model weights.We ablate over the three pooling head and two objective function choices. We perform alignment for 7680 steps and report the validation scores in Table 4 .For both separable and contrastive objectives the SwiGLU pooler performs best on the three benchmarks, and for both GLUE and RACE the separable objective performs best overall. However during these experiments we discovered that contrastive DPH was achieving higher scores than separable DPH for specifically the sentence completion style tasks like HellaSwag, WinoGrande and PIQA. We hypothesise this is caused by situations where multiple completions to a given prompt may be plausible even though there is only one 'gold' answer, and as such the model benefits from maximising the relative reward margin with the contrastive objective rather than optimising absolute rewards with the separable objective.By taking the DPH checkpoint and freezing all backbone parameters it is possible to learn task specific heads and pooling functions for different downstream tasks at the cost of only 19M parameters per task. We train new heads for the three task groups and plot the confusion matrix of each head for each task average in Table 5 . We further fine-tune for an additional 7680 steps on each task group using the same training setup as DPH alignment. Unsurprisingly the GLUE and GPT4All heads achieve the highest scores for GLUE and GPT4All benchmarks respectively, however the GPT4All head manages to outperform the RACE head on the RACE benchmark. We hypothesise this may be due to the inclusion of multiple choice QA and reading comprehension tasks in GPT4All which may prove better training signals than the RACE training data alone.Our final experiments involve exploring the behaviour of DPH when applied to frozen language models without further fine-tuning of the original model weights. We experiment using the Qwen 1.5 model family [2] and train only the pooler and reward head weights, reporting results in Table 6 .We use an identical training setup to DPH alignment but disable dropout due to the low number of trainable parameters.Because the model backbone and embeddings remain frozen during alignment the 'Log' scores represent the model's pre-trained (or fine-tuned) capabilities. When observing the difference between the Log scores of the 0.5B Qwen models it is evident that the fine-tuning and alignment used to transform the pre-trained model into the "chat" model resulted in degraded performance across the three tasks. This phenomenon is less apparent for the 1.8B models, and actually results in higher GLUE scores for the "chat" variant of the model. This further confirms the hypothesis that alignment can harm the reasoning capabilities of smaller language models. For all models DPH is consistently able to attain higher scores on the GLUE tasks compared to the log probabilities produced by the language modelling head, but the opposite is observed for RACE which suggests the hidden states produced by the frozen backbone do not contain rich enough features for long range modelling tasks such as reading comprehension. We also observe the "chat" variants produce higher task scores for DPH than the non-chat variants which we hypothesise is a result of the authors' fine-tuning with the Chat-ML format which lead to the models' greater understanding of message structure and therefore improved hidden state aggregation for the final end message token.When we combine these findings with those presented in Section 5.2.2, it becomes evident that the pooling function and reward head exhibit slower convergence when the model backbone is frozen. This observation further supports our hypothesis in Section 4.3, indicating that the hidden states generated by the models are are initially sub-optimal and that further fine-tuning is necessary to optimize these hidden states to achieve the best features for DPH.6 DiscussionQuestion: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: All data used for fine-tuning is publicly available on Hugging Face, the code used to train the models and perform the experiments is publicly available, and an additional repository containing the code and reproduction instructions will be made available.Guidelines:• The answer NA means that paper does not include experiments requiring code. 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). 
future work: 
 As shown in the results section, DPH is capable of learning to assign higher rewards to preferred outputs and lower rewards to dispreferred outputs which implies the pooling function learns rich features with respect to prompt-completion pairs. We believe that it would be possible to also extract additional information from the output of the pooling function to detect finer grained signals such as helpfulness, humour, creativity, toxic content, etc. Such an approach has been explored by ArmoRM [41] which was first trained with several reward heads to reflect different types of preference scores and then a gating system to chose the best combination of heads based on the context. We believe this technique can be integrated into the DPH process to facilitate context-dependent re-ranking.The main benefit of DPH being its ability to perform alignment without directly effecting the model's output distribution is also its main limitation: unlike other alignment techniques which can help prevent the model generating harmful outputs, DPH is only capable of detecting harmful outputs.Although we do include DPO alignment in our experiments to reduce the likelihood of harmful outputs, DPH does not require such model alignment to function, which shifts the responsibility of rejecting harmful outputs to the end user or service provider.From a computational perspective, DPH reformulates alignment as a re-ranking process, which necessitates sampling multiple candidate responses during inference. This approach requires more computational resources and memory compared to generating a single response. However, modern hardware used for LLM inference is often underutilized by smaller language models when generating singular responses in an autoregressive manner. By leveraging this unused compute capacity, we can generate multiple candidate responses in parallel when applying DPH to smaller models, without significantly increasing latency.In this paper we introduced Direct Preference Heads, a novel form of language model alignment which is performed at inference time to prune candidate completions for a given prompt. Unlike other alignment techniques which coerce the model into generating human preference aligned outputs, DPH instead produces reward scores for candidate outputs without affecting the actual generation process and therefore avoids the issue of RLHF leading to degraded performance when applied to smaller language models. We formulated two loss functions for DPH and find strong connections to Conservative DPO, implying that DPH is robust to label noise and can be tuned to a specific confidence margin. Finally, we evaluated our methods on a number of NLU, commonsense reasoning and reading Comprehension tasks and found that DPH is able to consistently outperform both our SFT baseline and multiple publicly available language model checkpoints of varying size and training volume.As with all language modelling systems we cannot guarantee all responses produced by our models are factually correct nor can we guarantee that they are safe and free from harmful content. Our work focuses on creating a system that helps filter out incorrect and harmful messages by scoring candidate outputs, but as with all alignment techniques our models may be susceptible to so-called 'jailbreaks' which can coerce the model into incorrectly assigning a higher score to less desirable content. To maximise safety DPH should be implemented alongside other safety guardrails such as Llama Guard [19] when used for publicly facing chat systems, and we intend for our provided model checkpoints to be used for reproduction of results and further research in the field of alignment.A.1 Full Proof of Theorem 1We can prove Theorem 1 by examining the partial gradients with respect to the rewards.∂ ∂rw L SepDPH (r w , r l ) = ϵ - 1 e rw + 1 (8a) ∂ ∂r l L SepDPH (r w , r l ) = 1 e -r l + 1 -ϵ (8b)From equations 8a and 8b we find that the partials gradients are both equal to zero at the points r w = log 1-ϵ ϵ and r l = log ϵ 1-ϵ respectively. It is also interesting to note that log 1-ϵ ϵ + log ϵ 1-ϵ = 0 which implies the positive and negative rewards will converge to an equal distance from 0.∂ 2 ∂r 2 w L SepDPH (r w , r l ) =e rw (e rw + 1) 2 (9a)∂ 2 ∂r 2 l L SepDPH (r w , r l ) = e r l (e r l + 1) 2 (9b)If we derive the second derivatives for the rewards, as shown in equations 9a and 9b, we find that they are both strictly positive for all values of r w and r l which implies that Separable DPH is convex with respect to the rewards.We can prove Theorem 2 by examining the partial gradients with respect to the rewards. From equations 10a and 10b we can see a symmetry emerge, where the partial gradients with respect to the preferred logits are equal and opposite to the partial gradients with respect to the dispreferred logits. If we reparameterise the loss function such that r ∆ = r w -r l we can derive the following partial derivativeEQUATIONwhich is equal to zero for ϵ ∈ (0, 0.5] at the point r ∆ = log 1-ϵ ϵ . If we derive the second derivative of the Contrastive DPH objective function with respect to the reward margin r ∆ we obtain the following formula∂ 2 ∂r 2 ∆ L ConDPH (r w , r l ) =e r∆ (e r∆ + 1) 2 (12) which is strictly positive for all values of r ∆ , and -with respect to the reward logits -frames Contrastive DPH as a convex optimization problem with the additional properties of guaranteed convergence to a fixed margin for all ϵ ∈ (0, 0.5].We provide an illustration of the loss landscapes to give a visual comparison of how our objective functions 'pull' rewards towards the optimal margin bounds. Note that three of the GLUE tasks have no license specified on their homepages nor within their publications: CoLA claims their dataset falls under "fair use," no concrete license can be found for RTE nor its pre-cursors, and SST-2 does not specify a license.For brevity, we only include the prompt templates of the tasks we use for evaluation. All other prompt templates are listed within the code repository.Below is an instruction that describes a task. Write a response that appropriately completes the request using the provided answer options.Given the following sentence, answer the question with "yes" or "no".Question: Does this sentence make sense?Answer:{{no | yes}}Below is an instruction that describes a task. Write a response that appropriately completes the request using the provided answer options.Given a premise statement and a hypothesis statment, respond with "True" if the premise entails the hypothesis, respond with "False" if the premise contradicts the hypothesis, or respond with "Neither" if the statements are neurtral.Premise: {{premise}} Hypothesis: {{hypothesis}} Question: True, False or Neither?Answer: {{True | Neither | False}} GLUE -MRPC System User AssistantBelow is an instruction that describes a task. Write a response that appropriately completes the request using the provided answer options.Given the following sentences, answer the question with "yes" or "no".Sentence 1: {{sentence1}}Sentence 2: {{sentence2}}Question: Do both sentences mean the same thing?Answer:{{no | yes}} GLUE -QNLI System User AssistantBelow is an instruction that describes a task. Write a response that appropriately completes the request using the provided answer options.Given the following sentences, answer the question with "yes" or "no". Below is an instruction that describes a task. Write a response that appropriately completes the request using the provided answer options.Given the following sentences, answer the question with "yes" or "no". User AssistantBelow is an instruction that describes a task. Write a response that appropriately completes the request using the provided answer options.Given the following sentences, answer the question with "yes" or "no". We initialise the embeddings from OPT-125m and use embedding tying for the language modelling head. Since our model dimension is 1536 while the embedding dimension is 768 the model contains an up-projection as the first layer of the backbone and a down-projection for the final layer. There are a total of 18 transformer blocks in the model backbone which use pre-layer norm in the attention and FFN residuals. The attention blocks have 24 attention heads and we use RoPE with a base frequency of 500,000 for positional embedding, and the FFN block uses SwiGLU activation with an intermediate dimension of 4096. The context window of the model is 2048 tokens and the Transformer XL recurrent memory contains 2048 tokens which allows the model to use a sliding window size of up to 4096 tokens at inference without any degradation.The model was trained for approximately 100 billion tokens on the first 24 shards of The Pile. Each batch is constructed of 480 sequences of 2048 tokens each which are continuously sampled from the datasets shards using queues for the Transformer XL style pre-training method.We use the LaProp optimizer [44] with β 1 = 0.9, β 2 = 0.95, a max learning rate of 6e-4 which warms up over 2000 steps and cosine decays down to 6e-5, LR-coupled weight decay of 0.1 and global gradient clipping with a max norm of 1.Each epoch of 256 steps takes 1 hour and 59 minutes on 4x RTX A4500 GPUs. For the full 398 epochs (or 101888 steps) this comes out to around 790 hours or just under 33 days of training time (ignoring time for validation in-between epochs and at the end of training).1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: We conduct a variety of experiments and ablations which validate the claims made in the abstract. We make use of standard benchmarks in the field of language models and compare our method with results from smaller, comparably sized and larger models. 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. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We include a broader impacts section at the end of our paper and describe potential pitfalls of using our 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., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Although we are publicly releasing our model checkpoints they have an incredibly low risk of misuse as other significantly more capable and more accessible models exist. Additionally, the Hugging Face checkpoints do not support 'AutoModel' loading, which means a level of technical knowledge is required to set the model up, making it less accessible those outside of the ML research sphere. 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? • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: The approach presented is interesting -- rather than directly fine-tuning a model, it proposes to train a reward model to be used at inference time only to rerank base model outputs. In practice, this reward model is the same as the base model, which has been fine-tuned to have a reward head, but not directly on generated output sequences.

The approach requires more computation at inference time (in contrast to simple fine-tuning) and its usefulness is also constrained by inference-time compute (required to sample candidates to rerank). 

The experiments focused only on several small models, but the authors argue that the approach itself is focused on improving small language models, where inference-time compute is less burdensome. I suggest that the authors refine their arguments about the approach being focused on small models, and why this justify experiments and baselines from small / older models only.

There are also not experiments that provide direct evidence that the model hasn't degraded through the fine-tuning process in training it to predict rewards for inference time.

There is no direct comparison with other standard RLHF approaches such as PPO, although there are comparisons to DPO, and authors argue that PPO would be computationally infeasible for their experiments.

Reviews: review: paper_summary: paper_summary: The authors propose an inference time method to align language models with human preferences without harming the model’s reasoning abilities. The method creates an auxiliary reward head that operates during inference to score potential outputs without changing the output distribution directly. The author validates approach by comprehensive evaluations on NLU, commonsense reasoning, and reading comprehension datasets.
strength_weakness: strength_weakness: 1. Inference Time Overhead: The Direct Preference Heads method involves using language models to generate multiple candidate responses, which must then be evaluated for selection. This process is less efficient compared to traditional fine-tuning methods due to the additional computational steps required.

2. Dependence on Initial Sample Quality: The effectiveness of the generation-then-reranking approach highly relies on the quality of the initial responses generated. If the sampled responses are of low quality, DPH is unable to enhance the output, as it can only rerank the given candidates.

3. Limited Evaluation Scope: The evaluation of the DPH method is restricted to a single model of 550M parameters. It is not clear how well the method would perform when scaled to larger models, as the results might not be consistent across different model sizes.
questions: questions: 1. The paper would be more convincing if the author applied the proposed method to a larger-sized model and demonstrated a scaling trend. The use of only one model makes the effectiveness of the approach a bit underdetermined

2. How does the proposed DPH method compare with other inference time algorithm such as rejection sampling? 

3. How do you choose baselines when reporting performance? For example, Table 1 reports the model performance on BERT, an early-stage masked language model, instead of other competitive baselines.
limitations: limitations: The authors addressed all the limitations in the conclusion section.
paper_summary: paper_summary: Paper proposes Direct Preference Heads, which learns a reward prediction head using a pretrained model, without affecting the model's output distribution.
strength_weakness: strength_weakness: Some missing/confusing parts about the proposed method:
- How are the learned reward predictions used? Some parts of the paper seems to suggest that multiple responses are sampled from the model, and the reward head is used to rank the samples? I think this is an important component of how the proposed method would be used, so it would be great to further elaborate on this procedure, e.g. how many samples?
- Some parts of the paper seems to suggest that the proposed approach only learns a reward prediction head on top of an existing model without modifying the original model (e.g. in abstract, "... without directly affecting the output distribution of the language modeling head ..."), but other parts of the paper stated that it is better to finetune the entire model when learning to predict rewards (e.g. Section 4.3). However, the stated reason for why DHP might be better than other RLHF techniques like DPO is that further alignment (past SFT) can hurt model performance. If DHP also requires further finetuning, what is the intuition for why DHP performs better than DPO?

Related works
- Paper introduces PPO, DPO and cDPO. However, there has been a large number of recent works in RLHF beyond these papers. It would be great to present a more complete related works section with a more comprehensive survey of existing work.
questions: questions: See weaknesses
limitations: limitations: Yes
paper_summary: paper_summary: This paper introduces Direct Preference Heads (DPH), a novel method for aligning language models (LLMs) with human preferences at inference time. DPH works by adding an auxiliary reward head to the LLM that learns to predict human preference scores for generated outputs. This allows the model to self-evaluate multiple candidate outputs and select the highest-scoring one, effectively pruning undesirable responses.

The authors argue that DPH offers several advantages over traditional Reinforcement Learning from Human Feedback (RLHF) methods like PPO and DPO:

* Inference-time alignment: DPH aligns the model at inference time, avoiding the potential degradation of reasoning abilities often observed with RLHF during training.
* Lightweight: DPH requires only a single model to produce both responses and rewards, unlike RLHF which typically involves multiple models.

The paper presents two objective functions for DPH, separable and contrastive, and demonstrates their connection to Conservative DPO. Experiments on GLUE, GPT4All, and RACE datasets show that DPH consistently outperforms both supervised fine-tuning (SFT) and DPO alone.
strength_weakness: strength_weakness: * Lack of statistical significance: The paper does not report error bars or statistical significance measures, making it difficult to assess the robustness of the results.

* Missing the comparison to standard RLHF methods: the paper does not include comparison to the PPO baseline, which is the most commonly used technique.

* The authors do not discuss other inference-time alignment techniques that exist in the literature.
questions: questions: * The paper uses different sampling strategies for SFT and DPH. It would be helpful to discuss the rationale behind these choices and explore the impact of different sampling strategies on DPH performance.

* Sampling multiple candidate outputs at inference time can be computationally expensive, especially for larger LLMs. The paper doesn't explicitly address this cost, which could be a practical concern for real-world applications. Can you elaborate more on this in the paper?
limitations: limitations: The paper focuses on a 551M parameter model. It would be valuable to see how DPH scales to larger models, as RLHF is known to be more effective for larger models. Comparing DPH to RLHF-aligned larger models would provide a more comprehensive evaluation.

Also RLHF baselines on top of DPO need to be added.
paper_summary: paper_summary: The paper introduces DPH, a new method for pre-trained language models that addresses the limitations of RLHF. Unlike traditional RLHF, which can compromise a model's reasoning abilities and cause hallucinations, DPH employs an auxiliary reward head to learn human preference signals without altering the LM's output distribution. The authors conduct a theoretical analysis linking their objective function to cDPO and demonstrate that DPH can be used in conjunction with existing alignment techniques to improve performance. The experiment results show that models fine-tuned with DPH achieve higher scores compared to those fine-tuned with SFT or DPO alone.
strength_weakness: strength_weakness: 1.	The claim that the auxiliary reward head avoids affecting the output distribution of the language modeling head is undermined by the practical implementation. Specifically, this claim does not hold when the backbone language model is updated to learn the preference distribution (line 180) and when the model is updated using a joint loss function (Eq. 7). These aspects of the implementation weaken the validity of their claim.
2.	Experiment setup: The baseline setup in the experiments is not reasonable. The authors compare their method with other language models that have distinct training settings and data, rather than comparing with other alignment methods. This makes it difficult to infer the superiority of the proposed method over existing alignment techniques, as the comparisons are not directly relevant.
3.	Clarity and organization: The writing is not easy to follow and requires improvement. Specifically, sections 4.3 and 4.4, which are parts of the methodology, are incorrectly placed in the experiment section. Additionally, the evaluation protocol (section 5.1) should be introduced in section 4 but is instead placed in the results section. This misplacement affects the clarity and logical flow of the paper, making it harder for readers to understand the methodology and its evaluation comprehensively.
questions: questions: 1. Why does DPH not require an SFT sampling and human labelling stage? In section 4.4, you mentioned SFT in your training pipeline.
2. What is the relationship between the proposed approach and the sentence "Would I Lie To You?" in your title?
limitations: limitations: The authors include a discussion regarding the limitations of the proposed method.

Comments: title: Response
comment: After reading the rebuttal and other reviews, here are my additional comments (hope other reviewers are also aware of following): 

TL;DR 
- The experimental setup is considerably flawed: 
    - Evaluating an **alignment/preference tuning approach** using baselines with only pretraining (Table 2, Pythia & TinyLlama; Table 3 Llama) (some baselines are SFT)
    - Not including other preference tuning approaches as baselines
- Motivation is disconnected the proposed approach
    - Claiming RLHF compromising an LLM (authors' motivation) but updating the parameters of the backbone model by DPH
    - Not showing DPH does not compromise an LLM

---   
***Experimental setup is considerably flawed***

The experimental design is deeply misguided. The proposed approach is a preference tuning (PT) algorithm that is conducted after SFT and trained by preference data. **A legitimate experimental setup would directly compare this with other preference tuning algorithms (DPO, KTO, IPO, etc.)**. However, the comparisons presented are like: 
- Proposed approach: Base_Model_A + Pre-training_A + SFT_A + Preposed_PT
- Others: Base_Model_B + Pre-training_B (+ SFT_B or no SFT_B + unknown PT or no PT)

It is **really confusing when you expect readers to evaluate a preference tuning appraoch by comparing a PT-ed model with a pre-trained model** (Table 2, Pythia & TinyLlama; Table 3 Llama). 

> To obtain baselines for similarly sized models using other alignment methods would have required us to train these models ourselves which requires significant time and compute. 

This excuse is unconvincing. I believe you only need to take SFT checkpoint of your model (551M) and conduct other PT baselines like DPO. The computational cost would not exceed your method (cDPO + DPH). Furthermore, the same comparison based on TinyLlama backbone is expected to make it "real" extensive evaluation. These are tiny models and can be done on a single A100 GPU.

---
***Motivation is disconnected the proposed approach***

The authors argue that RLHF compromises an LLM, and thus propose a method to update only the added reward heads, purportedly to prevent any impact on the model's output distribution (line 10). However, their practical implementations (section 4.3) are: 
- Updating the backbone model, which is shared by language modelling head, by DPH + prior regularization
- Updating the backbone model by DPH and cDPO loss

These steps clearly alter the model's output distribution, directly contradicting the claimed motivation. Thus, the superiority of the proposed approach should seek strong support from the empirical results, which is unfortunately missing.
Thank you for your response. It would indeed be valuable to include results from models such as Llama 2, however as stated in my previous response we elected to use verified results which were not collected by ourselves to reduce any potential bias in our testing pipeline (such as from prompt selection, or aggregation method for log-probabilities).

However, for some benchmarks (the GPT4All suite and RACE) it may be possible to use a standardised evaluation suite such as Eleuther AI's LM Eval Harness to obtain reasonably unbiased results, and we would have the compute necessary to do so for a future revision. Evaluating GLUE, however, would be a bit more tricky, especially since STSB is a regression tasks and we had to compromise our own model's evaluation by forcing it to chose from integer predictions.
Thanks a lot for the additional details you provided. However, I’m still not fully convinced about the choice of baselines included in the paper and the experimental setup. Specifically, I’m curious why larger-scale state-of-the-art models, such as LLaMA-2 7B, weren’t used for the main experimental results and compared with all baselines at that scale. The choice of smaller models makes it challenging to determine if the improvements would hold at a larger scale. I’m already leaning towards accepting the paper and will keep my score unchanged, but I hope future versions will include these considerations.
I thank the authors for the response! It would be great to incorporate the clarifications in the rebuttal into the paper.