Normalizing Flows with Iterative Denoising

Título: Normalizing Flows with Iterative Denoising (iTARFlow)

Autores: Tianrong Chen, Jiatao Gu, David Berthelot, Joshua Susskind, Shuangfei Zhai (Apple Machine Learning Research).

1. O Problema

O artigo aborda as limitações atuais dos modelos generativos de imagem, especificamente no contexto dos Fluxos de Normalização (Normalizing Flows - NFs):

• Dilema do Ruído (Noise Dilemma): Os NFs baseados em TARFlow (Transformers Autoregressive Flows) enfrentam um compromisso difícil durante o treinamento.
  • Se o nível de ruído injetado no treinamento for baixo, o modelo gera imagens com texturas locais ricas, mas com estrutura global pobre e instabilidade no treinamento.
  • Se o nível de ruído for alto, o modelo captura melhor a estrutura global, mas produz imagens desfocadas e com artefatos visíveis após o processo de "auto-denoising" (denoising automático), devido ao comportamento previsto pelo Lema de Tweedie.
• Limitações de Escalabilidade e Qualidade: Enquanto os modelos de Difusão alcançam qualidade de ponta (SOTA), eles exigem centenas de passos iterativos e têm comportamento de escalabilidade menos previsível. Os modelos autoregressivos discretos (como VQ-VAE + Transformer) sofrem com gargalos de informação devido à quantização. Os NFs contínuos anteriores não conseguiam equilibrar alta fidelidade visual com escalabilidade eficiente.

2. Metodologia: iTARFlow

Os autores propõem o iTARFlow (iterative TARFlow), uma extensão do TARFlow que integra um processo de denoising iterativo inspirado em modelos de difusão, mantendo um objetivo baseado em verossimilhança (likelihood-based) totalmente end-to-end.

Principais Componentes:

1. Treinamento Multi-Nível de Ruído:

   • Ao contrário do TARFlow original que usa um nível fixo de ruído, o iTARFlow é treinado em um contínuo de níveis de ruído $t \in [t_{min}, t_{max}]$.
   • O objetivo é otimizar a função de verossimilhança logarítmica $\log p_\theta(x_t, t)$, onde $x_t = x + t\epsilon$.
   • Isso permite que o modelo aprenda funções de pontuação (score functions) $\nabla_x \log p_\theta(x_t, t)$ para múltiplas escalas de ruído simultaneamente.

2. Geração Autoregressiva (Fase 1):

   • A amostragem começa gerando uma imagem ruidosa $x_{t_{max}}$ de forma autoregressiva (bloco a bloco/patch a patch), utilizando a arquitetura Transformer causal do TARFlow.
   • Utiliza-se Classifier-Free Guidance (CFG) apenas nesta fase para melhorar a coerência global.

3. Denoising Iterativo Paralelo (Fase 2):

   • Após a geração da imagem ruidosa, o modelo aplica um processo de denoising iterativo para refinar a imagem.
   • Diferente da geração autoregressiva, esta etapa é totalmente paralelizável.
   • O processo segue uma Equação Diferencial Ordinária (ODE) probabilística:
     $$dx_t = -t \nabla_x \log p_\theta(x_t, t) dt$$
   • O gradiente (score) é calculado diretamente via diferenciação automática da função de verossimilhança do modelo em relação à entrada.
   • Apenas 5 a 10 passos de ODE são necessários para atingir alta qualidade visual.

3. Contribuições Principais

• Identificação e Caracterização do "Dilema do Ruído": O trabalho demonstra sistematicamente como níveis de ruído diferentes afetam a estrutura global versus o detalhe local em NFs, propondo uma solução que combina os benefícios de ambos.
• Novo Paradigma de Treinamento e Inferência: Introdução de um esquema de treinamento multi-ruído combinado com um pipeline de inferência híbrido (autoregressivo para estrutura global + denoising iterativo paralelo para refinamento).
• Eficiência e Compatibilidade: O iTARFlow mantém a eficiência de treinamento dos NFs (convergência mais rápida que modelos de difusão) e demonstra compatibilidade surpreendente com backbones de difusão pré-treinados (como DiT), algo que variantes anteriores do TARFlow não conseguiam fazer bem.

4. Resultados Experimentais

O modelo foi avaliado no dataset ImageNet nas resoluções de 64, 128 e 256 pixels, utilizando a métrica FID (Fréchet Inception Distance).

• ImageNet-64: O iTARFlow-B (770M parâmetros) alcançou um FID de 1.68, superando o TARFlow original (2.66) e competindo diretamente com modelos de difusão como EDM-SDE (1.55), mas com menos parâmetros e arquitetura diferente.
• ImageNet-128: O iTARFlow-L alcançou FID de 2.44, superando significativamente o TARFlow anterior (5.03) e competindo com modelos de difusão como Simple Diff (1.94).
• ImageNet-256: O iTARFlow-XL alcançou FID de 3.32, reduzindo drasticamente a lacuna de desempenho entre NFs e modelos de difusão SOTA (como SiD2 patch 1 com 1.38).
• Eficiência de Treinamento: O modelo convergiu em cerca de 600 épocas, enquanto modelos comparáveis de difusão (DiT) exigiram 1400 épocas.
• Análise de Falhas: O estudo identificou dois artefatos principais: fundos colapsados (pretos) e imagens desfocadas (especialmente no canto inferior direito, correspondente ao primeiro token gerado). Isso sugere que a falta de contexto condicional no primeiro token autoregressivo é um gargalo.

5. Significado e Impacto

• Avanço nos Fluxos de Normalização: O trabalho revitaliza os NFs como uma alternativa viável e competitiva aos modelos de difusão e autoregressivos discretos, provando que é possível obter alta fidelidade visual sem quantização e com escalabilidade previsível.
• Unificação de Paradigmas: O iTARFlow demonstra que é possível integrar a estrutura autoregressiva (escalável) com a capacidade de refinamento iterativo (alta qualidade) de forma nativa, sem depender de modelos auxiliares (como VAEs finetuned) para o denoising.
• Limitações e Futuro: A principal limitação atual é o alto consumo de memória GPU durante a fase de denoising iterativo (devido à diferenciação automática em toda a rede). O artigo sugere que resolver os artefatos de "primeiro token" e otimizar a memória são os próximos passos para alcançar o estado da arte absoluto.

Em resumo, o iTARFlow representa um marco significativo ao resolver o dilema de ruído dos Fluxos de Normalização, oferecendo um modelo generativo contínuo que combina a velocidade de amostragem e escalabilidade dos autoregressivos com a qualidade visual dos modelos de difusão.