Ignoriere die KL-Strafe! Steigere die Exploration bei kritischen Tokens, um das Feinabstimmen von RL zu verbessern.

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

February 10, 2025

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

cs.AI

Zusammenfassung

Die Fähigkeit, langfristige Ziele zu erreichen, ist eine zentrale Herausforderung bei der aktuellen Entwicklung großer Sprachmodelle (LLMs). Um dies zu bewältigen, können vortrainierte LLMs mit reinforcement learning (RL) feinabgestimmt werden, um Lösungen zu erkunden, die ein gegebenes Ziel optimieren. Die Exploration mit LLMs ist jedoch schwierig, da ein Gleichgewicht zwischen der Entdeckung neuer Lösungen und der Nähe zum vortrainierten Modell gefunden werden muss, um grundlegende Fähigkeiten nicht zu beeinträchtigen. Dies wird in der Regel mit einer Kullback-Leibler (KL)-Strafe gesteuert. In diesem Paper untersuchen wir die Explorationsdynamik eines kleinen Sprachmodells bei einer einfachen arithmetischen Aufgabe. Wir zeigen, wie unterschiedliche Grade des Vortrainings die Exploration beeinflussen und die Bedeutung von "kritischen Tokens" aufzeigen, die einen dramatischen Einfluss auf das Endergebnis haben. Dementsprechend führen wir eine einfache Modifikation der KL-Strafe ein, die die Exploration bei kritischen Tokens begünstigt und die Effizienz der RL-Feinabstimmungsphase erhöht.

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.