Как мы готовим RL для Alignment в больших языковых моделях: опыт команды YandexGPT

Сегодня через API стала доступна новая модель YandexGPT 3 Lite. Одним из ключевых этапов её обучения, как и в случае с другими недавними моделями, стал этап Alignment (выравнивания), включающий в том числе стадию обучения с подкреплением (RL). Пожалуй, без этого этапа мы бы не смогли добиться такого роста в качестве, который был необходим для запуска новых возможностей и сервисов (например, Нейро). Поэтому эту статью мы полностью посвятим особенностям выравнивания моделей.

На тему Alignment и RL было написано уже немало статей. Кажется, любой ML‑инженер уже, так или иначе, сталкивался или читал о них. Поэтому мы хоть и напомним базовую информацию, но всё же сфокусируемся на тех деталях реализации, которые не на слуху.

Небольшое предисловие

Сегодня мы будем говорить про одну из стадий обучения языковых моделей‑ассистентов. Этот этап называется alignment (выравнивание) и идёт вслед за этапом предварительного обучения модели (pretrain).

Мы практически совсем не будем говорить об архитектуре нейронных сетей: все нейронные сети в этой статье — это трансформеры. Нам достаточно будет знать, что языковые модели получают на вход текст в токенизированном виде и выдают вероятность встретить тот или иной токен следующим. Если вы не знаете, что такое «токен», то просто считайте, что токены — это слова. То есть нейронная сеть смотрит на входной текст и выдаёт вероятности для всех возможных слов: с каким шансом какое слово будет следующим.

На стадии pretrain языковая модель учится на огромном количестве текстов из интернета: она изучает общие правила построения предложений, общие знания, которые можно найти в интернете. После этой стадии модель умеет продолжать тексты так, как это было в обучающей выборке.

Если мы хотим получить языковую модель‑помощник, использовать pretrain‑модель как есть не получится — она не сможет отвечать на запросы пользователя, хоть в ней и содержатся знания всего интернета. Вместо этого модель будет пытаться этот запрос продолжать.

Процесс превращения просто умной модели в модель‑ассистента и называется alignment — выравнивание. Мы будем пытаться «выравнивать» ответы модели с нашими человеческими ожиданиями. Команда Anthropic хорошо описала свойства, которыми должны обладать ответы такой модели. Речь о трёх H:

• Helpful — ответ должен решать задачу пользователя.

• Harmless — ответ не должен вредить пользователю.

• Honest — ответ должен быть фактически корректным.

Обычно подобного поведения от модели мы добиваемся в два этапа выравнивания:

1. Обучение с учителем на выборке, собранной людьми.

2. Дообучение модели с подкреплением (RLHF), которое позволяет не просто выучить «определённое поведение», но максимизирует удовлетворение пользователя от общения с моделью.

Об этих двух этапах мы и поговорим сегодня подробнее.

Подробнее о стадии SFT

На стадии обучения с учителем (Supervised FineTuning, SFT) мы будем дообучать pretrain‑модель на парах «запрос — ответ», чтобы модель начала своё превращение в ассистента. Основные сложности и тонкости этого этапа заключаются в сборе хорошей выборки для обучения. Для SFT‑стадии нам нужна разнообразная выборка запросов, которые могли бы поступить от пользователей, а также правильные ответы на каждый из этих запросов.

Ситуация с запросами несколько проще, чем с ответами: их можно как написать человеческими руками, так и собрать в интернете в полуавтоматическом режиме. Запросы для первых этапов обучения YandexGPT мы отчасти собирали с помощью внутренних конкурсов: просили энтузиастов придумать, как и откуда собрать большое количество возможных запросов, или даже придумать запросы самим. Получилось очень даже неплохо.

С правильными ответами всё куда прозаичнее — их мы собираем при помощи разметки людьми. Для этого нужно формализовать критерии helpful, harmless, honest в законченную инструкцию, какими именно характеристиками должен обладать ответ. Дальше мы разбиваем запросы по категориям и отдаём запросы каждой категории в руки опытных специалистов, чтобы они написали верный ответ.

На полученной таким образом выборке дообучаем pretrain‑модель и получаем так называемую SFT‑модель. Чтобы обучаемая модель могла понять, где в тексте заканчивается запрос пользователя и начинается желаемый ответ, мы разделяем запрос и ответ спецтокеном. На этапе использования модели в неё будет подаваться запрос + спецтокен и модель сразу поймёт, что следующим токеном ей нужно отвечать на запрос.

