Resumo Técnico: Amostragem de Posterior Baseada em Difusão Provável para Problemas Inversos Lineares via DDIM

Declaração do Problema

O artigo aborda o desafio de resolver problemas inversos lineares ruidosos da forma $Y = AX_0 + \varepsilon$, onde $X_0 \in \mathbb{R}^d$ é um sinal desconhecido, $A \in \mathbb{R}^{n \times d}$ é um operador de medição linear conhecido e $\varepsilon \sim \mathcal{N}(0, \sigma^2 I_n)$ é um ruído Gaussiano, sendo o objetivo gerar amostras da distribuição posterior Bayesiana $p_{X_0|Y}$, condicionada à observação $y$.

Embora os modelos de difusão tenham emergido como priors poderosos baseados em dados para capturar distribuições complexas de sinais, os métodos existentes para amostragem de posterior frequentemente sofrem de uma de duas limitações:

1. Falta de garantias teóricas: Muitas abordagens heurísticas (por exemplo, métodos baseados em projeção ou correções de escore aproximadas) carecem de prova rigorosa de que suas amostras geradas convergem para a verdadeira posterior.
2. Sobrecarga computacional: Métodos que tentam aproximar o escore da posterior ou que utilizam Monte Carlo Sequencial (SMC) frequentemente incorrem em custos computacionais proibitivos, como a necessidade de computações de gradiente de funções de escore ou um grande número de partículas.

Metodologia: Posterior-DDIM

Os autores propõem o Posterior-DDIM, um algoritmo simples e eficiente que realiza a amostragem de posterior adaptando o amostrador padrão de Modelos de Difusão Implícitos de Denoising (DDIM). A inovação central reside em um particionamento baseado na decomposição em valores singulares (SVD) do espaço do sinal, onde as atualizações são ajustadas de acordo com a força relativa da relação sinal-ruído (SNR) da observação versus a SNR da difusão.

1. Particionamento Guiado pela SNR

Utilizando a SVD da matriz de medição $A = U\Sigma V^\top$, o problema é decomposto em problemas de inferência escalar ao longo de cada direção singular $v_s$. Os autores definem um limiar de tempo $\tau_s$ para cada direção onde a SNR da difusão ($\alpha_t/\sigma_t$) é igual à SNR da observação ($\Sigma_{ss}/\sigma$). Com base nisso, as direções singulares são particionadas em qualquer tempo $t$ em três grupos:

• Direções dominadas pela medição ($S^{meas}_t$): Direções onde a SNR da observação excede a SNR da difusão ($t \ge \tau_s$).
• Direções observadas dominadas pelo prior ($S^{prior}_t$): Direções onde a SNR da observação é inferior à SNR da difusão ($t < \tau_s$), mas a direção ainda é observada.
• Direções não observadas ($S^c$): Direções no espaço nulo de $A$ onde não existe informação de medição.

2. Regras de Atualização por Coordenada

O algoritmo aplica regras de atualização distintas para esses grupos, exigindo apenas modificações leves ao padrão de atualização do DDIM:

• Para $S^{meas}_t$ (Dominadas pela medição): O algoritmo ignora o prior aprendido e injeta diretamente uma versão calibrada e ruidosa da observação. Especificamente, ele gera a coordenada $v_s^\top \hat{X}_{t_{i+1}}$ usando o processo direto auxiliar definido pela medição $y$, garantindo que a distribuição marginal corresponda ao processo de difusão direta projetado sobre essa direção.
• Para $S^{prior}_t$ (Observadas dominadas pelo prior): O algoritmo segue a dinâmica padrão do DDIM impulsionada pelo preditor de dados aprendido $\mu_t(\cdot)$. No caso geral onde os valores singulares não são idênticos, um termo de correção é adicionado ao preditor para contabilizar os variados valores singulares, garantindo que a atualização permaneça consistente com a posterior condicional.
• Para $S^c$ (Não observadas): Estas coordenadas são amostradas inteiramente a partir do prior de difusão aprendido, condicionado às coordenadas observadas. Para incentivar uma exploração suficiente do subespaço não observado, um parâmetro de estocasticidade $\eta_1$ maior é utilizado, tornando a atualização mais semelhante a uma Langevin.

