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

Este trabalho apresenta o modelo de consistência contínua regularizado por pontuação (rCM), uma abordagem escalável que supera as limitações de qualidade e diversidade dos métodos anteriores ao introduzir um regularizador de distilação de pontuação, permitindo a geração de imagens e vídeos de alta fidelidade em apenas 1 a 4 passos sem necessidade de ajuste GAN.

O essencial

Imagine que você tem um artista genial (o Modelo Professor) que consegue pintar quadros incríveis ou criar vídeos realistas, mas ele é extremamente lento. Para criar uma imagem, ele precisa fazer 50 ou 100 "rascunhos" e ajustes antes de chegar ao resultado final. Isso é ótimo para a qualidade, mas péssimo para quem quer ver o resultado rápido.

O objetivo deste trabalho é criar um Artista Aprendiz (o Modelo Aluno) que aprenda a fazer o mesmo trabalho do Professor, mas em apenas 1 ou 2 passos, mantendo a mesma qualidade.

Aqui está a explicação simplificada da descoberta deles, usando analogias do dia a dia:

1. O Problema: O "Aprendiz" que Perdeu o Foco

Antes deste trabalho, existia uma técnica chamada sCM (Modelo de Consistência Contínua). Pense nela como um aluno que tenta aprender a arte observando o Professor e tentando pular direto para o final do processo.

• O que funcionava: O aluno aprendia a ser muito rápido e criativo (gerava muitas variações diferentes).
• O problema: Quando tentava fazer coisas complexas (como escrever texto pequeno em um quadro ou manter um objeto estável em um vídeo), ele começava a "alucinar". As letras ficavam borradas, os objetos se fundiam ou a imagem perdia detalhes finos. Era como se o aluno tivesse pressa demais e esquecesse os detalhes importantes.

2. A Solução: O "Duplo Chefe" (rCM)

Os pesquisadores descobriram que o problema acontecia porque o aluno estava apenas tentando imitar o Professor de uma única maneira (chamada de "divergência direta"). Eles decidiram dar ao aluno dois tipos de instrução ao mesmo tempo:

1. O Chefe da Velocidade (Consistência): Continua dizendo: "Faça rápido! Pule direto para o resultado final!" (Isso mantém a velocidade e a diversidade).
2. O Chefe da Qualidade (Distilação de Pontuação): Um novo instrutor que diz: "Espere! Olhe para o que você mesmo criou e compare com o que o Professor faria. Se estiver feio, corrija!" (Isso foca na qualidade e nos detalhes).

Essa nova técnica se chama rCM (Modelo de Consistência Contínua Regularizado por Pontuação).

3. A Analogia da Cozinha

Imagine que o Professor é um Chef de cozinha famoso que leva 1 hora para fazer um prato perfeito.

• O método antigo (sCM): O aprendiz tenta copiar o prato do Chef em 1 minuto, apenas olhando para a foto final. O resultado? Um prato rápido, mas que parece uma "massa" sem sabor (borrado, sem detalhes).
• O novo método (rCM): O aprendiz tenta fazer o prato em 1 minuto, mas enquanto cozinha, ele tem um segundo Chefe ao lado.
  • O primeiro Chefe diz: "Mova-se rápido!"
  • O segundo Chefe diz: "Cheire a comida, prove, e se não estiver igual ao do Chef, ajuste o tempero agora mesmo."
  • Resultado: O prato sai em 1 minuto, mas tem o sabor e a apresentação de quem levou 1 hora.

4. O Desafio Técnico (O "Supercomputador")

Fazer isso em modelos gigantes (como os usados para criar vídeos de 5 segundos) era quase impossível antes.

• O Obstáculo: Calcular os ajustes finos exigia uma matemática complexa (chamada JVP) que travava os computadores modernos quando o modelo era muito grande (mais de 10 bilhões de parâmetros).
• A Invenção: Eles criaram um "atalho" de software (um novo kernel de FlashAttention) que permite fazer esses cálculos complexos sem travar o computador, mesmo em modelos gigantes. Foi como trocar um caminhão de carga lento por um trem de alta velocidade para transportar a matemática.