Здесь нужно сказать пару слов про выбор модели: обычно мы выбираем лучшую модель по целевой метрике на валидации. В задаче выравнивания моделей сложно придумать автоматизированную метрику — большинство из них непоказательны. Поэтому выбор модели обычно проводится тоже не без помощи человеческих рук: на не слишком большой отложенной выборке запросов мы генерируем ответы обучаемой моделью, а также моделью‑бейзлайном. Люди же решают, чей ответ был лучше, и вычисляют так называемое side‑by‑side превосходство (sbs). В силу использования не слишком больших выборок для сравнения обязательно необходимо убеждаться в статистической значимости прироста. Например, можно использовать биномиальный статистический тест.

И писать ответы, и оценивать качество ответов должны специалисты, так как различать нюансы в ответах модели на сложные темы очень непросто. В этом нам помогают AI‑тренеры, которые и хорошо разбираются в узких областях, и могут написать грамотный текст.

Считается, что на стадии выравнивания, в частности — SFT, языковая модель не получает новых знаний, а лишь учится использовать уже имеющиеся: совершать поиск в массиве уже выученной информации и правильно её доносить. В исследованиях показано, что использование небольшой выборки с высоким качеством разметки на этом этапе приводит к лучшим результатам, чем использование больших выборок сомнительного качества (LIMA: Less Is More for Alignment ). Более того, исследователи из Google Research показали, что попытка вложить новые знания в голову модели на этапе SFT может привести к галлюцинациям — поведению, когда модель с полной уверенностью говорит полную ерунду и иногда очень правдоподобную (Does Fine‑Tuning LLMs on New Knowledge Encourage Hallucinations? ).

Что дальше? Обычно после стадии SFT модель получается уже достаточно умной и «выровненной», чтобы быть ассистентом. Однако можно лучше: данных для обучения SFT всегда мало, а если вычищать «грязные», то ещё меньше.

Оказывается, можно ещё сильнее «выровнять» модель с нашими ожиданиями без использования прямой разметки — написания ответов людьми. Этот этап называется RL from Human Feedback.

Подробнее об RLHF

Есть такая область машинного обучения — обучение с подкреплением (Reinforcement Learning, RL). Это как обучение с учителем, но учитель тут — это не правильный ответ, а наказание или награда, сваливающаяся на голову студента.

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

Формально специалисты по RL любят говорить про марковские процессы принятия решений (Markov Decision Process, MDP), описывающие взаимодействие студента (агента) и среды, которая поставляет обучающий сигнал. В MDP агент наблюдает текущее состояние среды и выбирает действие в соответствии со своей стратегией поведения (в англоязычной литературе стратегия называется policy, и это слово часто переводят на русский, как «политика», что неверно, но общеупотребимо). Среда на действие реагирует изменением своего состояния на , а также некоторой наградой для агента — именно награда и является обучающим сигналом в MDP. Марковское свойство тут заключается в том, что, если агент знает текущее состояние среды, ему не нужно знать об истории взаимодействия, чтобы выбрать оптимальное действие.

Про взаимодействие пользователя с ИИ‑ассистентом можно тоже думать в контексте марковского процесса принятия решений. Состояние среды — это текущий текст в беседе, начинающийся с запроса пользователя (в начале взаимодействия это просто запрос пользователя). Действие агента — это, конечно же, ответ языковой модели. Интересно, что такое награда в этой задаче. Об этом подробно мы поговорим в следующей главе.

Из соображений простоты в задаче alignment обычно рассматривается одношаговое взаимодействие: агент видит запрос пользователя (или какой‑то диалог пользователя с ИИ, заканчивающийся репликой пользователя) и пытается ответить так, чтобы получить максимальное количество награды за этот ответ.

Существует множество алгоритмов RL, отличающихся теми или иными особенностями. Дальше в статье мы рассмотрим три алгоритма, наиболее подходящих для решения задачи выравнивания.

Если обучение с учителем больше похоже на «запихивание» знаний эксперта в голову модели, то на этапе обучения с подкреплением модель учится на собственном опыте генерации текстов. На этом этапе модель может «переусвоить» формулировки, внесённые на этапе SFT в более приятном для себя виде. Эффект галлюцинаций, которые могут появиться во время SFT, на этапе RL снижается, так как модель быстро понимает, что ответ «не знаю» ведёт к большей награде, чем ответ неправильный.

А судьи кто?

При наличии идеальной функции награды формализм RL позволяет находить оптимальные генеративные модели без использования дополнительной прямой разметки (написанных людьми ответов), нужен лишь пул хороших возможных запросов. Однако будет лукавством сказать, что никакой разметки для RL не нужно, — ведь у нас нет никакой идеальной награды!

Откуда её — награду — вообще взять? Первое, что приходит в голову, — вновь посадить людей размечать данные, но уже в куда более щадящем режиме. Давайте мы для большой выборки запросов сгенерируем по одному ответу с помощью нашей SFT‑модели с прошлой стадии. А людей попросим оценить, насколько был хорошим ответ: пусть ставят оценки от 1 до 5 в зависимости от крутости. Затем возьмём ту же SFT‑модель и заменим ей голову на необученный линейный слой. Нейронная сеть устроена так, что её последний слой выдаёт вероятностное распределение над токенами — его меняем на обычный линейный слой, который выдаёт одно число. Это и будет наша модель награды: будем скармливать ей запрос + ответ и обучать так, чтобы на выходе была та оценка, какую в среднем ставят люди такому ответу.

