Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Online Merging Optimizers for Boosting Rewards and Mitigating Tax in Alignment

Keming Lu, Bowen Yu, Fei Huang, Yang Fan, Runji Lin, Chang Zhou

TL;DR

This work tackles RLHF alignment tax by introducing Online Merging Optimizers that augment each RLHF update with the SFT reference delta $\tau_r$, via sparsified gradient consolidation (OnDARE) or sign-consensus (OnTIES). By online merging at every step, the method achieves higher alignment rewards while mitigating forgetting, and demonstrates strong, cross-backbone, cross-algorithm effectiveness across 14 benchmarks. The paper also links online/offline merging through Step-K variants and shows KL constraints can complement online merging to further balance reward and stability. Overall, the approach provides a practical optimizer to improve human-aligned behavior in LLMs while preserving foundational capabilities, with potential applicability to continual learning scenarios.

Abstract

Effectively aligning Large Language Models (LLMs) with human-centric values while preventing the degradation of abilities acquired through Pre-training and Supervised Fine-tuning (SFT) poses a central challenge in Reinforcement Learning from Human Feedback (RLHF). In this paper, we first discover that interpolating RLHF and SFT model parameters can adjust the trade-off between human preference and basic capabilities, thereby reducing the alignment tax at the cost of alignment reward. Inspired by this, we propose integrating the RL policy and SFT models at each optimization step in RLHF to continuously regulate the training direction, introducing the Online Merging Optimizer. Specifically, we merge gradients with the parameter differences between SFT and pretrained models, effectively steering the gradient towards maximizing rewards in the direction of SFT optimization. We demonstrate that our optimizer works well with different LLM families, such as Qwen and LLaMA, across various model sizes ranging from 1.8B to 8B, various RLHF algorithms like DPO and KTO, and existing model merging methods. It significantly enhances alignment reward while mitigating alignment tax, achieving higher overall performance across 14 benchmarks.

Online Merging Optimizers for Boosting Rewards and Mitigating Tax in Alignment

TL;DR

This work tackles RLHF alignment tax by introducing Online Merging Optimizers that augment each RLHF update with the SFT reference delta , via sparsified gradient consolidation (OnDARE) or sign-consensus (OnTIES). By online merging at every step, the method achieves higher alignment rewards while mitigating forgetting, and demonstrates strong, cross-backbone, cross-algorithm effectiveness across 14 benchmarks. The paper also links online/offline merging through Step-K variants and shows KL constraints can complement online merging to further balance reward and stability. Overall, the approach provides a practical optimizer to improve human-aligned behavior in LLMs while preserving foundational capabilities, with potential applicability to continual learning scenarios.

Abstract

Effectively aligning Large Language Models (LLMs) with human-centric values while preventing the degradation of abilities acquired through Pre-training and Supervised Fine-tuning (SFT) poses a central challenge in Reinforcement Learning from Human Feedback (RLHF). In this paper, we first discover that interpolating RLHF and SFT model parameters can adjust the trade-off between human preference and basic capabilities, thereby reducing the alignment tax at the cost of alignment reward. Inspired by this, we propose integrating the RL policy and SFT models at each optimization step in RLHF to continuously regulate the training direction, introducing the Online Merging Optimizer. Specifically, we merge gradients with the parameter differences between SFT and pretrained models, effectively steering the gradient towards maximizing rewards in the direction of SFT optimization. We demonstrate that our optimizer works well with different LLM families, such as Qwen and LLaMA, across various model sizes ranging from 1.8B to 8B, various RLHF algorithms like DPO and KTO, and existing model merging methods. It significantly enhances alignment reward while mitigating alignment tax, achieving higher overall performance across 14 benchmarks.
Paper Structure (19 sections, 8 equations, 4 figures, 7 tables, 2 algorithms)

This paper contains 19 sections, 8 equations, 4 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Preliminaries and Motivations
  4. Online Merging Optimizers
  5. From Gradient-based Optimizer to Online Merging Optimizer
  6. Implementations
  7. Experiments
  8. Experimental Setup
  9. Main Results
  10. Hyper-parameter Effects
  11. The Impact of RLHF Algorithms
  12. Discussion
  13. Bridging Online and Offline Merging via Step-K Online Merging Optimizers
  14. Complementary Effect of KL Constraints and Online Merging
  15. Conclusion
  16. ...and 4 more sections

Figures (4)

  • Figure 1: An illustrastion of RLHF with online merging optimizers described in \ref{['sec:methods']}. In each RLHF iteration, we first obtain the update weight $\Delta\theta^{(t)}$, and then sparcify it and make a consensus with the delta parameters of the reference model. We use this merged delta as the update of the policy model in this iteration. We also compare online merging with offline merging, shown in the lower part of the figure and further introduced in \ref{['sec:offline_merging']}.
  • Figure 2: Analysis of Parameter Reserve Rates
  • Figure 3: Analysis of Online Merging Gap Step
  • Figure 4: Analysis of KL constraints