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

Technical Summary: Transformers as Bayesian In-Context Experimenters

Problem Formulation
The paper addresses the challenge of designing adaptive experiments for estimating the Average Treatment Effect (ATE) in sequential, two-arm settings with covariates. The optimal strategy for such experiments is the "oracle" Neyman allocation, which assigns treatment probabilities based on the unknown arm-conditional residual variances to minimize the variance of the ATE estimator. Implementing this optimally requires a Bayesian experimenter that sequentially updates nonparametric beliefs over potential outcome moments (means and second moments) and translates these posterior uncertainties into treatment assignment probabilities.

A critical bottleneck in high-dimensional covariate spaces is the unknown smoothness of the potential outcome functions. Overly smooth priors induce bias, while overly rough priors converge slowly. Furthermore, deploying a full Bayesian update-and-design process in real-time is computationally intensive, requiring per-experiment engineering. The authors investigate whether this complex sequential inference and allocation process can be "amortized" using in-context learning, specifically via Transformer architectures, to create a policy that maps experimental histories directly to optimal assignment probabilities without explicit per-step Bayesian computation.

Methodology
The authors propose a framework centered on a "Bayesian in-context experimenter," implemented as a Transformer policy trained to imitate a "Bayesian posterior Neyman teacher."

1. The Bayesian Teacher:

   • Model: The teacher maintains nonparametric Bayesian beliefs over arm-conditional means ($\mu_w$) and second moments ($q_w$) using Gaussian-series priors. These priors are indexed by a smoothness parameter $\beta$.
   • Hierarchical Adaptation: To handle unknown smoothness, the teacher employs a hierarchical Bayesian extension where $\beta$ is treated as a random variable with a prior over a compact interval. This allows the posterior to concentrate on the appropriate smoothness class as data accumulates.
   • Allocation: The teacher updates beliefs using a simplified Gaussian quasi-likelihood (squared loss) and assigns the next treatment unit according to a posterior Neyman propensity: $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)})$, where $\bar{v}$ is the posterior mean of the residual variance.
   • Theoretical Guarantees: The paper proves that under mild regularity conditions, the teacher's posterior contracts at the minimax nonparametric rate (adapted to the true smoothness) and that the resulting design converges to the oracle Neyman allocation, ensuring semiparametric efficiency for the final ATE estimator (via Augmented Inverse Probability Weighting, AIPW).

2. Transformer Expressivity and Construction:

   • Architecture: The authors constructively demonstrate that a Transformer can represent the teacher's mapping.
     • Attention as Sufficient Statistics: Self-attention blocks aggregate historical tokens to compute empirical response statistics and covariance-vector products, effectively acting as sufficient statistics for the Gaussian-series model.
     • FFNs as Gradient Descent: Token-wise Feed-Forward Networks (FFNs) implement projected gradient descent steps to solve the finite-rank masked ridge regression problem corresponding to the Bayesian posterior mean update.
     • Dynamic Masking: The architecture uses a dynamic effective-dimension mask that activates a number of basis coordinates ($J$) dependent on the sample size ($N$) and the smoothness prior. This enforces the bias-variance tradeoff: $J$ is chosen such that the empirical Gram matrix is well-conditioned and the prior precision does not dominate.
     • Mixture-of-Experts (MoE): A MoE gate approximates the hierarchical posterior over smoothness classes. It weights different "expert" networks (each corresponding to a specific smoothness level $\beta_k$) based on approximate marginal likelihood scores derived from attention-based statistics.
   • Generalization via ERM: The paper analyzes the learning of this policy via supervised empirical risk minimization (ERM). The Transformer is trained on finite pretraining trajectories to imitate the teacher's input-output behavior (history-to-propensity mapping). The theory guarantees that the trained model can approximate the teacher's behavior up to approximation and estimation errors, without needing to explicitly recover the internal posterior state or execute the gradient steps during inference.

Key Contributions

1. Bayesian Design and Smoothness Adaptation: The authors define a Bayesian posterior Neyman teacher that adapts to unknown outcome smoothness via a hierarchical prior. They prove this design achieves near-oracle contraction rates and optimal ATE efficiency.
2. Transformer Expressivity: They provide a constructive proof that Transformers can represent the Bayesian update-and-design map. Specifically, they show how attention mechanisms and projected gradient descent blocks can realize finite-rank masked ridge regression and how an MoE gate can implement smoothness adaptation.
3. Generalization Theory: The paper establishes an ERM generalization bound, demonstrating that supervised pretraining over finite trajectories enables the Transformer to learn the adaptive experimentation map, imitating the teacher's history-to-propensity behavior.
4. Empirical Validation: Experiments on synthetic data show that:
   • A single Transformer, trained only on Bayesian mixture posterior labels (without an explicit MoE head), achieves smoothness-adaptive minimax rates across unseen smoothness levels.
   • The design Transformer reproduces Bayesian-Neyman allocation behavior, responding directionally to arm-specific residual variances, even though it does not explicitly learn the outcome moments in its latent states.
   • The learned policy improves ATE estimation precision (reducing MSE) compared to uniform randomization and approaches the performance of the oracle Neyman allocation, even under out-of-distribution variance ratios.

Results and Significance
The paper claims that Transformers can serve as "amortized Bayesian in-context experimenters." By learning the mapping from experimental history to treatment propensity directly, these models bypass the need for solving a full Bayesian design problem at deployment time.

The significance lies in the demonstration that:

• Amortization is feasible: Complex sequential Bayesian inference and optimal design can be distilled into a fixed neural network policy.
• Adaptivity emerges: The model can adapt to unknown function smoothness and variance structures without explicit engineering of the smoothness parameter or moment estimation pipelines.
• Efficiency is preserved: The learned policies achieve statistical efficiency comparable to the oracle Neyman allocation, validating the use of Transformers for statistically principled experimental design.

The authors note limitations, stating the current study is restricted to two-arm experiments and synthetic environments, and that extending the approach to richer action spaces, delayed outcomes, and real-world logs remains a future direction. The broader impact is framed as improved data efficiency in domains like digital health and online services.