Published on

Масштабирование вывода диффузионных моделей Новая парадигма

Авторы
  • avatar
    Имя
    Ajax
    Twitter

Введение

Недавние достижения в области больших языковых моделей (LLMs) продемонстрировали эффективность масштабирования во время вывода. Такие модели, как o1, o3, DeepSeek R1, QwQ и Step Reasoner mini, показали, что увеличение вычислительных ресурсов во время вывода может значительно улучшить производительность. Возникает вопрос: можно ли применить этот принцип и к диффузионным моделям?

Команда под руководством Се Санинга в Нью-Йоркском университете исследовала именно этот вопрос. В ходе систематического исследования с использованием общей поисковой структуры они обнаружили, что масштабирование во время вывода действительно эффективно для диффузионных моделей. Увеличение вычислительных ресурсов во время вывода приводит к значительному улучшению качества генерируемых образцов. Более того, сложность изображений позволяет использовать различные комбинации компонентов в рамках структуры, адаптированные к различным сценариям применения.

Основные выводы

  • Масштабирование во время вывода эффективно для диффузионных моделей: выделение большего количества вычислительных ресурсов во время вывода приводит к более высокому качеству образцов.
  • Гибкость в комбинациях компонентов: структура позволяет использовать различные конфигурации компонентов, адаптированные к различным приложениям.
  • За пределами шагов шумоподавления: исследование показывает, что поиск лучшего шума во время выборки является еще одним измерением для масштабирования NFE, помимо простого увеличения шагов шумоподавления.
  • Две оси проектирования: структура фокусируется на двух ключевых осях проектирования:
    • Верификаторы: Обеспечение обратной связи во время процесса поиска.
    • Алгоритмы: Поиск лучших кандидатов на шум.

Методология исследования

Команда исследовала три различных сценария для верификаторов, моделируя различные варианты использования:

  1. Сценарии, в которых доступна привилегированная информация об итоговой оценке.
  2. Сценарии, в которых доступна условная информация для управления генерацией.
  3. Сценарии без дополнительной информации.

Для алгоритмов они исследовали:

  1. Случайный поиск: Выбор лучшего из фиксированного набора кандидатов.
  2. Поиск нулевого порядка: Итеративное улучшение кандидатов на шум с использованием обратной связи от верификатора.
  3. Поиск пути: Итеративное улучшение траекторий выборки диффузии с использованием обратной связи от верификатора.

Первоначально исследование изучало эти конструкции в относительно простой настройке условной генерации классов ImageNet. Впоследствии они применили эти конструкции к более масштабной генерации текста в изображение и оценили предложенную ими структуру.

Масштабирование времени вывода

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

  • Верификаторы: Это предварительно обученные модели, которые оценивают качество сгенерированных образцов. Они принимают сгенерированные образцы вместе с дополнительными условиями и выводят скалярную оценку.
  • Алгоритмы: Эти алгоритмы используют оценки верификатора для поиска лучших кандидатов на образцы. Функция принимает верификатор (V), предварительно обученную диффузионную модель (_θ) и набор сгенерированных образцов и условий, выводя лучший начальный шум.

Общий бюджет вывода измеряется общим количеством вычислений функций (NFE), включая как шаги шумоподавления, так и затраты на поиск.

Поисковые верификаторы

Исследователи начали с верификатора Oracle, который имеет полную информацию об итоговой оценке выбранных образцов. Для ImageNet это включало такие метрики, как FID и IS. Затем они изучили более доступные предварительно обученные модели в качестве контролируемых верификаторов, такие как CLIP и DINO. Эти модели использовались для классификации образцов, выбирая образец с наивысшим логитом, соответствующим метке класса.

Однако они заметили, что эти классификаторы, работающие точечно, лишь частично соответствуют целям оценки FID. Это привело к снижению дисперсии образцов и коллапсу моды по мере увеличения вычислений. Это явление, названное "взломом верификатора", было ускорено неограниченным пространством поиска алгоритма случайного поиска.

Интересно, что исследование показало, что верификаторам не обязательно нужна условная информация для эффективного управления поиском. Они наблюдали сильную корреляцию между логитами классификаторов DINO/CLIP и косинусным сходством в пространстве признаков между x-прогнозом на низком уровне шума и итоговым чистым образцом. Это привело к использованию самообучающихся верификаторов, которые не требуют дополнительной условной информации и по-прежнему демонстрируют эффективное масштабирование.

