Google рассказала, как направлять генератор картинок, не меняя его веса

Исследователи Google Research показали Diffusion Controller — небольшую сеть-надстройку. По ходу генерации она подправляет картинку, чтобы та точнее следовала запросу, а сам генератор остаётся прежним.

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

Повод: препринт с описанием метода вышел 7 марта, а сейчас авторы подробно разобрали его в блоге Google. В работе участвовали также исследователи Google DeepMind и университетов.

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

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

• Основная модель заморожена. Маленькая сеть на каждом шаге смотрит на промежуточный результат. Основная модель в этот момент выдаёт указание, как дальше убирать шум, и сеть прибавляет к нему небольшую поправку.
• Поправка тянет картинку к заданной цели, например к стилю или к точности запроса. Штраф при этом не даёт далеко уйти от обычного поведения модели.
• Сеть учится по оценке готовой картинки двумя способами: алгоритмом PPO¹, который не даёт самой сети резко меняться за один шаг обучения, или взвешиванием примеров по награде.
• Силу вмешательства задаёт один параметр. Его можно менять прямо во время генерации, без переобучения.

Общая картина: сейчас генерацию направляют двумя разрозненными наборами приёмов. Одни усиливают влияние запроса прямо во время генерации, другие дообучают модель, например методом LoRA². Авторы описывают оба подхода на одном математическом языке. Так баланс между точностью и качеством картинки можно рассчитать, а не подбирать наугад.

В цифрах: метод проверили на Stable Diffusion v1.4 в трёх режимах обучения: обычном дообучении на примерах, взвешивании по награде и PPO. Качество оценивали по HPSv2³, считая долю сравнений, в которых картинка вышла лучше, чем у исходной модели:

• При обычном дообучении и взвешивании по награде надстройка без доступа к весам выигрывала у исходной модели чаще, чем LoRA. При этом она вмешивалась в заметно меньшее число слоёв.
• Версия с полным доступом к весам, которую обучали вместе с основной моделью, выиграла у исходной 90% сравнений.
• Люди оценивали картинки по сложным запросам с несколькими условиями. По качеству и соответствию запросу надстройка получила у них лучшие оценки.

Но: все замеры сделаны на Stable Diffusion v1.4, открытой модели 2022 года. Результатов на закрытых моделях вроде Nano Banana авторы не приводят. Кроме того, надстройке нужен доступ к промежуточным результатам на каждом шаге генерации: авторы называют такой режим «серым ящиком».

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