При всей прямолинейности этот подход не лишён недостатков:

• AI‑тренерам часто сложно выбрать правильную оценку из пяти опций.

• На этапе дизайна системы выставления оценок сложно определиться с тем, что такое «отлично», «хорошо» и так далее. Если прогадать, то может случиться так, что почти все оценки оказались, например, хорошими. Тогда модель не сможет выучить, что именно делает эти ответы хорошими. И разметку придётся пересобирать, а это дорого.

Хотя такой подход и используется, обычно применяется его более сложная модификация, упрощающая разметку для AI‑тренеров: куда проще сравнить два ответа и выбрать лучший, чем выставлять сухую оценку.

Будем делать так: на каждый запрос в выборке генерируем по два ответа и просим AI‑тренеров для каждой пары подписать, какой ответ был лучше. Так и размечать проще, и нет проблемы, что модель не сможет различать хорошие ответы. Если они одинаково хороши, то это тоже можно явно указать при разметке.

Если вы вдруг вздумали собирать разметку для собственной LLM‑компаньонки, то вам стоит знать про синдром утёнка. Утята готовы принять за свою маму первый подходящий предмет, который увидят. Так и AI‑тренеры, подобно утятам, больше любят первый прочитанный ответ при прочих равных. Обязательно перемешивайте порядок ответов случайным образом, прежде чем показывать AI‑тренеру, иначе можете уплыть не туда.

Более сложный вопрос: как на такой выборке обучить модель награды? Тут есть два наиболее популярных пути.

Путь первый. Модель Брэдли-Терри

Будем обучать всё ту же модель, которая на запрос + ответ выдаёт одно число (награду), но новым способом. Явной разметки, какая награда какому ответу соответствует, у нас теперь нет. Но мы можем попросить модель выдавать хорошему ответу число побольше, а тому ответу, который хуже, — поменьше. Для этого введём вероятностную модель, где вероятность (уверенность модели), что ответ лучше ответа выражается следующей формулой:

Здесь — это запрос, на который даны ответы, — это обучаемая нейронная модель награды с параметрами , а — функция, отображающая ничем не ограниченную разность в интервал (0, 1). Это нужно, так как вероятность должна жить в этом интервале:

Таким образом, уверенность модели, что ответ лучше ответа тем больше, чем больше разница в награде, которую модель даёт этим двум ответам.

Обучение модели проводится методом максимизации правдоподобия: мы будем подбирать параметры так, чтобы вероятность для собранной AI‑тренерами выборки была максимальной. С целью стабилизации всяких численных проблем вероятность обычно логарифмируют, что ведёт нас к следующей задаче оптимизации:

То есть обучение сводится к простому обучению с учителем, но со специфической функцией ошибки.

Экспериментальным путём мы пришли к пайплайну, в котором модель награды учится в два этапа:

1. Сначала учим на «грязных» данных, собранных в полуавтоматическом режиме.

2. Дообучаем награду на чистых данных, полученных от AI‑тренеров.

Если со вторым этапом всё на данный момент должно быть прозрачно, то вот про первый пока не понятно, на каких данных его проводить. Оказывается, что довольно много пар вида (запрос)+(лучший ответ, ответ похуже) можно бесплатно раздобыть в интернете.

Многие сервисы в интернете позволяют ранжировать ответы пользователей на какие‑то запросы. Например, здесь же на Хабре, хабровчане могут оценить ответы пользователей в комментариях, и мы буквально можем выбрать пару (хороший комментарий, плохой комментарий) в ответ на (запрос) — пост. Такая же логика может быть применена к небезызвестному StackOverflow.

А ещё можно найти много задачников, где для вопроса предлагается выбрать правильный ответ из списка и где правильный ответ заранее известен. Тут тоже несложно собрать выборку с двумя ответами на вопрос, где один — лучше.

Немного подумав, можно напарсить большую выборку сравнений сомнительного качества — вот она и пойдёт в первый этап обучения модели награды.

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

Путь второй. Модель, которая ест сразу два запроса

Модель Брэдли‑Терри вносит некоторые ограничения в то, какая может выучиться награда. Одно из основных ограничений — это предположение, что награда транзитивна. Говоря простыми словами, в модели Брэдли‑Терри не бывает циклов в предпочтениях: не может быть такого, что ответ , , но .

Может показаться абсурдным пожелание разрешить награде формировать циклы в предпочтениях, но теория социального выбора говорит, что мы, люди, именно так, нетранзитивно, выбор и делаем.

