Sample Efficient Alignment for LLMs

Authors: Zichen Liu, Changyu Chen, Chao Du, Wee Sun Lee, Min Lin
Primary area: foundation or frontier models, including LLMs
Keywords: rlhf, online dap, llm alignment, sample efficiency

Abstract:
We study methods for efficiently aligning large language models (LLMs) with human preferences given budgeted online feedback. We first formulate the LLM alignment problem in the frame of contextual dueling bandits. This formulation, subsuming recent paradigms such as online RLHF and online DPO, inherently quests for sample-efficient algorithms that incorporate online active exploration. Leveraging insights from bandit theory, we introduce a unified algorithm based on Thompson sampling and highlight its applications in two distinct LLM alignment scenarios. The practical agent that efficiently implements this algorithm, named SEA (Sample-Efficient Alignment), is empirically validated through extensive experiments across three model scales (1B, 2.8B, 6.9B) and three preference learning algorithms (DPO, IPO, SLiC). The results demonstrate that SEA achieves highly sample-efficient alignment with oracle's preferences, outperforming recent active exploration methods for LLMs. We will release our codebase to hopefully accelerate future research in this field.

Decision: reject
Reviewer scores: 5, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1730378370536

thread_type:
official_review

<thread_object>
creator: reviewer_fdal

<content>
summary:
This paper frames the problem of aligning LLMs with pairwise preferences as a contextual dueling bandit problem. The authors propose a method that adapts a well-know & effective bandit algorithm, Thompson sampling (TS), to LLMs. This is not trivial, because TS requires maintaining and sampling from a belief over reward functions. To address this, the authors provide tractable approximations. They evaluate their proposed approach empirically on LLMs of various sizes, and provide comparisons to related work.

strengths:
- The contextual dueling perspective is interesting, insightful, and to the best of my knowledge, new.
- The paper does a good job of covering the most relevant related literature, and how it relates to the proposed approach.
- The description is very limited with no details in the supplementary material, but it is possible that the testbed they develop to evaluate different approaches is a valuable contribution in its own right

weaknesses:
- The writing could be improved significantly; parts of the paper are unclear, imprecise, and many details are omitted and should at least be present in the supplementary materials. A non-exhaustive list of examples:
    - line 144: "based on a binary stochastic feedback $z$" -> z is never referred to anywhere outside of this line as far as I see
    - line 196: "offline RL" -> I think the authors might be conflating offline RL with offline preference collection. RLHF typically uses online RL. PPO is an online algorithm.
    - I struggle to make sense of Figure 3, for example I fail to understand how (d) is meant to differ from (b).
- The posterior approximation scheme (Eq. 9) is not well motivated. I can see it is computationally attractive, but why do we believe this is a reasonable way to approximate the posterior? A cursory glance at the papers mentioned around Eq. 9 did not help. Since this is one of the central components of the proposed approach, I would expect a more rigorous justification here.
- The experiments are comprehensive, but I struggle to take them at face value. Unfortunately, few details are provided, neither in the main text, nor in the supplementary materials.
    - There are design decisions involved in setting up baselines, and these are not discussed.
    - I can't reconcile any of the curves in Figure 6 with the curves in Figure 5.
- Some of the metrics are problematic, as I see it. The "online win rate" (line 387) would make a model that deliberately samples poor responses look great.

questions:
Concretely, in what ways does SEA differ from APL? How do you explain the large increase in sample efficiency?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730668708644

thread_type:
official_review

<thread_object>
creator: reviewer_rts4

<content>
summary:
This paper considers explorations by approximate TS in LLM alignment.

strengths:
This paper is nicely written and well-executed. The details are well-presented.

weaknesses:
This paper is a bit incremental and the idea is well expected. The contextual bandit formulation is well-known and the implementation of TS using ensemble is pretty straightforward.  

There are multiple key questions this paper fails to answer:

1. Why ensemble on the top of DPO makes sense and how good the uncertainty estimation is? This is not a standard supervise learning problem. This paper lacks a way to directly measure it. I do not think ensemble + X can quantify uncertainty in every problem.  

2. The reward model is too small (0.4B) and the base model Pythia family is outdated such that the empirical result may become useless in a practical setting. It is very common that the benefit of active exploration could disappear when the base model and reward model become much stronger. 

3. The baseline number is not comparable with other literatures since the author uses its own oracle reward model. There are tons of implementation details that could be hidden for other baselines. I think it would be good to have an apple-to-apple comparison with the same setting in the original DPO paper.

questions:
see above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692063302

thread_type:
official_review

<thread_object>
creator: reviewer_cf2f

<content>
summary:
This paper studies the online alignment for LLMs. It formulates the LLM alignment as a contextual dueling bandit problem. It proposes a Thompson sampling-like algorithm for both the regret minimization and best arm identification setting. For the challenge encountered when implementing this algorithm including the model constraints, the intractable operation to sample a posterior and argmax action, the paper proposes corresponding solutions to enable real-world application. Finally, the authors build a learning system for studying methods in online LLM alignment and compare the proposed method with available baselines in different scale levels of experiments. The proposed method shows superior sample efficiency compared with baselines.

