Resumo Técnico: Destilação de Decodificação Paralela para Geração Rápida de Imagem e Vídeo

1. Definição do Problema

Modelos de difusão e de correspondência de fluxo (flow matching) de larga escala alcançaram resultados de estado da arte (SOTA) em geração de texto para imagem, texto para vídeo e multimodalidade. No entanto, seus processos inerentes de amostragem iterativa exigem centenas de avaliações de rede (NFEs), criando custos computacionais e latência significativos. Este gargalo dificulta aplicações como geração de vídeo em tempo real, edição de conteúdo e modelagem de mundo interativa.

Os métodos de aceleração existentes geralmente se dividem em duas categorias:

• Métodos baseados em trajetória: Estes destilam o processo de amostragem sequencial do professor em menos passos. Embora eficazes para imagens, eles frequentemente degradam a qualidade do vídeo ou exigem algoritmos de treinamento custosos (ex: produtos de Jacobiano-vetor ou diferenças finitas) quando aplicados a modelos de vídeo.
• Métodos baseados em distribuição: Estes alinham as distribuições marginais do aluno e do professor em vez de seguir trajetórias específicas. Embora dominantes na destilação de vídeo (ex: usando Destilação de Score Variacional ou perdas adversariais), eles sofrem de colapso de modo (mode collapse), levando a uma perda de diversidade de vídeo e movimento estático. Eles também frequentemente exigem objetivos de treinamento alternados e são sensíveis a hiperparâmetros.

O artigo identifica a necessidade de um método de destilação baseado em trajetória que seja escalável, evitando as dificuldades de otimização dos métodos baseados em distribuição (colapso de modo) e o overhead computacional dos métodos baseados em trajetória existentes (JVPs, diferenças finitas), especificamente para geração de vídeo de alta resolução.

2. Metodologia: Destilação de Decodificação Paralela (PDD)

Os autores introduzem a Destilação de Decodificação Paralela (PDD), um framework de destilação baseado em trajetória projetado para acelerar a inferência de modelos de difusão e de correspondência de fluxo.

Conceito Central

Ao contrário dos métodos tradicionais que fundem múltiplos passos de denoising em um único passo grande, a PDD aprende um decodificador paralelo que prevê as velocidades médias para múltiplos intervalos de tempo consecutivos em uma única avaliação de rede.

• Discretização: O domínio do tempo $[0, 1]$ é discretizado em $N$ intervalos. Estes são agrupados em blocos de tamanho $L$.
• Predição Paralela: Dado um estado $X_n$ no tempo $t_n$, o decodificador paralelo $\bar{u}_\theta$ prevê as velocidades médias para todos os $L$ intervalos dentro do bloco $\{n, \dots, n+L-1\}$ simultaneamente:
  $$\bar{u}_\theta(n|X_n) \approx u_k(X_k), \quad k = n, \dots, n+L-1$$
  onde $u_k$ é a velocidade média do $k$-ésimo intervalo.
• Amostragem: Durante a inferência, o modelo avança $L$ intervalos em um único passo usando as velocidades previstas, reduzindo o total de NFEs para $N/L$. Ao variar o tamanho do bloco $L$ durante o treinamento, o modelo suporta NFEs variáveis no tempo de inferência sem necessidade de condicionamento de tempo adicional.

Objetivo de Treinamento

A PDD utiliza um objetivo de treinamento baseado em regressão simplificado que evita VSD, perdas adversariais, JVPs e diferenças finitas.

• Aproximação do Professor: As velocidades médias alvo são aproximadas usando um solver Runge-Kutta (Euler ou Midpoint) aplicado ao modelo professor pré-treinado.
• Treinamento On-Policy: A perda é computada sobre as próprias saídas do aluno (on-policy). O aluno prevê as velocidades, simula a trajetória dentro do bloco e a perda é calculada comparando a velocidade prevista pelo aluno em um passo amostrado $k$ contra a velocidade média aproximada do professor naquele estado.
• Função de Perda:
  $$\mathcal{L}_{PD}(\theta) = \mathbb{E} \left[ \| \bar{u}_\theta(n|X_n) - u_k(\text{sg}(\bar{X}_k)) \|^2 \right]$$
  onde $\text{sg}$ denota o operador de parada de gradiente (stop-gradient), e $\bar{X}_k$ é o estado alcançado seguindo as velocidades previstas pelo aluno.

Arquitetura

• Backbone: A PDD utiliza exatamente a mesma arquitetura de backbone do modelo professor pré-treinado.
• Camada de Saída: A camada linear final é expandida repetindo a dimensão de canal $N$ vezes (onde $N$ é o tamanho da grade). Isso permite que o modelo produza $N$ previsões de velocidade distintas correspondentes a cada passo de tempo.
• Fusão de Camadas: Durante a geração, em vez de aplicar $L$ camadas lineares separadas, os pesos são fundidos em uma única camada linear que produz a velocidade média ponderada do bloco. Isso garante que não haja custo computacional adicional durante a inferência em comparação com a avaliação de passo único do professor.

Treinamento Sem Dados (Data-Free)

