Resumo Técnico: FlowADMM: ADMM Plug-and-play com Priors de Renoise-Denoise Baseados em Fluxo

Declaração do Problema

Problemas inversos em imageamento, como remoção de ruído, desblur, super-resolução e inpainting, são tipicamente mal-postos e formulados como a recuperação de uma imagem $x$ a partir de medições degradadas $y = Ax + \nu$. Embora a aprendizagem profunda ofereça mapeamentos diretos poderosos, eles frequentemente carecem de interpretabilidade e garantias teóricas. Abordagens variacionais clássicas formulam isso como a minimização de um termo de fidelidade aos dados mais um termo de regularização, frequentemente resolvido via algoritmos proximais como a Divisão Forward-Backward (FBS).

Métodos Plug-and-Play (PnP) melhoram isso substituindo o operador proximal de um regularizador manual por um removedor de ruído aprendido. Recentemente, modelos de fluxo generativo e difusão forneceram famílias dependentes do tempo de removedores de ruído $\{D_t\}$ adaptados a níveis de ruído específicos. No entanto, surge um desafio chave: durante a reconstrução iterativa, as iterações intermediárias frequentemente se desviam da distribuição de ruído $p_t$ para a qual o removedor de ruído foi treinado. Métodos PnP baseados em fluxo existentes (por exemplo, PnP-Flow, Flower) abordam isso incorporando etapas explícitas de renoise-denoise, onde a iteração atual é perturbada em direção ao nível de ruído apropriado antes da remoção de ruído. Esses métodos frequentemente fazem a média sobre múltiplas amostras de ruído para estabilizar o processo.

Apesar de seu sucesso empírico, os métodos PnP de renoise-denoise existentes são formulados como algoritmos estocásticos, tornando seu comportamento de convergência difícil de analisar teoricamente. Especificamente, o operador determinístico implicitamente induzido pela média da saída do removedor de ruído sobre a distribuição de ruído latente não foi caracterizado, deixando as propriedades de convergência desses métodos pouco compreendidas.

Metodologia

O Operador Médio Renoise-Denoise

Os autores identificam e formalizam o operador determinístico subjacente aos métodos PnP estocásticos baseados em fluxo. Dada uma rede de velocidade de fluxo pré-treinada $v_\theta$ definindo um removedor de ruído MMSE dependente do tempo $D_t = I + (1-t)v_\theta^t$, o passo estocástico padrão envolve rebarulhar uma iteração $x$ com ruído $\epsilon \sim \mathcal{N}(0, I)$ para obter $tx + (1-t)\epsilon$, seguido de remoção de ruído.

O artigo define o operador médio renoise-denoise $\bar{S}_t(x)$ como a expectativa desse processo sobre a distribuição de ruído:
$$\bar{S}_t(x) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)} [D_t(tx + (1-t)\epsilon)]$$
Este operador representa o limite determinístico das atualizações estocásticas usadas por métodos como PnP-Flow. Os autores analisam sua forma residual $R_t(x) = \bar{S}_t(x) - x$, interpretando-a como um campo vetorial fornecendo correções locais em direção à variedade de dados. Sob suposições ideais de MMSE, pontos na variedade de dados são pontos fixos deste operador.

Algoritmo FlowADMM

Com base nessa percepção, os autores propõem o FlowADMM, que integra o operador médio renoise-denoise no framework do Método de Direção Alternada dos Multiplicadores (ADMM). Diferentemente do FBS, o ADMM permite múltiplos operadores de prioridade e lida com termos de dados não diferenciáveis sem restrições de passo dependentes do operador.

O algoritmo resolve o problema restrito $\min_{x,z} F_y(x) + R(z)$ sujeito a $x=z$ minimizando o Lagrangiano aumentado. As atualizações são:

1. Consistência dos Dados: $x^{k+1} = \text{prox}_{\tau F_y}(z^k - u^k)$
2. Atualização da Prioridade: $z^{k+1} = \bar{S}_{t_k}(x^{k+1} + u^k)$
3. Atualização Dual: $u^{k+1} = u^k + x^{k+1} - z^{k+1}$

Como a expectativa em $\bar{S}_t$ é intratável, ela é aproximada via amostragem de Monte Carlo. Os autores propõem uma estratégia de amostragem dependente do passo de tempo: usar menos amostras nas iterações iniciais (onde correções grosseiras são feitas e o operador é menos sensível ao ruído) e aumentar o número de amostras nas iterações posteriores (onde detalhes finos são refinados e a precisão é crítica).

Análise de Convergência

O artigo estabelece garantias de convergência para o FlowADMM sob condições de Lipschitz fracas na rede de fluxo subjacente:

