Почему сэмплеры 4+ порядка взрываются в диффузии, и как мы это починили: история создания Flow-WLS

Любой, кто когда-либо открывал список сэмплеров в ComfyUI или WebUI, наверняка замечал странное единообразие. В выпадающем меню присутствуют Euler, Heun, DPM++ 2M, изредка DPM++ 3M — и на этом список порядков заканчивается.

Возникает закономерный вопрос: почему численные методы дифференциальных уравнений, развивавшиеся в прикладной математике столетиями, в генеративных нейросетях застряли на втором и третьем порядке? Неужели нельзя взять метод Адамса-Башфорта 5-го, 6-го или 8-го порядка и получить сверхточную траекторию за меньшее число шагов?

Когда мы попытались напрямую реализовать многошаговые сэмплеры высших порядков для современных моделей Flow Matching (в частности, Anima DiT), генерация сразу же сломалась двумя разными способами: холст сначала выжигался в 90% белого цвета, а затем распадался на несвязный цифровой шум.

В этой статье подробно разберем:

1. Какая скрытая ловушка ждет разработчиков при переносе классических формул диффузии (EDM) на архитектуры Flow Matching.

2. Почему точная полиномиальная интерполяция Лагранжа порядка 5+ неизбежно разрушает латент.

3. Как мы применили идеи фильтрации Савицкого-Голея и взвешенных наименьших квадратов, создав метод Flow-WLS, генерирующий изображения за 10 шагов и 11 секунд на потребительской видеокарте.

1. Ловушка координат: почему экран становился белым

Первая же попытка запустить кастомный многошаговый сэмплер на модели Anima DiT привела к аномалии: VAE декодировал результат как равномерный белый прямоугольник с едва различимым силуэтом вверху.

Причина крылась в различии математических оснований между классическими моделями диффузии (Stable Diffusion 1.5, SDXL, EDM) и современными генераторами непрерывного потока (Flow Matching / DiT).

Классический EDM:

В моделях типа Karras EDM масштаб дисперсии данных , а параметр времени выражается через логарифмическое отношение сигнал/шум:

Flow Matching (comfy.model_sampling.CONST):

В моделях выпрямленного потока дисперсия распределена линейно:

Спрямленное время выражается через логит:

Во многих библиотеках для выпрямления траектории применяется замена переменной:

На старте диффузии () коэффициент . При делении начального тензора гауссова шума на амплитуда значений в латенте мгновенно умножалась в 10 000 раз. К 30-му шагу латенты уходили в глубокое насыщение, классификатор CFG усиливал дисбаланс, а VAE честно транслировал экстремальный клиппинг в белый цвет (255, 255, 255).

Решение:

Уравнение траектории Flow Matching уже изначально является прямолинейным в пространстве скоростей:

Оно не требует вспомогательных нормировок через . Достаточно интегрировать векторное поле напрямую:

def compute_velocity(x: torch.Tensor, denoised: torch.Tensor, sigma: float) -> torch.Tensor:
    safe_sigma: float = max(sigma, 1e-4)
    return (x - denoised) / safe_sigma

После перехода на прямое интегрирование поля скоростей пересвет полностью исчез.

2. Развал высоких порядков: почему Адамс-Башфорт 6-го порядка рухнул в шум

Устранив клиппинг, мы построили прямое многошаговое интегрирование Адамса-Башфорта на неравномерной сетке шагов для порядков от 2 до 6.

Казалось бы, чем выше порядок интерполяционного полинома, тем меньше локальная ошибка усечения (). Однако метрики сравнения с эталоном показали обратное:

Начиная с 5-го порядка генерация резко деградировала, а 6-й порядок разрушил изображение до состояния несвязного шума.

В чем причина коллапса?

1. Деление на степени шага: Разделенные разности -го порядка делятся на произведение расстояний между узлами сетки . При шаге знаменатель для порядка составляет порядка . Малейшая нелинейность отклика нейросети умножается в миллионы раз.

2. Знакопеременность весов Лагранжа: Веса квадратурной формулы для неравномерной сетки вычисляются через интегралы базисных полиномов:

При шаге веса для разных порядков ведут себя следующим образом:

• Порядок 2: (сумма модулей )

• Порядок 3: (сумма модулей )

• Порядок 6: веса начинают хаотично осциллировать со сменой знаков (), а сумма их модулей вырастает до .

Интегратор превращается в резонатор высокочастотного дребезга (классический феномен Рунге). Точная полиномиальная интерполяция на диффузионных траекториях имеет жесткое ограничение: порядки математически неустойчивы.

3. Метод Flow-WLS: регуляризация Савицкого-Голея

