RL for Diffusion modelsPart 1

Глава 1

Зачем reward, если есть DPO

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

Один из её методов — DPO. Собираем пары winner и loser и учим модель приближаться к первому и отдаляться от второго. Данные собирать легко: подойдут картинки любых моделей.

Но у DPO две слабые стороны. Модель легко уходит из распределения претренинга, и её приходится сдерживать — SFT на winner-картинках или методами вроде ThermoDPO (подробнее об этом — в предыдущем посте про preference optimization). А главное — пара сообщает только «эта картинка лучше той». Что именно исправить, пара не говорит.

Например, просим нарисовать «четыре чашки», а модель рисует три. Можно собрать пару «четыре чашки лучше трёх», но из неё модель не поймёт, что ошиблась именно в счёте. Да и такие пары сначала нужно где-то найти.

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

Если проблему можно оценить числом, её можно оптимизировать. Этим и занимается RL. Допустим, хорошая reward-модель (оценщик: смотрит на готовую картинку и возвращает это число) у нас уже есть.

Как обновить генератор, чтобы он чаще получал высокий reward? Об этом пост, первый из двух про RL для диффузии. Как строить и балансировать сами reward-модели — тема второго.

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

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

Глава 2

Один шаг в конце: ReFL

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

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

Но протащить градиент через десятки шагов дорого. Первым градиент от reward в генератор довёл ReFL — в той же работе, где появилась reward-модель ImageReward. Решение лобовое: раз через всю траекторию не пройти, возьмём один шаг.

Выбираем случайный поздний шаг , доходим генерацией до него без градиента и из  за один шаг предсказываем чистое изображение , считаем на нём reward и обновляем модель по его градиенту [ReFL, разд. 3]:

Какой шаг выбрать

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

Остаётся поздний, но не самый последний участок, где картинка уже читается, но ещё не готова [ReFL, рис. 4]. Например, при 50 шагах генерации — шаги примерно с 30-го по 40-й.

Отсюда и главное ограничение. Какой бы дифференцируемый reward мы ни взяли, к этому моменту композиция уже выбрана: если модель нарисовала три чашки, четвёртая не появится. ReFL может только сделать эти три аккуратнее.

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

А если сдвигать окно раньше?

Напрашивается идея. По мере обучения модель рисует всё более чёткие картинки уже на ранних шагах. Может, постепенно сдвигать окно обучения к началу траектории — и так добраться до композиции?

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

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

Глава 3

До конца, но только последние шаги: DRaFT

Через десятки шагов сразу — дорого. Но между «один шаг» и «все шаги» есть середина: пройти генерацию до конца, а градиент пустить только через последние шагов сэмплирования. Так делает DRaFT, вариант DRaFT-K. Reward считается на настоящем : мы проходим реальный путь до картинки, а не угадываем её за один шаг.

Но четвёртой чашки по-прежнему не будет. Мы всё так же обновляем только конец траектории, а сколько чашек в кадре, решилось задолго до него. Модель чуть лучше рисует текстуру или пальцы, и только.11 DRaFT-LV: усреднить reward. Простая модификация DRaFT, которую полезно знать для любого RL. Получив без градиента, зашумим его раз с разным шумом и снова расшумим за один шаг. Reward усредняем по всем предсказаниям. У такой оценки меньше дисперсия, и обновление получается чище: DRaFT-LV учится примерно вдвое быстрее ReFL [DRaFT, рис. 4].

Почему нельзя просто взять больше шагов

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

Если эти матрицы не близки к единичной, после десятка умножений градиенты сильно растут, и обновления перестают работать [DRaFT, рис. 8]. Добавьте к этому память на хранение всех промежуточных активаций. В самой работе DRaFT маленькие , вплоть до , оказываются лучшим компромиссом [DRaFT, рис. 9].

Ранние шаги, где решается число чашек, отделены от reward десятком таких умножений. И каждое из них добавляет к взрыву.

Глава 4

До ранних шагов без взрыва: DRTune, ShortFT

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

