Bridging data-driven priors via the score function for posterior sampling -- Comparative review and experimental study

Resumo Técnico: Unificando Priors Baseados em Dados via Função de Score para Amostragem de Posterior

Enunciado do Problema
Tarefas de restauração de imagem (IR), como denoising (redução de ruído), deblurring (desfoque), inpainting (preenchimento) e super-resolução, são tipicamente formuladas como problemas inversos lineares ($y = Ax + n$). Dentro de um arcabouço Bayesiano, resolver esses problemas mal postos requer a definição de uma distribuição a priori $p(x)$ para regularizar a solução. Embora os priors tradicionais baseados em modelos (ex: Variação Total) sejam bem compreendidos, eles frequentemente falham em capturar as estruturas sofisticadas de imagens naturais. Por outroônica, os priors baseados em dados aprendidos via deep learning (ex: Redes Adversárias Generativas, Autoencoders Variacionais, Fluxos de Normalização) oferecem expressividade superior, mas impõem desafios significativos para a amostragem de posterior.

Um desafio primário na IR Bayesiana é a alta dimensionalidade do espaço da imagem e a não suavidade das funções de regularização, o que dificulta a eficiência dos algoritmos de Cadeia de Markov Monte Carlo (MCMC). Métodos existentes frequentemente dependem do Algoritmo de Langevin Não Ajustado (ULA) para amostrar da posterior. No entanto, o ULA pode enfrentar dificuldades em problemas severamente mal condicionados devido às restrições de tamanho de passo vinculadas à constante de Lipschitz da verossimilhança. Além disso, muitos priors baseados em dados avançados (como aqueles baseados em denoisers ou modelos generativos baseados em score) são definidos implicitamente, tornando a amostragem direta difícil sem um arcabouço teórico unificado.

Metodologia
O artigo propõe um arcabouço unificado para amostrar distribuições posteriores definidas por diversos priors baseados em dados, desde que estes possam ser expressos via sua função de score ($s(x) = \nabla_x \log p(x)$). A metodologia central baseia-se no Langevin-within-Split Gibbs Sampler (LwSGS).

1. Aumentação de Dados Assintoticamente Exata (AXDA): O algoritmo introduz uma variável auxiliar $z$ para decompor a distribuição posterior em uma forma aumentada:
   $$\pi_\rho(x, z) \propto \exp\left[ -f(x, y) - g(z) - \frac{1}{2\rho^2}\|x - z\|^2 \right]$$
   onde $f(x,y)$ é o termo de fidelidade de dados e $g(z)$ é o potencial de regularização. À medida que o parâmetro de acoplamento $\rho \to 0$, a distribuição marginal de $x$ converge para a posterior alvo.

2. Split Gibbs Sampling (SGS): O algoritmo alterna entre a amostragem de $x$ e $z$:

   • Amostragem de $x$: Dado $z$, a distribuição condicional $p(x|y, z)$ é Gaussiana (para problemas inversos Gaussianos lineares), permitindo uma amostragem eficiente e em forma fechada. Crucialmente, este passo desacopla o tamanho do passo da constante de Lipschitz da verossimilhança, abordando uma limitação importante do ULA padrão.
   • Amostragem de $z$: Dado $x$, a distribuição condicional $p(z|x)$ é amostrada usando um passo de Langevin Monte Carlo (LMC). Este passo utiliza explicitamente a função de score do prior:
     $$z^{(t+1)} = z^{(t)} + \gamma \left( s(z^{(t)}) + \frac{1}{\rho^2}(x^{(t)} - z^{(t)}) \right) + \sqrt{2\gamma}\epsilon^{(t)}$$
     onde $s(z) = \nabla \log p(z)$ é a função de score do prior.

3. Representação Unificada de Priors: O artigo demonstra que quatro classes distintas de priors baseados em dados podem ser unificadas sob esta perspectiva de função de score:

   • Regularization-by-Denoising (RED): A função de score é derivada diretamente do resíduo de denoising ($s(x) \propto -(x - D(x))$).
   • Modelos Generativos Baseados em Score (SGM): A função de score é aprendida explicitamente via score matching (ou modelos probabilísticos de difusão de denoising).
   • Fluxos de Normalização (NF): A função de score é computada via diferenciação automática da densidade logarítmica derivada da fórmula de mudança de variáveis.
   • Regularizadores Convex-Ridge (CRR): A função de score corresponde à saída de uma CNN de camada única livre de viés.