strengths:
1.	The paper studies the LLM alignment problem from an online view. It is reasonable to formulate the problem as a contextual dueling bandit problem. 
2.	The paper is inspired by the classic Thompson sampling algorithm to actively explore the preference function and proposes a corresponding version suitable for real-world application.
3.	The paper builds a learning system for studying methods in online LLM alignment. This provides efficient tools for the community on this research line.

weaknesses:
1.	In Line 383, the authors claim they omit a baseline SELM since it shares a very similar algorithmic design as XPO. This is generally not an acceptable reason for excluding a baseline comparison.

questions:
Please see the last part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703989726

thread_type:
official_review

<thread_object>
creator: reviewer_z8j2

<content>
summary:
The paper considers the problem of aligning LLMs to humans via online interaction and in a sample efficient manner that minimizes the overall human annotations consumed. It proposes a practical algorithm that can be implemented and using the tldr dataset shows that it can achieve extremely high win-rates when compared against the baseline alignment algorithm.

strengths:
+ The paper builds upon the direction of aligning LLMs that considers the epistemic uncertainty of the reward function and utilizing the variance to guide the model to collect preference data in directions that maximally obtains information.
+ The algorithm is sound and theoretically solid. They build upon well performing methods for uncertainty estimates using the ensemble reward model method.
+ The results shown are promising (although see weakness for caveats).

weaknesses:
- The main weakenss of this paper is the lack of siginificant experimentation. The paper only considers the summarization task with the tldr dataset. It does not consider more diversity in tasks/datasets and does not utilize some of the most common tasks in the LLM literature (e.g., reasoning, code, general knowledge). Even within summarization, it does not consider > 1 datasets. Improvements in one dataset is typically very limiting.

questions:
- Do you have results that can showcase this method on a variety of tasks and many different datasets? How generalizable is this method (in terms of how well it works)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732042934498

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Paper Revision

comment:
We thank all reviewers for their constructive feedback, and we have responded to each reviewer individually. We have also uploaded a **Paper Revision** including additional results and illustrations:

* $\\textrm{\\color{red}Figure 8}$ & $\\textrm{\\color{red}Table 2}$ (Page 21): Evaluation results on AlpacaEval 2.0 of models trained with prompts from UltraFeedback dataset (Reviewers **z8J2**, **RtS4**, **fdaL**).
* $\\textrm{\\color{red}Appendix F}$ (Page 21): Reviewers **cF2f** and **fdaL** recognize the values of our online alignment training system (Section 5.1). We further conduct rigorous benchmarking to compare our codebase with huggingface/trl on the online DPO training and show ~2.5x latency improvement.
* $\\textrm{\\color{red}Appendix A}$ (Page 16): We add a detailed algorithm as complementary to Algorithm 1 in the main text, to enhance the clarity (Reviewer **fdaL**).
* $\\textrm{\\color{red}Appendix B}$ (Page 17): We discuss extensively how our bandit formulation is related to single-step preference-based reinforcement learning, explaining RLHF is an instance of offline model-based RL (Reviewer **fdaL**).
* $\\textrm{\\color{red}Appendices C, D}$ (Pages 17-20): We review all the baseline methods and provide details on the settings (Reviewer **fdaL**); we also clarify the reason why SELM is not compared (Reviewer **cF2f**).
* $\\textrm{\\color{red}Figure 2}$ (Page 3): We updated this figure to better illustrate how alignment as CDB differs from (offline) RLHF and DAP, to improve the clarity (Reviewer **fdaL**).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732332555185

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Looking forward to further feedback

comment:
Dear Reviewers,

Thank you again for your valuable comments and suggestions, which are really helpful for us. We have posted responses to the proposed concerns and included additional experiment results.

We totally understand that this is quite a busy period, so we deeply appreciate it if you could take some time to return further feedback on whether our responses solve your concerns. If there are any other comments, we will try our best to address them.

Best,

The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733111717311

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final Request for Discussion: Your Feedback Is Invaluable

comment:
Dear Reviewers,

We hope this message finds you well. As we approach the **final day** of the extended reviewing period, we wanted to kindly follow up on the rebuttals we provided to your thoughtful reviews.

We carefully addressed each of your questions within the first week to the best of our ability. We have been looking forward to your responses for the past two weeks, but unfortunately, we have not yet received any follow-up or further discussion.

We sincerely look forward to your additional feedback and will do our best to address any remaining concerns within the limited time available.

Thank you again for your consideration. We understand how busy this period can be, and your dedication to the review process means a great deal.

Best regards, \
The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces Sample-Efficient Alignment (SEA), a novel algorithm to align large language models (LLMs) with human preferences sample-efficiently. The problem is formulated as a contextual dueling bandit framework, and the SEA algorithm leverages Thompson Sampling with deep ensembles for uncertainty estimation. The authors present extensive experiments across three LLM scales and three preference learning algorithms, demonstrating the proposed method's superior sample efficiency compared to baseline methods. The work also provides a practical learning system for online alignment research.

Reviewers raised concerns about limited task diversity, inadequate baseline comparisons, particularly the omission of SELM, and incremental novelty. While the authors added experiments (e.g., UltraFeedback and AlpacaEval 2.0) and improved writing, reviewers found these changes insufficient to address broader generalization, weak comparability, and lack of rigor. Presentation issues were also noted, although some partial improvements were made in the revisions. Based on the common consensus, therefore, I would recommend the authors submit to the next suitable venue addressing all the concerns once all the necessary modifications are incorporated.