DRTune: стоп-градиент на входе сети

DRTune останавливает градиент на входе денойзера [DRTune, разд. 3.2]. Тогда градиент не проходит через сеть на каждом шаге и не перемножает якобианы. Он идёт по траектории только через линейную часть шага сэмплирования, и нормы перестают расти [DRTune, рис. 7].

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

Так мы обновляем модель на разумной, не слишком шумной картинке и дотягиваемся до ранних шагов, а значит — до композиции. Авторы показывают, что именно обучение ранних шагов важно для низкоуровневых reward, где методы с поверхностной супервизией не справляются [DRTune, табл. 1 и рис. 4].

ShortFT: короткая траектория через shortcut-модель

ShortFT решает задачу с другой стороны. Вместо траектории из пятидесяти шагов возьмём few-step модель на 4–10 шагов, которая сохраняет траекторию исходной модели. По такой короткой цепочке можно протащить градиент целиком, от reward до шума — и до шага, где выбирается число чашек.

Как перенести результат на полную модель? Shortcut можно сделать LoRA-дистилляцией, например Hyper-SD. Поверх неё совместно учим ещё одну LoRA уже на reward. После обучения дистилляционную LoRA убираем, а обученную оставляем — и она переносится на полную модель.

Насколько легко это работает, зависит от метода дистилляции. Разрыв между обучением и инференсом авторы отдельно закрывают прогрессивной схемой [ShortFT, разд. 3.3 и рис. 6].

Почему не дистиллировать после RL?

Логичный вопрос: может, сразу взять дистиллированную модель и сделать RL на ней? Проблема в разнообразии. Дистиллированная модель на четырёх шагах почти не даёт diversity. Без него не работает exploration: модели нечего сравнивать и не из чего выбирать.

Поэтому дистилляцию обычно делают после RL. К тому же модель после RL уже частично схлопнута, и её проще дистиллировать.

Что изменил flow matching

У flow matching траектории прямее. В диффузии на ранних шагах предсказанное изображение может быть очень далеко от итогового. У flow-моделей уже на раннем участке траектории видна почти готовая композиция: мелких деталей ещё нет, но сцена уже понятна. На таких предсказаниях можно считать reward.

Поэтому для flow-моделей полные rollouts (прогоны генерации от шума до готовой картинки) стали реальными и без shortcut-LoRA. Траектории изначально хорошие и прямые, можно генерировать на десяти шагах. Картинки получатся немного размытыми, но модель со временем учится делать их чётче сама.

Где этот подход остаётся полезным

Есть и практический вывод. Если задача — конкретные детали, полный RL может быть не нужен: попиксельные градиенты дифференцируемого reward очень плотные и хорошо работают на мелочах. Для композиции нужно что-то умнее.

Получить надёжный дифференцируемый reward для сложных глобальных свойств трудно; сейчас это получается в основном у больших VLM. Можно комбинировать reward: с полным rollout часть оценщиков считается только на готовой картинке, в конце траектории.

Но наш счётчик чашек — не сеть. Он возвращает «три» или «четыре», и градиента у него нет.

Глава 5

Если от reward нельзя взять градиент: Flow-GRPO

Градиент от reward и не нужен, если есть другое: число на готовом сэмпле и вероятности шагов, которыми мы к нему пришли. Тогда можно делать траектории с высоким reward вероятнее, а с низким — реже. Это policy gradient, и у него три причины вытеснить direct backpropagation.

Первая — многие оракулы бинарные: четыре чашки или нет, текстура хорошая или нет. Такие вещи не описываются дифференцируемо. Вторая — даже обученной reward-модели трудно доверять её градиенты. Модель легко находит направления, в которых reward растёт, а картинка не становится лучше, — reward hacking.

Третья — хочется online-контроля: генерировать rollouts текущей моделью и исправлять именно её проблемы.

Поэтому формулируем денойзинг как марковский процесс принятия решений. Состояние — текущий шумный сэмпл, действие — следующий шаг, policy (правило выбора действия; здесь это сама модель), reward выдаётся в конце:

