BOND: Het afstemmen van LLM's via Best-of-N-distillatie

BOND: Aligning LLMs with Best-of-N Distillation

July 19, 2024

Auteurs: Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amélie Héliou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev, Olivier Bachem

cs.AI

Samenvatting

Reinforcement learning from human feedback (RLHF) is een belangrijke drijvende kracht achter kwaliteit en veiligheid in state-of-the-art grote taalmodelen. Toch is Best-of-N sampling, een verrassend eenvoudige en krachtige strategie tijdens inferentie, een methode die de beste generatie selecteert uit N kandidaten. In dit artikel introduceren we Best-of-N Distillation (BOND), een nieuw RLHF-algoritme dat Best-of-N probeert na te bootsen, maar zonder de aanzienlijke rekenkosten tijdens inferentie. Specifiek is BOND een distributie-matching-algoritme dat ervoor zorgt dat de distributie van generaties van het beleid dichter bij de Best-of-N-distributie komt. We gebruiken de Jeffreys-divergentie (een lineaire combinatie van forward en backward KL) om een balans te vinden tussen mode-covering en mode-seeking gedrag, en leiden een iteratieve formulering af die gebruikmaakt van een bewegend anker voor efficiëntie. We demonstreren de effectiviteit van onze aanpak en verschillende ontwerpkeuzes via experimenten op abstractieve samenvatting en Gemma-modellen. Het uitlijnen van Gemma-beleid met BOND overtreft andere RLHF-algoritmen door betere resultaten te behalen op verschillende benchmarks.

English

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.