Почему «momentum»-сэмплер для ComfyUI рисует синюю кляксу на flow-модели и какой momentum на самом деле помогает

В ComfyUI есть набор сэмплеров ComfyUI-Extra-Samplers с «momentumized»-вариантами DPM-Solver++. Я запустил dpmpp_dualsde_momentumized на Anima (flow-модель, bf16, CFG 8, 30 шагов) и получил вместо рисунка бледно-голубой силуэт. Тот же промпт на dpmpp_2m_sde_gpu даёт чистую картинку. Дальше я разбирал, что именно сломано, а потом проверял, существует ли «momentum», который действительно уменьшает ошибку при том же числе вызовов сети.

Коротко:

• поломка не в momentum: через обычный KSampler его коэффициент равен 0;

• сэмплер написан под variance-exploding модели, и на flow-модели его шаг добавления шума имеет неверную дисперсию: при eta = 1 латент разрастается (std 3-4,5 вместо 1,2), при eta = 0 тот же код рисует кота;

• «momentum» там экспоненциальное скользящее среднее предсказания данных. С положительным коэффициентом это фильтр нижних частот: он добавляет запаздывание и портит результат (до +61 % ошибки);

• с отрицательным коэффициентом та же рекурсия становится устойчивой экстраполяцией и снижает ошибку на 11-21 % при том же числе шагов (в детерминированном случае);

• это скромное ускорение, примерно в 1,2 раза меньше шагов при той же ошибке, а не вдвое.

Идея не новая

Оговорка сразу. Momentum в сэмплировании диффузии уже исследовали. Ближайшая работа: Wizadwongsa и др., ICLR 2024 (arXiv:2307.11118), где Heavy Ball momentum встраивают в существующие численные солверы, чтобы расширить область их устойчивости и убрать артефакты расходимости при малом числе шагов. Есть также сэмплер с momentum в стиле Adam (Wang и др., arXiv:2308.11941), momentum для «странных» изображений (Jang и др., AAAI 2024), TADA (arXiv:2506.21757) и пакеты ComfyUI sonar и Seren Sampler (последний как раз для семейства Anima). Я прочитал только аннотации и README, код этих пакетов не запускал и сравнений не делал. Какие знаки коэффициента перебирают авторы Heavy Ball, я не проверял, поэтому мой вывод «положительный коэффициент только добавляет запаздывание» не опровергает их работу: модель, солвер и метрика у нас разные.

Чего я не нашёл: разбора, почему этот конкретный сэмплер ломается на flow-модели, и измерения при равном числе вызовов, как знак и величина такого momentum меняют ошибку. Это утверждение о моём поиске, а не о всей области.

Что именно сломано

Шаг метода DPM-Solver++ в форме предсказания данных:

x_t = (σ_t/σ_s)·e^(−ηh)·x_s + α_t·(1 − e^(−(1+η)h))·D_eff + шум

Все методы ниже отличаются только D_eff. dpmpp_dualsde_momentumized делает полный шаг 2S (два вызова сети) и затем добавляет сверху многошаговую поправку 3M. Шаг 2S уже содержит второй порядок, так что поправка считается дважды: сэмплер расходится даже без шума, в 2-35 раз хуже dpmpp_2m при том же числе вызовов. Ещё в цикле переменная r перезаписывается внутри функции, и середина стадии 2S уезжает с 1/2.

Главное, однако, другое: формулы написаны для моделей с x = x0 + σ·ε, а у flow-модели x = (1 − σ)·x0 + σ·ε. Дисперсия добавляемого шума неверна, и ошибка копится за 30 шагов. Доказательство косвенное: тот же алгоритм 2S SDE в flow-корректной форме из ядра ComfyUI (dpmpp_sde) даёт чистую картинку. Механизм я отдельным экспериментом не изолировал.

Я перенёс сэмплер в виде точного порта и сверил его с оригиналом до 1e-10 на игрушечной модели. Так можно включать и выключать каждый дефект.

Каким должен быть momentum

Тяжёлый шар Поляка, применённый к приращению x напрямую, не согласован: сумма весов приращений равна h + μ·h_prev, а не h. Согласованный вариант нормирует веса, и получается экспоненциальное скользящее среднее приращений; при μ = −1/2 и равномерных шагах это ровно метод Адамса-Башфорта 2. В рамке экспоненциального интегратора это даёт простое правило: D̄ = m·D̄_prev + (1 − m)·D. Положительный m сглаживает и запаздывает. Отрицательный экстраполирует. Рекурсия устойчива только при |m| < 1: при m = −0,9 выигрыш исчезает, при m = −1,2 расходится.

Измерения

Anima, bf16, 576×768, CFG 5, 15 юнитов (5 промптов x 3 сида), 8-24 вызова сети, парные бутстреп-интервалы, эталон — сошедшееся решение того же типа. Девять гипотез я зарегистрировал до измерений: четыре подтвердились, пять опровергнуты (включая мою первую гипотезу о том, что виновата поправка 3M).

• Детерминированный случай. ema с m = −0,6 даёт ошибку на 11-21 % ниже, чем dpmpp_2m, при том же числе шагов (интервалы не включают ноль ни в одной ячейке; на 8 шагах выигрыш в 14 юнитах из 15, на 16 шагах в 15 из 15). Примерно в 1,2 раза меньше шагов при той же ошибке. LPIPS до декодированного эталона лучше в 14 из 15 юнитов на 8 шагах. Форма с масштабированием по времени m = −0,5 не хуже лучшей константы (это находка post hoc).

• Положительный momentum. Вредит всегда: до +61 % ошибки при m = 0,6.

• Случай с шумом. Сильный отрицательный momentum перестаёт помогать: вносимый шум непредсказуем, и экстраполяция усиливает ошибку, которая от него происходит. Демпфированная экстраполяция (beta = 0,5) даёт на 7-10 % меньшую сильную ошибку, чем dpmpp_2m_sde, а мягкий ema с m = −0,3 на 6-15 %. Но на настройках исходной проблемы (CFG 8, 1216×832) мягкий ema давал синий оттенок, а демпфированная экстраполяция нет.

• Один вызов против двух. Лучший вариант с одним вызовом на шаг обходит двухвызовный dpmpp_sde при том же NFE на 16-29 %.

• Что не сработало. Статическая сетка шагов из моего прошлого исследования снижала ошибку вдвое на четырёх коротких промптах, на которых её подбирали, и переносилась между ними, но не помогла отложенному промпту на естественном языке и сломала 8-шаговую картинку при CFG 8. В пакете она лежит как экспериментальный узел.

Как попробовать

Поставьте пакет (ссылка ниже) и перезапустите ComfyUI. В списке сэмплеров появятся ddim_momentum (детерминированный, форма со временем) и dpmpp_2m_sde_damped. Узел Sampler Momentum позволяет задать режим и коэффициент вручную. Отдельно есть узел-фильтр, который оборачивает любой сэмплер с одним вызовом модели на шаг, и записи <base>_momentumized для 19 существующих сэмплеров (dpmpp_3m_sde_momentumized, er_sde_momentumized и так далее). Эту часть я добавил как инженерное расширение и на модели не оценивал. При одном дымовом запуске на 8 шагах lms_momentumized и ipndm_v_momentumized со значением силы по умолчанию дали латент со std около 10, поэтому для них берите силу мягче или не берите вовсе.

Код, 13 тестов на CPU, таблицы и README на английском, русском, китайском и марийском: https://github.com/YakovSava/momentum-cfg-samplers-research Запись на Zenodo: https://doi.org/10.5281/zenodo.24214606