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

Adaptive experiments for average treatment effect (ATE) estimation typically require complex, per-experiment tuning to balance statistical efficiency with valid inference. This paper investigates whether pretrained transformers can learn to perform this sequential variance-estimation and allocation process automatically, effectively acting as amortized Bayesian experimenters that map experimental histories directly to optimal treatment propensities.

The authors propose a transformer architecture that imitates a Bayesian posterior Neyman teacher. The model uses attention-based sufficient statistics to aggregate experimental history and projected gradient descent blocks to perform finite-rank masked ridge regression. To handle unknown outcome smoothness, the authors employ a mixture-of-experts (MoE) architecture where the gate approximates a hierarchical Bayesian posterior over smoothness classes. They prove that this policy can be learned via empirical risk minimization (ERM) using supervised pretraining, providing an algorithm-imitation guarantee that the learned model achieves near-oracle contraction rates.

The study demonstrates that a single transformer, trained on Bayesian mixture posterior labels, achieves smoothness-adaptive minimax rates across unseen smoothness levels without requiring explicit smoothness inputs. Empirical results show that the design transformer reproduces Bayesian-Neyman allocation behavior—responding directionally to arm-specific residual variance—without explicitly learning the underlying outcome moments. In online deployment, the learned policy improves ATE estimation precision relative to uniform randomization and approaches the performance of the oracle Neyman allocation, even under out-of-distribution variance ratios.

This approach reduces the need for manual engineering in adaptive experimental design, offering a scalable, automated system for clinical trials and online platforms. By distilling complex Bayesian updating into a direct history-to-propensity mapping, the method enables efficient, statistically principled experimentation that adapts to the complexity of the data in real-time.