5. Os Resultados: O Milagre da Velocidade

Com essa nova técnica (rCM), eles conseguiram:

• Velocidade: Gerar imagens e vídeos em 1 a 4 passos (antes eram 50+). Isso é um aumento de velocidade de 15 a 50 vezes.
• Qualidade: A qualidade visual é tão boa quanto a do modelo original lento, conseguindo escrever textos legíveis e manter objetos consistentes em vídeos.
• Diversidade: Ao contrário de outros métodos rápidos que tendem a criar sempre a mesma coisa (modo "colapso"), o rCM continua criando coisas variadas e criativas.

Em resumo: Eles ensinaram um robô super-rápido a não apenas correr, mas a correr olhando para os detalhes, garantindo que ele não tropece nos detalhes finos da imagem ou do vídeo. É como ter um Ferrari que dirige na velocidade da luz, mas com a precisão de um piloto de Fórmula 1.

Resumo técnico

Resumo Técnico: rCM (Score-Regularized Continuous-Time Consistency)

1. O Problema

Os modelos de consistência contínua (como sCM e MeanFlow) são teoricamente fundamentados e eficazes para acelerar a inferência de modelos de difusão em escala acadêmica. No entanto, sua aplicação em modelos de difusão de grande escala (para tarefas de imagem e vídeo de nível industrial, com bilhões de parâmetros) enfrenta barreiras críticas:

• Desafios de Infraestrutura: O cálculo do Produto Vetor-Jacobiano (JVP), essencial para o treinamento do sCM, é incompatível com as infraestruturas modernas de treinamento massivo (como precisão BF16, FlashAttention e paralelismo de contexto), gerando erros numéricos e instabilidade.
• Limitações de Qualidade: Em tarefas complexas (como renderização de texto ou consistência temporal em vídeos), o sCM puro sofre de degradação de qualidade, gerando artefatos, distorções e texturas borradas.
• Natureza do Objetivo: O sCM utiliza uma divergência "forward" (modo de cobertura), que tende a cobrir toda a distribuição de dados, resultando em alta diversidade, mas baixa qualidade de detalhes finos.

2. Metodologia

Os autores propõem o rCM (Score-Regularized Continuous-Time Consistency Model), uma nova estrutura de destilação que combina a consistência temporal com a regularização por pontuação (score distillation).

A. Infraestrutura Escalável (FlashAttention-2 JVP)
Para permitir o treinamento em modelos com mais de 10 bilhões de parâmetros e vídeos de alta dimensão, os autores desenvolveram:

• Um kernel FlashAttention-2 compatível com JVP, implementado em Triton, que integra o cálculo do JVP diretamente no passo forward da atenção, suportando tanto atenção auto quanto cruzada.
• Adaptação para FSDP (Fully Sharded Data Parallel) e CP (Context Parallelism), reestruturando as camadas da rede para aceitar entradas e saídas de tangentes, permitindo o cálculo do JVP dentro de cada camada sem violar a granularidade do paralelismo.

B. Arquitetura do rCM
O rCM resolve as limitações de qualidade do sCM através de uma combinação híbrida de objetivos:

1. Consistência Contínua (sCM): Mantém o objetivo original de consistência forward (divergência forward), que garante alta diversidade e estabilidade de treinamento.
2. Regularização por Pontuação (Score Distillation): Introduz um termo de regularização baseado em DMD (Distribution Matching Distillation). Este termo atua como um regularizador de "longo salto" (long-skip), utilizando uma divergência reversa (reverse divergence) que busca o modo (mode-seeking).
   • Mecanismo: O aluno gera amostras (rollouts) e um "fake score network" (uma rede auxiliar) aprende a estimar a pontuação dessas amostras. O objetivo de DMD penaliza a discrepância entre a distribuição do aluno e a do professor.
   • Equação de Perda: $L_{rCM} = L_{sCM} + \lambda L_{DMD}$. O peso $\lambda$ (encontrado empiricamente como 0.01) equilibra a diversidade (do sCM) e a qualidade (do DMD).

