Source-linked AI summary
Mitigating the Alignment Tax of RLHF
Yong Lin, Hangyu Lin, Wei Xiong, Shizhe Diao, Jianmeng Liu, Jipeng Zhang, Rui Pan, Haoxiang Wang, Wenbin Hu, Hanning Zhang, Hanze Dong, Renjie Pi, Han Zhao, Nan Jiang, Heng Ji, Yuan Yao, Tong Zhang
TL;DR
RLHF can cause LLMs to forget pretrained abilities, creating an alignment tax that trades off against alignment performance. This paper investigates the trade-off, analyzes model averaging, and introduces HMA to optimize layer-wise averaging ratios. Model averaging and HMA improve the alignment-forgetting Pareto front across models, benchmarks, and RLHF algorithms, though HMA does not eliminate the tax.
Problem
RLHF can cause forgetting of pretrained abilities, raising the question of how to reduce alignment tax without sacrificing alignment performance.
Method
The paper systematically evaluates mitigation methods, analyzes model averaging through feature diversity and layer effects, and introduces HMA for optimizing separate layer-part averaging ratios.
Results
Model averaging achieves a strong alignment-forgetting Pareto front, while HMA consistently improves that front across benchmarks, RLHF algorithms, and OpenLLaMA-3B and Mistral-7B.
Takeaways & Limitations
Averaging low-level Transformer layers and using heterogeneous layer ratios provide effective ways to improve alignment while retaining NLP task performance.
Takeaways & Limitations
HMA significantly alleviates but does not fully eliminate the alignment tax.
Abstract
from arXiv · showhide
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: https://github.com/avalonstrel/Mitigating-the-Alignment-Tax-of-RLHF.git.
1 Introduction
RLHF aligns LLMs with human preferences but can cause forgetting of pretrained abilities, creating an alignment tax. The paper investigates this trade-off and finds model averaging and HMA effective for retaining NLP performance while preserving alignment.
- Motivation: RLHF targets helpful, honest, and harmless behavior but can make LLMs forget abilities acquired during pre-training.This forgetting is termed the alignment tax.
- Investigation: The study evaluates alignment tax across common-sense QA, reading comprehension, and translation benchmarks.Experiments primarily align OpenLLaMA-3B using Rejection Sampling Finetuning, then extend to Mistral-7B and DPO.
- Investigation: As RLHF reward increases, alignment tax also increases, revealing an alignment-forgetting trade-off.The reported pattern is observed consistently across the evaluated NLP benchmarks.
- Model averaging: Model averaging, which interpolates pre- and post-RLHF weights, achieves the most efficient alignment-forgetting Pareto front among compared methods.The comparison spans methods from multiple research communities addressing forgetting or alignment tax.
- Analysis: The analysis links model averaging's effectiveness to increased feature diversity in layers where NLP and alignment tasks share feature spaces.Empirical evidence particularly supports averaging low-level Transformer layers.
- Heterogeneous Model Averaging: HMA optimizes distinct averaging ratios across model parts to maximize alignment reward while maintaining the overall alignment tax.It is evaluated across RLHF algorithms, benchmarks, and model types including OpenLLaMA-3B and Mistral-7B.
2 Discussion with existing works.
The paper distinguishes its contribution from prior work on model averaging, adaptive model merging, and language-model forgetting. It argues that task-specific layer weighting is insufficient for mitigating forgetting across many tasks simultaneously.
- Existing model averaging: Prior model-averaging studies address reward-model robustness, multi-objective RLHF, or pre-training, but not the alignment tax.The paper positions its investigation as addressing an unstudied aspect of model averaging.
- Adaptive combinations: Adaptive model-merging methods assign layer weights for a specific task and assume access to that task's data.The paper contrasts this setting with alignment-tax mitigation across a wide range of tasks.
- Adaptive combinations: Single-task optimization methods fail to optimize performance for multiple tasks simultaneously when mitigating broad forgetting.The paper reports that a method optimized for one task does not effectively address forgetting on other tasks.
- Novelty: The paper presents itself as the first to explain model averaging's effectiveness for alleviating forgetting and why heterogeneous combination ratios are useful.This claim is made in the discussion of existing works.
- Language-model forgetting: Much prior language-model forgetting research studies sequential pre-training or sequential task fine-tuning rather than RLHF alignment tax.These studies commonly evaluate forgetting by measuring performance on a task after subsequent training.
3 Experimental Settings
The experiments use an instruction-tuned OpenLLaMA-3B followed by RLHF, evaluate alignment tax across diverse NLP benchmarks, and verify findings with additional RLHF algorithms. The setting also extends to Mistral-7B.
- Basic Setting: OpenLLaMA-3B is selected as the primary model because it is computationally friendly and has an openly available pre-training dataset.The latter supports investigation of Experience Replay in the appendix.
- Basic Setting: The training pipeline instruction-tunes the pre-trained model on ShareGPT before applying RLHF to obtain the final model.The notation distinguishes the pre-trained model, instruction-tuned model, and post-RLHF model.
- Evaluation: Alignment tax is evaluated on common-sense QA, reading comprehension, and translation benchmarks.Common-sense QA includes ARC, RACE, and PIQA; reading comprehension includes SQuAD and DROP.
- RLHF Basics: The RLHF objective assumes a ground-truth reward function over prompts and outputs and seeks to maximize reward.The policy is induced by the language model, with outputs treated as actions.
- RLHF Algorithm: Rejection Sampling Finetuning is the main RLHF algorithm, with PPO and DPO used to verify the findings.RSF learns from a best-of-n policy that selects the highest-reward response among sampled responses.
4 Evaluating Existing Methods
RLHF increases alignment reward while pretrained NLP abilities can decline, creating an alignment-forgetting trade-off. Existing mitigation methods reduce forgetting but generally sacrifice reward, while model averaging offers the strongest reported Pareto front.
- As RLHF proceeds, reward increases while translation and reading-comprehension performance declines; common-sense performance rises initially before falling.
- The evaluated methods include early stopping, regularization toward θ0, LoRA, knowledge distillation, model averaging, and stochastic moving averaging.
- Model averaging, which interpolates between θ0 and θ, achieves the most efficient alignment-forgetting Pareto front among the compared methods.
- Each method alleviates alignment tax but reduces RLHF reward, revealing a clear trade-off between reward and retained pretrained abilities.
- Comparisons with Experience Replay and KL reward penalties produce similar conclusions for PPO.
5 Unravelling the Mysteries of Model Averaging for Alleviating Alignment Tax
The paper analyzes why model averaging can mitigate alignment tax and finds that its benefits depend on feature overlap across tasks and transformer layers. Motivated by layer-specific trade-offs, it studies averaging transformer parts, especially low-level layers, as a path to improving both alignment and NLP performance.
- The analysis compares model averaging across tasks with potentially overlapping feature and label spaces, rather than only across distributions for one task.
- The theoretical setup models two task-specific predictors using feature sets with n features and n_o overlapping features, with each weak feature failing independently with probability p.
- Theoretical implications: When tasks are more similar, averaging improves performance more, but the benefit decreases when their models rely on more overlapping features.
- Theoretical implications: Feature diversity can reduce simultaneous model failure when tasks share similar feature spaces, whereas averaging disjoint feature spaces does not improve either task.
- Layer-level analysis: Because low-level transformer layers may encode features shared by alignment and NLP tasks, averaging them could improve both alignment reward and NLP performance.
- Empirical validation: The 26 transformer layers are divided into input, middle, and output parts, and one part is averaged with its pre-RLHF counterpart while the others remain unchanged.
- Empirical validation: Averaging low-level parts produces a reported improvement in both NLP tasks and alignment rewards, with the effect also observed for DPO and PPO models.
6 Heterogeneous Model Averaging
Heterogeneous Model Averaging (HMA) assigns different averaging ratios to transformer parts to improve the alignment-forgetting trade-off beyond uniform model averaging. Experiments show consistent gains across RLHF algorithms and model evaluations, with performance depending on the number of parts and averaging ratio.
- Proof of Concept: Certain heterogeneous ratio combinations consistently outperform equal-ratio averaging across benchmarks.These combinations exceed the vanilla model-averaging trade-off curve.
- Heterogeneous Model Averaging: HMA replaces uniform averaging with optimized ratios for different transformer parts to maximize alignment performance.The transformer is divided into K parts, each assigned a ratio α_i.
- Results: HMA consistently advances the Pareto front of vanilla model averaging across different RLHF algorithms.The reported results include RSF and DPO settings.
- Ablation on K: Increasing K from 3 to 6 and 9 slightly decreases the trade-off curve but still outperforms vanilla model averaging.The decrease is attributed to possible overfitting; with larger K, alignment reward rises while reading-comprehension performance drops at the same mean ratio.
- Choosing the Averaging Ratio: An averaging ratio of α = 0.2 consistently alleviates alignment tax without hurting alignment performance in the reported experiments.This ratio is presented as a suitable choice for generalization to more tasks.
- Larger-Model Evaluation: On Zephyr-7B evaluations, HMA improves over model averaging with PairRM and consistently outperforms Zephyr-7B-β under GPT4 evaluation.The GPT4 comparison covers win rate and NLP-task metrics.
7 Conclusion
The paper highlights model averaging as effective for mitigating alignment tax and introduces Heterogeneous Model Averaging to further improve performance.
- The paper emphasizes model averaging's effectiveness and proposes Heterogeneous Model Averaging to further enhance performance.
Limitations
The paper reports that HMA significantly alleviates but does not fully eliminate alignment tax, leaving the optimal trade-off and its lower bound open for future work.
- HMA significantly alleviates alignment tax, but the alignment tax is not fully eliminated.The authors identify determining the theoretical lower bound and optimal trade-off as future work.
- The evaluation focuses on preserving pretrained abilities during RLHF across common-sense QA, reading comprehension, and translation benchmarks.These tasks use accuracy, F1, and BLEU-based evaluation respectively.
- Experience Replay is less practical when pretraining data are unavailable and introduces additional computational and implementation costs when used.The comparison also finds model averaging superior on two of three benchmark groups despite replaying substantially more tokens than the RLHF data.
- Larger KL penalties partially mitigate forgetting, but model averaging is more effective for the alignment-forgetting trade-off than reward penalties.
C.3 Consistency of different combination ratios among various tasks
Experiments varying three layer-wise averaging ratios show that unequal combinations can consistently improve the alignment-forgetting trade-off over vanilla model averaging.
- Three ratio patterns were tested across a base α ∈ {0.2, 0.3, 0.4} for the transformer’s three parts.The patterns included equal ratios and two configurations reducing later-part ratios by 0.1.
- Certain ratio combinations exceed vanilla model averaging’s trade-off curve and consistently outperform equal ratios across benchmarks.
C.4 Results of α = 0.2
At α = 0.2, model averaging and HMA consistently alleviate the alignment tax without sacrificing alignment performance. The section reports this result alongside implementation details for the compared RLHF methods.
- At α = 0.2, MA and HMA consistently alleviate the alignment tax without sacrificing alignment performance.
- Vanilla model averaging and HMA are illustrated using α = 0.2 in Figures 10 and 11.
- RSF uses iterative best-of-n sampling, generating n = 32 responses per prompt from the current model at each iteration.
- When combined with forgetting-mitigation methods, RSF collects best-of-32 samples during training and fine-tunes from the starting checkpoint with additional methods.
D.2 Implementation of PPO
The PPO implementation uses stabilization measures including preferred-sample warm-up and KL-regularized reward optimization, yet training can remain unstable. The section also describes tuning, restart-based recovery, and model-averaging implementation details.
- Implementation of PPO: PPO reward optimization is modified with a KL penalty controlled by η > 0.
- Implementation of PPO: PPO can collapse into refusals or incomplete responses after sudden reward drops, despite preferred-sample warm-up and KL regularization.
- Hyper-parameter tuning: The implementation tunes learning rate, KL coefficient, update epoch, and batch size by grid search; batch size 512 and update epoch 2 are fixed for further tuning.
- Implementation of PPO: Restarting PPO from a checkpoint three iterations before collapse still led to collapse after 10–20 iterations but achieved a much higher reward value.
- Model averaging: Model averaging interpolates linear-layer modules between the instruction-following model and the RLHF model while varying α from 0 to 1.
- Heterogeneous Model Averaging: HMA reparameterizes layer-wise averaging ratios through sigmoid-transformed variables and searches ε ∈ {0, 0.1, ..., 0.9}.
E.4 Results of AdaMerging (Yang et al., 2023)
AdaMerging is designed to optimize layer-wise merging ratios for a specific task, but the section contrasts this task-specific setup with HMA's use of only RLHF data. The accompanying analysis frames model averaging through shared and dissimilar feature spaces.
- Results of AdaMerging: AdaMerging uses labeled Reading Comprehension data to optimize 26 layer-wise merging ratios, treating Reading Comprehension as task A.
- Results of AdaMerging: AdaMerging optimized for Reading Comprehension cannot preserve performance on Common Sense when the training data omit that task.
- Results of AdaMerging: HMA requires only RLHF data and does not use data from the tasks whose abilities it aims to preserve.
- Results of AdaMerging: HMA alleviates the alignment tax across Reading Comprehension, Commonsense QA, and translation, using RSF, PPO, and DPO.
- Theoretical analysis: For disjoint tasks, the averaged model preserves each task's original accuracy in the analyzed case, yielding ξ(2) = 0.
F.3 Discussion on the Effect of Task Similarity on Model Averaging
The analysis explains that model averaging offers limited improvement when tasks are dissimilar because one model's features do not affect the other task's predictions. The section provides supporting functional definitions and experiment hyper-parameter references.
- When task label spaces are disjoint, model averaging provides little improvement because the second model's classifier has zero weights in the first task's class dimensions.
- The cited analysis compares model averaging under similar and independent task conditions using feature-space overlap.
- For dissimilar tasks, the averaged model's prediction on task a equals the original prediction of model fa because wb(k) = 0.
- The section references separate hyper-parameter tables for RLHF, auxiliary, and HMA experiments.