• Programação de Tempo Fixa: Para um passo de tempo constante $t$, a convergência é garantida se a constante de Lipschitz da rede de fluxo $L_v$ satisfizer $L_v < \frac{1}{1-t}$. Esta condição é significativamente mais fraca que a contração estrita ($L < 1$) exigida por muitos outros métodos PnP, especialmente quando $t \to 1$.
• Programações de Tempo Não Estacionárias: A análise é estendida a programações de tempo $(t_k)_k$ que convergem para um limite $t_{max} < 1$. Os autores provam a convergência para um ponto fixo do operador limitante $T_{t_{max}}$, desde que a sequência de operadores se aproxime do limite suficientemente rápido e a rede de fluxo seja contínua de Lipschitz em relação ao tempo $t$.

Contribuições Principais

1. Caracterização Teórica: Os autores identificam e analisam o operador médio renoise-denoise determinístico subjacente aos métodos PnP estocásticos baseados em fluxo, caracterizando sua estrutura residual e comportamento de Lipschitz.
2. Proposta de Algoritmo: Eles propõem o FlowADMM, um algoritmo PnP que integra diretamente este operador médio no framework ADMM.
3. Garantias de Convergência: Eles derivam critérios de convergência para o FlowADMM sob suposições fracas de Lipschitz, cobrindo tanto programações de passo de tempo constantes quanto não constantes (variáveis no tempo).
4. Estratégia de Amostragem: Eles propõem uma estratégia de amostragem de Monte Carlo de grosso para fino que aumenta o número de amostras nas iterações posteriores, motivada pelo comportamento do operador médio.
5. Desempenho Empírico: Eles demonstram desempenho de última geração em múltiplos problemas inversos, exigindo menos avaliações de consistência de dados do que abordagens anteriores de média estocástica.

Resultados Experimentais

Os autores avaliaram o FlowADMM nos conjuntos de dados CelebA e AFHQ-Cat em cinco tarefas: remoção de ruído, desblur, super-resolução, inpainting aleatório e inpainting em caixa.

• Desempenho: O FlowADMM alcançou resultados de última geração em métricas de distorção (PSNR e SSIM) em 7 de 10 tarefas em comparação com métodos PnP recentes baseados em fluxo e difusão (incluindo PnP-Flow, Flower, OT-ODE, D-Flow e DiffPIR). Melhorias notáveis foram observadas nas tarefas de super-resolução e inpainting.
• Eficiência: Embora igualando o número total de avaliações de fluxo (passagens forward da rede) do principal competidor Flower5-OT, o FlowADMM exigiu significativamente menos avaliações do operador proximal do termo de dados. Isso ocorre porque o Flower5-OT faz a média de trajetórias completas de reconstrução, enquanto o FlowADMM faz a média apenas do passo de prioridade dentro de cada iteração.
• Estudos de Ablação: Experimentos confirmaram que a programação de amostragem "late-heavy" (aumentando amostras nas iterações posteriores) melhora significativamente o desempenho sobre a amostragem constante. Além disso, aplicar esta programação ao PnP-Flow melhorou seus resultados, mas o FlowADMM ainda o superou, sugerindo que a própria formulação ADMM fornece uma vantagem.
• Limitações: A análise de convergência depende das constantes de Lipschitz da rede pré-treinada, que não são explicitamente penalizadas durante o treinamento. Adicionalmente, as garantias teóricas aplicam-se ao operador médio determinístico, enquanto a implementação prática usa aproximações estocásticas de Monte Carlo.

Significado e Alegações

O artigo alega que, ao formalizar explicitamente o operador determinístico subjacente aos métodos estocásticos de renoise-denoise, ele preenche a lacuna entre o sucesso empírico e a compreensão teórica em PnP baseado em fluxo. O trabalho fornece:

• Convergência Provável: Diferentemente de métodos PnP estocásticos anteriores, o FlowADMM vem com garantias de convergência sob condições de Lipschitz realistas na rede de fluxo.
• Eficiência Algorítmica: A integração no ADMM permite um uso mais eficiente de recursos computacionais (menos passos de consistência de dados) em comparação com métodos de média de trajetória.
• Insight Teórico de Operadores: A análise sugere que o comportamento "de grosso para fino" do operador médio pode ser explorado através de programações de amostragem adaptativas, um princípio que beneficia tanto abordagens baseadas em ADMM quanto em FBS.

Os autores concluem que o FlowADMM oferece um framework robusto e teoricamente fundamentado para resolver problemas inversos com priores baseados em fluxo, embora reconheçam que trabalhos futuros poderiam focar em impor restrições de Lipschitz durante o treinamento para fortalecer as garantias de estabilidade.