BOND: Aligning LLMs with Best-of-N Distillation
Authors: name: Pier Giuseppe Sessa; name: Robert Dadashi; name: Léonard Hussenot; name: Johan Ferret; name: Nino Vieillard; name: Alexandre Ramé; name: Bobak Shariari; name: Sarah Perrin; name: Abe Friesen; name: Geoffrey Cideron; name: Sertan Girgin; name: Piotr Stanczyk; name: Andrea Michi; name: Danila Sinopalnikov; name: Sabela Ramos; name: Amélie Héliou; name: Aliaksei Severyn; name: Matt Hoffman; name: Nikola Momchev; name: Olivier Bachem
Published: 2024-07-19 18:38:25

Abstract: Reinforcement learning from human feedback (RLHF) is a key driver of quality
and safety in state-of-the-art large language models. Yet, a surprisingly
simple and strong inference-time strategy is Best-of-N sampling that selects
the best generation among N candidates. In this paper, we propose Best-of-N
Distillation (BOND), a novel RLHF algorithm that seeks to emulate Best-of-N but
without its significant computational overhead at inference time. Specifically,
BOND is a distribution matching algorithm that forces the distribution of
generations from the policy to get closer to the Best-of-N distribution. We use
the Jeffreys divergence (a linear combination of forward and backward KL) to
balance between mode-covering and mode-seeking behavior, and derive an
iterative formulation that utilizes a moving anchor for efficiency. We
demonstrate the effectiveness of our approach and several design choices
through experiments on abstractive summarization and Gemma models. Aligning
Gemma policies with BOND outperforms other RLHF algorithms by improving results
on several benchmarks.