CORTEXA
← Browse
arxivcs.LGcs.AI2026-06-30

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

Jiachun Li, David Simchi-Levi

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.

View free PDFSource page

Related papers

arxivcs.CLcs.AIcs.LG2026-06-29

MATCH: Modulating Attention via In-Context Retrieval for Long-Context Transformers

Linrui Ma, Chun Hei Lo, Xinyu Wang, Peng Lu, Xihao Yuan, Hanting Chen, et al.

The quadratic computational cost of traditional attention mechanisms poses a major bottleneck to the scalability and practical deployment of large language models (LLMs), particularly in long-context scenarios. To improve efficiency, existing approaches often enforce rigid struct…

View free PDFSource page
arxivcs.LGcs.AI2026-07-17

In-context learning of closed form solution to simple linear regression task using transformer with linear self-attention

Katsuyuki Hagiwara

In-context learning is a remarkable property of transformers and has recently received a lot of interest. In many studies of in-context learning, it has been shown that transformers are capable of implementing solver for linear and non-linear regression problems, in which the mos…

View free PDFSource page
arxivcs.LGcs.AI2026-07-20

Censoring-Aware In-Context Learning for Generalized Supplier Lead Time Estimation in Supply Chain Planning

Christopher Wang, Sebastien Ouellet, Behrouz Haji Soleimani, Ali Etemad

Supplier lead time forecasting is a central input to material requirements planning, inventory optimization, and supply chain risk management. However, many industrial lead time datasets are naturally right-censored: at the time forecasts are required, some orders have not yet ar…

View free PDFSource page
arxivcs.LGcs.AI2026-07-24

Optimization of time-consuming experimental conditions using pseudo-experimental data guided by adaptive polynomial regression

Hirotaka Sugawara, Yujin Taguchi, Kei Minagawa, Yusuke Hiki, Takashi Morikura, Akira Funahashi

Bayesian optimization (BO) is an optimization method that sequentially proposes the next candidate explainable variables for optimizing target variables by balancing exploration and exploitation. BO is often used under a limited evaluation budget, such as hyperparameter tuning of…

View free PDFSource page
arxivcs.LGcs.AIcs.CL2026-07-15

Data-Efficient Adaptation of LLMs via Attention Head Reweighting

Tuomas Oikarinen, Zixiao Chen, Charlotte Siska, Tsui-Wei Weng, Chandan Singh, Jianfeng Gao

Learning effectively from limited data is critical in domains like security where labeled examples are scarce. Large language models (LLMs) have demonstrated some capabilities for data-efficient learning, especially through parameter-efficient adaptation methods, but continue to…

View free PDFSource page