Online Preference Alignment for Language Models via Count-based Exploration
Chenjia Bai, Yang Zhang, Shuang Qiu, Qiaosheng Zhang, Kang Xu, Xuelong Li
TL;DR
COPO introduces a theoretically grounded, count-based exploration term into online RLHF to address data-coverage and OOD reward generalization. By replacing the reward-model-centered objective with a Direct Preference Optimization backbone augmented with an optimistic, count-based exploration bonus, COPO achieves a $ ilde{O}(\sqrt{T})$ regret bound under a linear reward assumption. The practical CFN module provides scalable pseudo-count estimates that drive exploration in large prompt–response spaces, and experiments on Zephyr and Llama-3 demonstrate superior instruction-following and benchmark performance compared with online DPO and SELM. This approach expands LLM exploration while maintaining stable preference optimization, offering a principled path toward broader data coverage and robust policy improvement in online RLHF.
Abstract
Reinforcement Learning from Human Feedback (RLHF) has shown great potential in fine-tuning Large Language Models (LLMs) to align with human preferences. Existing methods perform preference alignment from a fixed dataset, which can be limited in data coverage, and the resulting reward model is hard to generalize in out-of-distribution responses. Thus, online RLHF is more desirable to empower the LLM to explore outside the support of the initial dataset by iteratively collecting the prompt-response pairs. In this paper, we study the fundamental problem in online RLHF, i.e. \emph{how to explore} for LLM. We give a theoretical motivation in linear reward assumption to show that an optimistic reward with an upper confidence bound (UCB) term leads to a provably efficient RLHF policy. Then, we reformulate our objective to direct preference optimization with an exploration term, where the UCB-term can be converted to a count-based exploration bonus. We further propose a practical algorithm, named \emph{Count-based Online Preference Optimization (COPO)}, which leverages a simple coin-flip counting module to estimate the pseudo-count of a prompt-response pair in previously collected data. COPO encourages LLMs to balance exploration and preference optimization in an iterative manner, which enlarges the exploration space and the entire data coverage of iterative LLM policies. We conduct online RLHF experiments on Zephyr and Llama-3 models. The results on instruction-following and standard academic benchmarks show that COPO significantly increases performance.