JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

Resumo Técnico: JADAI – Amortização Conjunta de Design Adaptativo e Inferência Bayesiana

Definição do Problema
Muitos domínios científicos e de engenharia envolvem a inferência de parâmetros não observados $\theta$ a partir de observações $x$ geradas por simuladores complexos. Embora a Inferência Baseada em Simulação (SBI) tenha avançado a capacidade de aproximar distribuições posteriores $p(\theta|x)$ sem avaliações explícitas de verossimilhança, esses métodos tipicamente assumem um design experimental fixo. Inversamente, o Design Experimental Bayesiano (BED) e o Design Adaptativo (BAD) focam na seleção de variáveis de design $\xi$ para maximizar o ganho de informação (ex: Ganho de Informação Esperada, EIG), mas frequentemente tratam a inferência como uma etapa separada e computacionalmente cara realizada após a coleta de dados. As abordagens existentes frequentemente desacoplam essas tarefas: métodos de BAD frequentemente dependem de inferência não amortizada e lenta ou de famílias paramétricas simples, enquanto métodos de SBI não otimizam a aquisição de dados. Essa separação leva a um ganho de informação subótimo e à ineficiência computacional, particularmente em configurações de alta dimensão ou multimodais.

Metodologia
Os autores introduzem o JADAI (Jointly Amortizing Adaptive Design and Bayesian Inference), um framework que unifica o design adaptativo e a inferência posterior em um único sistema treinável de ponta a ponta. A inovação central é a amortização conjunta de três componentes:

1. Uma Rede de Política ($\pi_\phi$): Propõe o próximo design $\xi_t$ com base no histórico experimental atual.
2. Uma Rede de Histórico/Sumário ($\eta_\omega$): Codifica a sequência de pares design-observação passados $\{(\xi_k, x_k)\}_{k=1}^t$ em uma estatística de sumário de dimensão fixa $h_t$.
3. Uma Rede de Inferência ($q_\psi$): Aproxima a distribuição posterior $p(\theta|h_t)$ em cada etapa da sequência experimental.

Objetivo de Otimização
O JADAI minimiza uma função de perda genérica que agrega reduções incrementais no erro da posterior ao longo de sequências experimentais. Em vez de maximizar o Ganho de Informação Esperada (EIG) diretamente — o que exige avaliações de densidade — o framework utiliza uma soma telescópica de perdas da posterior.

• O objetivo de treinamento é derivado de um limite inferior variacional no Ganho de Informação Total Esperado (TEIG).
• A perda no passo $t$, $\ell_t$, é definida pela perda do estimador da posterior (ex: perda de difusão ou flow matching).
• A função de utilidade $u_T$ é construída como uma soma de diferenças entre perdas consecutivas: $u_T = \sum (\ell^*_{t-1} - \ell_t)$, onde $\ell^*$ denota uma versão desconectada (stop-gradient) da perda anterior.
• Esta formulação garante que a rede de política e as redes de sumário sejam treinadas para maximizar a redução na perda da posterior em cada passo, efetivamente maximizando o ganho de informação sem exigir avaliações explícitas de densidade da posterior.

Escolhas Arquiteturais Chave

• Modelos de Difusão: A rede de inferência $q_\psi$ é instanciada usando modelos de difusão (ou flow matching), permitindo que o framework aproxime posteriors de alta dimensão e multimodais em cada etapa experimental, uma capacidade frequentemente ausente em métodos de BAD anteriores restritos a fluxos de normalização (normalizing flows).
• Estabilidade de Treinamento: Para evitar o backpropagação através do tempo (BPTT) aninhado e garantir um treinamento estável, o framework emprega:
  • Estados de Histórico Desconectados: Os designs são gerados a partir de um estado de sumário desconectado para evitar caminhos de gradiente cíclicos entre a rede de política e a rede de histórico.
  • Aprendizado por Currículo: O comprimento do rollout é aumentado gradualmente durante o treinamento, e uma estratégia de "mistura de prior de design" é usada no início do treinamento para expor as redes a designs diversos antes de depender exclusivamente da política aprendida.

Resultados
O JADAI foi avaliado em três benchmarks representando complexidade crescente em estrutura de política e posterior:

1. Localização de Fonte (LF): Uma tarefa que requer a inferência de localizações de origem a partir de medições de intensidade ruidosas. O JADAI alcançou escores sPCE (Estimativa Contrastiva de Prior Sequencial) superiores ou competitivos em relação às baselines de estado da arte (ex: DAD, iDAD, RL-BOED), particularmente para horizontes longos ($T > 10$). A política aprendida equilibrou efetivamente exploração e explotação, e a posterior baseada em difusão capturou corretamente a natureza multimodal e intercambiável das fontes.
2. Elasticidade de Substituição Constante (CES): Um problema de design desafiador envolvendo inferência de preferência onde os designs informativos residem em um estreito "ponto ideal" (sweet spot). O JADAI superou todas as baselines existentes no horizonte padrão $T=10$, demonstrando robustez em cenários onde designs aleatórios são não informativos.
3. Descoberta de Imagem MNIST: Uma tarefa de alta dimensão onde o objetivo é reconstruir uma imagem a partir de medições parciais. O JADAI alcançou escores de Índice de Similaridade Estrutural (SSIM) e Erro Quadrático Médio Normalizado (NRMSE) de estado da arte, superando a baseline CoDiff. O método identificou com sucesso classes de dígitos precocemente e refinou detalhes locais, mantendo o desempenho mesmo sob ruído de medição aditivo.

Significância e Alegações
O artigo afirma que o JADAI representa um passo significativo ao unificar design e inferência, eliminando assim a necessidade de etapas de inferência separadas e custosas durante o processo adaptativo. As principais contribuições incluem:

• Amortização Conjunta: É o primeiro framework a amortizar conjuntamente o design adaptativo e a inferência da posterior usando uma utilidade geral que agrega reduções incrementais de erro da posterior, evitando a necessidade de avaliações de densidade.
• Escalabilidade para Posteriors Complexas: Ao utilizar modelos de difusão, o JADAI pode lidar com posteriors de alta dimensão e multimodais em cada etapa experimental, superando as limitações de expressividade de métodos anteriores que dependiam de fluxos de normalização ou misturas gaussianas.
• Desempenho: O framework alcança desempenho superior ou competitivo em diversos benchmarks, incluindo reconstrução de imagens de alta dimensão, mantendo a eficiência computacional durante a implantação (rollouts completados em milissegundos).

Os autores observam que, embora o JADAI atualmente dependa de simuladores diferenciáveis, a abordagem é bem adequada para o crescente cenário de frameworks de simulação diferenciável. Eles reconhecem que a amostragem da posterior continua sendo um gargalo para a implantação de sub-segundo em casos de alta dimensão, mas sugerem a destilação post-hoc como uma direção futura potencial. O trabalho não pretende resolver a otimização de simuladores de caixa preta sem gradientes, mas se posiciona como uma solução robusta para configurações diferenciáveis onde pipelines clássicos são computacionalmente impraticáveis.