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

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.

• The paper introduces transformers as amortized Bayesian experimenters that adapt experimental designs for semiparametrically efficient average treatment effect estimation.
• It maps Bayesian updates to transformer primitives—attention, feedforward networks, and MoE gating—achieving minimax nonparametric contraction rates.
• Empirical analysis and theoretical guarantees confirm that the approach adapts to unknown smoothness, reducing mean squared error in treatment effect estimation.

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

Problem Formulation

The paper investigates whether amortized in-context learning with pretrained transformer architectures can efficiently support adaptive experimental design for average treatment effect (ATE) estimation. The classical Neyman allocation rule for optimal ATE estimation depends on covariate-specific residual outcome variances in each arm, which are unknown and must be estimated from past experiment data. Bayesian adaptive designs sequentially update beliefs about these variances and assign treatments accordingly. The challenge lies in the nonparametric, high-dimensional estimation of outcome moments, where unknown function smoothness influences statistical precision and convergence rates, and in automating optimal allocation without custom per-deployment engineering.

Bayesian Posterior Neyman Teacher and Smoothness Adaptation

The authors formalize the Bayesian "posterior Neyman teacher" as a sequential process updating beliefs over arm-specific conditional means and second moments using orthonormal basis expansion, with Gaussian-series priors parametrized by a smoothness index $\beta$. The variance belief is projected onto bounded intervals, and adaptive assignment probabilities are computed as posterior Neyman propensities, converging to the oracle allocation. They prove that with correct prior smoothness, the posterior contracts to the oracle outcome functions at the minimax nonparametric rate $\epsilon_{t} \asymp N_{t}^{-\beta/(2\beta+d)}$, yielding semiparametrically efficient ATE inference via AIPW.

When smoothness is unknown, a hierarchical Bayesian hyperprior on $\beta$ achieves smoothness adaptation: the posterior automatically concentrates on the near-oracle smoothness slab, maintaining rates within logarithmic factors of optimal. Theoretical results guarantee that this hierarchical teacher delivers valid inference and smoothness-adaptive efficiency under general sub-exponential noise.

Figure 1: Smoothness-adaptive estimation rates for increasing smoothness indices $\beta$, demonstrating the transformer’s empirical scaling matches the Bayesian-minimax contraction.

Transformer Realization: Computational and Statistical Mapping

The key technical advance is to constructively show that transformer architectures can explicitly represent the Bayesian experimenter through:

• Attention-based sufficient statistics: Linear self-attention executes normalized aggregation of history-level features—empirical averages, covariance products, and dot products—directly mapping to Bayesian statistics.
• Tokenwise feedforward networks (FFNs): These implement coordinate masking, dynamic active dimension selection, diagonal prior shrinkage, and projected gradient descent for masked ridge regression updates needed for posterior mean estimation.
• Mixture-of-experts (MoE) gating: The MoE head approximates hierarchical Bayesian smoothness selection, concentrating on experts with penalized empirical risk and properly scaling complexity penalty.

By decomposing the Bayesian update into these primitives, the paper proves that a ReLU-attention transformer of width $O(J_n)$ and depth $O(\log n)$ contains a parameter configuration that approximates the smoothness-adaptive Bayesian estimator at minimax rates uniformly over all arms, moments, experts, and prefixes. The architecture is parameterized to hide poly-logarithmic factors, and the MoE gate is shown to suppress suboptimal experts exponentially as data accumulate.

Supervised Pretraining and Generalization Guarantees

Rather than hand-engineering the transformer to mimic the Bayesian pipeline, practical models are trained by empirical risk minimization (ERM) over finite pretraining trajectories, with history-prefix-to-propensity labels supplied by the teacher. The paper establishes generalization bounds: if a transformer class contains a comparator with risk $\epsilon_n^2$, the ERM-trained model achieves population risk within $O(\sqrt{P_n/N_{pre}})$ of optimal, where $P_n$ is parameter complexity. The proof deals with trajectory-level independence (not token-level) and polynomial hidden-state Lipschitz constants, ensuring stability and uniform deviation. The guarantee applies both to estimation transformers and to design transformers that output allocation propensities.

Empirical Analysis

Synthetic experiments across seven unseen smoothness levels demonstrate that the trained estimation transformer, even without an explicit MoE head, achieves empirical log–log slopes matching the Bayesian-minimax contraction rates. The design transformer, trained to imitate the Neyman teacher allocation rule, is shown to respond correctly to arm-specific residual variance changes, confirming that it learns a variance-sensitive allocation policy rather than simplistic shortcuts. Layerwise probes reveal that while propensities become linearly decodable, outcome moments are not explicitly learned—indicating the transformer directly learns the mapping from history to allocation.

Online deployment experiments show that the transformer closely tracks the teacher’s allocation convergence to oracle Neyman propensities, with matching step-level fluctuations and only minor gaps in late-phase imitation. Downstream ATE estimation out-of-distribution demonstrates improved mean squared error (MSE) relative to uniform randomization, closing in on oracle Neyman allocation results, including tasks with extreme variance ratios.

Theoretical and Practical Implications

The results establish that pretrained transformers can serve as amortized Bayesian experimenters, learning complex adaptive design rules directly from history, and efficiently implementing smoothness-adaptive experimental policies. The constructive mapping of Bayesian statistical operations to transformer primitives provides theoretical insight into the expressivity limits of modern sequence models for causal and experimental tasks. Practically, this enables automated deployment of principled allocation in clinical and online experiments, improving data efficiency without reengineering for each use case.

Theoretical implications include formalizing attention and FFNs as modules for nonparametric Bayesian estimation, proving near-oracle contraction and efficiency, and showing hierarchical adaptation and mixture selection are tractable with deep learning architectures. Future directions include extending to multi-arm settings, delayed outcomes, real experimental logs, and pretraining on misspecified distributions.

Conclusion

This paper rigorously demonstrates that transformer architectures can be trained, via in-context learning and empirical risk minimization, to act as amortized Bayesian experimenters for adaptive experiments. The explicit connection to nonparametric Bayesian updating, proven minimax rates, smoothness adaptation, and semiparametric efficiency, together with empirical validation, push forward the understanding of transformer expressivity in causal inference and experimental design. Broader applications include improved data efficiency in digital health, education, and online services through automated statistically principled allocation policies.