C. Estabilidade Numérica
Para evitar colapsos durante o treinamento de longas durações, o rCM emprega técnicas para calcular a derivada temporal de forma estável:

• Tempo Semi-Contínuo: Usa diferenças finitas para a derivada parcial em modelos menores.
• Tempo de Alta Precisão: Para modelos grandes (>10B) e vídeo, utiliza o cálculo contínuo nativo via JVP, mas força precisão FP32 nas camadas de embedding de tempo para evitar erros numéricos acumulados no BF16.

3. Principais Contribuições

1. Primeira Escala Industrial: É o primeiro trabalho a escalar a destilação de consistência contínua para modelos de difusão de texto-para-imagem (T2I) e texto-para-vídeo (T2V) com até 14 bilhões de parâmetros e vídeos de 5 segundos.
2. Solução de Infraestrutura: Desenvolvimento de kernels de JVP compatíveis com FlashAttention-2 e paralelismo massivo, resolvendo o gargalo computacional que impedia o uso do sCM em grandes modelos.
3. Novo Paradigma de Destilação: Demonstra que a combinação de divergências forward (consistência) e reversa (score distillation) supera as limitações individuais de cada método, eliminando a necessidade de ajuste fino complexo de GANs ou múltiplos estágios de treinamento.
4. Aceleração Extrema: Modelos destilados geram amostras de alta fidelidade em 1 a 4 passos, acelerando a inferência em 15x a 50x em comparação aos modelos professores.

4. Resultados Experimentais

Os autores validaram o rCM em modelos de ponta: Cosmos-Predict2 (T2I) e Wan2.1 (T2V).

• Qualidade vs. Diversidade:
  • O sCM puro sofre com distorções e falta de detalhes finos (ex: texto ilegível).
  • O DMD2 (estado da arte atual) tem alta qualidade, mas sofre de colapso de modo (baixa diversidade, objetos com posições/orientações muito similares).
  • O rCM atinge a qualidade do DMD2 (ou superior em métricas como GenEval e VBench) enquanto mantém a alta diversidade do sCM, evitando o colapso de modo.
• Métricas:
  • Em T2I (Cosmos-Predict2 14B), o rCM alcançou uma pontuação geral de 0.83 no GenEval em apenas 4 passos, superando modelos pré-treinados e competindo com o FLUX.1-schnell.
  • Em T2V (Wan2.1 14B), o rCM superou o modelo professor de 14B no VBench (pontuação total de 85), gerando vídeos de 5 segundos com consistência temporal e detalhes nítidos.
• Eficiência: O modelo consegue gerar vídeos de alta qualidade em 2 passos e imagens em 1 passo, com taxas de quadros (FPS) significativamente maiores que os métodos de difusão tradicionais.

5. Significado e Impacto

Este trabalho estabelece o rCM como um framework prático e teoricamente fundamentado para o futuro da destilação de difusão em grande escala.

• Viabilidade Prática: Remove a barreira de infraestrutura que impedia o uso de modelos de consistência contínua em aplicações reais de vídeo e imagem de alta resolução.
• Unificação Teórica: Demonstra que a união de princípios de divergência forward e reversa oferece um caminho superior para obter simultaneamente alta qualidade, alta diversidade e velocidade de inferência, superando as limitações de métodos puramente baseados em GANs ou apenas em consistência.
• Futuro: Abre caminho para o desenvolvimento de "modelos de mundo" interativos e geradores de vídeo em tempo real, onde a velocidade e a consistência são críticas.

Em suma, o rCM resolve o dilema clássico entre velocidade, qualidade e diversidade na geração de conteúdo, permitindo que modelos massivos de IA generativa sejam utilizados de forma eficiente em cenários do mundo real.