Difusão em Grande Escala via Distilação de Pontuação Regularizada em Consistência de Tempo Contínuo

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

October 9, 2025

Autores: Kaiwen Zheng, Yuji Wang, Qianli Ma, Huayu Chen, Jintao Zhang, Yogesh Balaji, Jianfei Chen, Ming-Yu Liu, Jun Zhu, Qinsheng Zhang

cs.AI

Resumo

Este trabalho representa o primeiro esforço para escalar a destilação de consistência em tempo contínuo para modelos de difusão de imagem e vídeo em nível de aplicação geral. Embora o modelo de consistência em tempo contínuo (sCM) seja teoricamente fundamentado e empiricamente poderoso para acelerar a difusão em escala acadêmica, sua aplicabilidade em tarefas de grande escala, como geração de texto para imagem e vídeo, permanece incerta devido a desafios de infraestrutura no cálculo de produtos Jacobiano-vetor (JVP) e às limitações dos benchmarks de avaliação padrão. Primeiro, desenvolvemos um kernel JVP compatível com paralelismo baseado no FlashAttention-2, permitindo o treinamento de sCM em modelos com mais de 10 bilhões de parâmetros e tarefas de vídeo de alta dimensionalidade. Nossa investigação revela limitações fundamentais de qualidade do sCM na geração de detalhes finos, que atribuímos ao acúmulo de erros e à natureza "mode-covering" de seu objetivo de divergência direta. Para remediar isso, propomos o modelo de consistência em tempo contínuo regularizado por pontuação (rCM), que incorpora a destilação de pontuação como um regularizador de salto longo. Essa integração complementa o sCM com a divergência reversa "mode-seeking", melhorando efetivamente a qualidade visual enquanto mantém alta diversidade na geração. Validado em modelos de grande escala (Cosmos-Predict2, Wan2.1) com até 14 bilhões de parâmetros e vídeos de 5 segundos, o rCM iguala ou supera o método de destilação state-of-the-art DMD2 em métricas de qualidade, oferecendo vantagens notáveis em diversidade, tudo sem ajustes de GAN ou extensas buscas de hiperparâmetros. Os modelos destilados geram amostras de alta fidelidade em apenas 1 a 4 passos, acelerando a amostragem de difusão em 15 a 50 vezes. Esses resultados posicionam o rCM como uma estrutura prática e teoricamente fundamentada para avançar a destilação de difusão em grande escala.

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.