Denoising, Fast and Slow: Difficulty-Aware Adaptive Sampling for Image Generation

1. Problema Identificado

Os modelos modernos de geração de imagem baseados em difusão e flow matching (como DiT e SiT) geralmente alocam recursos computacionais de forma uniforme em todo o espaço da imagem. Em cada passo de desnoising (remoção de ruído), todas as regiões da imagem são atualizadas com o mesmo nível de ruído global e o mesmo número de avaliações de função (NFE).

O artigo argumenta que essa abordagem ignora a heterogeneidade das imagens naturais:

• Regiões Fáceis: Grandes áreas de fundo, cores saturadas ou texturas simples são fáceis de desnoising e não necessitam de muitos refinamentos.
• Regiões Difíceis: Estruturas finas, bordas de objetos, texto pequeno ou oclusões permanecem ambíguas até os estágios finais do processo e beneficiam-se de mais refinamento e contexto adicional.

Tratar todas as regiões de forma idêntica força o modelo a desperdiçar computação em áreas fáceis, enquanto subaloca recursos para áreas que necessitam de mais contexto para resolver ambiguidades.

2. Metodologia: Patch Forcing (PF)

Os autores propõem o Patch Forcing (PF), um framework que permite que diferentes regiões de uma imagem sigam trajetórias de desnoising distintas, baseadas na dificuldade local. A metodologia baseia-se em três pilares principais:

A. Escalas de Ruído por Patch (Patch-level Noise Scales)

Em vez de um timestep global $t$, o modelo utiliza um timestep específico para cada patch (bloco de tokens) da imagem. Isso permite que certas regiões "avançem" mais rápido no processo de desnoising do que outras.

B. Desafio do Treinamento vs. Inferência (O "Train-Test Gap")

O artigo identifica um problema crítico ao tentar aplicar Diffusion Forcing (técnicas anteriores para vídeo) em imagens:

• Se os timesteps dos patches forem amostrados independentemente e uniformemente, a média dos timesteps por amostra tende a se concentrar em $t \approx 0.5$. Isso significa que, durante o treinamento, o modelo frequentemente vê patches já quase limpos (com pouco ruído), criando um "vazamento de contexto" que não ocorre na inferência (que começa com ruído puro).
• Solução (Amostrador LTG): Os autores introduzem um amostrador de timesteps chamado Logit-Normal Truncated Gaussian (LTG).
  1. Eles primeiro amostram um $t_{max}$ (o timestep máximo de ruído permitido para aquela amostra) usando uma distribuição Logit-Normal.
  2. Em seguida, amostram os timesteps individuais dos patches ($t_i$) a partir de uma metade inferior de uma distribuição Gaussiana truncada em $t_{max}$.
  • Resultado: Isso garante que o modelo seja treinado sem patches excessivamente limpos, alinhando a distribuição de treinamento com a inferência (que começa do ruído total), enquanto ainda permite heterogeneidade espacial.

C. Cabeça de Dificuldade (Uncertainty Head)

Para saber quais patches devem avançar mais rápido, o modelo é equipado com uma cabeça leve que prevê a incerteza (ou dificuldade) de cada patch.

• A incerteza é modelada como a variância prevista da estimativa de velocidade do modelo.
• Correlação: Patches com alta incerteza correspondem a regiões difíceis (bordas, detalhes finos) e patches com baixa incerteza correspondem a regiões fáceis.
• Mecanismo de Contexto: Regiões "confiantes" (baixa incerteza) avançam mais rápido no tempo de desnoising e fornecem contexto limpo para as regiões "incertas" (alta incerteza), ajudando a resolver suas ambiguidades.

D. Estratégias de Amostragem Adaptativa

Com base na previsão de incerteza, o PF introduz duas estratégias de inferência adaptativa:

1. Dual-Loop (Duplo Loop):
   • Loop Externo: Avança os patches confiantes com passos de timestep maiores.
   • Loop Interno: Usa o estado atualizado dos patches confiantes como contexto para refinar os patches incertos com passos menores.
   • O processo se repete, ajustando as prioridades dinamicamente.
2. Look-Ahead (Antecipação):
   • Projeta os patches confiantes para um futuro timestep ($t + \delta$) para gerar uma representação "limpa" antecipada.
   • Essa representação futura é usada como condição contextual para guiar a geração dos patches incertos no timestep atual.

3. Contribuições Principais

1. Patch Forcing (PF): Um novo framework para síntese de imagem que utiliza timesteps por patch e uma cabeça de dificuldade aprendida para permitir a alocação adaptativa de recursos computacionais.
2. Amostrador LTG: Uma estratégia de amostragem de timesteps durante o treinamento que resolve a discrepância entre treinamento e inferência em modelos de desnoising heterogêneo, superando abordagens anteriores como a do Spatial Reasoning Models (SRM).
3. Amostradores Conscientes de Dificuldade: Demonstração de que usar a incerteza prevista para guiar a ordem de desnoising (avançando regiões fáceis para ajudar as difíceis) melhora significativamente a qualidade da geração.
4. Generalização: O método é ortogonal a técnicas de alinhamento de representação (como REPA) e de guidance (como CFG), podendo ser combinado com elas para ganhos adicionais. Funciona tanto para geração condicional por classe (ImageNet) quanto para text-to-image.

4. Resultados Experimentais

Os autores avaliaram o método no benchmark ImageNet 256x256 e em tarefas de Text-to-Image (T2I):

• ImageNet (Class-Conditional):
  • O modelo PFT (Patch Forcing Transformer) supera consistentemente o baseline SiT (Scalable Interpolant Transformers) mesmo com o mesmo número de parâmetros e passos de amostragem (NFE).
  • Exemplo (SiT-XL/2): FID de 17.2. Com PFT + Look-Ahead: FID de 9.8.
  • A combinação com REPA (Representation Alignment) e Look-Ahead atingiu um FID de 6.7, superando o estado da arte anterior.
• Text-to-Image:
  • Um modelo PFT de 1.2B parâmetros foi treinado no conjunto de dados COYO.
  • O modelo demonstrou superioridade em benchmarks de composição (T2I-CompBench++) e avaliação de objetos (GenEval).
  • Qualidade de Texto: O PF produziu texto mais legível e claro em imagens geradas em comparação com modelos treinados com Flow Matching padrão, sugerindo que a heterogeneidade de ruído ajuda na renderização de detalhes finos.
• Eficiência: As estratégias adaptativas (Dual-Loop e Look-Ahead) oferecem melhor qualidade (FID mais baixo) para o mesmo orçamento computacional (NFE) em comparação com amostragem Euler paralela padrão.

5. Significado e Impacto

O trabalho propõe uma mudança de paradigma na geração de imagens: em vez de tratar a imagem como um bloco monolítico com um ritmo de desnoising único, o Patch Forcing reconhece que a complexidade visual é local.

• Alocação Inteligente de Computação: O método direciona a capacidade computacional para onde é mais necessária (regiões difíceis), permitindo que regiões fáceis forneçam contexto para as difíceis.
• Fundação para Amostragem Adaptativa: Estabelece que agendamentos de desnoising em nível de patch são uma base promissora para uma nova classe de samplers adaptativos.
• Futuro: Os autores sugerem que essa abordagem é particularmente relevante para modelos de poucos passos (few-step) e modelos destilados, onde saltos de tempo maiores podem amplificar os benefícios de avançar progressivamente as regiões confiantes.

Em resumo, o artigo demonstra que a heterogeneidade controlada no processo de desnoising, guiada pela incerteza do modelo, leva a imagens de maior qualidade e mais detalhadas sem aumentar o custo computacional total.