Усовершенствованные методы обучения моделей согласованности20.06.2024 00:00

Модели согласованnosti (consistency models) — это новое семейство генеративных моделей, способных создавать высококачественные данные за один шаг без необходимости состязательного обучения. Текущие модели согласованности достигают оптимального качества выборки путем дистилляции из предварительно обученных диффузионных моделей и использования заученных метрик, таких как LPIPS. Однако дистилляция ограничивает качество моделей согласованности уровнем предварительно обученной диффузионной модели, а LPIPS вносит нежелательное смещение при оценке. Для решения этих задач мы представляем усовершенствованные методы обучения согласованности, при которых модели согласованности учатся напрямую на данных без дистилляции. Мы углубляемся в теорию, лежащую в основе обучения согласованности, и выявляем ранее упущений недостаток, который мы устраняем путем исключения экспоненциального скользящего среднего (EMA) из учительской модели согласованности. Чтобы заменить такие заученные метрики, как LPIPS, мы используем функции потерь Псевдо-Хьюбера из робастной статистики. Кроме того, мы внедряем логнормальное расписание шума для целевой функции обучения согласованности и предлагаем удваивать общее число шагов дискретизации каждые заданные итерации обучения. В сочетании с более качественной настройкой гиперпараметров эти модификации позволяют моделям согласованности достигать показателей FID 2.51 и 3.25 на наборах данных CIFAR-10 и ImageNet 64×64 соответственно всего за один шаг семплирования. Эти результаты демонстрируют улучшение в 3.5 и 4 раза по сравнению с предыдущими подходами к обучению согласованности. Используя двухшаговое семплирование, мы дополнительно снижаем показатели FID до 2.24 и 2.77 для этих двух наборов данных, превосходя результаты, полученные с помощью дистилляции как в одношаговом, так и в двухшаговом режимах, и при этом сужая разрыв между моделями согласованности и другими передовыми генеративными моделями.