Mitigating the Alignment Tax of RLHF (2024)

Abstract: LLMs acquire a wide range of abilities during pre-training, but aligning LLMs under Reinforcement Learning with Human Feedback (RLHF) can lead to forgetting pretrained abilities, which is also known as the alignment tax. To investigate alignment tax, we conducted experiments with existing RLHF algorithms using OpenLLaMA-3B, which revealed a pronounced alignment tax in NLP tasks. Whereas, despite various techniques to mitigate forgetting, they are often at odds with the RLHF performance, leading to a trade-off between alignment performance and forgetting mitigation, leading to an alignment-forgetting trade-off. In this paper we show that model averaging, which simply interpolates between pre and post RLHF model weights, surprisingly achieves the most strongest alignment-forgetting Pareto front among a wide range of competing methods. To understand its effectiveness, we offer theoretical insights into model averaging, revealing that it enhances performance Pareto front by increasing feature diversity on the layers where tasks share overlapped feature spaces. Empirical evidence corroborates our analysis by showing the benefits of averaging low-level transformer layers. Building on the analysis and the observation that averaging different layers of the transformer leads to significantly different alignment-forgetting trade-offs, we propose Heterogeneous Model Averaging (HMA) to Heterogeneously find various combination ratios of model layers. HMA seeks to maximize the alignment performance while incurring minimal alignment tax. Moreover, we validate HMA’s performance across a range of RLHF algorithms over OpenLLaMA-3B and further extend our findings to Mistral-7B which is evaluated by open-sourced preference model and GPT4. Code available here.

