WARP : Sur les avantages des politiques récompensées par moyenne pondérée

WARP: On the Benefits of Weight Averaged Rewarded Policies

June 24, 2024

Auteurs: Alexandre Ramé, Johan Ferret, Nino Vieillard, Robert Dadashi, Léonard Hussenot, Pierre-Louis Cedoz, Pier Giuseppe Sessa, Sertan Girgin, Arthur Douillard, Olivier Bachem

cs.AI

Résumé

L'apprentissage par renforcement à partir de retours humains (RLHF) aligne les grands modèles de langage (LLMs) en encourageant leurs générations à obtenir des récompenses élevées, grâce à un modèle de récompense entraîné sur les préférences humaines. Pour éviter l'oubli des connaissances pré-entraînées, le RLHF intègre généralement une régularisation KL ; celle-ci contraint la politique à rester proche de son initialisation par fine-tuning supervisé, bien qu'elle entrave l'optimisation des récompenses. Pour résoudre ce compromis entre KL et récompense, nous introduisons dans cet article une nouvelle stratégie d'alignement nommée Weight Averaged Rewarded Policies (WARP). WARP fusionne les politiques dans l'espace des poids à trois étapes distinctes. Premièrement, elle utilise la moyenne mobile exponentielle de la politique comme ancre dynamique dans la régularisation KL. Deuxièmement, elle applique une interpolation sphérique pour fusionner des politiques indépendamment fine-tunées en une nouvelle politique améliorée. Troisièmement, elle effectue une interpolation linéaire entre ce modèle fusionné et l'initialisation, afin de récupérer des caractéristiques issues du pré-entraînement. Cette procédure est ensuite appliquée de manière itérative, le modèle final de chaque itération servant d'initialisation avancée pour la suivante, affinant progressivement le front de Pareto KL-récompense et obtenant des récompenses supérieures à KL fixe. Des expériences avec des politiques GEMMA valident que WARP améliore leur qualité et leur alignement, surpassant d'autres LLMs open-source.

English

Reinforcement learning from human feedback (RLHF) aligns large language models (LLMs) by encouraging their generations to have high rewards, using a reward model trained on human preferences. To prevent the forgetting of pre-trained knowledge, RLHF usually incorporates a KL regularization; this forces the policy to remain close to its supervised fine-tuned initialization, though it hinders the reward optimization. To tackle the trade-off between KL and reward, in this paper we introduce a novel alignment strategy named Weight Averaged Rewarded Policies (WARP). WARP merges policies in the weight space at three distinct stages. First, it uses the exponential moving average of the policy as a dynamic anchor in the KL regularization. Second, it applies spherical interpolation to merge independently fine-tuned policies into a new enhanced one. Third, it linearly interpolates between this merged model and the initialization, to recover features from pre-training. This procedure is then applied iteratively, with each iteration's final model used as an advanced initialization for the next, progressively refining the KL-reward Pareto front, achieving superior rewards at fixed KL. Experiments with GEMMA policies validate that WARP improves their quality and alignment, outperforming other open-source LLMs.