Collaborative Few-Step Distillation and Low-Bit Quantization for Wan2.2 Dual-Expert Video Diffusion Models

Resumo Técnico: Destilação Colaborativa de Poucos Passos e Quantização de Baixo Bit para Modelos de Difusão de Vídeo de Especialista Duplo Wan2.2

1. Definição do Problema

Grandes modelos de difusão de vídeo, como o Wan2.2-T2V-A14B, alcançam alta qualidade visual, mas enfrentam barreiras significativas de implantação devido a dois gargalos primários:

1. Custo Computacional: O processo iterativo de denoising exige um grande número de passos de amostragem, tornando a inferência cara em termos de latência.
2. Pegada de Memória: Os modelos possuem uma contagem de parâmetros residentes elevada (aprox. 27B no total, com ~14B ativos por passo na arquitetura de especialista duplo), excedendo os orçamentos de memória de muitos cenários de implantação.

As estratégias de compressão existentes frequentemente falham quando aplicadas diretamente a esses modelos. Simplesmente reduzir os passos de amostragem sem retreinamento leva a uma baixa qualidade, enquanto a quantização da trajetória original de passos longos pode introduzir desajustes na distribuição de ativação durante a inferência, causando propagação de erro temporal e degradação da fidelidade do vídeo. Além disso, o Wan2.2 utiliza uma arquitetura de Mistura de Especialistas (MoE) de especialista duplo (especialista de alto ruído e especialista de baixo ruído), o que complica as abordagens de quantização padrão que tratam o modelo como um tronco denso.

2. Metodologia

O artigo propõe um pipeline de compressão orientado à implantação que integra destilação de poucos passos e quantização de baixo bit de maneira co-projetada. O pipeline segue uma sequência específica: adaptação do modelo $\rightarrow$ destilação de poucos passos $\rightarrow$ quantização de baixo bit $\rightarrow$ construção de calibração $\rightarrow$ geração $\rightarrow$ avaliação.

Componentes Técnicos Principais:

• Tratamento de Arquitetura Consciente de Especialista Duplo:

  • O pipeline respeita o mecanismo de roteamento específico do Wan2.2, onde um especialista de alto ruído lida com o layout global/movimento e um especialista de baixo ruído refina a textura/detalhes.
  • Calibração por Ramificação: Em vez de agrupar ativações, o método captura e calibra estatísticas separadamente para as ramificações de alto e baixo ruído. Isso evita que o quantizador misture distribuições de ativação incompatíveis.
  • Proteção de Camadas de Entrada: As camadas iniciais e blocos sensíveis ao roteamento são mantidos em maior precisão ou calibrados de forma mais agressiva. Isso é crítico porque erros na fase inicial de alto ruído (inicialização de layout/movimento) propagam-se por toda a sequência, enquanto erros posteriores são frequentemente locais.

• Destilação de Poucos Passos via Correspondência de Distribuição:

  • O método emprega Destilação de Correspondência de Distribuição (DMD) para aprender uma trajetória de denoising mais curta, em vez de meramente truncar o amostrador.
  • O modelo estudante é otimizado para corresponder à distribuição do professor ao longo do número reduzido de passos, com um objetivo consciente de estágio ($L_{distill} = \lambda_1 L_{high} + \lambda_2 L_{low}$) que respeita os papéis distintos das fases de alto e baixo ruído.

• Quantização Co-projetada:

  • Calibração Alinhada à Trajetória: Uma inovação central é realizar a calibração de quantização no estudante destilado de poucos passos em vez do modelo original de passos longos. Isso garante que o quantizador observe as estatísticas de ativação que encontrará durante a inferência real, reduzindo o desajuste de distribuição.
  • Representação HiF4 (HiFloat4): O pipeline utiliza um formato de ponto flutuante em bloco de 4 bits (HiF4) projetado para inferência de baixo bit. Este formato apresenta uma hierarquia de escala de três níveis (Global E6M2, Nível-2 E1, Nível-3 E1) e amplitude assinada S1P2. Este design fornece um amplo alcance dinâmico global (aprox. $2^{-50}$ a $2^{18}$) para lidar com intervalos de ativação heterogêneos entre as fases de denoising, enquanto preserva o detalhe local através de micro-expoentes.
  • Granularidade: Os pesos são quantizados por canal, e as ativações são quantizadas por token para rastrear variações de prompt e de frame.