Пусть существует некоторая группа AI‑тренеров, а также некоторое множество различных ответов на запрос. Каждый AI‑тренер может отранжировать ответы по своему предпочтению (например, ) — для каждого отдельно взятого AI‑тренера транзитивность соблюдена. Но как только мы ответы разных AI‑тренеров агрегируем, вполне может сложиться цикл , , но Это явление называется парадоксом Кондорсе и широко обсуждается обычно в контексте демократических выборов.

На самом деле, позволить модели награды быть достаточно гибкой, чтобы она могла воспроизвести возможные нетранзитивности в выборе AI‑тренеров, несложно. Достаточно обучить нейронную сеть , которая вместе с запросом принимает на вход не один, а сразу два ответа и возвращает уверенность модели, что . Для обучения такой модели у нас как раз имеется выборка с разметкой, кто из ответов был лучше. Обучаем решать классическую задачу классификации.

У такой модели, правда, есть существенный недостаток: с помощью неё можно лишь сравнивать разные модели — абсолютного значения награды она не выдаёт. А ещё, во время обучения такой модели нужно как‑то решать проблему несимметричности: . Кроме этого, такую модель награды просто дороже инферить, так как в неё подаётся сразу два ответа вместо одного.

В процессе разработки YandexGPT мы проводили эксперименты с потенциально нетранзитивной моделью награды. Однако показать её превосходство по сравнению с моделью Брэдли‑Терри не смогли — они работали одинаково хорошо. В силу сложности использования нетранзитивной модели, мы остановились на классической модели Брэдли‑Терри.

Для чего (ещё) нужна модель награды?

Вот, наконец, у нас на руках есть мерило успеха — оценщик качества ответов генеративной модели. Конечно, основная идея, что с этой моделью делать, тривиальна: давайте забабахаем дообучение с подкреплением SFT‑модели, чтобы максимизировать эту самую награду. И мы, конечно же, обсудим этот подход далее.

Однако применение модели награды не ограничивается одним лишь RL. Пусть у вас есть модель награды и какая‑то SFT‑модель, и вот вы насобирали новых данных и планируете переобучить SFT на новом датасете. Хорошая идея — смотреть во время обучения не только на графики лосса на валидации, но ещё и на награду. Для этого во время валидации нужно на запросы из валидационной выборки генерировать ответы обучаемой моделью и считать среднюю награду для этих ответов.

Часто оказывается так, что ранний останов по награде случается не на той итерации, что ранний останов по лоссу. Мы используем этот нехитрый трюк в YandexGPT и стабильно получаем SFT‑модели более высокого качества.

Здесь и далее мы не раз будем говорить про генерацию ответов моделью с целью валидации или для обучения. Эффективная по скорости генерация — залог успеха. Существует целый ряд библиотек, позволяющих скомпилировать вашу модель в формат, пригодный для быстрого исполнения. Например, TensorRT, поддержанный в Torch. В Яндексе мы используем собственный внутренний фреймворк для быстрого исполнения моделей.

Не всегда конвертация модели в «быстрый» формат стоит свеч. Например, если вы планируете генерировать тексты на каждой итерации валидации, временные затраты на конвертацию могут превзойти затраты на не очень эффективную генерацию, так как конвертировать модель придётся каждый раз заново. Здесь можно воспользоваться правилом: если веса модели меняются редко или вовсе не меняются, можно конвертировать, а если меняются часто — лучше инферить на обычном Torch. По этой логике генерацию для валидации во время обучения мы делаем без конвертации, а вот модель награды держим в «быстром» формате, потому что она не меняется.

«RL для бедных» — Cross-Entropy Method

Наконец, можно поговорить и об основном способе использования модели награды — обучения генеративной модели, которая бы набирала много награды.

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

• генерируем SFT‑моделью для каждого запроса по N ответов,

• выбираем лучший ответ по награде для каждого запроса,

• дообучаем SFT‑модель на полученных парах запрос + лучший ответ.

При всей простоте из этого метода можно выжать довольно неплохой профит. Здесь для генерации и оценки как раз осмысленно использовать сконвертированные в «быстрый» формат генеративную модель и модель награды.

Среди RL‑специалистов этот метод больше известен под именем Cross‑Entropy Method (CEM / CE‑RL). Применение CEM к RL — это, на самом деле, частный случай. CEM — это общий алгоритм оптимизации, относящийся к семейству генетических алгоритмов. В двух словах CEM можно описать так: генерируем множество точек и выбираем из них лучшие, затем двигаем генератор точек в сторону лучших.

Важно понимать, что CEM может быть эффективен лишь для моделей, обладающих ненулевым разнообразием ответов. Если вы отвечаете одно и то же, то выбирать лучший ответ бессмысленно.

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

