Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps (
сайт проекта)
Последнее время активно исследуется тема масштабирования вычислений во время инференса (применения модели). В LLM ярким событием стал анонс o1 от OpenAI, где модель могла исписать 50 страниц рассуждений вместо 5, что привело к улучшению качества в нешироком наборе задач. Авторы из DeepMind решили попробовать масштабировать вычисления на инференсе для диффузионных моделей генерации картинок по текстовому запросу.
Диффузионные модели, если упрощать, работают так:
1) Сначала создается случайный шум — просто хаотичная картинка, никак не связанная с запросом (может выглядеть
так, ткните картинку чтобы понять о чем речь) и моделью
2) Этот шум пропускают через модель, получают на выходе параметры распределения (например, среднее и дисперсию) другого шума, который нужно вычесть из текущего, чтобы получить более качественное изображение, соответствующее вашему запросу.
3) берут семпл из этого распределения (то есть случайным образом выбирают какое-то одно значение с учётом предсказанных выше параметров; более просто: случайный выбор значения из предсказанного моделью диапазона)
4) из текущего зашумленного изображения (на первой итерации это то, что получилось в пункте 1) вычитают то, что получилось в пункте (3); модель как бы предсказала, какой шум нужно вычесть, чтобы «очистить» изображение (поэтому называется denoising, убирание шума). Это делает картинку чуть более четкой, но пока она всё ещё далека от финального результата. На этом этапе могут применяться специальные алгоритмы, связанные с дифференциальными уравнениями, но об этом как нибудь в другой раз.
5) Обновленное изображение снова пропускают через модель, повторяя процесс. Постепенно шум убирается шаг за шагом, и через множество итераций модель выдает готовую картинку.
Прочитав это, легко сходу придумать, как именно масштабировать вычисления во время предсказания: нужно просто увеличить количество шагов! К сожалению, эта мера не так эффективна, и после относительно небольшого увеличения прирост качества генераций прекращается. Условно между 20 и 50 шагами (итерациями) вы увидите разницу, а между 100 и 200 почти наверняка нет (в некоторых случаях это и вовсе портит картинку). То есть этот метод масштабирования очень ограничен.
Поэтому авторы рассматривают альтернативные способы. Они подмечают, что существует такая вещь как черри-пикинг — это когда для одного и того же запроса одна и та же модель генерирует много картинок, а после этого для демонстрации выбирается лучшая, например, чтобы похвастаться в статье или на сайте. То есть в среднем генерации могут быть просто хорошими, но вот иногда появляется картинка красивее и качественнее — хотя казалось бы ничего не меняется (кроме случайного шума в первом пункте из списка выше).
Значит, какие-то исходные шумы более удачны, какие-то менее. Это и будет первый метод поиска для масштабирования вычислений: давайте сгенерируем N картинок из N разных шумов, затем пропустим их через отдельную модель, которая даёт оценки, и выберем лучшую. «Отдельная модель» будет называться verifier (верификатор?), она принимает на вход картинку и, опционально, текстовый запрос и выдаёт какую-то цифру, по которой и судим.
Верификаторы могут быть разные — это может быть и одна модель, натренированная оценивать эстетику изображения (такие давно есть) и не опирающаяся на текст запроса, и LLM, которой дали промпт «ну чё ты, оцени по десятибальной», и ансамбль моделей, где несколько разных независимых оценок суммируются в одну. В статье верификаторам уделяется много внимания, но я про них писать не буду — важно то, что они есть, и это существенно отличает подобный метод от, например, о1, где модель генерирует сама без опоры на внешнюю валидацию.