Градиент ожидаемого reward записывается по трюку REINFORCE:

насколько сэмпл лучше среднего по запросу ( — baseline)сделать каждый шаг этой траектории вероятнее

Ни одной производной от  здесь нет: reward входит как число.

Flow-GRPO

Flow-GRPO — сейчас один из основных методов RL для flow-моделей. Делаем полный rollout на небольшом числе шагов, считаем reward и через policy gradient обновляем модель целиком. Поверх этого — GRPO: advantage (насколько сэмпл лучше среднего по группе) считается внутри группы rollouts на один запрос.

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

Используем известный приём: у каждого ODE есть эквивалентное SDE с теми же маргинальными распределениями [Flow-GRPO, разд. 4]. Тогда из одного шума можно получать разные траектории, а шаги становятся гауссовыми — у них есть посчитанный .

Уровень шума задаётся так:

Параметр регулирует разнообразие траекторий. Если он слишком маленький, примерно меньше 0,5, rollouts почти не отличаются. Если слишком большой, сэмплы начинают портиться [Flow-GRPO, рис. 7b]. На практике я беру порядка 0,7.

Вторая проблема — стоимость rollouts. Здесь помогает то же наблюдение, что и в ShortFT: траектории flow matching прямые. В Flow-GRPO rollouts для обучения делают на 10 шагах, а на инференсе остаются исходные 40 [Flow-GRPO, разд. 4 и рис. 7a]. Грубых rollouts хватает, чтобы ранжировать сэмплы внутри группы.

Дальше — обычный GRPO. Для запроса «четыре чашки» генерируем группу rollouts: детектор ставит 1 там, где чашек четыре, и 0 там, где три. Нормируем reward внутри группы и получаем advantage : у четырёх чашек он положительный, у трёх — отрицательный. Каждый шаг траектории получает тот же advantage, включая ранние, где число чашек и решалось.

Механически всё выглядит так. Делаем полный rollout и сохраняем шаги. Считаем reward на итоговой картинке. Затем идём по траектории, на каждом шаге считаем , накапливаем градиенты и делаем один шаг оптимизатора. Градиент через цепочку шагов не протаскивается, поэтому проблемы со взрывом нет.

Результат заметный. Для SD3.5-M Flow-GRPO поднимает точность на GenEval с 63% до 95%, а на рендеринге текста — с 59% до 92% [Flow-GRPO, табл. 2]. Авторы также добавляют KL к reference-модели (замороженной копии исходной модели): он сдерживает reward hacking [Flow-GRPO, рис. 6].

Где Flow-GRPO неудобен

У метода две проблемы. Первая: мы оптимизируем обратный процесс, а не прямой, хотя модель предобучалась на прямом. Из-за этого сильные ограничения на сэмплер: работает только эйлеровская дискретизация SDE, иначе не сходится математика [Flow-GRPO, разд. 4].

Вторая: rollouts делаются с CFG. Значит, в обучении участвуют и conditional-, и unconditional-ветки. Вероятность шага считается по скорости после CFG, и нужно хранить и обновлять обе. Это в два раза больше памяти и в два раза больше вычислений.

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

Глава 6

Без вероятностей и сэмплера: DiffusionNFT, AWM

Можно, если учиться на прямом процессе, а не на обратном. Хранить только чистые сэмплы и их reward, зашумлять их и учить модель тем же loss, что на претренинге, — только с весом от reward. Так делают DiffusionNFT и AWM, и в итоге diffusion RL оказывается взвешенным flow matching.

DiffusionNFT: учиться на прямом процессе

DiffusionNFT отказывается от likelihood обратного процесса. Rollouts генерируются любым сэмплером, хоть ODE высокого порядка. Сохраняем только чистые сэмплы и их reward. Затем зашумляем эти сэмплы и учим модель обычным flow matching loss (функция ошибки претренинга: насколько предсказанная скорость отличается от настоящей).

