Mutual-Taught for Co-adapting Policy and Reward Models

TL;DR

The paper introduces Mutual-Taught, a self-training iterative method that enhances both policy and reward models in large language models without extra human labels, addressing distribution shift issues during preference optimization.

Contribution

It presents a novel EM-like self-training approach for co-adapting policy and reward models, improving their performance through mutual updates without additional annotations.

Findings

Improved policy model achieves 54.1% win rate on AlpacaEval-2.

Reward model performs comparably to GPT-4 on RewardBench.

Method demonstrates consistent improvements in both models.

Abstract

During the preference optimization of large language models (LLMs), distribution shifts may arise between newly generated model samples and the data used to train the reward model (RM). This shift reduces the efficacy of the RM, which in turn negatively impacts the performance of the policy model (PM). To address this challenge, we propose Mutual-Taught, a self-training method that iteratively improves both the PM and RM without requiring additional human annotation. Our approach mirrors the expectation-maximization (EM) algorithm. In the E-step, the PM is updated using feedback from the current RM, guiding the PM toward a better approximation of the latent optimal preference distribution. In the M-step, we update the RM by constructing training data from the outputs of the PM before and after the E-step update. This process ensures that the RM adapts to the evolving policy distribution. Experimental results demonstrate that this iterative approach leads to consistent improvements in both models. Specifically, our 8B policy model, LLaMA-3-8B-Instruct-MT, achieves a length-controlled win rate of 54.1\% on AlpacaEval-2, while our 8B reward model, FsfairX-LLaMA3-RM-MT, performs on par with GPT-4o-2024-08-06 on RewardBench.