Online Preference Alignment for Language Models via Count-based Exploration

Authors: Chenjia Bai, Yang Zhang, Shuang Qiu, Qiaosheng Zhang, Kang Xu, Xuelong Li
Primary area: foundation or frontier models, including LLMs
Keywords: reinforcement learning from human feedback, rlhf, preference alignment, exploration, llms

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., 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 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.

Decision: accept (spotlight)
Reviewer scores: 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730311316248

thread_type:
official_review

<thread_object>
creator: reviewer_pkvp

<content>
summary:
This paper introduces an online RLHF approach via the Count-based Online Preference Optimization (COPO) algorithm, which integrates count-based exploration and direct preference optimization to improve performance. Experiments on Zephyr and Llama-3 models demonstrate COPO’s effectiveness in boosting instruction-following capabilities and general academic benchmarks.

strengths:
The paper rigorously establishes the theoretical underpinnings of COPO, providing a regret-bound analysis that demonstrates the efficiency of the UCB-based exploration strategy.

weaknesses:
It seems there’s no notable improvement over SELM for the Llama-3-8B Instruct model in Table 1, especially when considering the potential variance from using GPT-4 as the evaluator. And can authors provide more baselines, for example comparison with Best-of-N and PPO?

questions:
Please see the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730449865852

thread_type:
official_review

<thread_object>
creator: reviewer_j1hg

<content>
summary:
This paper identifies a central problem in current online RLHF methods, that is, efficiently exploring the prompt-response space in each iteration. To address this, the paper designs a novel RLHF algorithm named Count-based Online Preference Optimization (COPO), which re-formulates the original DPO objective to balance space exploration and preference optimization. The authors adopt a simple coin-flip counting network approach to support COPO's pseudo-count, achieving promising experimental results.

strengths:
1. Solid theoretical analysis. This work is built upon a solid theoretical foundation, starting from a relatively reasonable linear reward assumption and ultimately analyzing the regret upper bound of the optimization objective, providing a more credible motivation.

2. The authors demonstrate good insights into the existing problems in online RLHF, specifically the focus on preference optimization while lacking effective space exploration. Around this issue, they propose an algorithm suitable for LLMs with vast token sequence spaces.

3. The proposed COPO algorithm shows notable improvements after multiple iterations.

weaknesses:
1. The proposed COPO algorithm is only implemented based on the DPO optimization objective, without verification of its adaptability to improve other online RLHF algorithms.

2. The visit count item in the COPO algorithm relies solely on a pseudo-count mechanism implemented by a coin flipping network (CFN) with accompanying random labels, which seems too simplistic. More convincing visit count algorithms need to be explored.

3. As mentioned by the authors in the experimental section, a better reward model could improve COPO's performance. However, this point wasn't explored in the experiments. In fact, the choice of different scoring RMs has a significant impact on the experiment and should be given more attention.

questions:
1. What would be the experimental results if COPO's main ideas were applied to other online RLHF algorithms?

2. Are there other alternative or explorable counting algorithms?

3. How would choosing a stronger Reward model (or Judge model) affect algorithm performance, and can this be demonstrated with specific experimental results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730694971726

thread_type:
official_review

<thread_object>
creator: reviewer_4qfj

<content>
summary:
This paper introduces the COPO algorithm for better exploration of online RLHF formulation. The authors provide a theoretical background of the proposed method and explain the practical implementation of it. To be specific, it assumes that the reward function could be expressed as a dot product between a parameter vector and a feature vector. Based on it, the model incorporates a USB term for exploration. Also, they employ a lightweight Config Flipping Network (CFN) to obtain pseudo-visit counts for the large state space of language generation. Experimental results show the efficacy of the proposed COPO algorithm. It performs better than the other exploration-aware online RLHF baseline, SELM, in several benchmarks. When they adjust the coefficient for the exploration term, it supports the importance of the exploration term in their experiments.

strengths:
It tackles the important problem innate in RLHF-based alignment learning, i.e., exploration, with formal theoretical motivation and practical algorithm.

weaknesses:
* Their explanation and selection of baseline methods are not enough. Especially, I think other iterative/online methods should have been 
considered in addition to the SELM. There are many variants to implement online DPO. For example, [1], [2], and [3].

