Amortized Factor Inference Networks for Posterior Inference

Resumo Técnico: Redes de Inferência de Fatores Amortizados (AFINs)

1. Declaração do Problema

A inferência amortizada visa aprender uma única rede neural que mapeia observações para distribuições posteriores, oferecendo inferência rápida no momento do teste após um custo de treinamento único. No entanto, os métodos amortizados existentes estão tipicamente vinculados a modelos probabilísticos fixos. Se o prior, a verossimilhança, a dimensionalidade latente ($d$) ou o número de observações ($N$) mudarem, a rede de inferência geralmente requer retreinamento ou ajuste fino no momento do teste, o que é custoso. Embora alguns trabalhos recentes generalizem para $N$ variável ou mudanças específicas de parâmetros, eles frequentemente carecem da capacidade de lidar com combinações arbitrárias de famílias de prior/verossimilhança variáveis, dimensionalidade e contagens de observações simultaneamente, sem sacrificar a precisão ou exigir adaptação específica ao modelo.

A questão central colocada por este trabalho é: Uma única rede de inferência pode generalizar para toda uma família de problemas distintos de inferência bayesiana, lidando com priors, verossimilhanças, dimensões e contagens de observações variáveis, enquanto iguala a precisão de métodos não amortizados como o NUTS?

2. Metodologia: Redes de Inferência de Fatores Amortizados (AFINs)

Os autores introduzem as Redes de Inferência de Fatores Amortizados (AFINs), uma arquitetura de codificar-fundir-decodificar projetada para mapear uma representação tipada de uma especificação de modelo e suas observações diretamente para os parâmetros de uma posterior variacional.

A. Representação de Fatores Tipados
Em vez de tratar o modelo como um bloco monolítico, a AFIN representa a densidade conjunta $p(z, y_{1:N})$ como uma coleção de $N+1$ fatores: um fator de prior $p_0(z)$ e $N$ fatores de verossimilhança $p_n(y_n|z)$. Cada fator é representado por:

• Tipo ($t_n$): Um rótulo indicando a família de distribuição (ex: "verossimilhança gaussiana", "verossimilhança t-Student").
• Parâmetros ($\theta_n$): Parâmetros específicos do fator (ex: escala, graus de liberdade) necessários para instanciar o fator.
• Observações ($y_n$): Os dados associados aos fatores de verossimilhança.

Esta representação permite que a rede lide com modelos heterogêneos onde diferentes observações podem seguir famílias de verossimilhança distintas.

B. Arquitetura Independente de Dimensão
Para suportar dimensões latentes variáveis ($d$) e números de fatores ($N$), a AFIN emprega módulos que são independentes de $d$ e equivariantes a permutações de coordenadas latentes:

1. BoxMLP: Uma Rede Neural Perceptron Multicamada com pesos acoplados, onde os pesos aprendidos não dependem de $d$. Ela processa características coordenada a coordenada e resumos agrupados (ex: média sobre dimensões) para garantir que a saída seja equivariante a permutações do vetor latente $z \in \mathbb{R}^d$.
2. BoxTransformer: Um análogo independente de dimensão do Transformer padrão. Ele opera sobre uma sequência de incorporações de fatores. Crucialmente, os mapas de projeção (Query, Key, Value, Output, FFN) são implementados como BoxMLPs aplicados a tokens indexados por coordenada. Isso permite que a rede modele interações entre fatores sem achatar a dimensão latente, preservando a estrutura $O(d)$ ou $O(d^2)$ dos parâmetros.

C. O Pipeline Codificar-Fundir-Decodificar

1. Codificação: Cada fator é processado independentemente. Um adaptador específico do tipo mapeia os parâmetros e observações do fator para um espaço de incorporação comum, seguido por um codificador BoxMLP compartilhado. A incorporação inclui tanto características de nó (coordenada a coordenada) quanto características de par (interações de pares de coordenadas, úteis para matrizes de covariância/precisão).
2. Fusão: A sequência de incorporações de fatores é processada por $M$ blocos BoxTransformer. Esses blocos usam atenção ao longo do eixo dos fatores (tratando os fatores como a dimensão da sequência) para agregar informações. O mecanismo de atenção combina pontuações de nó e de par para atualizar as representações.
3. Decodificação: As representações fundidas são agrupadas por soma sobre o eixo dos fatores para produzir uma única incorporação no nível da tarefa. Um decodificador BoxMLP final mapeia essa incorporação para os parâmetros variacionais $\phi$ (ex: média e precisão para uma Gaussiana, ou parâmetros de fluxo para um fluxo normalizante).