Идея «давайте сдистиллируем хорошие ответы модели назад в модель» на самом деле довольно общая. При этом крутые ответы не обязательно должны быть получены случайной генерацией разных ответов. Можно получить хорошие ответы, применив метод Chain‑of‑Thought где модель посредством подводок генерирует ответ в несколько шагов, «проговаривая свои мысли вслух». А можно применить тяжёлую артиллерию — Monte‑Carlo Tree Search (MCTS) — для направленного поиска по дереву ответов с высокой наградой.

Почти в любом курсе по RL вам скажут, что в RL переобучения не бывает. Имеется в виду, что если агент научился набирать много награды для какой‑то задачи, то специальным образом его валидировать не нужно — он уже решает задачу максимизации награды. В языковом моделировании это правило не работает, так как агент видит не все возможные запросы, а лишь какое‑то их подмножество. Поэтому критически важно валидировать обучаемую с помощью RL языковую модель на отложенной выборке запросов.

«RL для богатых» — Proximal Policy Optimization

От простого к сложному: «RL для бедных» при всей простоте имплементации довольно сложно заставить найти оптимальную в смысле награды модель. Основных причин две:

• генерировать N гипотез на каждый запрос для сравнительно большой выборки весьма дорого по времени и ресурсам;

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

Здесь на помощь приходят более теоретически совершенные алгоритмы обучения с подкреплением. Для задачи alignment обычно принято использовать алгоритмы из семейства policy gradients.

Идея простая: мы хотим максимизировать среднюю награду, которую набирает наш обучаемый студент. Давайте посчитаем градиент этой средней награды и изменим веса модели в направлении этого вектора, то есть применим самый обычный градиентный подъём (это как градиентный спуск, но для максимизации, а не минимизации функций).

Формально, среднюю награду агента можно записать так:

Да, здесь — это матожидание, но не нужно пугаться! Сегодня нам достаточно про него думать как про простое усреднение по большому (бесконечному) количеству примеров. Например, — это усреднение по большому количеству запросов из выборки . А — это усреднение по большому количеству ответов модели на запрос . Здесь за мы обозначили собственно обучаемую генеративную языковую модель с параметрами . Получается, — это вероятность ответа на запрос . Именно такие вероятности нам выдают языковые модели.

Сложная часть метода — как, собственно, найти производную этой штуки по . Неприятно то, что сидит под матожиданием. Однако нерешаемых проблем почти не бывает, и, если напрячь матапарат, можно получить не слишком неприятную формулу для производной:

Чтобы понять, откуда взялась такая формула для градиента, вспомнить, что такое математическое ожидание, всё же придётся. Распишем матожидание по ответам в функционале, который мы собираемся дифференцировать:

Суммирование здесь происходит по множеству всех возможных ответов, коих бесконечно (счётно) много. Попробуем взять производную по :

С полученным выражением есть проблема: раньше у нас было матожидание по ответам, а теперь — сумма. Эта сумма не является математическим ожиданием, потому что под суммой должны быть вероятности, а тут — градиенты вероятностей. Матожидания лучше сумм, так как матожидания мы умеем оценивать методом Monte‑Carlo (усреднением, по сути), а вот суммы — нет. Чтобы вернуть желаемое матожидание, предлагается воспользоваться так называемым log‑derivative трюком (трюк с производной логарифма):

Эту формулу несложно получить, если расписать формулу для производной логарифма. Собственно, её и предлагается подставить в формулу производной оптимизируемого функционала:

Теперь под суммой есть вероятности, и можно выделить матожидание назад, как мы и хотели:

Ч.Т.Д.

В целом, уже из этой формулы можно сообразить рабочий алгоритм для максимизации средней награды. Но давайте остановимся и подумаем, что эта формула нам говорит. Если долго и пристально на неё смотреть, можно увидеть, что этот градиент нам говорит изменять вероятность ответа пропорционально его награде .Если награда отрицательная, то вероятность уменьшаем, если положительная — увеличиваем, а если сильно положительная — сильно увеличиваем.

Пока мы используем матожидания — усреденения по бесконечному числу примеров — всё работает как положено. Но на практике мы неизбежно перейдём к усреднениям по конечному, небольшому количеству ответов модели. И тут могут быть проблемы: пусть награда выучилась так, что она везде положительная: для плохих примеров она поменьше, для хороших — побольше. Тогда этот градиент будет нас толкать увеличивать вероятности для любых ответов, какие бы модель ни сгенерировала.

Для решения этой проблемы на практике обычно используют видоизменённую формулу (тоже теоретически обоснованную):

Здесь — это средняя награда, которую агент набирает, если генерирует ответы для запроса (так называемая функция ценности для агента). Обычно под эту нужду обучается ещё одна нейронная сеть, которая принимает запрос и возвращает одно число. То есть идея изменения в том, что мы будем повышать вероятности только тех ответов, которые лучше среднего. Учится эта нейронная сеть минимизировать среднее квадратичное отклонение своих выходов от награды за ответ:

