Online Merging Optimizers for Boosting Rewards and Mitigating Tax in Alignment

TL;DR

This paper introduces the Online Merging Optimizer, a method that dynamically combines RLHF and SFT models during training to improve alignment rewards and reduce the alignment tax in large language models.

Contribution

It proposes a novel optimizer that merges gradients from SFT and pretrained models during RLHF, enhancing alignment performance across multiple models and benchmarks.

Findings

Improves alignment reward in various LLMs.

Reduces alignment tax while maintaining capabilities.

Works effectively with different RLHF algorithms.

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.