DyWeight: Dynamic Gradient Weighting for Few-Step Diffusion Sampling

Resumo Técnico: DyWeight

1. O Problema

Os Modelos de Difusão (DMs) alcançaram o estado da arte na geração de dados, mas sofrem de um gargalo crítico: a lentidão na amostragem. Gerar uma única imagem de alta qualidade tradicionalmente exige centenas ou milhares de avaliações da rede neural (NFEs - Number of Function Evaluations).

Embora solvers de equações diferenciais ordinárias (ODE) de múltiplos passos (como DPM-Solver++ e iPNDM) tenham melhorado a eficiência ao reutilizar gradientes históricos, eles apresentam limitações fundamentais em cenários de poucos passos (few-step):

• Coeficientes Artesanais (Handcrafted): Métodos existentes utilizam coeficientes fixos derivados de suposições de estacionariedade e trajetórias suaves de baixa ordem.
• Dinâmica Não Estacionária: A dinâmica de difusão é altamente não estacionária; gradientes ruidosos de passos iniciais podem entrar em conflito com gradientes mais confiáveis de passos posteriores quando os coeficientes são fixos.
• Erros de Truncamento: Em passos grandes (poucos NFEs), as suposições clássicas de solvers numéricos falham, gerando grandes erros e degradação da qualidade da imagem.
• Otimização Complexa: Solvers baseados em aprendizado anteriores tentaram otimizar coeficientes e agendamentos de tempo, mas frequentemente adotam uma abordagem desacoplada (otimizando pesos e tempos separadamente), o que cria um landscape de otimização complexo e instável.

2. Metodologia: DyWeight

O DyWeight propõe um solver de múltiplos passos baseado em aprendizado que introduz um paradigma de acoplamento implícito simplificado. A abordagem central é relaxar as restrições numéricas clássicas para permitir que o solver aprenda dinamicamente como agregar gradientes e ajustar o tempo simultaneamente.

Principais Componentes:

• Pesagem de Gradiente Dinâmica (Relaxação de Restrições):

  • Em solvers clássicos, os pesos dos gradientes históricos devem somar 1 ($\sum w_i = 1$). O DyWeight remove essa restrição, permitindo pesos não normalizados e aprendíveis.
  • Isso cria dois mecanismos emergentes:
    1. Parte Normalizada: Aprende a combinação ótima de gradientes históricos.
    2. Parte Não Normalizada (Soma dos Pesos): Atua implicitamente como um mecanismo de escalonamento do tamanho do passo efetivo. A soma dos pesos ajusta a magnitude da atualização, permitindo que o solver adapte o "tamanho do passo" dinamicamente sem precisar prever explicitamente novos tempos.

• Calibração de Tempo (Time Calibration):

  • Para corrigir a discrepância entre o tamanho do passo numérico e as expectativas internas de ruído do modelo de difusão, o método emprega duas estratégias:
    1. Deslocamento de Tempo Implícito (Time Shifting): Utiliza a soma não normalizada dos pesos aprendidos para calcular o próximo passo de tempo ($\tilde{t}_{n-1}$), ajustando a trajetória numérica.
    2. Escala de Tempo (Time Scaling): Introduz um escalar aprendível por passo ($s_n$) que é aplicado diretamente ao tempo $t$ alimentado na rede neural ($\epsilon_\theta(x_t, s_n t)$). Isso compensa o viés de exposição e garante que a rede seja consultada no nível de ruído correto.

• Treinamento por Distilação (Teacher-Student):

  • O solver é treinado sem dados rotulados, utilizando um framework de distilação de ponta a ponta (endpoint supervision).
  • Um solver "professor" de alta fidelidade (ex: 35 NFEs) gera trajetórias de referência.
  • O solver "aluno" (DyWeight) gera a imagem final em poucos passos (ex: 3-7 NFEs).
  • A perda é calculada apenas na imagem final (distância entre a saída do aluno e a do professor), permitindo que o aluno descubra trajetórias eficientes que podem desviar do professor no meio do caminho, mas convergem para o mesmo manifold de dados.

3. Contribuições Chave

1. Novo Paradigma de Acoplamento: Substitui a parametrização complexa e desacoplada de solvents anteriores por uma formulação de pesos relaxada que acopla implicitamente a agregação de gradientes com o ajuste de tamanho de passo.
2. Mecanismo de Calibração de Tempo: Combina deslocamento e escala de tempo para manter o alinhamento temporal com a dinâmica interna do modelo de difusão, essencial para a estabilidade em poucos passos.
3. Eficiência e Desempenho: O método é leve, não requer modificações no modelo de difusão subjacente e possui um custo de treinamento mínimo (distilação rápida).

4. Resultados Experimentais

Os autores avaliaram o DyWeight em diversos benchmarks, incluindo CIFAR-10, FFHQ, ImageNet, LSUN-Bedroom, Stable Diffusion e o modelo de ponta FLUX.1-dev.

• FID (Fréchet Inception Distance): O DyWeight estabeleceu um novo estado da arte (SOTA) em quase todos os conjuntos de dados e configurações de NFE.
  • Exemplo: No CIFAR-10 com 3 NFEs, o DyWeight alcançou FID de 8.16, superando significativamente o iPNDM (24.55) e o EPD-Solver (10.40).
  • Exemplo: No FLUX.1-dev (MS-COCO), o DyWeight superou consistentemente DPM-Solver++ e iPNDM tanto em FID quanto em pontuações CLIP e preferências humanas (HPSv2.1, ImageReward).
• Qualidade Visual: As imagens geradas exibem maior fidelidade visual, melhor alinhamento com o prompt e coerência estrutural, especialmente em cenários de 3 a 7 passos.
• Estabilidade: A ablação mostrou que tanto o deslocamento quanto a escala de tempo são essenciais; remover qualquer um degrada o desempenho drasticamente.
• Eficiência de Treinamento: O modelo converge rapidamente (em ~600 imagens de treinamento para CIFAR-10) e requer apenas alguns minutos de treinamento em GPUs modernas.

5. Significado e Impacto

O DyWeight representa um avanço significativo na aceleração de modelos de difusão. Ao demonstrar que coeficientes fixos são subótimos para a dinâmica complexa e não estacionária da difusão, o trabalho valida a necessidade de solvers adaptativos aprendidos.

Sua principal inovação é a simplificação da otimização: ao acoplar implicitamente o ajuste de tempo à pesagem de gradientes, elimina a necessidade de otimizações bi-níveis complexas ou previsões explícitas de agendamento de tempo. Isso torna o método escalável e aplicável a modelos de difusão modernos e complexos (como FLUX e Stable Diffusion), permitindo a geração de imagens de alta qualidade em tempo real ou com latência extremamente baixa, removendo uma barreira crítica para aplicações interativas.