Improving Diffusion Posterior Samplers with Lagged Temporal Corrections for Image Restoration

Resumo Técnico: Melhoria de Amostradores de Posterior Difusivo com Correções Temporais Atrasadas (LAMP)

Declaração do Problema

A amostragem de posterior baseada em difusão (PS) emergiu como um quadro líder para resolver problemas inversos de imageamento (por exemplo, desfocagem, super-resolução, inpainting). Esses métodos combinam um prior de difusão aprendido com restrições de medição para amostrar a distribuição posterior $p(x_0|y)$, onde $y = Kx_0 + e$. Formulações padrão de PS, como DiffPIR e DDRM, operam atualizando iterativamente a estimativa atual $x_t$ usando uma etapa de remoção de ruído do modelo de difusão e uma correção de consistência de dados $D_t$.

No entanto, os autores identificam uma limitação estrutural fundamental nessas abordagens padrão: elas dependem de estimativas de consistência de dados instantâneas e passo a passo ($D_t$) que são recalculadas a cada iteração reversa. Esse design introduz variabilidade temporal e instabilidade na dinâmica reversa. Sob uma perspectiva de sistemas dinâmicos, as atualizações padrão de PS correspondem a uma discretização de primeira ordem da dinâmica de difusão mais um termo de correção residual. Esse termo residual atua como um sinal de forçamento variante no tempo que perturba a trajetória, fazendo com que ela se desvie da dinâmica contínua suave e bem definida assumida por solucionadores de alta ordem (como DPM-Solver) usados na geração incondicional. Consequentemente, a falta de memória temporal nos métodos padrão de PS leva a atualizações inconsistentes e estabilidade subótima, particularmente em regimes mal-postos.

Metodologia: LAMP

O artigo propõe LAMP (LAgged Multistep Posterior), um esquema modular de plug-in que aborda essas instabilidades introduzindo uma correção temporal atrasada.

1. Reinterpretação Dinâmica: Os autores reinterpretem a atualização padrão de PS como uma discretização de primeira ordem de uma dinâmica de difusão forçada. Eles mostram que a atualização $x^{PS}_{t-\Delta t}$ pode ser decomposta em uma etapa exponencial de primeira ordem impulsionada por um alvo consciente da medição $D_t$ e um termo de correção residual proporcional ao desajuste $(D_t - \hat{x}_{0|t})$.
2. Discretização de Segunda Ordem: Para melhorar a estabilidade, o LAMP emprega uma discretização de segunda ordem da dinâmica subjacente. Ao contrário de solucionadores de alta ordem padrão que assumem uma EDO suave, o LAMP leva explicitamente em conta o termo de forçamento dependente do residual inerente a problemas inversos.
3. Mecanismo de Correção Atrasada: A inovação central é a introdução de uma correção temporal baseada na variação de estimativas consecutivas. A atualização do LAMP combina a discretização de segunda ordem com a correção residual:
   $$x^{LAMP}_{t-\Delta t} = x^{PS}_{t-\Delta t} + \alpha_{t-\Delta t} A_1(h) \gamma \frac{D_t - D_{t+\Delta t}}{h_{prev}}$$
   Aqui, o termo $(D_t - D_{t+\Delta t})/h_{prev}$ atua como uma derivada temporal discreta. Ao escolher um coeficiente negativo $\gamma$, o método aplica uma correção atrasada, puxando a estimativa atual ligeiramente em direção às anteriores em vez de extrapolar. Isso suaviza a trajetória e contrabalança a variabilidade induzida por correções instantâneas.
4. Preservação Estrutural: Crucialmente, os autores provam (Proposição 1) que o LAMP preserva a estrutura padrão de amostragem de posterior. A atualização pode ser reescrita na forma canônica $x_{t-\Delta t} = \alpha_{t-\Delta t} \tilde{D}_t + \sigma_{t-\Delta t} \epsilon_\Theta(x_t, t)$, onde $\tilde{D}_t$ é uma estimativa filtrada temporalmente definida como uma combinação convexa das estimativas de consistência de dados atuais e anteriores: $\tilde{D}_t = (1-\beta_t)D_t + \beta_t D_{t+\Delta t}$. Isso garante que o LAMP seja uma extensão modular compatível com qualquer espinha dorsal de PS existente (por exemplo, DiffPIR, DDRM) sem modificar o removedor de ruído ou o procedimento de treinamento.