Поисковые алгоритмы

Чтобы смягчить взлом верификатора, исследователи изучили более точные алгоритмы поиска, которые постепенно оптимизируют образцы-кандидаты. Это включало метод поиска нулевого порядка:

  1. Начало со случайного гауссовского шума в качестве опорной точки.
  2. Нахождение N кандидатов в окрестности опорной точки.
  3. Пропуск кандидатов через решатель ODE для получения образцов и оценок верификатора.
  4. Обновление опорной точки лучшим кандидатом и повторение шагов 1-3.

Они также исследовали алгоритм поиска пути, который изучает возможность поиска вдоль траектории выборки:

  1. Выборка N начальных шумовых образцов и пропуск решателя ODE до уровня шума σ.
  2. Добавление шума к каждому образцу и моделирование процесса прямого зашумления.
  3. Запуск решателя ODE на каждом зашумленном образце и сохранение N лучших кандидатов на основе оценок верификатора, повторяя до тех пор, пока решатель ODE не достигнет σ=0.
  4. Случайный поиск оставшихся N образцов и сохранение лучшего.

Оба алгоритма поиска нулевого порядка и поиска пути поддерживают сильную локальность по сравнению со случайным поиском.

Масштабирование в сценариях преобразования текста в изображение

Команда изучила возможности масштабирования поисковой структуры в более крупных задачах преобразования текста в изображение. Для оценки они использовали наборы данных DrawBench и T2I-CompBench, а в качестве основы - модель FLUX.1-dev. Они также расширили выбор контролируемых верификаторов, включая Aesthetic Score Predictor, CLIPScore и ImageReward. Кроме того, они создали Verifier Ensemble, объединив эти три верификатора.

Анализ: соответствие верификатора задаче

В исследовании сравнивались результаты различных комбинаций верификаторов и алгоритмов на разных наборах данных. На DrawBench они обнаружили, что использование всех верификаторов в целом улучшает качество образцов. Однако они заметили, что использование верификаторов Aesthetic и CLIP по отдельности может привести к переобучению их смещениям, что приведет к негативным последствиям друг для друга. Это связано с несоответствием их фокуса оценки: Aesthetic Score фокусируется на визуальном качестве, часто отдавая предпочтение сильно стилизованным изображениям, в то время как CLIP отдает приоритет визуально-текстовому выравниванию, иногда жертвуя визуальным качеством.

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

Производительность алгоритма

Все три алгоритма поиска (случайный, нулевого порядка и пути) эффективно улучшили качество выборки на DrawBench. Однако случайный поиск в некоторых аспектах показал лучшие результаты из-за локального характера двух других методов. Случайный поиск быстрее сходился к смещению верификатора, в то время как двум другим алгоритмам требуется улучшение на менее чем оптимальных кандидатах.

Совместимость с тонкой настройкой

Команда исследовала совместимость своего метода поиска с моделями с тонкой настройкой. Они использовали модель Stable Diffusion XL с тонкой настройкой DPO и обнаружили, что метод поиска можно обобщить на разные модели и улучшить производительность уже выровненных моделей.

Влияние различных измерений вычислений вывода

В исследовании изучалось, как различные аспекты вычислений вывода влияют на результаты:

  • Количество итераций поиска: Увеличение количества итераций приближает шум к оптимуму.
  • Вычисления за итерацию поиска: Настройка количества шагов шумоподавления за итерацию выявляет различные вычислительно оптимальные области.
  • Вычисления окончательной генерации: Команда использовала оптимальные настройки для окончательных шагов шумоподавления, чтобы обеспечить наивысшее качество окончательного образца.

Эффективность инвестиций в вычисления

Исследователи изучили эффективность масштабирования времени вывода на меньших диффузионных моделях. Они обнаружили, что для ImageNet масштабирование меньших моделей может быть очень эффективным. В некоторых случаях поиск на меньшей модели может превзойти более крупные модели без поиска. Однако эффективность зависит от базовой производительности меньшей модели.

В текстовых настройках PixArt-Σ, используя лишь часть вычислений, превзошел FLUX-1.dev. Эти результаты показывают, что значительные вычислительные ресурсы, затраченные во время обучения, могут быть компенсированы меньшими объемами вычислений во время генерации, что приводит к более качественным образцам более эффективно.