Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation

Resumo Técnico: Transformers como Experimentadores Bayesianos em Contexto

Formulação do Problema
O artigo aborda o desafio de projetar experimentos adaptativos para estimar o Efeito Médio do Tratamento (ATE) em configurações sequenciais de dois braços com covariáveis. A estratégia ótima para tais experimentos é a alocação de Neyman "oracle", que atribui probabilidades de tratamento baseadas nas variâncias residuais condicionais aos braços, desconhecidas, para minimizar a variância do estimador do ATE. Implementar isso de forma ótima requer um experimentador bayesiano que atualize sequencialmente crenças não paramétricas sobre os momentos (médias e segundos momentos) dos resultados potenciais e traduza essas incertezas posteriores em probabilidades de atribuição de tratamento.

Um gargalo crítico em espaços de covariáveis de alta dimensão é a suavidade desconhecida das funções de resultados potenciais. Priores excessivamente suaves induzem viés, enquanto priores excessivamente rugosas convergem lentamente. Além disso, implantar um processo completo de atualização e design bayesiano em tempo real é computacionalmente intensivo, exigindo engenharia por experimento. Os autores investigam se esse complexo processo de inferência e design sequencial pode ser "amortizado" usando aprendizado em contexto, especificamente via arquiteturas Transformer, para criar uma política que mapeie históricos experimentais diretamente para probabilidades de atribuição ótimas sem a necessidade de computação bayesiana explícita a cada etapa.

Metodologia
Os autores propõem um framework centrado em um "experimentador bayesiano em contexto", implementado como uma política Transformer treinada para imitar um "professor de Neyman de posterior bayesiana".

1. O Professor Bayesiano:

   • Modelo: O professor mantém crenças bayesianas não paramétricas sobre médias condicionais ao braço ($\mu_w$) e segundos momentos ($q_w$) usando priors de série gaussiana. Essas priors são indexadas por um parâmetro de suavidade $\beta$.
   • Adaptação Hierárquica: Para lidar com a suavidade desconhecida, o professor emprega uma extensão bayesiana hierárquica onde $\beta$ é tratado como uma variável aleatória com uma prior sobre um intervalo compacto. Isso permite que a posterior se concentre na classe de suavidade apropriada conforme os dados se acumulam.
   • Alocação: O professor atualiza as crenças usando uma quasi-verossimilhança gaussiana simplificada (perda ao quadrado) e atribui a próxima unidade de tratamento de acordo com uma propensão de Neyman posterior: $e_t(x) \propto \sqrt{\bar{v}_{1,t-1}(x)} / (\sqrt{\bar{v}_{0,t-1}(x)} + \sqrt{\bar{v}_{1,t-1}(x)})$, onde $\bar{v}$ é a média da variância residual posterior.
   • Garantias Teóricas: O artigo prova que, sob condições de regularidade brandas, a posterior do professor contrai na taxa não paramétrica minimax (adaptada à suavidade real) e que o design resultante converge para a alocação de Neyman oracle, garantindo eficiência semiparamétrica para o estimador final do ATE (via Pesagem Inversa de Probabilidade Aumentada, AIPW).