## introduction large language models (llms), such as gpt4 @xcite , bard @xcite , and claude @xcite , have attracted widespread attention due to their remarkable achievements. llms are pre-trained on vast datasets, which equip them with the ability to effectively handle diverse tasks, e.g., gpt-3 showcases its prowess in various tasks such as reasoning, common sense questionanswering (qa), translation, and so on. while llms exhibit strong abilities among various benchmarks, they still require alignment with human preferences, including the principles of being helpful, honest, and harmless as outlined by @xcite . the goal is to ensure that llms are designed to assist users in completing tasks, provide truthful information without deception, and avoid causing harm, whether physical, psychological, or social, to individuals or the environment. the process of aligning llms with human preferences often involves the application of reinforcement learning with human feedback (rlhf) @xcite , as shown in figure 1 . although rlhf allows llms to align with human expectations, prior studies @xcite @xcite have found that this approach can lead to forgetting in the diverse abilities that the llms have already acquired, as illustrated in figure 1 . this phenomenon, also known as the "alignment tax" in the literature, has accumulated substantial attention from both academia and industry @xcite @xcite @xcite . investigating alignment tax. in this paper, we first conduct a comprehensive investigation on alignment tax and develop methods to reduce alignment tax while maintaining the alignment performance. in particular, we followed the approach presented by @xcite and evaluated alignment tax using multiple nlp benchmarks from common sense qa, such as arc easy and challenge @xcite , race @xcite , and piqa @xcite , reading comprehension benchmarks including squad @xcite and drop @xcite , and translation tasks, including wmt 2014 french to english translation @xcite (c.f. section 3). our primary focus is on aligning the openllama-3b on the helpfulness and harmlessness dataset @xcite using rejection sampling finetuning methods @xcite (also known as the best-of-n algorithm). in the later part, we extend our experiments to mistral-7b and direct preference optimization (dpo, @xcite ). we mainly focus on rsf and dpo since they are popular and nearly all of the latest opensourced llms on the leaderboards are aligned by these two methods foot_0 . indeed, we observed a substantial alignment tax on these benchmarks consistently, confirming the findings of @xcite . specifically, as we gained a higher reward during rlhf, indicating better alignment with human preference, the alignment tax also increased simultaneously, clearly inducing a alignment-forgetting trade-off. surprising effectiveness of model averaging over. we then compare various methods developed in different communities as potential rescues to alleviate the alignment tax. this includes the model averaging method @xcite from out-of-distribution (ood) generalization literature, regularization-based techniques from the continual learning literature @xcite @xcite , low-rank adaptation (lora) @xcite from the parameterefficient fine-tuning literature, as well as the utilization of reward penalty from the reinforcement learning literature @xcite @xcite . interestingly, we found that model averaging, which simply interpolates between the weights of models before and after rlhf, achieves the most efficient alignment-forgetting pareto front. in appendix c.1, we further show and discuss the in-effectiveness of experience reply (rebuffi et al.) method compared with ma. understanding the effectiveness of model averaging. to understand the effectiveness of model averaging, we provide theoretical insights based on the framework of @xcite . in particular, we show that the method can enhance pareto front by increasing feature diversity on layers where two tasks share similar feature spaces. empirical evidence also indicates that averaging the low-level layers of transformers consistently improves both alignment reward and nlp task performance. this aligns with our theoretical insights, as tasks could share similar lower-level features, e.g., better word representation on low-level layers benefits both nlp and alignment tasks. heterogeneous model averaging. we noticed that averaging different layers of the transformers unveiled notably distinct patterns of alignmentforgetting trade-off, aligning with our earlier analysis that tasks may exhibit varying overlapping feature spaces in different layers. motivated by this observation, we propose heterogeneous model averaging (hma), which adaptively averages different parts of the models during model averaging. we start by dividing the transformer into k parts and assigning unique averaging ratios for each part, represented as α @xmath0, 1] for the ith part. hma aims to maximize alignment reward by optimizing the averaging ratios (α 1 , . . . , α k ) while maintaining the overall alignment tax, thus consistently improve the alignment-forgetting pareto front. to demonstrate the efficiency of hma, we also contrasted our method with other rlhf techniques, including direct preference optimization (dpo). @xcite we further substantiate our findings on mistral-7b where evaluations conducted by open sourced perference model and gpt4, which further corroborates our empirical findings on openllama-3b. we summarize our contributions as follows: the paper is structured as follows: we conduct a systematic investigation of existing methods in section 3-4. in section 5, we provide insights into the effectiveness of model averaging. subsequently, we propose heterogeneous model averaging in section 6. we conclude the paper in section 7. in this section, we provide comparison of this work with existing works to highlight the novelty of our findings. we defer more comprehensive related works to appendix a. existing works of model averaging for llms. previous research has covered certain aspects of model averaging. @xcite demonstrate the utilization of model averaging to construct a more resilient reward model for reinforcement learning with human feedback (rlhf). in a similar vein, @xcite employ model averaging to merge policy models trained for distinct objectives, facilitating multi-objective rlhf. @xcite introduce the integration of moving averaging to enhance pre-training. however, none of these studies investigate the alignment tax, and their findings are independent of our research. existing works on finding adaptive combinations for model merging. previous studies @xcite have also discussed the idea of dynamically assigning different weights to different layers when merging models, aiming to maximize performance on a specific task (e.g., t i ). these approaches assume access to the taskspecific data t i . however, considering the nature of alleviating alignment tax, which aims to mitigate forgetting across a extremely wide range of tasks (t j 1 ...t j k ), these methods fail to effectively optimize performance for multiple tasks simultaneously. in the appendix e.4, we demonstrate that using the method proposed by @xcite , which optimizes for a single task, does not effectively address forgetting on the other tasks. furthermore, our work is the first to provide an explanation for the surprising effectiveness of model averaging in alleviating forgetting, as well why we should assign heterogeneous combination ratios. existing works on the forgetting of language models. most research on forgetting in language models focuses on sequentially pre-training @xcite @xcite @xcite or fine-tuning tasks @xcite @xcite @xcite , e.g., sequentially training on task t i and then task t j . they evaluate forgetting by measuring the model's performance on a task (e.g., task t i ) after training it on another task (e.g., task t j ). however, these methods have not explored the effectiveness of model averaging. in our case, we demonstrate the significant power of model averaging which outperform a wide range of existing methods. furthermore, existing works assume that the data size of each task is comparable (i.e., the dataset size of t i and t j is similar), allowing for a subset (e.g., 10%) of old task data replay, which is shown to effective alleviate the forgetting without excessive computation overhead in their settings. however, in our alignment tax situation, we aim to preserve a wide range of abilities gained during pre-training, which is challenging since pretraining datasets are often not publicly available. in appendix c.1, we show that even when we have access to the pre-training data and replay a subset up to four times larger than the rlhf data (which costs significant computation overhead), experience replay still under-performs model averaging in two out of three benchmarks. this is likely due to the vast size of the pre-training data, where the subset only covers a small fraction of it (e.g., only covers ~0.01% of the pre-training data). so replay methods are less practical for alleviating alignment tax. ## experimental settings basic setting. we chose the openllama-3b model @xcite because (1) it is computational friendly compared with 7b models (2) it has openly available pre-training dataset, which is convenient to investigate experience replay in appendix. c.1. furthermore, we extend the experiments to mistral-7b in sec. 6. following the standard procedure outlined in @xcite , we initially conducted instruction tuning, followed by rlhf. here, θ represents an llm with parameters θ, with the pre-trained model denoted as θ pre . we commenced with instruction fine-tuning for θ pre on sharegpt 3 , which yielded θ 0 . subsequently, rlhf was performed on θ 0 to obtain θ. similar to the methodology proposed in @xcite , the alignment tax was evaluated by comparing the performance regression of θ with θ 0 across various nlp tasks. the whole procedure and notations are illustrated in fig. 1 . datasets for evaluating alignment tax. following the approach in @xcite , our evaluation of alignment tax encompasses various nlp benchmarks: (a) common sense qa: this includes arc easy and challenge @xcite , race @xcite , and piqa @xcite , with the performance being assessed using accuracy. (b) reading comprehension: we employ squad @xcite and drop @xcite to gauge reading comprehension ability, with evaluation based on the f1 score for both datasets. (c) translation: our evaluation utilizes wmt 2014 french to english translation @xcite , with performance measured using bleu @xcite scoring. ## evaluating existing methods in figure 12 of appendix e.1, we visualize the training procedure in terms of the alignmentforgetting trade-off during rlhf. specifically, we can clearly see that as the rlhf proceeds, the reward begins to increase while the translation and reading comprehension ability continues to drop. interestingly, we observe that the performance of common sense increases first and then drops. given that alignment tax is inherently a catastrophic forgetting issue, we then proceed to explore methods to reduce alignment tax. research focused on reducing forgetting is mainly classified into two main categories, depending on the availability of the pretraining dataset. we also investigate the reward penalty method developed in rl community in appendix c.2. ## unravelling the mysteries of model averaging for alleviating alignment tax given the promising performance of model averaging, we try to understand the efficacy of model averaging in this section and motivate our method to improve it. we utilize the theoretical framework proposed by @xcite to gain insights into its effectiveness in alignment tax. while the framework addresses classification problems, the insights derived can aid our understanding of model averaging. we also conduct empirical analysis using a generative model (openllama-3b) to verify these theoretical insights. analyzing the performance of model averaging in alignment tax is more intricate compared to the work of the study by @xcite focuses on out-of-distribution (ood) scenarios, where the same task is performed under different distributions. in contrast, our focus in alignment tax is to comprehend the performance trade-offs among different tasks. to illustrate, consider the entire feature space y and two tasks with label spaces y a ⊂ y and y b ⊂ y, with the simplifying assumption that @xcite only considers the case where y @xmath1, we extend these results to encompass the case where theoretical settings. suppose we have many features s @xmath2, ..., x d . following @xcite , we adopt a simplified model f (x) = wφ(x) where @xmath3, φ(x) = foot_1 furthermore, even if the tasks are very similar, fitting two models on them can rely on different features due to randomness in data or training procedures @xcite ). we will investigate the performance of model averaging in case (1) and ( 2 ) to gain insights on when it works. following @xcite , we assume each feature is weak, failing with probability p. the effectiveness of model averaging is given by implications. proposition 5.1 demonstrates that when t a and t b are more similar, the averaging of models (f a and f b ) yields greater improvement. however, this improvement is reduced if f a and f b use more overlapping features. recall that each weak feature can fail with probability p. if t a and t b are similar, the feature utilized by the two models would be projected into a shared space, allowing model averaging to take advantage of a more diverse set of features. this diversity reduces the probability of model failure because a diverse set of features is less likely to fail together simultaneously @xcite . however, if t a and t b are dissimilar, for example, if |y a ∩ y b | = 0 and the feature spaces corresponding to y a and y b are disjoint, then the features in the space of y a would not provide any information for predicting y b . therefore, averaging f a and f b would not improve the prediction of either task in this case. refer to appendix f.3 for a detailed discussion. notably, the model θ 0 excels in nlp abilities before rlhf, while the model θ excels in alignment reward after rlhf. using an analogy, we can equate nlp tasks with t a , alignment with t b , θ 0 to f a , and θ to f b . recall that we adopt a simplified model for theoretical analysis by considering only one layer feature learner, although, in practice, we average a deep transformer with 26 layers. research has shown that different layers in deep neural networks capture varying levels of features @xcite @xcite . for instance, low-level layers capture low-level features. furthermore, tasks share similar feature space at a low level (alternatively, from the perspective of low-level layers, tasks look more similar). for example, improving the low-level features such as better word representation could enhance both rlhf reward and nlp tasks. therefore, according to proposition 5.1, averaging the low-level layers could potentially elicit more improvements in both t a (nlp tasks) and t b (alignment reward) than higher layers. ## heterogeneous model averaging we have already shown that averaging different layers results in diverse patterns of alignmentforgetting trade-off @xcite . therefore, different layers should not be equally treated during averaging. this leads to a natural question: can we enhance the alignmentforgetting trade-off by using adaptive weights for different layers? consequently, we conduct proofof-concept experiments to provide affirmative answers to this question and subsequently propose a practical algorithm. proof of concept. the following proof of concept experiments provide insights into average different layers with various ratios. we use different averaging ratio, i.e., α 1 , α 2 , α 3 , for the three parts. specifically, the ith part of the averaged model is simply α i θ [i] + (1α i )θ heterogeneous model averaging. upon dividing the transformer into k parts, our objective is to adaptively determine a combination ratio for different layers that consistently perform well across an extensive range of tasks. the conventional averaging method uses a shared α for all layers, playing a crucial role in defining the trade-off between re-ward and tax. we aim to identify an optimized combination of (α 1 , ..., α k ) to replace a uniform α. let θ(k) represent the model merged by (α 1 , ..., α k ). in particular, the kth component of the merged model θ(k) is given by θ [k] (k) := α k θ [k] + (1α k )θ to optimize the pareto-front influenced by α, we identify combination ratios corresponding to each α. subsequently, we establish the mean of (α 1 , ..., α k ) as α and ascertain the best combination of (α 1 , ..., α k ) to maximize the reward. specifically, denoting ω := 1 k k α @xmath4, α 1 , ..., α @xmath5, 1] , we solve: the intuition behind hma is outlined as follows: (1) when maintaining the mean, i.e., 1 k k α k , as α, we can compare hma performance with the performance of vanilla model averaging with the same α. (b) we only optimize k parameters, where k is typically small. for example, we adopt @xmath6. this helps to ensure that the forgetting level of (α 1 , ..., α k ) remains close to α. intuitively, if we optimize a large number of parameters, it could easily lead to over-fitting in the in-domain (rlhf reward) and may also result in more significant forgetting. the whole algorithm is summarized algorithm 1 in appendix. results. the results of hma are shown in figure 5 . we can see that hma can consistently push forward the perato-front of the vanilla model averaging. furthermore, such improvement is consistent over various rlhf algorithms. more detailed results (e.g., on commonsense qa and translation with different rlhf algorithms) of hma can be found in appendix e.5. ablation results on different k. we tested different values of k with α = 0.2, 0.4, 0.6 as illustrated in figure 5 (right). the trade-off curve shows a slight decrease as we increase k from 3 to 6 and 9, but still consistently improves over the vanilla model averaging. this decrease is likely due to overfitting. specifically, comparing the performance of hma with different k for the same mean ratio, we observe that as the alignment reward increases with an increase in k from 3 to 9, the reading comprehension performance drops. how to choose the averaging ratio. in practice, we determine the averaging ratio α for adopt- other models results. to further validate our method on larger llms, e.g., mistral-7b @xcite based models, we apply model averaging (ma) and heterogeneous model averaging (hma) on zephyr-7b-β foot_2 @xcite which is trained with dpo on the sft version, mistral-7b-sft-β foot_3 . we also apply hma on zephyr-7b-gemma foot_4 which is aligned based on gemma-7b foot_5 model. here we use the the publicly available preference model pairrm @xcite to judge the helpfulness and evaluate models on alpacaeval 2.0 @xcite . we reports the win rates of each model. figure 6 (top) shows that the trends of averaging different layers evaluated by pairrm are similar with the results evaluated by our own reward model. the results range across α = 0, 0.2, . . . , 1.0 depicted in figure 6 (bottom) demonstrate that ma effectively achieves a strong pareto front to mitigate forgetting in the mistral-7b models. additionally, our hma algorithm shows further improvement compared to the ma method. ## conclusion in this paper, we highlight the surprisingly effectiveness of model averaging and propose the heterogeneous model averaging (hma) framework to further enhance the performance.