Análise Teórica

O artigo fornece uma análise de risco de um passo (Proposição 2) que caracteriza o compromisso viés-variância da correção atrasada. A análise mostra que o LAMP melhora a transição reversa quando a redução na variância do erro de estimativa (devido ao suavização temporal) supera o viés introduzido pelo atraso (devido à deriva do alvo posterior). Os autores demonstram que essa condição é tipicamente satisfeita nos estágios posteriores do processo de difusão, motivando uma estratégia híbrida onde o LAMP é ativado após uma breve fase de aquecimento.

Contribuições Principais

• Identificação de Incompatibilidade Estrutural: Os autores destacam que interpretações padrão baseadas em EDOs de solucionadores de difusão falham em capturar os termos de forçamento dependentes do residual presentes na amostragem de posterior de problemas inversos, levando à instabilidade temporal.
• Framework LAMP: Eles propõem o LAMP, um esquema corrigido de segunda ordem que integra uma correção temporal atrasada na atualização de amostragem de posterior.
• Modularidade e Compatibilidade: O LAMP é provado ser um amostrador de posterior válido que pode ser implementado como um plug-in sobre qualquer método de PS existente (por exemplo, DiffPIR, DDRM) sem retreinamento ou alteração do modelo direto.
• Justificativa Teórica: Uma análise de erro local estabelece as condições sob as quais a correção atrasada reduz o risco, especificamente quando os alvos posteriores adjacentes evoluem suavemente em relação ao ruído de estimativa.

Resultados Experimentais

Os autores avaliam o LAMP em problemas inversos lineares padrão (desfocagem gaussiana, desfocagem de movimento e super-resolução) em três conjuntos de dados: FFHQ, ImageNet e CelebA. Eles comparam versões aprimoradas com LAMP (DiffPIR-LAMP, DDRM-LAMP) contra suas linhas de base e variantes ingênuas de alta ordem (1M, 2M).

• Desempenho: O LAMP melhora consistentemente a qualidade de reconstrução sobre linhas de base fortes. No cenário ruidoso ($\sigma_y = 0.05$), o DDRM-LAMP alcança o melhor desempenho geral, mostrando ganhos consistentes em PSNR, SSIM e LPIPS em todas as tarefas. O DiffPIR-LAMP também mostra melhorias, particularmente em desfocagem e super-resolução.
• Eficiência: Essas melhorias são alcançadas sem aumentar o número de avaliações de remoção de ruído (NFE), mantendo o custo computacional dos métodos base (por exemplo, 20 NFE para DDRM, 100 NFE para DiffPIR).
• Comparação com Alta Ordem Ingênua: Discretizações ingênuas de EDO de primeira e segunda ordem (sem a correção residual específica) degradam consistentemente o desempenho, confirmando que a correção residual atrasada proposta é o fator chave para a melhoria.
• Resultados Qualitativos: Comparações visuais mostram que o LAMP produz texturas mais suaves, limites mais nítidos e melhor preservação de detalhes finos (por exemplo, textura de cabelo em CelebA, escamas em imagens de Iguana) em comparação com as linhas de base, reduzindo granulação e artefatos de desfocagem.

Significado e Afirmações

O artigo afirma que o filtragem temporal de estimativas conscientes da medição é um mecanismo crítico, anteriormente pouco explorado, para estabilizar solucionadores inversos baseados em difusão. Ao reformular a amostragem de posterior como um sistema dinâmico com um termo de forçamento dependente do residual, os autores demonstram que uma simples correção atrasada pode melhorar significativamente a qualidade e a estabilidade da reconstrução.

O significado reside na modularidade da abordagem: o LAMP não requer novos priors, retreinamento ou estratégias complexas de condicionamento. Em vez disso, oferece um aprimoramento de propósito geral e baixo custo para amostradores de posterior existentes de última geração. Os autores observam modestamente que o LAMP herda as limitações de seus métodos de espinha dorsal (por exemplo, não pode compensar incompatibilidades fundamentais do modelo) e que sua eficácia depende do comportamento temporal da sequência de correção do amostrador subjacente. No entanto, os resultados sugerem que, para linhas de base estáveis, o LAMP fornece uma melhoria consistente e perceptualmente significativa na fidelidade dos dados, particularmente para estruturas de pequena escala.