* The improvements seem marginal despite the complexity of the proposed method. The performance gaps among the baselines could be changed with different random seeds, especially for the LLM-as-a-Judge tasks.

* It would be good if more benefits of the exploration were suggested. For example, it might be more robust in some adversarial cases.

[1] Self-Rewarding Language Models (Yuan et al., 2024)

[2] Aligning Large Language Models by On-Policy Self-Judgment (Lee et al., 2024)

[3] D2PO: Discriminator-Guided DPO with Response Evaluation Models (Singhal et al., 2024)

questions:
* How do we measure/believe the CFN works properly? (before testing the model on the final benchmarks.)
  * (Declaimer) I do not have any idea about the original CFN work. However, the model (CFN module) seems too small to represent the hidden state of the modern LLM.
  * If we can improve the performance of the module, will the final alignment performances be also improved?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733108885944

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear reviewers and AC,

We sincerely appreciate your valuable time and effort spent reviewing our manuscript. We propose Count-based Online Preference Optimization (COPO) for LLM alignment that leverages coin-flip counting to encourage exploration in online RLHF with theoretical guarantees. As reviewers highlighted, our method has solid theoretical motivation (4QFJ, j1HG, pkVP), good insights and practical algorithm (j1HG, 4QFJ), and notable improvements (j1HG, pkVP). We appreciate your constructive feedback on our manuscript. In response to the comments, we have carefully revised and enhanced the manuscript as follows: 
- we add experiments on D2PO, an automatic ranking method compared to the off-the-shelf PairRM reward model.
- we add experiments on an adversarial case that restricts the data coverage of the initial preference dataset to show the benefit of our exploration objective. 
- we add experiments on online KTO with count-based exploration, which serves as a new preference optimization objective in addition to the DPO-based method.
- we add experiments on a more powerful reward model, *ArmoRM-Llama3-8B*, which shows that a better reward model can enhance the online RLHF algorithms.
- we add experiments on both offline and online methods to integrate the Best-of-N strategy into the RLHF process. 

Meanwhile, we give more clarifications of the algorithmic details and experiments as follows:
- we add an explanation of the measurement of CFN in the training process and the estimation error of CFN after coverage.
- we add an explanation of the marginal improvement on some specific tasks in academic benchmarks.
- we add an explanation of other exploration algorithms to show why CFN is more suitable for online RLHF algorithms.
- we add an explanation of the consistency results evaluated by AlpacaEval 2.0.

Meanwhile, we wish to highlight that Reviewer pkVP with negative evaluations did not give any feedback on our response in the rebuttal period. The major concern of Reviewers pkVP is the potential variance from using GPT-4 as the evaluator and the lack of baselines of Best-of-N and PPO. In the rebuttal period, we add a detailed analysis of the consistency of AlpacaEval 2.0 which has a diverse evaluation set and a good alignment to human evaluators. Meanwhile, we add Best-of-N as a baseline that combines offline RLHF and online RLHF methods. We also give explanations on the PPO-based method. We believe our response has successfully addressed the concerns raised by the reviewers. We kindly request that the AC dedicate more attention to the evaluation of our submission.

In the revised manuscript, these updates are temporarily highlighted in brown for your convenience to check. We sincerely believe that these updates may help us better deliver the benefits of the proposed COPO to the ICLR community.

Thank you very much,

Authors.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper addresses the practical challenge of improving online preference alignment. The authors begin by identifying limitations in existing methods and then provide a theoretical foundation for their proposed approach. Building on this theoretical framework, they introduce a novel alignment method based on a simple count-based approach.  The experimental results demonstrate clear and consistent performance gains.

A key strength of this work is the strong theoretical grounding for the proposed method, which effectively addresses the fundamental exploration-exploitation trade-off in reinforcement learning systems.  Furthermore, the authors translate this theoretical foundation into a novel yet simple alignment method that proves effective across various datasets.

While the paper presents compelling results, there is potential for further improvement. Exploring more sophisticated reward models could enhance the method's performance.  Additionally, expanding the empirical evaluation could provide a more comprehensive understanding of the approach's capabilities and limitations.