• Engenharia e Avaliação:

  • Gerenciamento de Memória: Para evitar erros de Out-Of-Memory (OOM) em GPUs únicas (ex: H100 de 80GB), a calibração e a substituição são realizadas camada por camada, com buffers auxiliares carregados sob demanda.
  • Avaliação de Ciclo Fechado: O sistema utiliza o VBench para avaliar cinco dimensões específicas: Consistência de Assunto (SC), Qualidade Estética (AQ), Qualidade de Imagem (IQ), Consistência Geral (OC) e Suavidade de Movimento (MS). O protocolo de avaliação garante uma comparação justa utilizando prompts, resoluções e cronogramas de amostragem idênticos entre as linhas de base.

3. Principais Contribuições

1. Co-Design de Destilação e Quantização: O artigo demonstra que calibrar a quantização na trajetória destilada de poucos passos (em vez do modelo original de passos longos) reduz significativamente o desajuste da distribuição de ativação, permitindo que o modelo comprimido mantenha uma qualidade próxima à do estudante de precisão total.
2. Estratégia de Compressão Consciente de MoE: O método introduz um esquema de calibração por ramificação e proteção de camadas de entrada especificamente adaptado para a estrutura de especialista duplo do Wan2.2, evitando a mistura de estatísticas de ativação de alto e baixo ruído.
3. Integração de HiF4 para Difusão de Vídeo: A aplicação da escala hierárquica do HiF4 mostra-se eficaz para modelos de difusão de vídeo, equilibrando a necessidade de um amplo alcance global (para evitar overflow nas fases de alto ruído) com precisão local (para preservar a textura nas fases de baixo ruído).
4. Validação Empírica de Contagem de Passos: O estudo avalia sistematicamente as configurações de 4, 8, 20 e 40 passos, identificando o ponto de equilíbrio ideal entre qualidade e eficiência.

4. Resultados Experimentais

Os experimentos foram conduzidos no modelo Wan2.2-T2V-A14B usando o conjunto de dados OpenS2V para calibração e o VBench para avaliação.

• Desempenho vs. Linhas de Base:

  • Em 4 passos, o modelo quantizado mantém um pequeno erro em relação ao modelo de precisão total de mesmo número de passos, mas fica atrás da linha de base original de passos longos.
  • Em 8 passos, o modelo quantizado (Pontuação Média: 0.7002) começa a exceder a linha de base original de precisão total (0.6976).
  • Em 20 passos, o modelo quantizado atinge o melhor desempenho (Pontuação Média: 0.7074), superando tanto a linha de base original de precisão total quanto o modelo destilado de precisão total de mesmo número de passos (0.7051).
  • Em 40 passos, o desempenho estabiliza, indicando retornos decrescentes ao aumentar ainda mais a extensão da trajetória.

• Análise de Métricas (Configuração de 20 Passos):

  • O modelo quantizado melhorou em relação à linha de base original em Consistência de Assunto (0.9629 vs. 0.9495), Qualidade Estética, Qualidade de Imagem e Suavidade de Movimento.
  • A Consistência Geral permaneceu comparável à do modelo de mesmo número de passos.
  • Os resultados indicam que a calibração alinhada à distribuição recuperou com sucesso as perdas de baixo bit sem sacrificar a estabilidade temporal.

• Sensibilidade ao Erro:

  • A análise revela que perturbações de quadro único afetam principalmente a Qualidade de Imagem e Estética, enquanto erros acumulados ao longo dos passos se manifestam como Deriva de Assunto (Subject Drift) e Tremor de Movimento (Motion Jitter). A proteção das camadas de entrada proposta mitiga efetivamente o último caso.

5. Significância e Alegações

O artigo alega que o pipeline proposto oferece um caminho prático para implantar modelos de difusão de vídeo de alta qualidade sob orçamentos limitados de memória de GPU e latência.

• Compromisso Qualidade-Eficiência: Os autores identificam a configuração de 20 passos como a configuração ideal, fornecendo capacidade de denoising suficiente para recuperar a qualidade temporal e visual, enquanto requer significativamente menos passos do que os amostradores convencionais.
• Mudança Estratégica: O trabalho argumenta que a quantização não deve ser um "adendo com perda" (lossy afterthought), mas sim uma parte integrante da estratégia de aceleração. Ao co-projetar a destilação e a quantização, o método transforma a conversão de baixo bit em um componente estável do pipeline de inferência.
• Prontidão para Implantação: O uso de proteção de camadas de entrada e calibração por ramificação aborda os modos de falha específicos de modelos de vídeo de especialista duplo (ex: deriva de assunto e instabilidade de movimento), tornando o modelo comprimido adequado para geração interativa e processamento em lote com recursos restritos.

O artigo conclui que, embora o pipeline atual foque na quantização pós-treinamento, o framework estabelece um loop de verificação onde falhas específicas de métricas (ex: quedas na Consistência de Assunto) fornecem sinais de depuração direta para refinar a calibração e a contagem de passos.