BOND: Alineando LLMs con Destilación de lo Mejor-de-N

BOND: Aligning LLMs with Best-of-N Distillation

July 19, 2024

Autores: 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

Resumen

El aprendizaje por refuerzo a partir de la retroalimentación humana (RLHF) es un impulsor clave de la calidad y la seguridad en los modelos de lenguaje grandes de última generación. Sin embargo, una estrategia sorprendentemente simple y sólida en tiempo de inferencia es la selección de la Mejor de N muestras que elige la mejor generación entre N candidatos. En este documento, proponemos Destilación Mejor-de-N (BOND), un algoritmo novedoso de RLHF que busca emular Mejor-de-N pero sin su significativo costo computacional en tiempo de inferencia. Específicamente, BOND es un algoritmo de emparejamiento de distribuciones que obliga a la distribución de generaciones de la política a acercarse a la distribución Mejor-de-N. Utilizamos la divergencia de Jeffreys (una combinación lineal de KL directa e inversa) para equilibrar entre la cobertura de modo y el comportamiento de búsqueda de modo, y derivamos una formulación iterativa que utiliza un anclaje móvil para eficiencia. Demostramos la efectividad de nuestro enfoque y varias elecciones de diseño a través de experimentos en resumen abtractivo y modelos Gemma. Alinear las políticas de Gemma con BOND supera a otros algoritmos de RLHF al mejorar los resultados en varios 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.