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

2026-06-30

AI Summary

This study addresses the challenge of achieving both statistical efficiency and valid inference in average treatment effect (ATE) estimation through adaptive randomization. To this end, the authors propose a Bayesian contextual experimenter framework that employs a mixture-of-experts Transformer policy to imitate a Bayesian posterior Neyman oracle. Leveraging experimental history, the method nonparametrically updates beliefs about potential outcomes and dynamically adjusts treatment assignment probabilities. The approach innovatively integrates attention mechanisms to construct sufficient statistics, Bayesian updating under Gaussian series priors, projected gradient descent, and smoothness-adaptive posterior gating to accommodate unknown smoothness of the outcome functions. Empirical results demonstrate that the proposed strategy effectively learns and accurately replicates the oracle assignment rule, significantly outperforming existing baselines in ATE estimation accuracy.

Abstract

Adaptive experiments for average treatment effects (ATE) require randomized allocations balancing valid inference with statistical efficiency. The oracle design is a covariate-dependent Neyman rule governed by unknown arm-conditional outcome variances. We investigate whether this sequential variance-estimation and allocation process can be amortized via in-context learning. We introduce Bayesian in-context experimenters: transformer policies trained to imitate a Bayesian posterior Neyman teacher. The teacher updates nonparametric beliefs over potential outcomes using experimental history to assign posterior Neyman treatment probabilities. This design converges to the oracle rule, supporting efficient ATE inference. Transformers constructively implement this mapping through attention-based sufficient statistics and projected gradient descent, imitating Bayesian updating for Gaussian-series priors. To address unknown outcome smoothness, we combine smoothness-indexed experimenters using a mixture-of-experts transformer. The gate acts as a hierarchical posterior over smoothness classes, concentrating on near-oracle experts. By bounding the complexity of the transformer class, we prove this amortized policy can be learned via empirical risk minimization using supervised pretraining. Experiments confirm accurate teacher imitation, adaptive allocation, and improved ATE precision over baselines.

Problem

• average treatment effect
• adaptive experiment
• in-context learning
• Bayesian inference
• smoothness adaptation

Innovation

• Bayesian in-context learning
• Transformer-based experimental design
• Adaptive ATE estimation
• Mixture-of-experts
• Smoothness-adaptive inference