3. Parâmetros de Estocasticidade

O método utiliza dois parâmetros de estocasticidade, $\eta_0$ e $\eta_1$, controlando a injeção de ruído nas atualizações DDIM para $S^{prior}_t$ e $S^c$, respectivamente.

• $\eta_0$ é mantido pequeno (frequentemente 0) para manter uma atualização de tipo determinístico para as direções observadas onde o prior é o principal guia.
• $\eta_1$ é definido para ser suficientemente grande para garantir que o subespaço não observado seja explorado efetivamente, já que sua distribuição é determinada unicamente pelo prior.

Principais Contribuições

1. Simplicidade Algorítmica: O método proposto reduz a amostragem de posterior a uma sequência de atualizações DDIM simples por coordenada. Ele incorpora explicitamente o modelo de medição sem exigir etapas de projeção complexas, computações de gradiente de verossimilhança ou retreinamento do modelo de difusão.
2. Garantias Teóricas: O artigo estabelece a correção assintótica da posterior. Sob a suposição de um prior de difusão exato e uma discretização de tempo suficientemente fina, os autores provam que a distribuição das amostras geradas pelo Posterior-DDIM converge para a verdadeira posterior $p_{X_0|Y}$ em distância de variação total. Isso se aplica tanto ao caso especial de valores singulares idênticos quanto ao caso geral com valores singulares arbitrários.
3. Desempenho Empírico: Experimentos extensos em tarefas de restauração de imagem (inpainting, super-resolução, desfoque gaussiano e compressão de sensores) demonstram que o Posterior-DDIM supera os amostradores de difusão existentes (incluindo DPS, DDRM, DDNM+ e ProjDiff) na maioria das métricas de avaliação (PSNR, SSIM, LPIPS e FID).

Resultados Experimentais

Os autores avaliaram o algoritmo nos conjuntos de dados CelebA e ImageNet. Principais descobertas incluem:

• Superioridade: O Posterior-DDIM alcançou o melhor desempenho em pelo menos três das quatro métricas na maioria das tarefas e em pelo menos duas métricas em cada tarefa considerada.
• Robustez: O método superou consistentemente os baselines através de diferentes níveis de ruído e tipos de degradação.
• Eficiência: O algoritmo requer apenas 100 avaliações de função (NFE) para atingir um desempenho estável, sendo comparável ou mais rápido que outros métodos que frequentemente exigem significativamente mais passos (ex: 1000 NFE para o DPS).
• Sensibilidade a Hiperparâmetros: Os experimentos mostraram que definir $\eta_0 = 0$ (atualização determinística para direções dominadas pelo prior) e um $\eta_1$ suficientemente grande (ex: 16 ou 32) produz resultados ótimos.

Significância e Alegações

O artigo afirma fornecer um arcabouço principiado e provável para integrar priors de difusão com modelos de medição linear. Ao converter o problema complexo de amostragem de posterior em atualizações DDIM simples por direção, os autores unem o rigor teórico à eficiência prática.

A significância deste trabalho é destacada por sua capacidade de:

• Oferecer suporte teórico rigoroso para amostradores do tipo DDIM em problemas inversos, abordando a falta de garantias de convergência em muitos métodos heurísticos.
• Fornecer uma solução de fácil implementação que preserva a eficiência computacional dos amostradores de difusão padrão enquanto lida explicitamente com o modelo de medição.
• Demonstrar que a consistência da posterior pode ser alcançada sem as proibições computacionais de Monte Carlo Sequencial ou as aproximações heurísticas de métodos baseados em projeção.

Os autores observam que, embora o presente trabalho foque em problemas inversos lineares, estender este arcabouço para cenários não lineares e caracterizar as taxas de convergência são direções naturais para pesquisas futuras.