Resumo Técnico: Guia Amortizado para Inpainting de Imagem com Modelos de Difusão Pré-treinados (AID)

Declaração do Problema

O inpainting de imagem é uma tarefa de geração condicional cujo objetivo é reconstruir regiões ausentes de uma imagem, mantendo a consistência com os pixels visíveis. Embora modelos de difusão pré-treinados tenham avançado significativamente essa capacidade, as abordagens existentes enfrentam uma dicotomia entre eficiência e adaptabilidade:

1. Treinamento Específico para Tarefa: Métodos como LaMa, MAT e BrushNet treinam modelos dedicados ou grandes módulos de adaptação. Embora isso amortize o custo de inferência, exige tempo substancial de treinamento e adiciona complexidade significativa ao modelo.
2. Otimização por Instância: Métodos como RePaint, DPS e MCG mantêm a base de difusão pré-treinada fixa, mas a adaptam separadamente para cada imagem mascarada no momento do teste. Isso preserva o modelo pré-treinado, mas incorre em custos de otimização repetidos durante a implantação, tratando cada imagem como um novo problema de inferência.

Existe uma lacuna para uma solução de "meio-termo" que amortize a inferência sem treinar um grande modelo específico para a tarefa ou realizar otimização por instância no momento do teste.

Metodologia: Inpainting Amortizado com Difusão (AID)

Os autores propõem o AID, um framework que mantém a base de difusão pré-treinada fixa e treina um módulo de guia pequeno e reutilizável offline. Este módulo é então reutilizado em todas as imagens mascaradas no momento do teste, sem otimização adicional.

1. Formulação de Controle

O problema é formulado como um problema de controle determinístico.

• Dinâmica: O processo de difusão reverso é governado por uma EDO de fluxo de probabilidade impulsionada por uma função de pontuação pré-treinada fixa $\hat{S}$.
• Guia: Um termo de controle dependente da tarefa $u(t, x; \xi)$ é adicionado à deriva, onde $\xi = (M, y)$ representa a máscara observável e os pixels visíveis. O controle não depende da imagem limpa oculta $x^\dagger$.
• Objetivo: O objetivo é minimizar uma perda terminal $\Psi$ (equilibrando a fidelidade aos pixels visíveis e à região ausente) mais um custo de execução que penaliza a magnitude do campo de guia. Isso leva a uma equação de Hamilton-Jacobi-Bellman (HJB) onde o controle ótimo $u^*$ é determinístico.

2. Ponte Teórica: Determinístico para Estocástico

Resolver diretamente a equação HJB determinística de alta dimensão é intratável. Os autores introduzem um problema de controle aleatorizado auxiliar usando políticas Gaussianas para tornar o problema aprendível via Aprendizado por Reforço de Tempo Contínuo (CT-RL).

• Problema Auxiliar: Em vez de um controle determinístico, o sistema usa uma política Gaussiana $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$, onde a média $\mu$ é a lei de guia candidata.
• Ponte Preservadora de Otimizador: O artigo prova dois teoremas críticos:
  1. Suficiência: O guia determinístico ótimo $u^*$ pode ser elevado a uma política Gaussiana ótima definindo a média da política como $u^*$.
  2. Necessidade: Qualquer política Gaussiana ótima para o problema auxiliar deve ter uma média que recupera exatamente o guia determinístico ótimo $u^*$.
• Significado: Esta ponte garante que aprender uma política estocástica (o que permite atualizações baseadas em gradiente) não altera o guia determinístico alvo usado na implantação. A variância da política Gaussiana é fixa e serve apenas como dispositivo de aprendizado, não como regularizador de entropia.

3. Algoritmo Atores-Críticos Amortizado

Com base na ponte teórica, os autores desenvolvem um algoritmo totalmente orientado a dados de atores-críticos:

• Parametrização: O ator (média de guia $\hat{\mu}_\phi$) e o crítico (função de valor $\hat{V}_\theta$) são parametrizados por redes neurais.
• Aprendizado: O algoritmo usa condições de ortogonalidade de martingale derivadas da equação HJB auxiliar. Ele atualiza o ator e o crítico usando aproximação estocástica baseada em trajetórias amostradas da distribuição da tarefa.
• Implantação: Uma vez treinado, o módulo de guia é implantado deterministicamente. Para uma nova imagem mascarada, o sistema usa a média de guia aprendida $\hat{\mu}_\phi$ dentro do solucionador de tempo reverso, descartando o ruído da política.

Contribuições Principais

1. Metodologia: O AID introduz um framework de guia controlado que amortiza a inferência treinando um módulo de guia leve offline, que é então reutilizado em tarefas sem otimização por instância.
2. Teoria: O artigo estabelece uma rigorosa "ponte de equivalência de política" entre o problema de guia determinístico e uma formulação Gaussiana auxiliar. Ele prova que o problema auxiliar é preservador de otimizador, justificando o uso de métodos estocásticos de atores-críticos para aprender um campo de controle determinístico.
3. Eficiência: A abordagem adiciona menos de 1% de parâmetros treináveis em relação à base pré-treinada, enquanto melhora significativamente o compromisso entre qualidade e velocidade.

Resultados Empíricos

Os autores avaliaram o AID no AFHQv2, FFHQ (usando o pipeline EDM no espaço de pixels) e ImageNet (usando o pipeline EDM2 no espaço latente).

• Desempenho vs. Baselines de Base Fixa: O AID supera consistentemente fortes baselines de base fixa (Sem Guia, Substituição, MCG, DPS) em múltiplas métricas (PSNR, SSIM, LPIPS) e tipos de máscara (livre-forma, central, faixa).
• Desempenho vs. RePaint: O AID supera o RePaint (uma forte baseline de otimização iterativa) em qualidade, usando apenas ~10% do orçamento de amostragem (por exemplo, 35 passos vs. 350 passos para o RePaint).
• Desempenho vs. LatentPaint: No ImageNet, o AID supera o LatentPaint (um método de adaptação reutilizável no espaço latente) em todas as métricas, usando um módulo de guia com menos de 45% dos parâmetros treináveis.
• Latência: O AID suporta um ponto de operação de baixa latência (12 passos) que reduz ainda mais o tempo de inferência, mantendo desempenho superior às baselines que não são RePaint e permanecendo competitivo com o RePaint.
• Generalização: O mesmo módulo de guia treinado transfere-se efetivamente para tipos de máscara não vistos e diferentes orçamentos de amostragem sem retreinamento.

Significado e Alegações

O artigo afirma que o AID é o primeiro framework de guia controlado amortizado para inpainting de imagem que:

1. Mantém a base de difusão pré-treinada congelada.
2. Treina apenas um pequeno módulo de guia reutilizável.
3. Fornece uma ponte exata preservadora de otimizador do aprendizado estocástico de atores-críticos ao campo de guia determinístico usado na implantação.

Os autores posicionam este trabalho como um passo em direção ao uso de frameworks baseados em controle e RL para fornecer métodos fundamentados teoricamente para estudar e melhorar modelos de difusão. Eles enfatizam que o RL de tempo contínuo fundamenta tanto a análise rigorosa quanto o design prático do algoritmo, oferecendo uma alternativa principiante à otimização heurística por instância ou ao pesado retreinamento específico para tarefa.