Из‑за свойств среднего квадратичного отклонения выучится как раз нужное нам среднее. Тут прежних проблем с градиентами нет — параметры под самим матожиданием не фигурируют.

Итак, что же у нас получилось? Давайте оформим это в виде псевдокода:

Алгоритм 1.

1. Вход: множество запросов для обучения

2. Инициализируем политику SFT-моделью:

3. Инициализируем ценность моделью награды:

4. Повторять до сходимости:

4.1. Выбираем батч запросов

4.2. Вычисляем ценность для каждого запроса из батча

4.3. Генерируем по одному ответу на каждый запрос

Важно генерировать именно актуальной обучаемой моделью

4.4. Вычисляем награду для всех пар

4.5. Вычисляем лосс для агента

4.6. Вычисляем лосс для функции ценности

4.7.

4.8.

Алгоритм, который мы записали выше, называется Advantage Actor Critic (A2C). Это уже очень неплохой алгоритм для максимизации награды. Прежде чем мы нырнём в обсуждение нюансов, давайте проговорим его основной минус — большую дороговизну обучения (в смысле времени). Алгоритм требует от вас генерировать по ответу на каждый элемент батча обязательно моделью с актуальной версией весов — нельзя заранее эффективно сгенерировать по ответу на каждый запрос выборки, как это было в «бедном RL». Тут вам наверняка придётся совершать куда менее эффективную генерацию обычной торчовой моделью.

Скорее всего, если вы где‑то читали про alignment, вы слышали, что на практике применяется некто PPO (Proximal Policy Optimization). Это ни что иное, как A2C, к которому приделали костыли (importance sampling и обрезку градиентов), чтобы можно было делать больше одного шага оптимизатора на одних и тех же ответах. Скорее всего, это всё, что вам нужно знать про PPO, но если вы хотите узнать больше, приходите на курс по обучению с подкреплением в ШАД (гитхаб курса ).

PPO частично решает одну из проблем A2C — необходимость генерировать очень много текстов — но оставляет и привносит другие сложности:

• вам нужно хранить три модели в памяти: политику, ценность и награду;

• всё же нужно написать эффективный код генерации на торче;

• у PPO очень много гиперпараметров, и он к ним достаточно чувствителен — вас ждёт поиск гиперпараметров по сетке.

В литературе часто советуют не заводить отдельную модель под V‑функцию, а обойтись двухголовой моделью, где из общего тела торчат две головы: политика и V‑модель. В YandexGPT мы заметили, что такую архитектуру очень сложно обучать: лоссы политики и ценности нужно замешивать с каким‑то коэффициентом, который сложно подобрать (а ещё хотелось бы его менять в процессе обучения). Неудачно подобранный коэффициент ведёт либо к тому, что V не учится, либо к тому, что она перетягивает на себя все силы оптимизатора, и агент выходит не торт.

Существует ряд работ, предлагающих отказаться от V‑модели совсем, — использовать заранее посчитанную константу вместо неё. Мы провели этот эксперимент с YandexGPT и увидели, что это работает. Но не всегда. В 50% запусков, при неидеально подобранных гиперпараметрах, обучение жесточайше разваливается. Хотя избавление от V‑модели в памяти и ускоряет жизнь, нестабильность такого подхода не позволяет на него положиться.

Если вы хотите узнать какую‑нибудь The‑One‑Киллер‑Фичу, которая заставит ваш PPO работать, наверняка — это нормализация advantage. Advantage в RL принято называть разницу награды за ответ и средней награды , которую мы используем при обучении агента. Если эту штуку отнормировать по батчу на нулевое среднее и единичную дисперсию, то вы, с одной стороны, потеряете теоретические гарантии к алгоритму, а с другой — получите намного более стабильную сходимость. Рекомендуем.

Доменный сдвиг, он же — Гудхартинг

Вы научились максимизировать награду. Всегда ли это круто? Казалось бы, мы этого и хотели — хорошие ответы, которые набирают много награды. К сожалению, высокая награда — не всегда показатель хорошего ответа. Ведь награду нам выдаёт нейронная сеть, обученная на конечной выборке, — у неё есть слабые места.

Хуже ситуация становится от понимания, что RL изменяет ответы генератора именно в том направлении, где модель награды работает плохо. Проще всего этот эффект увидеть на визуальном примере ниже.

Представьте себе, что ваш агент стоит у подножия горы, и он получает награду пропорционально его координате выше — лучше. Агент изучает свои окрестности и выучивает закономерность, которой следует награда: если идти вправо по картинке, то награда растёт. Далее запускается алгоритм поиска оптимальной политики — RL. Этот алгоритм, конечно, говорит, что нужно просто всегда идти вправо — это оптимально с точки зрения нашей оценки награды. Рано или поздно, агент, следующий этой функции награды, свалится с горы.