Contribuições Principais

• Unificação de Priors: O artigo estabelece que os priors RED, SGM, NF e CRR, apesar de suas diferentes origens e motivações, podem todos ser integrados em um único arcabouço de amostragem via suas funções de score.
• Versatilidade Algorítmica: Demonstra que o LwSGS é um algoritmo genérico capaz de lidar com qualquer log-prior diferenciável com uma função de score explícita ou eficientemente computável, superando as limitações dos esquemas simples de ULA.
• Desacoplamento de Restrições: Ao dividir o problema, o método permite que o passo de amostragem para o prior seja independente do condicionamento do termo de fidelidade de dados, potencialmente melhorando as propriedades de mistura em problemas mal postos.

Resultados Experimentais
Os autores avaliaram o arcabouço LwSGS contra contrapartes baseadas em ULA (PnP-ULA, NF-ULA e os novos SGM-ULA e CRR-ULA) em duas tarefas: inpainting de imagem e super-resolução de imagem única.

• Datasets: Experimentos foram conduzidos em 30 imagens RGB dos datasets FFHQ e ImageNet.
• Métricas de Desempenho: O estudo utilizou PSNR, SSIM, LPIPS (similaridade perceptual) e Tempo de Autocorrelação Integrada (IAT) para medir a qualidade da reconstrução e a eficiência da amostragem.
• Descobertas:
  • Qualidade de Reconstrução: O LwSGS consistentemente produziu reconstruções de alta qualidade que correspondiam de perto aos detalhes do ground-truth. Em tarefas de inpainting, preencheu efetivamente regiões ausentes; em super-resolução, preservou detalhes intrincados.
  • Comparação com ULA: O LwSGS geralmente superou os métodos baseados em ULA, particularmente ao usar priors RED e SGM, produzindo detalhes mais nítidos e maior fidelidade. Mesmo com priors NF e CRR, o LwSGS mostrou desempenho competitivo ou superior.
  • Eficiência: O algoritmo LwSGS demonstrou convergência mais rápida, indicada por valores de IAT mais baixos e funções de autocorrelação de decaimento mais rápido em comparação ao ULA.
  • Quantificação de Incerteza: Como um método de amostragem, o LwSGS permitiu a estimativa da incerteza pixel a pixel (desvio padrão), mostrando maior incerteza em regiões inpainted, conforme esperado.
  • Aplicação no Mundo Real: O arcabouço foi aplicado à restauração de amostras geológicas reais (imagens de desfoque de fora de foco). Usando um prior RED, o método recuperou com sucesso detalhes nítidos de imagens borradas, demonstrando utilidade prática em geociências.

Significância e Alegações
O artigo afirma que o arcabouço LwSGS fornece uma alternativa versátil e sistemática para o ULA em problemas inversos Bayesianos envolvendo priors baseados em funções de score. Sua significância reside em:

1. Ampla Aplicabilidade: Não se limita a modelos generativos específicos, mas pode lidar com qualquer prior onde a função de score seja acessível, unindo diversos paradigmas de regularização baseados em dados.
2. Eficiência Prática: A estratégia de divisão mitiga efetivamente os problemas de condicionamento inerentes ao ULA, levando a uma melhor mistura e convergência mais rápida em espaços de alta dimensão.
3. Utilidade no Mundo Real: A restauração bem-sucedida de amostras geológicas demonstra que o método é robusto o suficiente para pipelines operacionais onde os protocolos de aquisição podem ser subótimos.

Os autores concluem que, embora o trabalho atual foque em priors específicos, o arcabouço unificado é suficientemente geral para lidar com qualquer distribuição posterior definida por uma ampla classe de priors baseados em função de score, pavimentando o caminho para a futura integração de diversos modelos de deep learning na inferência Bayesiana.