Few-Step Diffusion Sampling Through Instance-Aware Discretizations

1. O Problema

Os Modelos de Probabilidade de Difusão (DPMs) e modelos de Flow Matching geram dados de alta fidelidade simulando trajetórias definidas por Equações Diferenciais Ordinárias (ODEs) ou Estocásticas (SDEs). Para acelerar a geração, utiliza-se a formulação de ODE de fluxo de probabilidade, permitindo o uso de solvers numéricos avançados.

No entanto, a eficiência da amostragem depende criticamente da estratégia de discretização (como o tempo é dividido em passos).

• Limitação Atual: A maioria dos métodos existentes, sejam heurísticos (como Uniform ou LogSNR) ou baseados em otimização, impõe um cronograma de passos de tempo globalmente compartilhado. Isso significa que a mesma sequência de tempos é aplicada a todas as amostras, independentemente da complexidade específica de cada instância de dados.
• A Lacuna: Diferentes entradas (instâncias) podem exigir trajetórias de amostragem distintas com níveis variados de complexidade. Um cronograma global ótimo em média pode ser subótimo para instâncias específicas, limitando a qualidade da geração, especialmente em cenários de poucos passos (few-step).

2. Metodologia: INDIS

Os autores propõem o INDIS (Instance-Specific Discretization), um framework que aprende a adaptar a alocação de passos de tempo com base nas características de cada entrada (o ruído inicial e as condições).

Principais Componentes:

1. Observação Empírica (Dados Sintéticos):

   • Experimentos controlados em dados sintéticos 2D demonstraram que uma estratégia de discretização específica para cada instância (instance-specific) reduz significativamente o erro de ponto final em comparação com cronogramas globais otimizados ou heurísticos.
   • A perda esperada sob um cronograma global serve como um limite superior para a perda alcançável por uma abordagem específica da instância.

2. Rede de Condicionamento Prévio (Prior Conditioning Network):

   • O método introduz uma rede leve, $\phi(\cdot)$, que toma como entrada o ruído inicial $x_T$ e a informação condicional $c$ (como rótulos de classe ou prompts de texto).
   • A saída da rede é um conjunto de parâmetros de discretização personalizados ($\xi_\phi$) para aquela amostra específica.
   • A rede é treinada para minimizar a distância entre a amostra gerada pelo solver com passos personalizados e uma amostra "professora" (gerada por um solver com muitos passos e alta precisão).

3. Adaptações para Escala e Viés:

   • Condicionamento Condicional: A rede aceita tanto rótulos de classe quanto embeddings de texto (ex: CLIP, T5) para guiar a geração.
   • Fatores de Deslocamento e Escala (Shift & Scale Factors): Para mitigar o problema de exposure bias (mismatch entre a distribuição de ruído no treinamento e na amostragem), a rede aprende fatores de deslocamento temporal ($\Delta\tau$) e escala ($\gamma$). Isso permite ajustar a função de avaliação do modelo ($\hat{\epsilon}_\theta$) dinamicamente para cada instância.
   • Formulação de Otimização: O objetivo é minimizar a distância (ex: MSE, LPIPS) entre a saída do solver estudante (com $\xi_\phi$) e a saída do solver professor, condicionada a $x_T$ e $c$.

3. Contribuições Chave

• Identificação da Subotimalidade Global: Demonstração quantitativa de que cronogramas de tempo globais são inerentemente subótimos para instâncias individuais, estabelecendo a motivação teórica para a discretização adaptativa.
• Framework INDIS: Proposta de um método escalável que generaliza a busca de discretização baseada em gradiente para o contexto generativo condicional, aprendendo a mapear o estado inicial e a condição para uma estratégia de tempo ótima.
• Mitigação de Viés de Exposição: Integração de fatores de deslocamento e escala aprendíveis para alinhar melhor a dinâmica de amostragem com a distribuição de dados, superando limitações de métodos anteriores que usam heurísticas fixas.
• Validação Abrangente: Testes extensivos em diversos domínios (dados sintéticos, espaço de pixels, espaço latente) e modalidades (imagens estáticas e vídeo).

4. Resultados Experimentais

O método foi avaliado em diversos modelos pré-treinados (EDM, Stable Diffusion, FLUX.1-dev, LTX-Video) e datasets (CIFAR-10, ImageNet, FFHQ, LSUN, MS-COCO).

• Desempenho em Poucos Passos (Few-Step):
  • Em configurações de baixo NFE (Number of Function Evaluations, ex: 3 a 7 passos), o INDIS superou consistentemente os melhores heurísticos e métodos de otimização global (como DMN, GITS, LD3).
  • Exemplos de Melhoria:
    • CIFAR-10 (Pixel-space): Redução de FID de ~16.52 (LD3) para 9.26 (INDIS) em 3 passos.
    • LSUN-Bedroom (Latent-space): Redução de FID de ~14.62 (LD3) para 12.44 (INDIS) em 3 passos.
    • FLUX.1-dev: Melhoria significativa no FID (44.35 vs 56.82 do baseline global otimizado) e nos escores CLIP/CMMD.
• Eficiência:
  • O custo de inferência adicional é mínimo (apenas uma passagem forward de uma rede leve $\phi$), representando menos de 2.5% do tempo total de amostragem.
  • Não requer re-treinamento do modelo de difusão base, apenas o ajuste fino da rede de condicionamento.
• Qualidade Visual: Resultados qualitativos mostram maior coerência, detalhe e consistência de sujeitos em comparação com métodos globais, especialmente em tarefas complexas de geração de vídeo e imagens de alta resolução.

5. Significado e Conclusão

O trabalho INDIS representa uma mudança de paradigma na aceleração de modelos de difusão. Ao abandonar a suposição de que uma única estratégia de tempo serve para todas as amostras, o método reconhece e explora a variabilidade intrínseca na complexidade de cada trajetória de geração.

• Impacto: Permite gerar amostras de alta qualidade com menos passos computacionais, tornando a geração de imagens e vídeos mais rápida e eficiente sem sacrificar a fidelidade.
• Escalabilidade: O método é leve e pode ser aplicado a modelos pré-treinados existentes (como FLUX e Stable Diffusion) com custo de treinamento de ajuste fino negligenciável.
• Futuro: Abre caminho para solvers que se adaptam dinamicamente não apenas ao tempo, mas à complexidade específica do conteúdo sendo gerado, potencialmente integrando-se com técnicas de distillation e inference-time scaling.

Em resumo, o INDIS demonstra que a "inteligência" na escolha dos passos de tempo, adaptada a cada instância, é um fator crítico para desbloquear o verdadeiro potencial de geração rápida de modelos de difusão.