Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Авторы из Harbin Institute of Technology представили LeanGRPO — способ убрать повторные вычисления из обучения диффузионных моделей с подкреплением; работа опубликована как препринт, не проходивший рецензирования. По их замерам, перестройка выполнения ускоряет полный шаг обучения до 1,83 раза без изменения исходной цели оптимизации. Это важно командам, которые дообучают генераторы изображений и видео через FlowGRPO, DanceGRPO или похожие методы и уже упираются во время GPU.
Что сделали
В рассматриваемых методах процесс удаления шума считается траекторией политики. Модель последовательно генерирует скрытые состояния, получает итоговое изображение или видео, после чего система вычисляет награду и преимущество — оценку того, насколько конкретный результат лучше или хуже других результатов для того же запроса.
Проблема возникает между генерацией траектории и обновлением параметров. Во время генерации модель уже выполняет прямой проход на каждом шаге, но обычно делает это без отслеживания градиентов. Когда награда становится известна, выбранные шаги приходится выполнить повторно, чтобы построить граф вычислений для обратного распространения ошибки.
При обучении на текущей политике, когда параметры между этими стадиями не меняются и используется одна вычислительная реализация, оба прохода дают одинаковые предсказания. Повтор нужен не для нового результата, а для восстановления ранее отброшенного графа.
LeanGRPO переносит построение градиентов в стадию генерации и предлагает два расписания выполнения. LeanGRPO-Retain сохраняет графы и промежуточные активации выбранных шагов до получения награды, а затем сразу запускает обратное распространение. Это простая схема, но её расход памяти растёт вместе с числом выбранных шагов.
LeanGRPO-Reweight выполняет предварительное обратное распространение сразу после каждого выбранного шага. Полученный градиент временно считается с единичным преимуществом, а активации освобождаются. После расчёта настоящего преимущества система масштабирует накопленный градиент и только затем синхронизирует его между GPU. По выводу авторов, в точной арифметике это даёт тот же градиент, что и исходная схема.
Чтобы не хранить на одном GPU состояния нескольких результатов для одного запроса, авторы также меняют распределение работы: все процессы получают один запрос, но генерируют разные варианты. Награды собираются после завершения группы, и обновление выполняется до перехода к следующему запросу.
Что показали
Авторы проверили LeanGRPO с несколькими вариантами FlowGRPO и DanceGRPO на генераторах изображений и видео. Выигрыш увеличивался, когда для обучения выбирали большую долю шагов удаления шума: именно в этом случае повторный проход занимал большую часть общего времени.
В разборе полного дообучения LeanGRPO-Reweight ускорил шаг в 1,81 раза, тогда как Retain — в 1,29 раза. Разница связана не только с устранением повторного прохода: Reweight объединяет несколько операций распределённой синхронизации градиентов в одну.
При дообучении через LoRA результаты вариантов сблизились: соответственно 1,45 и 1,42 раза. Авторы объясняют это меньшим объёмом обучаемых градиентов, из-за которого объединение синхронизаций почти не влияет на время. В таком режиме основную экономию даёт именно отказ от повторного вычисления.
Кривые награды для Native, Retain и Reweight сохраняли сопоставимую динамику. Однако Reweight не был численно идентичен исходному выполнению: порядок операций и распределённое суммирование в конечной точности вносили небольшое расхождение, хотя направление градиента оставалось близким.
Ограничения
Эквивалентность показана для обучения на текущей политике с одним обновлением, неизменными параметрами между генерацией и обратным проходом и общей вычислительной реализацией. Работа не доказывает применимость той же схемы к повторному использованию траекторий, нескольким обновлениям по одной выборке, асинхронному обучению или разделению генерации и обновления по разным процессам.
Основные замеры выполнены на одном узле с восемью NVIDIA RTX A6000; масштабирование проверялось в диапазоне от четырёх до 32 GPU. Эксперименты охватывают несколько открытых алгоритмов и генераторов изображений и видео, но в предоставленном тексте не указано, из какого набора взяты обучающие запросы. Качество отслеживали через HPSv2 и PickScore, поэтому работа не показывает сохранение качества на прикладных наградах конкретного продукта.
Сравнение проведено прежде всего с исходным расписанием тех же алгоритмов. Авторы не сопоставляют LeanGRPO по стоимости и времени внедрения с системами, где генерация вынесена в отдельный специализированный сервис. Также нет независимых воспроизведений результатов.
Что это значит
Для команды, уже строящей обучение на FlowGRPO, DanceGRPO или близкой схеме, работа может изменить план оптимизации инфраструктуры. Перед заменой GPU или сокращением числа обучающих шагов стоит измерить, какую долю времени занимает повторный прямой проход: LeanGRPO устраняет конкретный системный расход, не меняя награду, выбор шагов и целевую функцию.
Выбор варианта зависит от профиля памяти. Retain подходит, если выбранных шагов немного и сохранённые активации помещаются рядом с состояниями оптимизатора. Reweight лучше ограничивает рост памяти на длинных видеотраекториях, но требует хранить полный предварительный градиент для каждого локального результата и откладывать распределённую синхронизацию.
Это не готовая универсальная замена контура обучения. Потребуются совместимые процессы генерации и обновления, доступ к графу автоматического дифференцирования и перераспределение запросов между GPU. Если существующая архитектура разделяет эти стадии или повторно обучается на сохранённых траекториях, выводы работы напрямую планы не меняют. Если стадии уже совмещены, LeanGRPO задаёт проверяемое направление: сначала убрать повторное вычисление, затем решать, выгоднее ли хранить активации или предварительные градиенты.
Источники
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
