ShuffleFlow: Scalable Posterior Inference for Bayesian Inverse Imaging

Resumo Técnico: ShuffleFlow

Definição do Problema
A imagem reversa Bayesiana visa recuperar imagens latentes $x$ a partir de medições ruidosas e incompletas $y = A(x) + n$ através da estimação da distribuição posterior $p(x|y)$. Embora a Inferência Variacional (VI) ofereça uma abordagem fundamentada para a inferência posterior que pode gerar amostras de forma eficiente após a otimização, sua aplicação à reconstrução de imagens em larga escala é dificultada pela baixa escalabilidade. Métodos tradicionais de VI que dependem de redes baseadas em fluxos tipicamente modelam a distribuição completa da imagem diretamente, resultando em custos computacionais que escalam mal com a dimensão da imagem (por exemplo, $O(N^4)$ para $N \times N$ pixels). Isso leva a orçamentos de memória e tempo proibitivos durante a otimização. Tentativas existentes para melhorar a escalabilidade, como estratégias de patching ingênuas ou de pixelização, frequentemente sacrificam a qualidade da imagem e a expressividade da distribuição posterior. Por outro lado, modelos de difusão, embora poderosos, requerem grandes conjuntos de dados para treinamento e incorrem em altos custos computacionais para gerar posteriors de alta resolução e alta contagem de amostras devido à sua natureza de amostragem iterativa.

Metodologia: ShuffleFlow
Os autores propõem o ShuffleFlow, um framework de VI escalável projetado para decompor a tarefa de modelagem da posterior em três componentes modulares: amostragem de coordenadas, codificação de características e estimativa da posterior.

1. Amostragem de Coordenadas (Pixel-Unshuffling): Em vez de modelar a imagem completa de $N \times N$ ou patches independentes, o ShuffleFlow particiona a imagem em uma pilha de $S^2$ sub-imagens correlacionadas usando uma operação de pixel-unshuffling com um fator de subamostragem $S$. Cada sub-imagem possui dimensões $(N/S) \times (N/S)$. Isso reduz a dimensionalidade do problema de $N^2$ para $N^2/S^2$.
2. Codificação de Características (Campo Neural): Para capturar estruturas espaciais sem sobrecarregar a rede de fluxo com tarefas de reconstrução de imagem completa, o framework emprega um codificador de Campo Neural (NF) (especificamente um MLP baseado em coordenadas, SIREN). Os deslocamentos de coordenadas 2D das localizações de amostragem são inseridos no NF para gerar embeddings de características ($h_i$) para cada sub-imagem.
3. Estimativa da Posterior (Fluxo Normalizante Condicional): Um Fluxo Normalizante Condicional (CNF) compartilhado modela a distribuição conjunta da pilha de sub-imagens. O CNF é condicionado aos embeddings de características ($h_i$) correspondentes à localização de cada sub-imagem. Crucialmente, o framework adota uma abordagem inspirada em De Finetti, compartilhando uma única variável latente base $z_0$ entre todas as sub-imagens. Isso permite que a distribuição conjunta seja fatorizada como:
   $$q(x) = \left( \prod_{i=1}^{S^2} q(x_i | z_0; c_i) \right) q(z_0)$$
   onde $z_0$ é extraído de uma distribuição normal padrão. Este mecanismo de compartilhamento correlaciona as sub-imagens, evitando artefatos de ruído associados à independência de pixels, enquanto reduz significativamente o tamanho da rede de fluxo de $O(N^4)$ para $O(N^4/S^4)$.

Otimização e Inferência
O CNF e o NF são treinados conjuntamente minimizando a divergência de Kullback-Leibler (KL) entre a distribuição variacional e a posterior Bayesiana. A função de perda compreende três termos: verossimilhança dos dados, um prior (seja um prior clássico de Variação Total ou um prior baseado em score com pesos congelados) e um termo de entropia negativa para incentivar a diversidade. Durante a inferência, o NF gera embeddings para as coordenadas de amostragem, e o CNF gera amostras condicionadas a esses embeddings e a uma variável base compartilhada, que são então remontadas em uma imagem completa via pixel-shuffling.

Principais Contribuições

1. Decomposição do Problema: Os autores introduzem um novo framework de VI modular que decompõe a modelagem de posterior de imagem grande em amostragem de coordenadas, representação parcial de imagem e quantificação de incerteza leve. Esta estrutura é agnóstica a arquiteturas de rede específicas, permitindo a integração de avanços em campos neurais e priors generativos.
2. Escalabilidade e Eficiência: Ao alavancar o pixel-unshuffling e variáveis latentes compartilhadas, o ShuffleFlow alcança uma redução massiva no tamanho do modelo e nos recursos computacionais. Análises teóricas e empíricas mostram que o método escala quase tão eficientemente quanto a inferência por pixel, oferecendo uma redução de 4096$\times$ no tamanho da rede de fluxo para imagens de $256 \times 256$ (com $S=8$) em comparação com modelos de fluxo de imagem completa.
3. Geração de Posterior de Alta Resolução: O método demonstra a capacidade de gerar rapidamente posteriors de alta contagem de amostras. Ao contrário dos amostradores de difusão, cujo tempo de amostragem cresce linearmente com o número de amostras, o ShuffleFlow pode gerar milhares de amostras de forma eficiente após uma fase curta de treinamento.

Resultados Experimentais
Os autores avaliam o ShuffleFlow em problemas inversos lineares (desfoque por movimento, Ressonância Magnética de compressão de dados) e não lineares (recuperação de fase de Fourier).

• Problemas Lineares: O ShuffleFlow atinge o Estado da Arte (SOTA) entre os métodos de VI em termos de qualidade de imagem (PSNR, SSIM, LPIPS) e calibração de incerteza (NLL, ECE), superando métodos anteriores de VI baseados em fluxo (DPI) e estratégias de patching ingênuas. Ele apresenta desempenho comparável aos principais amostradores de difusão (ex: DAPS) em qualidade de imagem, mas com overhead computacional significativamente menor.
• Problemas Não Lineares: Na desafiadora tarefa de recuperação de fase de Fourier, que exibe uma posterior bimodal (imagens verticais vs. rotacionadas), o ShuffleFlow captura com sucesso a bimodalidade. Em contraste, métodos como CF-NeRF e patching ingênuo falham em modelar essas correlações globais complexas.
• Eficiência: Em um cenário que exige 10.000 amostras da posterior, o ShuffleFlow gera uma posterior bimodal de alta resolução em 16 minutos. No mesmo intervalo de tempo, amostradores de difusão (DAPS, DPS) geram muito menos amostras e falham em caracterizar claramente o segundo modo, exigindo aproximadamente 20$\times$ mais tempo para atingir resolução semelhante.

Significância e Alegações
O artigo afirma que o ShuffleFlow aborda o gargalo crítico de escalabilidade na Inferência Variacional para imagens científicas sem sacrificar a precisão da posterior. Ele posiciona a VI como uma alternativa vital aos modelos de difusão, particularmente em regimes de escassez de dados onde modelos de difusão não podem ser treinados, e em aplicações que exigem a exploração rápida de posteriors complexas e de alta resolução (ex: identificar distribuições bimodais em problemas inversos não lineares). Os autores observam que, embora o método seja altamente eficiente, ele atualmente exibe certa excessiva confiança (overconfidence) em configurações não lineares devido à instabilidade de treinamento, o que identificam como um alvo para trabalhos futuros. O framework é apresentado como uma ferramenta exploratória que pode determinar quando a complexidade da posterior justifica uma geração extensiva de amostras.