Где диффузионной модели тратить вычисления при генерации

Дополнительные вычисления при генерации диффузионной моделью полезнее тратить на отдельные важные шаги, а не поровну на всю траекторию. В работе Peking University и UC Berkeley алгоритм GAINS достиг качества равномерного поиска при сокращении числа проходов через модель на 20–50%; поскольку препринт не рецензирован, эти числа получили сами авторы. Команды могут снизить стоимость уже обученной модели, если отделят поиск кандидатов от распределения бюджета.

Почему один дополнительный кандидат полезнее другого

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

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

Авторы описывают полезность шага через чувствительность. Она объединяет масштаб шума на этом участке траектории и то, насколько критерий качества реагирует на изменение результата. Профиль чувствительности зависит от модели и сэмплера: у Stable Diffusion полезные участки распределились ближе к началу траектории, а у EDM сконцентрировались в её средней части.

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

Как water-filling превращается в расписание проходов

Убывающая отдача позволяет распределять бюджет методом water-filling, то есть выравнивать предельную пользу дополнительных проходов. Планировщик сравнивает, где следующий кандидат даст наибольшую ожидаемую прибавку, назначает его этому шагу и повторяет процедуру, пока бюджет не закончится.

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

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

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

GAINS поэтому сочетает два уровня. Сначала небольшой калибровочный запуск даёт усреднённый профиль чувствительности и базовое расписание. Затем контроллер корректирует его для конкретного начального шума или условия генерации, но не строит распределение заново.

Когда результат меняет план внедрения

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

На Stable Diffusion при одинаковом бюджете GAINS поднял балл яркости с 0,7025 до 0,7248, а балл сжимаемости — с 0,8833 до 0,8946. Отдельные опыты также показали улучшение при случайном поиске кандидатов и при локальном поиске возмущениями: глобальное расписание не привязано к одному способу выбирать шум.

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

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

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

Источники

• Leveraging Inference-Time Compute for Diffusion Models via Global Scheduling of Denoising Trajectories