Описанная проблема возникает из‑за того, что оптимизация метрики уводит нас из того домена, где эта метрика хорошо работает. Этот эффект называется законом Гудхарта: если у вас есть хорошая метрика, она перестанет быть хорошей, как только вы начнёте её оптимизировать. Это касается любых KPI в компаниях, макроэкономических показателей целых государств, и, конечно, нашей прокси‑награды в alignment.

При обучении награды для YandexGPT на одном из этапов мы наблюдали забавный эффект: награда предпочитала плохой ответ с красивым форматированием хорошему ответу без форматирования. Так происходило из‑за того, что в выборке для награды были примеры, где ответ с форматированием лучше, чем ответ без него, а обратных примеров не было совсем. Вот и случился классический гудхартинг.

Способов побороть эту беду по большей части два.

Первый способ — не давать модели уходить далеко от начальной инициализации (то есть от SFT‑модели). Если ваш агент не уходит далеко от подножия, он, скорее всего, не свалится с горы. Для этого мы добавляем к награде KL‑штраф за отклонение PPO‑модели от SFT‑модели:

Здесь — это коэффициент, определяющий силу штрафа.

Классически штраф за отклонение от SFT‑модели вводится через KL‑дивергенцию — меру удалённости двух вероятностных распределений друг от друга. По определению:

Честный KL‑штраф посчитать довольно сложно, потому что необходимо вычислить матожидание. На практике используют Monte‑Carlo‑оценку той или иной степени грубости. Самый распространённый вариант — сгенерировать один ответ из обучаемой модели и оценить KL по нему (если речь про PPO, то нужно использовать тот пример, который вы и так генерируете для алгоритма):

Поскольку трансформеры выдают вероятность не для всего ответа, а для каждого токена, итоговая формула будет выглядеть как‑то так:

В YandexGPT мы используем более точную оценку матожидания — вычисляем сумму потокенных KL‑штрафов:

То есть, для каждого токена считаем не просто логарифм отношения вероятностей, а честное матожидание на уровне токенов — усредняем по всем токенам из словаря с соответствующими вероятностями.

Второй способ — постоянно дообучать модель награды на ответах, к которым ведёт RL. Обучаете PPO, генерируете этой моделью по ответу на каждый запрос какой‑нибудь выборки и отдаёте AI‑тренерам сравнить, у кого ответ лучше: у SFT или у PPO. Полученную разметку добавляете в выборку для обучения награды. Переобучаете награду.

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

Существуют и более экзотические способы решить проблему, ещё не успевшие обрести большую популярность. Например, использование эпистемической неопределённости в оценке награды — неопределённости, связанной с недостаточным объёмом обучающих данных. Один из способов оценить эту неопределённость — обучить ансамбль моделей наград на немного разных данных и на немного разной начальной инициализации весов. Если члены ансамбля соглашаются в оценке награды, это значит, что они уверены в ней. Если расходятся, то такой оценке нужно доверять с опаской. Можно, например, variance прогноза ансамбля вычитать из награды агента.

А можно чуть дешевле? Direct Preference Optimization

PPO — это дорогой алгоритм: его сложно реализовать, он нестабилен в обучении, долго учится, требует подбора гиперпараметров и щепотку удачи. Но если все звёзды сошлись, он достигает высоких наград. А можно ли то же самое, но дешевле? Да, можно! Может, конечно, и не совсем то же самое, но очень близко по качеству: где‑то лучше, где‑то хуже.

Direct Preference Optimization (DPO) — сравнительно новый алгоритм, решающий задачу максимизации награды в RLHF способом, который не требует ни генерировать очень много текстов, ни обучать модель награды. Высокоуровнево идея проста: давайте соберем датасет с парами ответов и сравнениями, кто лучше, как мы делали для награды. А потом сразу как‑то обучим генератор на какой‑нибудь контрастный лосс, чтобы генератор порождал хорошие ответы чаще, а плохие — реже.

DPO от многих других контрастивных методов выгодно отличается строгим теоретическим обоснованием, и мы попробуем с ним познакомиться поближе. Существует забавный, довольно давно известный факт про то, как связана оптимальная политика и награда в случае, если используется KL‑штраф из прошлой главы. Мы буквально можем записать аналитическую формулу для оптимальной политики через награду:

Загвоздка тут в том, что формулу эту практически невозможно применить на практике в таком виде: здесь — это нормировочная константа вероятностного распределения. Чтобы её вычислить, нужно просуммировать награду по всем возможным ответам, которых бесконечно много. Это невозможно. Ну и ладно, мы от этой избавимся чуть позже.

