Saved in:
Bibliographic Details
Main Authors: Vassoyan, Jean, Beau, Nathanaël, Plaud, Roman
Format: Preprint
Published: 2025
Subjects:
Online Access:https://arxiv.org/abs/2502.06533
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866916606323458048
author Vassoyan, Jean
Beau, Nathanaël
Plaud, Roman
author_facet Vassoyan, Jean
Beau, Nathanaël
Plaud, Roman
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2502_06533
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Ignore the KL Penalty! Boosting Exploration on Critical Tokens to Enhance RL Fine-Tuning
Vassoyan, Jean
Beau, Nathanaël
Plaud, Roman
Computation and Language
Machine Learning
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.
title Ignore the KL Penalty! Boosting Exploration on Critical Tokens to Enhance RL Fine-Tuning
topic Computation and Language
Machine Learning
url https://arxiv.org/abs/2502.06533