Post #39

Controlled Decoding from Language Models Сегодня разбираем недавнюю статью от авторов из Google (один из которых уже успел перейти в OpenAI 😅), посвященную новому алгоритму алайнмента. Давайте разбираться, но прежде немного контекста для лучшего понимания. 1. Речь идет про заключительный этап файнтюнинга LLM, когда мы хотим настроить модель под helpful/harmless/respective/… ответы. То есть альтернативами здесь являются, например, RLHF (PPO) и DPO (не RL, подробнее писал здесь). 2. PPO/DPO алгоритмы меняют веса основной модели. 3. При использовании наиболее популярного сейчас алгоритма PPO, мы применяем on-policy RL, то есть для обновления весов модели каждый раз требуются свежие данные, полученные от самой этой модели. В таком случае схема выглядит как “генерируем ответы моделью → считаем лосс и делаем шаг оптимизатором → генерируем ответы уже обновленной моделью → считаем лосс …”. PPO позволяет делать несколько оптимизационных шагов за одно наблюдение, но все равно это накладывает свои ограничения. В данной статье авторы предлагают другой подход к той же задаче: 1. Обучение отдельной модели Prefix Scorer (PS) в off-policy режиме, то есть нам не требуется генерировать ответы основной моделью. Мы можем использовать любые доступные данные и обучаться на них. 2. Так как PS представляет из себя отдельную модель, веса основной LLM остаются неизменны. Основная идея метода — обучить PS для аппроксимации Value функции текущего состояния. Value функция по определению — средняя награда, которая может быть получена из состояния s (в данном случае s — любой момент на этапе генерации ответа). Делается это путем минимизации MSE ошибки между предсказанием PS и наградой, полученной из Reward Model (с оговоркой, посмотрите, как обрабатывается кейс при y ≠ EOS на картинке). Имея аппроксимацию Value функции, можно корректировать финальные предсказания модели в сторону бОльших наград. Авторы предлагают два способа: 1. Token-wise. На этапе генерации каждого токена прогоняем текущее выходное распределение через Prefix Scorer и изменяем предсказания модели. 2. Block-wise. Семплируем несколько продолжений ответа длины M (например, семплируем 5 вариантов, каждый из которых состоит из 4 токенов) и далее оставляем то, у которого оценка Value функции максимальна. Также можно использовать несколько разных Prefix Scorer для оптимизации сразу нескольких наград. Хотелось бы видеть больше наглядных примеров работы алгоритма, а не просто графиков типа win rate, но все равно работа интересная.