Мы поставили вопрос: как задействовать информацию за 5–8 предыдущих шагов, не создавая осциллирующий полином высокой степени?

Решение пришло из спектрального анализа и цифровой обработки сигналов — фильтр Савицкого-Голея (Savitzky-Golay).

Вместо интерполяционного полинома степени , который обязан строго пройти через каждую точку истории, Flow-WLS аппроксимирует траекторию жестко стабилизированной параболой 2-й степени методом взвешенных наименьших квадратов (WLS) по предыдущим точкам с коэффициентом затухания памяти :

Что дает этот переход:

1. Физическая невозможность осцилляций: Парабола () не имеет дополнительных степеней свободы и не способна изгибаться волнами.

2. Сжатие нормы весов: При росте истории с 3 до 6 шагов сумма модулей весов у Flow-WLS падает с до (у классического полинома она возрастала до ).

3. Фильтрация шума: Увеличение глубины памяти теперь не накапливает ошибки экстраполяции, а сглаживает скачки скорости нейросети на жестких участках траектории.

Реализация расчета эффективных весов на произвольной сетке:

def compute_wls_weights(
    dt: float,
    past_dts: list[float],
    deg: int = 2,
    gamma: float = 0.85,
) -> list[float]:
    length: int = 1 + len(past_dts)
    degree: int = min(deg, length - 1)
    
    nodes: list[float] = [0.0]
    accumulated_time: float = 0.0
    for step_size in past_dts:
        accumulated_time -= step_size
        nodes.append(accumulated_time)
        
    nodes_array: np.ndarray = np.array(nodes)
    weights_diagonal: np.ndarray = np.array([gamma**idx for idx in range(length)])
    weight_matrix: np.ndarray = np.diag(weights_diagonal)
    
    design_matrix: np.ndarray = np.column_stack(
        [nodes_array**power for power in range(degree + 1)]
    )
    normal_matrix: np.ndarray = design_matrix.T @ weight_matrix @ design_matrix
    pseudoinverse_matrix: np.ndarray = np.linalg.pinv(normal_matrix)
    projection_matrix: np.ndarray = pseudoinverse_matrix @ design_matrix.T @ weight_matrix
    
    integral_vector: np.ndarray = np.array(
        [(dt ** (power + 1)) / (power + 1) for power in range(degree + 1)]
    )
    effective_weights: np.ndarray = integral_vector @ projection_matrix
    return [float(weight) for weight in effective_weights]

4. Результаты экспериментов

Тестирование проводилось на модели Anima DiT (anima_aestheticV11) в разрешении на видеокарте Nvidia RTX 5060 без компиляторов и сторонних оптимизаций.

4.1. Стандартный бюджет (30 шагов)

Там, где классический полином cab_6 рухнул в шум с SSIM 0.1341, сэмплер flow_wls_8 превзошел встроенный uni_pc по всем трем метрикам.

4.2. Экстремальный режим: 10, 12 и 15 шагов

Главная ценность численного интегратора в генеративном пайплайне — способность удерживать геометрию при малом числе обращений к нейросети (Low-NFE):

При сокращении шагов до 10 метрика SSIM выросла до 0.8003, а время генерации составило 11 секунд.

5. Синергия трех уровней: как получить 8 секунд на кадр

Сэмплер Flow-WLS ускоряет генерацию во временном измерении (сокращая число шагов с 30 до 10). Однако его можно объединить с оптимизациями в других измерениях:

1. Временная ось: flow_wls_4 на 10 шагах.

2. Пространственно-спектральная ось: нода SPEED (Spectral Progressive Diffusion). Первые 50% шагов DiT-трансформер вычисляется на половинном пространственном разрешении с последующим DCT-расширением, что ускоряет внимание Attention в несколько раз.

3. Архитектурная ось: TeaCache, кэширующий медленно меняющиеся слои внимания между соседними шагами.

Итоговый замер:

• Базовый пайплайн (30 шагов, стандартный сэмплер): ~34 секунды.

• Связка Flow-WLS (10 шагов) + SPEED + TeaCache: 8 секунд на изображение 720×1280 на RTX 5060 при сохранении полноценного CFG 5.0 и работе любых сторонних LoRA.

Заключение

Эксперименты подтвердили два важных принципа:

1. Для генеративных моделей прямого потока (Flow Matching) опасна высокая степень интерполяционного полинома, но полезна глубокая ретроспективная выборка при низкой степени аппроксимации.

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

Код кастомной ноды открыт и доступен на GitHub, а препринт зарегистрирован в международной научной базе Zenodo:

• Репозиторий: github.com/YakovSava/research-unipc-dpmpp

• DOI: 10.5281/zenodo.23337579