Ignora la Penalità di KL! Potenziare l'Esplorazione sui Token Critici per Migliorare il Fine-Tuning di RL

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

February 10, 2025

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

cs.AI

Abstract

La capacità di raggiungere obiettivi a lungo termine è una sfida chiave nello sviluppo attuale dei grandi modelli linguistici (LLM). Per affrontare questo problema, i LLM pre-addestrati possono essere ottimizzati ulteriormente con apprendimento per rinforzo (RL) per esplorare soluzioni che ottimizzano un obiettivo specifico. Tuttavia, l'esplorazione con i LLM è complessa, poiché è necessario trovare un equilibrio tra la scoperta di nuove soluzioni e il mantenimento della coerenza con il modello pre-addestrato, al fine di non compromettere le capacità di base. Questo equilibrio è tipicamente controllato con una penalità di Kullback-Leibler (KL). In questo articolo, esaminiamo la dinamica dell'esplorazione di un piccolo modello linguistico su un semplice compito aritmetico. Mostreremo come diversi gradi di pre-addestramento influenzino l'esplorazione e dimostreremo l'importanza dei "token critici" che hanno un impatto significativo sull'esito finale. Di conseguenza, introduciamo una semplice modifica alla penalità KL che favorisce l'esplorazione sui token critici, aumentando l'efficienza della fase di ottimizzazione ulteriore con 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.