Что такое KL penalty в PPO? Зачем ограничивать новую модель от старой?

Краткий тезис

В RLHF при дообучении языковой модели с помощью PPO наивная максимизация наград от reward model приводит к катастрофическому отклонению политики от исходной модели. Чтобы этого избежать, в целевую функцию вводят KL penalty — штраф за расхождение KL-дивергенции между обновляемой политикой $\pi_{[text](/wiki/text){new}}$ и референсной моделью $\pi_{[text](/wiki/text){ref}}$. Итоговая награда вычисляется как $R = [text](/wiki/text){RM}(x,y) - [beta](/wiki/beta) \cdot [text](/wiki/text){KL}(\pi_{[text](/wiki/text){new}} | \pi_{[text](/wiki/text){ref}})$, где $[beta](/wiki/beta)$ — коэффициент, управляющий компромиссом между следованием награде и сохранением языковых навыков базовой модели.

Без этого штрафа модель быстро вырождается: она начинает «обманывать» reward model, генерируя неестественные или грамматически неправильные токены, которые приносят высокий RM-скор, но разрушают полезные свойства языка (катастрофическое забывание). KL penalty действует как регуляризатор, удерживая $\pi_{[text](/wiki/text){new}}$ в окрестности $\pi_{[text](/wiki/text){ref}}$ и гарантируя, что улучшения происходят только в рамках доверительной области.

2. KL penalty предотвращает катастрофическое забывание

Основная причина введения штрафа — катастрофическое забывание (catastrophic forgetting). Языковая модель, обученная на масштабном корпусе текстов, обладает богатыми знаниями о грамматике, фактах, стилях. Если запустить PPO напрямую с $[beta](/wiki/beta) = 0$, модель начнёт оптимизировать RM-скор любой ценой:

• Reward model может быть «сломана» и выдавать высокие оценки за бессмысленные, но патернированные последовательности (adversarial hacking).
• Модель может забыть правильный синтаксис, потому что «выгодно» генерировать токены, которых reward model не видела в обучении (out-of-distribution).
• Снижается разнообразие генераций (mode collapse).

KL penalty действует как «якорь»: он удерживает новую политику около референсной. Если политика начинает отклоняться (например, ради высокого RM-скора), KL быстро растёт, и награда падает. Таким образом, штраф гарантирует, что все улучшения достигаются в пределах доверительного региона (trust region), где модель ещё «помнит» язык.

Это напрямую связано с концепцией TRPO — ограничение шага обновления политики. В PPO вместо явного ограничения KL используется penalty для сглаживания адаптации.

3. Коэффициент β балансирует exploration vs stability

Выбор $[beta](/wiki/beta)$ критически влияет на поведение обучения. Слишком малое $[beta](/wiki/beta)$ делает штраф незначительным — модель перестраивается ради RM-скора, возникают описанные выше проблемы. Слишком большое $[beta](/wiki/beta)$ замораживает политику, и PPO не может улучшить модель (exploration подавлено).

На практике $[beta](/wiki/beta)$ часто динамически адаптируется:

• В оригинальном InstructGPT использовался коэффициент KL, который изменялся так, чтобы фактическое значение KL-дивергенции за эпоху обучения оставалось в окрестности целевого значения (например, 0.1 nats). Если KL слишком мал — $[beta](/wiki/beta)$ уменьшают; если слишком велик — $[beta](/wiki/beta)$ увеличивают.
• Такая адаптация реализуется с помощью пропорционального контроллера (proportional controller) и позволяет автоматически балансировать exploration (попытки отклониться) и stability (удержание около $\pi_{[text](/wiki/text){ref}}$).

Коэффициент $[beta](/wiki/beta)$ также может быть разным для разных токенов или этапов обучения, но общая идея — жёсткое ограничение на суммарное отклонение политики.

4. Без KL: модель может сломать язык

Рассмотрим гипотетический сценарий обучения без KL penalty ($[beta](/wiki/beta)=0$):

• Reward model обучена на небольшой выборке предпочтений. Она может «выучить» нежелательную корреляцию: например, ставить высокий скор за предложения, заканчивающиеся на «...». Модель быстро научится генерировать длинные строки точек, получая максимальный RM-скор, но полностью теряя смысл.
• Даже если reward model честна, модель может войти в область, где её генерации грамматически некорректны, но reward model по ошибке их поощряет. В результате метрики (perplexity, BLEU) резко падают.
• На практике наблюдалось, что KL-дивергенция без штрафа может вырасти до десятков nats, а сгенерированные тексты становятся бессвязными.

Именно поэтому KL penalty — не «приятный бонус», а обязательный компонент PPO в RLHF. Он превращает задачу максимизации награды в регуляризованную оптимизацию, сохраняя языковые способности базовой модели.

5. Пет-проект для закрепления

Задача: реализовать простой симулятор обучения языковой модели с PPO и визуализировать эффект KL penalty.

Инструменты:

• Python + PyTorch / TensorFlow
• Библиотека для RL (Stable‑Baselines3 не подходит, нужен собственный PPO)
• Простая reward model (бинарный классификатор на небольшом датасете)
• Референсная модель: условная LSTM (можно загрузить предобученный GPT-2 с Hugging Face)

Шаги:

1. Создать toy environment: запрос (prompt) → генерация 20 токенов из небольшого словаря.
2. Обучить reward model на синтетических предпочтениях (например, «предложения длиннее 10 токенов»).
3. Реализовать PPO с двумя вариантами награды: $R = [text](/wiki/text){RM}$ и $R = [text](/wiki/text){RM} - [beta](/wiki/beta) \cdot [text](/wiki/text){KL}$.
4. Запустить обучение с $[beta](/wiki/beta)=0$, $[beta](/wiki/beta)=0.01$, $[beta](/wiki/beta)=0.1$. Для каждого варианта записывать средний RM-скор, KL-дивергенцию, perplexity на валидационном корпусе.
5. Построить графики — убедиться, что без штрафа KL расходится, а с ним остаётся в разумных пределах.

Ожидаемый результат:

• Диаграмма, показывающая, что при $[beta](/wiki/beta)=0$ средний RM-скор быстро растёт, но perplexity ухудшается (модель «ломает язык»).
• При $[beta](/wiki/beta)=0.1$ RM-скор растёт медленнее, но perplexity остаётся близкой к исходной.

Связь с другими вопросами

Навигация

• Предыдущий: 986
• Следующий: 988
• Индекс: 00. Индекс разборов