¡Ignora la penalización de KL! Impulso a la exploración en tokens críticos para mejorar el ajuste fino de RL.

Ignore the KL Penalty! Boosting Exploration on Critical Tokens to Enhance RL Fine-Tuning

February 10, 2025

Autores: Jean Vassoyan, Nathanaël Beau, Roman Plaud

cs.AI

Resumen

La capacidad de lograr objetivos a largo plazo es un desafío clave en el desarrollo actual de grandes modelos de lenguaje (LLMs). Para abordar esto, los LLMs pre-entrenados pueden ser ajustados con aprendizaje por refuerzo (RL) para explorar soluciones que optimicen un objetivo dado. Sin embargo, la exploración con LLMs es difícil, ya que se debe encontrar un equilibrio entre descubrir nuevas soluciones y mantenerse lo suficientemente cerca del modelo pre-entrenado, para no degradar las capacidades básicas. Esto se controla típicamente con una penalización de Kullback-Leibler (KL). En este documento, investigamos la dinámica de exploración de un pequeño modelo de lenguaje en una tarea aritmética simple. Mostramos cómo diferentes grados de pre-entrenamiento influyen en la exploración y demostramos la importancia de los "tokens críticos" que tienen un impacto dramático en el resultado final. En consecuencia, introducimos una modificación simple a la penalización de KL que favorece la exploración en los tokens críticos, aumentando la eficiencia de la etapa de ajuste fino de RL.

English

The ability to achieve long-term goals is a key challenge in the current development of large language models (LLMs). To address this, pre-trained LLMs can be fine-tuned with reinforcement learning (RL) to explore solutions that optimize a given goal. However, exploration with LLMs is difficult, as a balance has to be struck between discovering new solutions and staying close enough to the pre-trained model, so as not to degrade basic capabilities. This is typically controlled with a Kullback-Leibler (KL) penalty. In this paper, we investigate the exploration dynamics of a small language model on a simple arithmetic task. We show how varying degrees of pre-training influence exploration and demonstrate the importance of "critical tokens" which have a dramatic impact on the final outcome. Consequently, we introduce a simple modification to the KL penalty that favors exploration on critical tokens, increasing the efficiency of the RL fine-tuning stage.