DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

Resumo Técnico: DIAMOND – Inferência Direcionada para Mitigação de Artefatos em Modelos de Flow Matching

Declaração do Problema
Apesar da alta fidelidade dos recentes modelos de texto-para-imagem, particularmente aqueles baseados em arquiteturas de Rectified Flow (ex: FLUX.2) e Modelos de Difusão (ex: SDXL), as imagens geradas frequentemente exibem artefatos estruturais e anatômicos. Estes incluem mãos distorcidas, membros extras, rostos malformados e texto incorreto. As estratégias de mitigação existentes sofrem de limitações significativas:

• Refinamento post-hoc: Métodos que corrigem imagens após a geração frequentemente falham em intervir durante o processo central de formação.
• Modificações invasivas: Técnicas que exigem ajuste fino (fine-tuning) do modelo ou pesos LoRA são computacionalmente caras, difíceis de escalar e muitas vezes carecem de pesos disponíveis publicamente para arquiteturas modernas.
• Limitações heurísticas: Métodos de detecção em tempo real dependem frequentemente de limiares heurísticos ou substituições aleatórias de latentes, levando a uma localização imprecisa ou correções subótimas.

Metodologia
Os autores propõem o DIAMOND (Directed Inference for Artifact Mitigation in Flow-Matching Models), um framework zero-shot, livre de treinamento, que corrige artefatos durante a trajetória de inferência sem modificar os pesos do modelo. O método opera sob o princípio de reconstruir uma estimativa da amostra limpa em cada passo para guiar a geração para longe de estados latentes propensos a artefatos.

1. Estimativa de Latente Limpo:
   Ao contrário das abordagens padrão que operam diretamente sobre latentes ruidosos ($x_t$), o DIAMOND estima o latente de dados limpos ($\hat{x}_{0,t}$) em cada etapa.

• Para modelos de Rectified Flow (prevendo a velocidade $v_\theta$): $\hat{x}_{0,t} = x_t - t \cdot v_\theta(x_t, t)$.
• Para modelos de Difusão (prevendo o ruído $\epsilon_\theta$): $\hat{x}_{0,t} = x_t - \sigma_t \cdot \epsilon_\theta(x_t, t)$.
  Este latente estimado é decodificado para o espaço de pixels ($\hat{I}_t = D(\hat{x}_{0,t})$) para garantir que o detector de artefatos opere em seu domínio ideal (dados limpos) em vez de em estados intermediários ruidosos.

2. Detecção de Artefato e Perda:
   Um Detector de Artefatos (AD) diferenciável e pré-treinado (ex: do DiffDoctor) analisa $\hat{I}_t$ para gerar uma máscara de probabilidade de artefato $M_t$. Uma perda de artefato ($L_a$) é calculada como o erro quadrático médio entre esta máscara e uma máscara zero (representando uma imagem livre de artefatos):
   $$L_a = \frac{1}{HW} \sum_{i,j} (M_{i,j}^t)^2$$

3. Correção de Trajetória Baseada em Gradiente:
   O método calcula o gradiente da perda em relação ao latente atual $x_t$ ($\nabla_{x_t} L_a$). Para garantir estabilidade e evitar que a magnitude da correção seja dominada pelo tamanho da região do artefato, o gradiente é normalizado:
   $$\delta_t = \lambda_t \frac{\nabla_{x_t} L_a}{\|\nabla_{x_t} L_a\|^2 + \epsilon}$$
   Este vetor de deslocamento $\delta_t$ é subtraído do passo de atualização do solver padrão, efetivamente direcionando a trajetória para longe dos artefatos.

4. Agendamento Dinâmico ($\lambda_t$):
   Um agendamento de lei de potência controla a intensidade da correção ($\lambda_t$). Com base na observação de que artefatos estruturais aparecem em estágios iniciais, dominados por ruído, enquanto detalhes de alta frequência são refinados posteriormente, o agendamento aplica uma correção forte no início do processo de geração e ela decai em direção ao final.

5. Preservação de Identidade (Opcional):
   Para evitar que a correção altere a estrutura global da imagem ou o alinhamento com o prompt, um termo de regularização opcional ($L_{rec}$) pode ser aplicado. Este termo minimiza o Erro Absollyuto Médio (MAE) entre a imagem corrigida e a imagem base em regiões não-artefatos.

Principais Contribuições

• Framework Livre de Treinamento: O DIAMOND é o primeiro método zero-shot que integra o direcionamento da trajetória geométrica diretamente no processo de inferência de modelos de Rectified Flow e Difusão sem exigir treinamento adicional ou modificações de pesos.
• Estimativa On-the-Fly: O método introduz um mecanismo inovador para estimar latentes limpos ($\hat{x}_{0,t}$) durante a inferência, permitindo a detecção confiável de artefatos em etapas intermediárias onde detectores padrão tipicamente falham.
• Generalização: Demonstra-se que a abordagem funciona efetivamente tanto para modelos de Rectified Flow (FLUX.1, FLUX.2) quanto para modelos de Difusão padrão (SDXL).
• Normalização de Gradiente: O artigo identifica e implementa a normalização de gradiente para desacoplar a força de atualização do tamanho espacial do artefato, garantindo uma correção consistente independentemente da extensão do artefato.

Resultados Experimentais
Avaliações quantitativas foram conduzidas em três conjuntos de dados (animais, palavras, pessoas) usando modelos como FLUX.1 [dev], FLUX.1 [schnell], FLUX.2 [dev] e SDXL.

• Redução de Artefatos: No conjunto de dados "words" com FLUX.1 [dev], o DIAMOND reduziu a Frequência Média de Artefatos de uma linha de base de 100% para 9,85%, superando significativamente as bases como DiffDoctor (46,97%) e HPSv2 (56,82%).
• Razão de Pixels: A Razão de Pixels de Artefato (APR) foi reduzida a níveis próximos de zero (ex: 0,009% no conjunto de dados de palavras) em comparação com as bases.
• Preservação de Qualidade: O método manteve pontuações CLIP-T (alinhamento texto-imagem) e ImageReward (preferência humana) elevadas, indicando que a remoção de artefatos não degradou a consistência semântica ou a qualidade visual.
• Preservação de Identidade: Estudos de ablação mostraram que, embora o DIAMOND não regularizado remova artefatos de forma eficaz, adicionar o termo $L_{rec}$ (com $\alpha=0,1$) equilibra de forma ideal a supressão de artefatos com a preservação da identidade do modelo base.

Significância e Alegações
Os autores alegam que o DIAMOND fornece um caminho robusto e zero-shot para a síntese de imagens de alta fidelidade e livres de artefatos. Ao intervir diretamente durante o processo de amostragem através da correção de trajetória, o método supera os custos computacionais do ajuste fino e as imprecisões de heurísticas post-hoc. O artigo enfatiza que esta abordagem permite a supervisão espacial detalhada e a correção precisa de artefatos localizados (como mãos ou texto), mantendo a estrutura global e a aderência ao prompt da imagem gerada. O método é apresentado como uma solução prática para aplicações profissionais onde a correção anatômica é crítica, como em visuais de e-commerce e marketing.