Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Faster WIND: Accelerating Iterative Best-of-$N$ Distillation for LLM Alignment

Tong Yang, Jincheng Mei, Hanjun Dai, Zixin Wen, Shicong Cen, Dale Schuurmans, Yuejie Chi, Bo Dai

TL;DR

This work addresses the high cost of iterative best-of-$N$ distillation (BoN) for LLM alignment by revealing a unified game-theoretic connection between iterative BoN and self-play. It introduces WIND, a win-rate-dominance framework with a regularized two-player game that approximates iterative BoN in the parameter space and provides convergence and sample-efficiency guarantees. The authors develop an exact last-iterate wind optimizer with linear convergence and a family of practical, sample-efficient estimators based on squared risk, KL, or NCE objectives. Empirical results on contextual bandits and LLM alignment benchmarks show WIND achieves competitive or superior performance with reduced sampling and training costs compared to state-of-the-art methods such as J-BOND and SPPO, highlighting its potential for scalable and efficient alignment.

Abstract

Recent advances in aligning large language models with human preferences have corroborated the growing importance of best-of-N distillation (BOND). However, the iterative BOND algorithm is prohibitively expensive in practice due to the sample and computation inefficiency. This paper addresses the problem by revealing a unified game-theoretic connection between iterative BOND and self-play alignment, which unifies seemingly disparate algorithmic paradigms. Based on the connection, we establish a novel framework, WIN rate Dominance (WIND), with a series of efficient algorithms for regularized win rate dominance optimization that approximates iterative BOND in the parameter space. We provides provable sample efficiency guarantee for one of the WIND variant with the square loss objective. The experimental results confirm that our algorithm not only accelerates the computation, but also achieves superior sample efficiency compared to existing methods.

Faster WIND: Accelerating Iterative Best-of-$N$ Distillation for LLM Alignment

TL;DR

This work addresses the high cost of iterative best-of- distillation (BoN) for LLM alignment by revealing a unified game-theoretic connection between iterative BoN and self-play. It introduces WIND, a win-rate-dominance framework with a regularized two-player game that approximates iterative BoN in the parameter space and provides convergence and sample-efficiency guarantees. The authors develop an exact last-iterate wind optimizer with linear convergence and a family of practical, sample-efficient estimators based on squared risk, KL, or NCE objectives. Empirical results on contextual bandits and LLM alignment benchmarks show WIND achieves competitive or superior performance with reduced sampling and training costs compared to state-of-the-art methods such as J-BOND and SPPO, highlighting its potential for scalable and efficient alignment.

Abstract

Recent advances in aligning large language models with human preferences have corroborated the growing importance of best-of-N distillation (BOND). However, the iterative BOND algorithm is prohibitively expensive in practice due to the sample and computation inefficiency. This paper addresses the problem by revealing a unified game-theoretic connection between iterative BOND and self-play alignment, which unifies seemingly disparate algorithmic paradigms. Based on the connection, we establish a novel framework, WIN rate Dominance (WIND), with a series of efficient algorithms for regularized win rate dominance optimization that approximates iterative BOND in the parameter space. We provides provable sample efficiency guarantee for one of the WIND variant with the square loss objective. The experimental results confirm that our algorithm not only accelerates the computation, but also achieves superior sample efficiency compared to existing methods.

Paper Structure

This paper contains 50 sections, 11 theorems, 153 equations, 2 figures, 1 table, 3 algorithms.

Table of Contents

  1. Introduction
  2. Contributions
  3. Related work
  4. RLHF and LLM alignment.
  5. RLHF via self-play.
  6. Best-of-$N$ and BOND.
  7. Notation.
  8. Preliminaries
  9. RLHF: reward versus win rate
  10. Reward maximization.
  11. Win rate maximization.
  12. Best-of-$N$ distillation
  13. A Unified Game-Theoretic View
  14. Iterative BoN as game solving
  15. Iterative BoN.
  16. ...and 35 more sections

Key Result

Theorem 1

Let $\pi_{\textnormal{ref}}\in \text{relint}\left(\Delta_{\mathcal{Y}}^{\mathcal{X}}\right)$ and $n\geq 2$ in Algorithm alg:iter_BoN. Then $\pi_\infty\coloneqq\lim_{T\rightarrow\infty}\pi_T$ exists, and $(\pi_\infty,\pi_\infty)$ is a Nash equilibrium of the log-win-rate game eq:log_game when:

Figures (2)

  • Figure 1: Empirical validation of Theorem \ref{['thm:wind_approx_relation']} on contextual bandit experiments. For (a) the no-mixing case, we show the convergence of both iterative BoN (c.f. Algorithm \ref{['alg:iter_BoN']}) and exact WIND (c.f. Algorithm \ref{['alg:iter_PMDA']}) to $\overline\pi^\star_0$; for (b) the mixing case, we show $\overline{\pi}^\star_{\beta}$ and $\pi^\star_\beta$ are very close to each other.
  • Figure 2: Run time (seconds) of different methods.

Theorems & Definitions (18)

  • Theorem 1: Iterative BoN solves the log-win-rate game \ref{['eq:log_game']}
  • Proposition 1: existence of $\pi^\star_\beta$
  • Theorem 2: relationship between two games (informal)
  • Theorem 3: Linear last-iterate convergence of Algorithm \ref{['alg:iter_PMDA']}, sokota2023unifiedapproachreinforcementlearning
  • Remark 1
  • Lemma 1: Conditional mean minimizes the square loss
  • Remark 2: Assumption \ref{['asmp:PL']} is satisfied with linear function approximation
  • Theorem 4: Convergence of Algorithm \ref{['alg:practice']}
  • Theorem 5: solution to iterative BoN (formal)
  • Remark 3
  • ...and 8 more