Overcoming Reward Overoptimization via Adversarial Policy Optimization with Lightweight Uncertainty Estimation

TL;DR

This paper presents AdvPO, a lightweight adversarial policy optimization method that mitigates reward over-optimization in RLHF for LLMs by estimating reward uncertainty from model embeddings, leading to improved human-aligned performance.

Contribution

The paper introduces a novel lightweight uncertainty estimation method and an adversarial optimization framework to address reward over-optimization in RLHF for large language models.

Findings

AdvPO effectively reduces reward over-optimization in LLM training.

Experimental results show improved human-aligned performance.

Uncertainty estimation from last layer embeddings is computationally efficient.

Abstract

We introduce Adversarial Policy Optimization (AdvPO), a novel solution to the pervasive issue of reward over-optimization in Reinforcement Learning from Human Feedback (RLHF) for Large Language Models (LLMs). Over-optimization occurs when a reward model serves as an imperfect proxy for human preference, and RL-driven policy optimization erroneously exploits reward inaccuracies. In this paper, we begin by introducing a lightweight way to quantify uncertainties in rewards, relying solely on the last layer embeddings of the reward model, without the need for computationally expensive reward ensembles. AdvPO then addresses a distributionally robust optimization problem centred around the confidence interval of the reward model's predictions for policy improvement. Through comprehensive experiments on the Anthropic HH and TL;DR summarization datasets, we illustrate the efficacy of AdvPO in mitigating the overoptimization issue, consequently resulting in enhanced performance as evaluated through human-assisted evaluation.