Para modelos de larga escala onde os dados não estão disponíveis, a PDD suporta um esquema de treinamento sem dados. O modelo amostra um ruído inicial, alternando entre passos de solver (avançando o estado usando as próprias previsões do aluno) e passos de treinamento, efetivamente aprendendo com suas próprias trajetórias geradas.

3. Principais Contribuições

1. Formulação da PDD: Um método de destilação baseado em trajetória escalável que prevê velocidades médias para múltiplos intervalos em uma única avaliação, permitendo inferência rápida para modelos de correspondência de fluxo e difusão.
2. Treinamento Simplificado: Um objetivo de regressão único que elimina a necessidade de JVPs, diferenças finitas, treinamento de múltiplos estágios, VSD ou perdas GAN, mantendo a produção de amostras de alta qualidade e diversidade.
3. Arquitetura Flexível: Um design compatível com qualquer modelo pré-treinado que suporte NFEs variáveis no tempo de inferência sem exigir condicionamento de tempo adicional ou uma segunda coordenada de tempo.
4. Validação Empírica: Demonstração de desempenho SOTA em ImageNet-256, Qwen-Image (Texto-para-Imagem), Wan2.1 (Texto-para-Vídeo 1.3B/14B) e LTX-2.3 (Texto-para-Vídeo/Áudio), com melhorias significativas na diversidade de geração em comparação com baselines baseados em distribuição.

4. Resultados Experimentais

Os autores avaliaram a PDD em três tarefas:

• Geração de Imagem Condicional por Classe (ImageNet-256):

  • A PDD alcançou pontuações FID competitivas no cenário de passo único (FID 2.69 com Midpoint), superando ou igualando métodos SOTA como FreeFlow e Pi-Flow, enquanto suporta NFEs variáveis (1, 2, 4, 8).
  • O método mostrou convergência estável através de diferentes hiperparâmetros.

• Texto-para-Imagem (Qwen-Image 20B):

  • Avaliado nos benchmarks OneIG, DPG-Bench e GenEval com 2, 4 e 8 NFEs.
  • A PDD (Midpoint) alcançou as melhores pontuações gerais no OneIG e DPG-Bench com 4 e 8 NFEs, superando baselines como DMD2, TwinFlow e Pi-Flow.
  • Crucialmente, a PDD preservou uma diversidade significativamente maior (medida pelas métricas de diversidade OneIG) em comparação ao DMD2, que sofreu de colapso de modo.

• Texto-para-Vídeo (Wan2.1 1.3B & 14B):

  • Avaliado no VBench com 4 e 8 NFEs.
  • A PDD alcançou SOTA em qualidade de vídeo e pontuações gerais, ocupando o primeiro lugar no modelo 1.3B e o segundo lugar no modelo 14B (atrás do AnyFlow em métricas gerais, mas com melhor diversidade).
  • Diversidade: A PDD demonstrou pontuações de diversidade superiores em comparação com baselines baseados em distribuição (DMD2, AnyFlow).
  • Movimento: A análise visual indicou que os vídeos gerados pela PDD exibiram graus mais elevados de movimento e consistência dinâmica em comparação aos baselines, que frequentemente produziam vídeos estáticos ou repetitivos.

• Texto-para-Vídeo/Áudio (LTX-2.3):

  • A PDD destilou o modelo LTX-2.3 de 22B para geração de vídeo 720p de 10s com áudio usando apenas 8 NFEs.
  • Após apenas 250 iterações de treinamento, a PDD teve desempenho equivalente ou superior ao modelo destilado oficial de 8 passos, apesar de não ter acesso aos dados de treinamento.
  • A avaliação de preferência humana (via Gemini 3.1 Pro) mostrou a PDD vencendo ou empatando com o modelo destilado oficial em alinhamento de prompt, qualidade visual, qualidade de movimento e qualidade de áudio.

5. Significância e Alegações

O artigo afirma que a PDD representa o primeiro método puro de destilação baseado em trajetória capaz de alcançar geração de vídeo de alta resolução e poucos passos com diversidade aprimorada.

• Superação do Colapso de Modo: Ao aderir a uma abordagem baseada em trajetória sem depender de perdas baseadas em distribuição (VSD/GANs), a PDD evita o colapso de modo e a falta de movimento frequentemente observados nos métodos atuais de destilação de vídeo.
• Eficiência Computacional: O método remove a necessidade de JVPs caros e diferenças finitas, tornando-o escalável para grandes modelos (até 22 bilhões de parâmetros) e compatível com treinamento sem dados.
• Flexibilidade: A capacidade de suportar NFEs variáveis no tempo de inferência através da seleção do tamanho do bloco oferece uma vantagem prática para diversos cenários de aplicação.
• Movimento e Diversidade: Os autores enfatizam que a PDD não apenas acelera a geração, mas também preserva a "riqueza" da saída do modelo professor, especificamente no que diz respeito à dinâmica de movimento e diversidade de amostras, que são críticos para a geração de vídeo.

Os autores concluem que, embora a PDD seja atualmente limitada a modelos de fluxo e difusão, o princípio da decodificação paralela poderia potencialmente ser generalizado para modelos autorregressivos discretos em trabalhos futuros. Eles também observam que investigar a PDD em configurações dependentes de dados além do ImageNet-256 permanece como uma área aberta para pesquisas futuras.