InfAlign: алайнмент языковых моделей с учётом процедуры инференса

Метод RLHF (Reinforcement Learning from Human Feedback) доказал эффективность в задаче алайнмента языковых моделей. Однако у него есть существенный недостаток: на практике возникает расхождение между процессом обучения и реальным использованием модели.

Например, после RLHF модель обычно старается избегать неверных ответов. Но при использовании стратегии генерации Best-of-N (выбор лучшего из нескольких сгенерированных ответов) такое жёсткое ограничение становится неоптимальным — модель могла бы давать лучшие ответы, разреши мы ей экспериментировать более агрессивно за счёт небольшой доли неверных ответов.

Для решения этого несоответствия авторы статьи разработали метод InfAlign, адаптирующий процесс обучения к конкретным процедурам генерации, используемым на практике.

Рассмотрим проблему детальнее. Классический подход RLHF с учётом KL-регуляризации гарантирует оптимальность модели по средней награде, если ответы генерируются сэмплированием. На практике, однако, нам интересна не столько средняя награда, сколько доля запросов, на которых новая модель лучше старой. И уже для такой метрики (при фиксированной модели, по отношению к которой мы считаем винрейт) RLHF даёт субоптимальные результаты даже для простого сэмплирования — что уж говорить о более продвинутых методах.

К счастью, авторам статьи удалось доказать, что оптимизация винрейта для некоторых процедур генерации, включая Best-of-N, Worst-of-N и сэмплирование, эквивалентна применению RLHF с модифицированной функцией награды.

Предложенный подход состоит из трёх основных этапов.

1. Калибровка награды. На этом этапе исходные награды преобразуются в значения от 0 до 1 таким образом, чтобы распределение наград ответов модели стало равномерным на каждом запросе. Это эквивалентно применению обусловленной на запрос функции распределения награды к самой награде. Забавно, что в первой версии статьи авторы предложили использовать медианную аппроксимацию функции распределения, однако спустя месяц удалили все упоминания об этом методе и перешли к использованию эмпирической функции распределения.

2. Трансформация награды. На следующем этапе откалиброванная награда адаптируется под конкретную процедуру генерации. Например, для стратегии Best-of-N применяется экспоненциальное преобразование, усиливающее различия между отличными и посредственными ответами, а для сэмплирования — логарифм, штрафующий за плохие ответы. Заметим, что на самом деле логарифм и экспонента — это лишь хорошие приближения оптимального преобразования. Но, как показывают эксперименты, погрешностью можно пренебречь ради простоты реализации.

3. Обучение с модифицированной наградой. Модель обучается при помощи классического RLHF, используя модифицированную награду, адаптированную под конкретную процедуру генерации.

Проведённые авторами эксперименты показали, что предложенный подход демонстрирует значительное улучшение качества генерации с точки зрения винрейта даже для процедуры сэмплирования.

Отметим, что сейчас метод InfAlign применим к весьма ограниченному набору реально используемых процедур генерации, таких как Best-of-N, Worst-of-N и сэмплирования.

Разбор подготовил ❣ Федор Лебедь

Душный NLP