Как отделить хорошие сэмплы от плохих? Сначала reward в группе нормируют и переводят в вероятность оптимальности . Можно было бы просто дообучаться на хороших примерах. Но rejection fine-tuning без негативов быстро схлопывает модель [DiffusionNFT, разд. 4.4]. От трёх чашек нужно отталкиваться так же, как тянуться к четырём.

Для этого обучаемую скорость представляют как скорость старой модели плюс добавку . Добавка с плюсом задаёт неявную позитивную policy, с минусом — негативную:

Позитивную ветку учим на сэмплах с весом , негативную — с весом :

Посмотрим, куда двигается добавка в начале обучения:

Для хорошего сэмпла , и модель тянется к его flow matching target (тому, к чему loss тянет предсказание). Для плохого — отталкивается. Old-модель остаётся опорой: обучение отталкивается от неё, а не от пустого места. Поэтому оно стабильнее, чем DPO по тем же сэмплам и чем дообучение только на позитивных.

Old-модель, которая генерирует rollouts, обновляют мягко, через EMA [DiffusionNFT, разд. 3.3]. Если сразу копировать в неё текущие веса, reward поначалу растёт быстро, но потом обучение разваливается [DiffusionNFT, рис. 8].

Плюсы видны сразу. Нет ограничений на сэмплер. Нет проблем с CFG: всё обучение идёт без него, учится только conditional-ветка. По сути, дистилляция CFG происходит по ходу.

Разнообразие rollouts при этом берётся из самой модели, а не из искусственного шума SDE. Шаги для loss — те же, что при сэмплировании. DiffusionNFT до 25 раз эффективнее Flow-GRPO и поднимает GenEval с 0,24 до 0,98 за тысячу шагов без CFG [DiffusionNFT, рис. 1 и 6].

Advantage Weighted Matching: просто взвешенный flow matching

Следующий шаг делает связь явной. Advantage Weighted Matching показывает, что DDPO-подобные методы, включая Flow-GRPO, по сути делают flow matching с шумными targets [AWM, разд. 2]. Отсюда большая дисперсия и медленная сходимость.

Рецепт AWM: в policy gradient заменить недоступный на его диффузионную ELBO-оценку. Эта оценка — и есть loss претренинга. Получается обычный flow matching, где каждый сэмпл взвешен своим advantage.

При модель тянется к target этого сэмпла, при — отталкивается. Как и в Flow-GRPO, к loss добавляют KL к reference-модели. Сэмплер может быть любым: ODE, SDE, few-step. Вес — тот же приём, которым перевзвешивают шаги в обычном диффузионном loss.

На SD3.5-M и FLUX это до 34 раз быстрее Flow-GRPO при том же качестве генерации [AWM, табл. 2 и рис. 5].

Если это просто SFT, зачем такие сложности?

Если сэмплы с высоким reward — это просто SFT на позитивных, нельзя ли сделать это напрямую? Можно: AWM это и показывает. Достаточно выписать ELBO-оценку policy gradient, и останется такая схема: отталкиваемся от примеров, которые не нравятся, и двигаемся к тем, что нравятся, ориентируясь только на reward.

Конструкция DiffusionNFT с позитивной и негативной ветками — скорее интуитивно найденный способ это стабилизировать. При этом сохраняется главный плюс online policy gradient: reward может быть недифференцируемым, он задаёт только направление движения.

А если генерировать сэмплы заранее?

Насколько часто нужно обновлять сэмплы с разметкой reward? Можно ли, как в DPO, заготовить их заранее? Самый надёжный вариант — всегда online. Можно попробовать отставание на один шаг, но это сложно сделать стабильным: даже такая задержка может схлопнуть модель.

Поэтому, если есть возможность, генерируем rollouts текущей моделью — точнее, её медленной EMA-копией, как в DiffusionNFT. От коллапса на собственной синтетике защищает KL к reference-модели: модель не уходит от неё слишком далеко. Эта регуляризация здесь работает намного лучше, чем в DPO.

Доставка reward закончилась весом на loss претренинга. Значит, и следить нужно за тем же, за чем на претренинге: за данными и за тем, чему учит reward.