D. Objetivo de Treinamento
A AFIN é treinada via inferência baseada em simulação. A rede é treinada para maximizar a densidade logarítmica atribuída a amostras latentes $z$ amostradas do modelo conjunto $p(z, y)$, condicionada às observações geradas $y$. Isso é equivalente a minimizar a divergência KL forward esperada ($KL(p(z|y) || q_\phi(z))$) entre a posterior verdadeira e a aproximação variacional. O objetivo da KL forward é notado como benéfico para o uso da aproximação como proposta para amostragem por importância.

3. Contribuições Principais

• Generalização: Uma única AFIN treinada pode realizar inferência posterior em famílias de priors variáveis, famílias de verossimilhança, dimensões latentes ($d$) e números de observações ($N$) sem retreinamento.
• Independência de Dimensão: A introdução dos módulos BoxMLP e BoxTransformer permite que a rede opere em espaços latentes de dimensão arbitrária $d$ enquanto mantém uma contagem fixa de parâmetros, um desvio significativo das redes amortizadas padrão que exigem tamanhos de entrada/saída fixos.
• Abstração no Nível de Fatores: Ao operar em fatores tipados em vez de grafos de modelos fixos, a AFIN suporta verossimilhanças heterogêneas (ex: misturar observações Gaussianas e Bernoulli) dentro de uma única execução de inferência.
• Eficiência: O método alcança alta precisão com custo computacional no momento do teste significativamente menor em comparação com métodos iterativos.

4. Resultados Experimentais

Os autores avaliaram um único checkpoint de AFIN com 5 milhões de parâmetros em tarefas sintéticas e do mundo real sem ajuste fino.

• Desempenho Sintético: Em 48 configurações (16 combinações prior-verossimilhança $\times$ 3 níveis de dificuldade), a AFIN alcançou precisão posterior comparável ao NUTS e várias bases de Inferência Variacional (IV) (Gaussiana Full-Rank, IAF, MAF, NSF).
  • Velocidade: A inferência única da AFIN é 2 a 4 ordens de magnitude mais rápida que o NUTS e métodos de IV baseados em fluxo.
  • Refinamento: Quando usada como proposta para Amostragem por Importância Auto-Normalizada (SNIS), a AFIN+SNIS igualou ou superou a precisão do NUTS com tempo de relógio substancialmente menor.
• Extrapolação: A AFIN demonstrou generalização zero-shot para configurações fora da distribuição (OOD), incluindo dimensões latentes ($d=32$) e contagens de observações ($N=512$) maiores do que as vistas durante o treinamento ($d \le 16, N \le 256$).
• Transferência para o Mundo Real: Em 12 conjuntos de dados UCI (incluindo configurações de verossimilhança heterogêneas), a AFIN+SNIS alcançou precisão competitiva com referências NUTS de longa execução quase imediatamente, enquanto as bases iterativas exigiram tempo significativo de otimização para convergir.
• Variantes de Decodificador: Tanto decodificadores Gaussianos quanto de Fluxo foram testados. O decodificador de Fluxo ofereceu ligeiras melhorias de precisão sobre o decodificador Gaussiano, mantendo a vantagem de velocidade da amortização.

5. Significado e Afirmações

O artigo afirma que a AFIN conecta com sucesso a lacuna entre a velocidade da inferência amortizada e a flexibilidade dos métodos não amortizados. Ao desacoplar a rede de inferência de especificações de modelo fixas através de uma interface de fatores tipados e módulos independentes de dimensão, a AFIN habilita inferência amortizada reutilizável.

Os autores posicionam a AFIN como um passo em direção a um "modelo de fundação" para inferência bayesiana, capaz de lidar com programas probabilísticos diversos. Eles reconhecem limitações, como a exigência de um conjunto pré-definido de tipos de fatores e a escalabilidade $O(d^2)$ das incorporações de componentes de par, mas enfatizam que a abordagem permite que uma única rede resolva uma ampla família de problemas de inferência que anteriormente exigiam modelos distintos e retreinados. O trabalho sugere que a amortização pode igualar a qualidade dos métodos de amostragem de última geração (NUTS) enquanto oferece economias computacionais massivas no momento do teste.