BOND : Alignement des LLM par distillation Best-of-N

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

Résumé

L'apprentissage par renforcement à partir de retours humains (RLHF) est un facteur clé de qualité et de sécurité dans les modèles de langage à la pointe de la technologie. Pourtant, une stratégie d'inférence étonnamment simple et efficace est l'échantillonnage Best-of-N, qui sélectionne la meilleure génération parmi N candidats. Dans cet article, nous proposons Best-of-N Distillation (BOND), un nouvel algorithme RLHF qui vise à imiter Best-of-N sans son coût computationnel significatif au moment de l'inférence. Plus précisément, BOND est un algorithme d'appariement de distributions qui force la distribution des générations de la politique à se rapprocher de la distribution Best-of-N. Nous utilisons la divergence de Jeffreys (une combinaison linéaire des KL avant et arrière) pour équilibrer entre un comportement couvrant les modes et un comportement recherchant les modes, et nous dérivons une formulation itérative qui utilise une ancre mobile pour plus d'efficacité. Nous démontrons l'efficacité de notre approche et de plusieurs choix de conception à travers des expériences sur la synthèse abstraite et les modèles Gemma. L'alignement des politiques Gemma avec BOND surpasse d'autres algorithmes RLHF en améliorant les résultats sur plusieurs 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.