2. Expressividade e Construção do Transformer:

   • Arquitetura: Os autores demonstram construtivamente que um Transformer pode representar o mapeamento do professor.
     • Atenção como Estatísticas Suficientes: Blocos de autoatenção agregam tokens históricos para computar estatísticas de resposta empírica e produtos de covariância-vetor, atuando efetivamente como estatísticas suficientes para o modelo de série gaussiana.
     • FFNs como Gradiente Descendente: Redes Feed-Forward (FFNs) por token implementam etapas de gradiente descendente projetado para resolver o problema de regressão de ridge mascarada de rank finito correspondente à atualização da média da posterior bayesiana.
     • Mascaramento Dinâmico: A arquitetura utiliza uma máscara de dimensão efetiva dinâmica que ativa um número de coordenadas de base ($J$) dependente do tamanho da amostra ($N$) e da prior de suavidade. Isso impõe o tradeoff viés-variância: $J$ é escolhido de tal forma que a matriz Gram empírica seja bem condicionada e a precisão da prior não domine.
     • Mistura de Especialistas (MoE): Um portão (gate) de MoE aproxima a posterior hierárquica sobre classes de suavidade. Ele pondera diferentes redes "especialistas" (cada uma correspondendo a um nível de suavidade $\beta_k$) com base em pontuações de verossimilhança marginal aproximadas derivadas de estatísticas de atenção.
   • Generalização via ERM: O artigo analisa o aprendizado desta política via minimização de risco empírico supervisionado (ERM). O Transformer é treinado em trajetórias de pré-treinamento finitas para imitar o comportamento de entrada-saída do professor (mapeamento histórico-para-propensão). A teoria garante que o modelo treinado pode aproximar o comportamento do professor até erros de aproximação e estimação, sem a necessidade de recuperar explicitamente o estado interno da posterior ou executar os passos de gradiente durante a inferência.

Principais Contribuições

1. Design Bayesiano e Adaptação de Suavidade: Os autores definem um professor de Neyman posterior bayesiano que se adapta à suavidade desconhecida do resultado via uma prior hierárquica. Eles provam que este design alcança taxas de contração próximas ao oracle e eficiência de ATE ótima.
2. Expressividade do Transformer: Eles fornecem uma prova construtiva de que Transformers podem representar o mapa de atualização e design bayesiano. Especificamente, mostram como mecanismos de atenção e blocos de gradiente descendente projetado podem realizar regressão de ridge mascarada de rank finito e como um portão MoE pode implementar a adaptação de suavidade.
3. Teoria de Generalização: O artigo estabelece um limite de generalização de ERM, demonstrando que o pré-treinamento supervisionado sobre trajetórias finitas permite que o Transformer aprenda o mapa de experimentação adaptativo, imitando o comportamento de história-para-propensão do professor.
4. Validação Empírica: Experimentos em dados sintéticos mostram que:
   • Um único Transformer, treinado apenas com rótulos de mistura bayesiana (sem um cabeçote MoE explícito), alcança taxas minimax adaptativas à suavidade em níveis de suavidade não vistos.
   • O Transformer de design reproduz o comportamento de alocação de Neyman bayesiano, respondendo direcionamente às variâncias residuais específicas de cada braço, mesmo que não aprenda explicitamente os momentos do resultado em seus estados latentes.
   • A política aprendida melhora a precisão da estimativa do ATE (reduzindo o MSE) em comparação com a randomização uniforme e aproxima-se do desempenho da alocação de Neyman oracle, mesmo sob razões de variância fora da distribuição.

Resultos e Significância
O artigo afirma que Transformers podem servir como "experimentadores bayesianos em contexto amortizados". Ao aprender o mapeamento do histórico experimental para a propensão de tratamento diretamente, esses modelos evitam a necessidade de resolver um problema completo de design bayesiano no momento do deployment.

A significância reside na demonstração de que:

• A amortização é viável: Inferência bayesiana sequencial complexa e design ótimo podem ser destilados em uma política de rede neural fixa.
• A adaptatividade emerge: O modelo pode se adaptar à suavidade da função e às estruturas de variância desconhecidas sem engenharia explícita do parâmetro de suavidade ou pipelines de estimativa de momentos.
• A eficiência é preservada: As políticas aprendidas alcançam eficiência estatística comparável à alocação de Neyman oracle, validando o uso de Transformers para design experimental estatisticamente fundamentado.

Os autores observam limitações, afirmando que o estudo atual é restrito a experimentos de dois braços e ambientes sintéticos, e que estender essa abordagem para espaços de ação mais ricos, resultados atrasados e logs do mundo real permanece uma direção futura. O impacto mais amplo é enquadrado como melhoria na eficiência de dados em domínios como saúde digital e serviços online.