Tokenised Flow Matching for Hierarchical Simulation Based Inference

Resumo Técnico: Tokenised Flow Matching for Hierarchical Simulation Based Inference (TFMPE)

1. O Problema

A Inferência Baseada em Simulação (SBI) é uma técnica fundamental para a estimativa bayesiana de parâmetros quando a função de verossimilhança é intratável, mas um simulador computacional está disponível. No entanto, em cenários hierárquicos (comuns em epidemiologia, física e ciências sociais), onde existem parâmetros globais compartilhados e parâmetros locais específicos para cada "sítio" (ex: países, pacientes, estações), a inferência enfrenta um gargalo crítico:

• Custo Computacional: A estimativa tradicional requer a simulação de todos os sítios para cada amostra de treinamento. À medida que o número de sítios ($n_s$) aumenta, o custo de simulação cresce linearmente, tornando a inferência inviável para grandes conjuntos de dados.
• Limitações das Abordagens Existentes: Métodos atuais de SBI hierárquico geralmente fatorizam a posteriori, mas ainda exigem múltiplas simulações por amostra de treinamento (ex: simular $n_{max}$ sítios). Isso limita a eficiência de amostragem, especialmente quando os simuladores são caros (ex: dinâmica de fluidos computacional).

2. Metodologia Proposta (TFMPE)

Os autores propõem o TFMPE (Tokenised Flow Matching for Posterior Estimation), uma abordagem que combina três conceitos inovadores para resolver o problema de escalabilidade:

A. Fatorização da Verossimilhança (Likelihood Factorisation - LF)
Em vez de fatorizar a distribuição posterior (como fazem métodos anteriores), o TFMPE propõe fatorizar a verossimilhança.

• Estratégia de Duas Etapas:
  1. Surrogado Neural por Sítio: Treina-se um modelo neural ($q_{\phi_l}$) para aprender a distribuição condicional de observações de um único sítio ($p(y_s | \theta_g, \eta_s)$) usando apenas uma simulação por amostra.
  2. Montagem Sintética: Utiliza-se esse surrogado para gerar observações sintéticas multi-sítio ($\hat{y}$) que são usadas para treinar um estimador da posterior completa ($q_{\phi_p}$).
• Vantagem: Reduz drasticamente o número de chamadas ao simulador real por amostra de treinamento (de $n_s$ para 1), amortizando o custo computacional.

B. Tokenização para Observações Funcionais
O modelo utiliza um esquema de tokenização baseado em Transformers (Encoder-only) para representar os dados.

• Flexibilidade: Os parâmetros e observações são achatados em uma sequência de tokens. Cada token contém: valor, identificador da variável, posição intra-variável, identificador do grupo (sítio) e, crucialmente, entradas funcionais (ex: tempo ou coordenadas espaciais).
• Benefício: Permite lidar com observações de tamanho variável, irregulares ou com dados faltantes (comuns em dados de vigilância epidemiológica ou sensores), algo que métodos tradicionais de SBI não suportam bem.

C. Flow Matching (Correspondência de Fluxo)
Para estimar a densidade posterior, o método utiliza Flow Matching, uma técnica de modelagem generativa que aprende um campo vetorial para transportar uma distribuição base para a distribuição alvo.

• Estabilidade: Oferece dinâmicas de treinamento mais estáveis e caminhos de probabilidade mais simples em comparação com métodos baseados em score matching (como Diffusion Models).
• Arquitetura: Um Transformer processa os tokens e um cabeçalho linear simples (readout) extrai o campo vetorial para os parâmetros.

3. Contribuições Principais

1. Fatorização da Verossimilhança (LF Sampling): Uma nova estratégia de treinamento que reduz o custo de simulação para um único sítio por amostra, permitindo a inferência em escalas antes proibitivas.
2. TFMPE: A primeira aplicação de Flow Matching tokenizado para SBI hierárquico, capaz de lidar com observações funcionais e estruturas complexas de dados.
3. Novo Benchmark Hierárquico: Os autores desenvolveram o primeiro conjunto de testes (benchmark) específico para SBI hierárquico, adaptado do benchmark SBI padrão, incluindo tarefas com diferentes estruturas de agrupamento (separados e partial pooling).
4. Validação em Cenários Realistas: Demonstração da eficácia do método em modelos complexos do mundo real, superando as limitações de métodos MCMC e SBI não hierárquicos.

4. Resultados e Avaliação

A. Benchmark Hierárquico

• O TFMPE com amostragem LF foi comparado a baselines não hierárquicos (NPE, SNPE, Simformer, FMPE) e ao método hierárquico mais recente (Posterior Factorisation - PF).
• Eficiência de Amostragem: O TFMPE demonstrou a melhor eficiência de amostragem na maioria das tarefas, especialmente quando o número de sítios ($n_s$) é grande em relação ao orçamento de simulação.
• Trade-off: Em tarefas onde o modelo de observação é altamente não-linear e difícil de aproximar (ex: "Two Moons"), a fatorização da posterior (PF) pode ser superior devido ao erro de aproximação do surrogado. No entanto, para a maioria dos casos (ex: SEIR, CFD), a LF superou as outras abordagens.

B. Casos de Uso Realistas

1. Modelagem de Doenças Infecciosas (SEIR Sazonal):
   • Cenário: Vigilância de doenças com dados observados em centenas de unidades administrativas com frequências temporais irregulares.
   • Resultado: O TFMPE conseguiu inferir parâmetros globais e locais com alta calibração, enquanto o MCMC (NUTS) falhou em convergir para mais de 5 sítios. O método lidou nativamente com a irregularidade temporal dos dados.
2. Dinâmica de Fluidos (Hemodinâmica):
   • Cenário: Calibração de modelos de fluxo sanguíneo 1D para pacientes específicos, onde cada simulação leva ~12 segundos.
   • Resultado: O TFMPE reduziu o custo computacional de treinamento em 2x para 2 pacientes e mostrou um ganho de velocidade de ~1850x na geração de amostras sintéticas em comparação com simulações diretas. A calibração da posterior foi mantida com 16 pacientes, algo computacionalmente proibitivo para métodos diretos.

5. Significado e Impacto

O trabalho TFMPE representa um avanço significativo na viabilidade prática da Inferência Baseada em Simulação para problemas científicos complexos e de grande escala:

• Escalabilidade: Permite a calibração de modelos hierárquicos com milhares de sítios, algo que era impossível com métodos de SBI tradicionais devido ao custo de simulação.
• Flexibilidade de Dados: A capacidade de processar observações funcionais e irregulares (comuns em dados reais) sem necessidade de pré-processamento agressivo ou agregação de dados.
• Eficiência: A estratégia de fatorização da verossimilhança oferece um caminho para reduzir o custo computacional de ordens de magnitude em simulações caras, tornando a inferência bayesiana acessível para aplicações em tempo real ou de grande porte.

Em resumo, o TFMPE combina a eficiência de aprendizado de surrogados locais com a flexibilidade de arquiteturas de Transformers e a estabilidade do Flow Matching, estabelecendo um novo estado da arte para a inferência hierárquica baseada em simulação.