Откуда вообще взялась формула, связывающая оптимальную политику и награду? Её несложно получить из задачи максимизации награды с KL-штрафом. Запишем эту задачу для произвольного пользовательского запроса :

Раскроем KL‑дивергенцию по определению:

Подставим в наш функционал и объединим математические ожидания:

Идея наших манипуляций в том, чтобы увидеть в этой задаче максимизации задачу минимизации уже другой KL‑дивергенции. Чтобы переход был прозрачен, давайте поделим выражение на . Поскольку мы делим на отрицательное число, задачу максимизации следует заменить на минимизацию:

Последний шаг, необходимый, чтобы увидеть в этом выражении KL‑дивергенцию: применим тождественное преобразование:

Подставим в оптимизируемый функционал и сразу же воспользуемся свойством логарифма: разность логарифмов — это логарифм отношения:

Полученное выражение — это ни что иное, как KL‑дивергенция между обучаемой политикой и неотнормированным распределением :

Известно, что минимум KL‑дивергенции достигается в случае, когда распределения одинаковы. С учётом нормировочной константы, получаем:

Ч.Т.Д.

Нам в этой формуле интересна её перевёрнутая запись: выразим награду через политику.

Никакой магии не произошло — просто выразили награду. Что интересно: а что будет, если в эту формулу подставить не оптимальную политику , а произвольную . Политика, неоптимальная для вашей задачи, может быть оптимальной для другой. По этой формуле мы могли бы вычислить функцию награды, в которой оптимальна. Давайте это даже явно запишем:

Изменение в том, что теперь политика не оптимальная, а произвольная, а награда записана с индексом , так как она теперь явно параметризуется через параметры политики. Это та награда, в которой оптимальна.

У такой записи есть ещё одно интересное свойство: вообще говоря, оптимальная политика не должна меняться от добавления к награде константы, не зависящей от ответа. То есть если из формулы просто выкинуть , то полученной функции награды будет соответствовать всё та же оптимальная политика :

Сейчас будет сложный финт ушами: давайте подберём параметры так, чтобы награда была максимально правдоподобной в смысле модели Брэдли‑Терри (см. часть про модель награды):

При этом никакой отдельной нейронной сети для награды у нас не будет — мы ведь умеем выражать награду через политику. Нейронная сеть будет только для политики. Чтобы избежать упоминания награды в оптимизируемом функционале, подставим в него запись награды через политику:

Этот функционал можно оптимизировать по параметрам нейронной сети вашим любимым способом.

Итого, что мы сделали:

• выразили функцию награды через политику,

• подставили это выражение в функцию ошибки награды.

Теперь, решая задачу обучения модели награды, мы одновременно решаем и задачу обучения генератора. Получается так, что — это как раз та модель, которая максимизирует . И обходимся мы при этом только одной нейронной сетью — политикой . В случае если мы по какой‑то причине хотим вычислить награду, у нас для этого есть аналитическая формула.

DPO требует от вас сбора такой же обучающей выборки, как для модели награды. Обучается политика обычным обучением с учителем, но на довольно специфический лосс. В отличии от PPO, в DPO нет необходимости генерировать данные во время обучения. Как и нет необходимости учить отдельную модель награды. По сложности вычислений DPO всё же дороже ванильного обучения с учителем, потому что в функции ошибки участвуют вероятности не только обучаемой модели, но и SFT — её тоже нужно держать в памяти. В ходе работы над YandexGPT мы заметили, что DPO слабо чувствителен к подбору гиперпараметров (в частности, ). Это позволяет делать alignment малой кровью, и больше времени тратить на эксперименты с данными, чем на поиск гиперпараметров, с которыми обучение работает.

Если вы обучаете DPO и у вас всё же есть уже обученная модель награды, её тоже можно использовать. Нередко DPO запускают поверх синтетической выборки сравнений: генерируем по два ответа на запрос и выбираем лучший с помощью модели награды. На полученной выборке сравнений обучаем DPO. Конечно, если нормальные данные имеются, лучше их тоже добавить или попробовать обучить модель в две стадии. Казалось бы, масло масляное, но подход показал себя полезным на практике.

Так какой же метод лучше?

На текущий момент научное сообщество не пришло к однозначному решению, какой метод лучше. Обычно сравнивают DPO и PPO, говоря, что CEM достигает худших результатов. Однако серьёзные исследования, где бы обучение с помощью CEM проводилось бы в несколько этапов с контролем энтропии, авторам этой статьи неизвестны.

Даже в отношении DPO и PPO никакой ясности нет. Существует витающее в воздухе ощущение, что PPO сложнее заставить работать, но можно получить более хорошую модель, чем от DPO. Но однозначно этот тезис пока что подтвердить не представляется возможным.

Что касается нашей команды, мы ведём исследования по всем направлениям.