Крупномасштабная дистилляция диффузии через непрерывную согласованность с регуляризацией по скорингу

Аннотация

Данная работа представляет собой первую попытку масштабирования непрерывного согласованного дистилляции (continuous-time consistency distillation) для общих моделей диффузии изображений и видео на уровне приложений. Хотя непрерывная согласованная модель (sCM) теоретически обоснована и эмпирически эффективна для ускорения академических моделей диффузии, её применимость к крупномасштабным задачам генерации изображений и видео из текста остаётся неясной из-за инфраструктурных сложностей в вычислении произведения Якобиана на вектор (JVP) и ограничений стандартных бенчмарков для оценки. Мы разработали совместимое с параллельными вычислениями ядро FlashAttention-2 для JVP, что позволило обучать sCM на моделях с более чем 10 миллиардами параметров и задачах с высокоразмерными видео. Наше исследование выявило фундаментальные ограничения sCM в генерации мелких деталей, что мы связываем с накоплением ошибок и "покрытием мод" (mode-covering) в её целевой функции прямого расхождения. Для устранения этого мы предлагаем регуляризованную непрерывную согласованную модель (rCM), которая включает дистилляцию оценок (score distillation) в качестве регуляризатора с длинным пропуском. Такая интеграция дополняет sCM "поиском мод" (mode-seeking) через обратное расхождение, эффективно улучшая визуальное качество при сохранении высокой разнородности генерации. Проверенная на крупномасштабных моделях (Cosmos-Predict2, Wan2.1) с до 14 миллиардами параметров и 5-секундными видео, rCM соответствует или превосходит современный метод дистилляции DMD2 по метрикам качества, предлагая заметные преимущества в разнообразии, всё это без настройки GAN или обширного поиска гиперпараметров. Дистиллированные модели генерируют высококачественные образцы всего за 1–4 шага, ускоряя выборку диффузии в 15–50 раз. Эти результаты позиционируют rCM как практичный и теоретически обоснованный фреймворк для продвижения крупномасштабной дистилляции диффузии.

English

This work represents the first effort to scale up continuous-time consistency distillation to general application-level image and video diffusion models. Although continuous-time consistency model (sCM) is theoretically principled and empirically powerful for accelerating academic-scale diffusion, its applicability to large-scale text-to-image and video tasks remains unclear due to infrastructure challenges in Jacobian-vector product (JVP) computation and the limitations of standard evaluation benchmarks. We first develop a parallelism-compatible FlashAttention-2 JVP kernel, enabling sCM training on models with over 10 billion parameters and high-dimensional video tasks. Our investigation reveals fundamental quality limitations of sCM in fine-detail generation, which we attribute to error accumulation and the "mode-covering" nature of its forward-divergence objective. To remedy this, we propose the score-regularized continuous-time consistency model (rCM), which incorporates score distillation as a long-skip regularizer. This integration complements sCM with the "mode-seeking" reverse divergence, effectively improving visual quality while maintaining high generation diversity. Validated on large-scale models (Cosmos-Predict2, Wan2.1) up to 14B parameters and 5-second videos, rCM matches or surpasses the state-of-the-art distillation method DMD2 on quality metrics while offering notable advantages in diversity, all without GAN tuning or extensive hyperparameter searches. The distilled models generate high-fidelity samples in only 1sim4 steps, accelerating diffusion sampling by 15timessim50times. These results position rCM as a practical and theoretically grounded framework for advancing large-scale diffusion distillation.

Крупномасштабная дистилляция диффузии через непрерывную согласованность с регуляризацией по скорингу

Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency

Аннотация

Support