Ignore a Penalidade de KL! Impulsionando a Exploração em Tokens Críticos para Aprimorar o 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

Resumo

A capacidade de alcançar metas de longo prazo é um desafio fundamental no desenvolvimento atual de grandes modelos de linguagem (LLMs). Para lidar com isso, os LLMs pré-treinados podem ser ajustados com aprendizado por reforço (RL) para explorar soluções que otimizam um objetivo dado. No entanto, a exploração com LLMs é difícil, pois é necessário encontrar um equilíbrio entre descobrir novas soluções e permanecer próximo o suficiente do modelo pré-treinado, de modo a não degradar as capacidades básicas. Isso é tipicamente controlado com uma penalidade de Kullback-Leibler (KL). Neste artigo, investigamos a dinâmica de exploração de um pequeno modelo de linguagem em uma tarefa simples de aritmética. Mostramos como diferentes graus de pré-treinamento influenciam a exploração e demonstramos a importância dos "tokens críticos", que têm um impacto dramático no resultado final. Consequentemente, introduzimos uma modificação simples na penalidade de KL que favorece a exploração nos tokens críticos, aumentando a eficiência da 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.