Игнорируйте штраф KL! Повышение исследования критических токенов для улучшения тонкой настройки RL.

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

February 10, 2025

Авторы: Jean Vassoyan, Nathanaël Beau, Roman Plaud

cs.AI

Аннотация

Способность достижения долгосрочных целей является ключевой проблемой в текущем развитии больших моделей языка (LLM). Для решения этой проблемы предварительно обученные LLM могут быть донастроены с помощью обучения с подкреплением (RL) для исследования решений, оптимизирующих заданную цель. Однако исследование с LLM затруднено, поскольку необходимо найти баланс между обнаружением новых решений и сохранением достаточной близости к предварительно обученной модели, чтобы не ухудшить базовые возможности. Обычно это контролируется с помощью штрафа Кульбака-Лейблера (KL). В данной статье мы исследуем динамику исследования небольшой модели языка на простой арифметической задаче. Мы показываем, как различные степени предварительного обучения влияют на исследование и демонстрируем важность "критических токенов", которые имеют решающее значение для конечного результата. В результате мы предлагаем простое изменение штрафа KL, которое способствует исследованию критических токенов, повышая эффективность этапа донастройки 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.