An invertible generative model for forward and inverse problems

Resumo Técnico: Um Modelo Generativo Invertível para Problemas Diretos e Inversos

Enunciado do Problema
Problemas inversos, tipicamente formulados como encontrar $u$ tal que $Ku = f$ para um operador direto $K$, são frequentemente mal postos (ill-posed). Embora métodos variacionais (por exemplo, minimizando a fidelidade aos dados mais a regularização) ofereçam uma estrutura para reconstrução, eles podem ser computacionalmente dispendiosos. Por outro via, abordagens baseadas em dados aprendem um mapeamento inverso direto, mas frequentemente lutam com a natureza mal posta inerente e com a falta de quantificação de incerteza. Estruturas Bayesianas fornecem uma solução rigorosa ao caracterizar a distribuição posterior, contudo, a amostragem de posteriors de alta dimensão é computacionalmente exigente (por exemplo, via MCMC).

Abordagens atuais de modelagem generativa frequentemente tratam a simulação (amostragem da verossimilhança $F|U=u$) e a inferência (amostragem da posterior $U|F=f$) como tarefas separadas usando modelos condicionais distintos. Essa separação falha em impor consistência entre os processos direto e inverso e exige o treinamento de dois modelos separados, aumentando os custos computacionais.

Metodologia: O Simulador Reversível
Os autores propõem um framework unificado centrado em um único mapa invertível, denominado Simulador Reversível ($S: \mathbb{R}^{n+m} \to \mathbb{R}^{n+m}$). Este mapa é projetado para realizar tanto a simulação quanto a inferência, dependendo da direção de aplicação:

1. Direção Direta (Simulação): Dada uma variável latente $Y$ e um parâmetro $u$, $S$ mapeia $(u, Y)$ para $(X, F)$, onde $F$ é distribuído de acordo com a verossimilhança $\mu_{F|U=u}$.
2. Direção Reversa (Inferência): Dada uma variável latente $X$ e uma observação $f$, o mapa inverso $R = S^{-1}$ mapeia $(X, f)$ para $(U, Y)$, onde $U$ é distribuído de acordo com a posterior $\mu_{U|F=f}$.

Construção via Mapas de Transporte Triangulares
O artigo fornece uma construção explícita de $S$ combinando dois mapas de transporte triangulares (rearranjos de Knothe-Rosenblatt):

• Um mapa inferior-triangular ($\tilde{F}$) que realiza o push-forward de uma medida de referência para a distribuição conjunta $\mu_{U,F}$, facilitando a amostragem condicional da verossimilhança.
• Um mapa superior-triangular ($\hat{F}$) que similarmente facilita a amostragem condicional da posterior.

O Simulador Reversível $S$ é construído como a composição desses mapas e de seus inversos. Especificamente, $S = U \circ L$, onde $L$ utiliza a estrutura inferior-triangular para condicionar em $u$, e $U$ utiliza a estrutura superior-triangular para condicionar em $f$. Esta construção garante que $S$ seja um mapa de transporte válido entre o produto da prior e o ruído latente ($\mu_U \otimes \mu_Y$) e o produto do ruído latente e a marginal dos dados ($\mu_X \otimes \mu_F$).

Formulação Variacional e Treinamento
Para permitir o treinamento com arquiteturas de redes neurais invertíveis genéricas (em vez de depender apenas de construções triangulares explícitas), os autores derivam uma formulação variacional. Eles definem quatro funções de perda baseadas em uma métrica de distância $M$ (por exemplo, distância de Wasserstein) entre medidas:

1. $J_1(S)$: Mede a discrepância entre o push-forward do produto prior-ruído e a distribuição conjunta (via componente de verossimilhança).
2. $J_2(R)$: Mede a discrepância para o componente da posterior.
3. $J_3(S)$ e $J_4(R)$: Medem a discrepância entre os push-forwards das distribuições conjuntas nas direções direta e reversa.

Os autores provam que minimizar essas perdas (especificamente $J_1$ e $J_3$) garante que o mapa aprendido satisfaça as propriedades de um Simulador Reversível. Isso permite que o modelo seja treinado diretamente a partir de amostras pareadas $\{(u_i, f_i)\}_{i=1}^N$ sem exigir conhecimento explícito das densidades subjacentes.

Principais Contribuições

1. Framework Unificado: O artigo introduz o "Simulador Reversível", um único mapa invertível que unifica simulação e inferência, impondo consistência entre os modelos direto e inverso.
2. Construção Explícita: Uma instância explícita de tal simulador é construída combinando fluxos de normalização (normalizing flows) inferiores-triangulares e superiores-triangulares.
3. Análise Teórica: Os autores investigam a existência e a não-unicidade de Simuladores Reversíveis, mostrando que existem infinitos tais mapas via composição com transformações preservadoras de medida.
4. Objetivo de Treinamento: Uma formulação variacional é derivada que permite o treinamento direto de dados pareados usando arquiteturas padrão de modelos generativos.
5. Validação Numérica: O framework é demonstrado em problemas Gaussianos lineares analiticamente tratáveis, exemplos sintéticos não-lineares e uma tarefa estilizada de inpainting de imagem usando o dataset MNIST.

Resultos

• Caso Gaussiano Linear: No caso de problemas inversos lineares com priors e verossimilhanças Gaussianas, o Simulador Reversível construído mostra-se bem definido mesmo quando o modelo direto é determinístico (ou seja, conforme a variância do ruído se aproxima de zero), enquanto os mapas triangulares padrão podem degenerar.
• Exemplo Não-Linear: A construção lida com sucesso com dependências não-lineares, visualizando planos de transporte distintos para os mapas inferior-triangular, superior-triangular e o combinado reversível.
• Inpainting de Imagem (MNIST): Usando uma arquitetura baseada em GLOW adaptada para o Simulador Reversível, o modelo aprende com sucesso a:
  • Gerar amostras de verossimilhança que refletem com precisão o operador direto (adicionando ruído e máscara) enquanto preservam a estrutura do dígito original.
  • Gerar amostras da posterior que reconstroem com precisão as partes ausentes dos dígitos (inpainting) e localizam corretamente a incerteza nas regiões mascaradas.
  • O modelo demonstra generalização para dados fora da distribuição (letras 'A' e 'F'), embora as amostras de verossimilhança para essas letras indiquem que o operador direto aprendido pelo modelo é influenciado pelos valores de pixels específicos da distribuição de treinamento (dígitos), indicando que o operador direto não é perfeitamente aprendido para classes não vistas.

Significância e Alegações
O artigo alega que o Simulador Reversível oferece uma abordagem principada e unificada para modelagem generativa condicional tanto para problemas diretos quanto inversos. Ao usar um único mapa invertível, o framework reduz o ônus computacional de treinar modelos separados e garante consistência entre os modelos de simulação e inferência. Os autores destacam que a construção pode oferecer um melhor condicionamento do mapeamento aprendido em comparação com fluxos triangulares padrão, particularmente em limites determinísticos. Eles posicionam este trabalho como uma ponte entre redes neurais invertíveis genéricas e fluxos de normalização para inferência Bayesiana amortizada, fornecendo uma base teórica para aprender tais mapeamentos diretamente dos dados. Os autores notam modestamente que, embora os resultados de unicidade atuais tenham sido estabelecidos, a caracterização adicional de todos os possíveis Simuladores Reversíveis poderia levar a critérios de otimalidade para o treinamento. Eles também sugerem extensões futuras para outros modelos generativos (ex: modelos de difusão) e aplicações em design experimental.