Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Randomized Experiments Using Foundation Models

Piersilvio De Bartolomeis, Javier Abad, Guanbo Wang, Konstantin Donhauser, Raymond M. Duch, Fanny Yang, Issa J. Dahabreh

TL;DR

This work tackles the high cost and imprecision of randomized experiments by introducing Hybrid Augmented Inverse Probability Weighting (H-Aipw), which safely integrates multiple foundation-model predictions with experimental data to estimate the average treatment effect. The method forms a convex ensemble of AIPW estimators, weighting them to minimize asymptotic variance while preserving valid inference even when external predictions are biased, and it remains consistent and asymptotically normal with variance no larger than the standard AIPW. The authors provide a step-by-step MLOps-ready recipe for implementing H-Aipw with large language models, and they demonstrate substantial precision gains across eight social-science experiments, equivalent to saving up to about 20% in sample size in some settings. They also show that increasing LLM scale and inference-time compute improves prediction accuracy, further reducing estimator variance, and discuss practical considerations such as model selection, covariance estimation, and cross-fitting. Overall, H-Aipw offers a principled, scalable pathway to leverage foundation-model predictions for efficient, valid causal inference in randomized trials, with potential implications for medicine and other data-constrained domains.

Abstract

Randomized experiments are the preferred approach for evaluating the effects of interventions, but they are costly and often yield estimates with substantial uncertainty. On the other hand, in silico experiments leveraging foundation models offer a cost-effective alternative that can potentially attain higher statistical precision. However, the benefits of in silico experiments come with a significant risk: statistical inferences are not valid if the models fail to accurately predict experimental responses to interventions. In this paper, we propose a novel approach that integrates the predictions from multiple foundation models with experimental data while preserving valid statistical inference. Our estimator is consistent and asymptotically normal, with asymptotic variance no larger than the standard estimator based on experimental data alone. Importantly, these statistical properties hold even when model predictions are arbitrarily biased. Empirical results across several randomized experiments show that our estimator offers substantial precision gains, equivalent to a reduction of up to 20% in the sample size needed to match the same precision as the standard estimator based on experimental data alone.

Efficient Randomized Experiments Using Foundation Models

TL;DR

This work tackles the high cost and imprecision of randomized experiments by introducing Hybrid Augmented Inverse Probability Weighting (H-Aipw), which safely integrates multiple foundation-model predictions with experimental data to estimate the average treatment effect. The method forms a convex ensemble of AIPW estimators, weighting them to minimize asymptotic variance while preserving valid inference even when external predictions are biased, and it remains consistent and asymptotically normal with variance no larger than the standard AIPW. The authors provide a step-by-step MLOps-ready recipe for implementing H-Aipw with large language models, and they demonstrate substantial precision gains across eight social-science experiments, equivalent to saving up to about 20% in sample size in some settings. They also show that increasing LLM scale and inference-time compute improves prediction accuracy, further reducing estimator variance, and discuss practical considerations such as model selection, covariance estimation, and cross-fitting. Overall, H-Aipw offers a principled, scalable pathway to leverage foundation-model predictions for efficient, valid causal inference in randomized trials, with potential implications for medicine and other data-constrained domains.

Abstract

Randomized experiments are the preferred approach for evaluating the effects of interventions, but they are costly and often yield estimates with substantial uncertainty. On the other hand, in silico experiments leveraging foundation models offer a cost-effective alternative that can potentially attain higher statistical precision. However, the benefits of in silico experiments come with a significant risk: statistical inferences are not valid if the models fail to accurately predict experimental responses to interventions. In this paper, we propose a novel approach that integrates the predictions from multiple foundation models with experimental data while preserving valid statistical inference. Our estimator is consistent and asymptotically normal, with asymptotic variance no larger than the standard estimator based on experimental data alone. Importantly, these statistical properties hold even when model predictions are arbitrarily biased. Empirical results across several randomized experiments show that our estimator offers substantial precision gains, equivalent to a reduction of up to 20% in the sample size needed to match the same precision as the standard estimator based on experimental data alone.

Paper Structure

This paper contains 45 sections, 3 theorems, 47 equations, 6 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Integrating foundation models
  4. Integrating observational data
  5. Background on randomized experiments
  6. A class of valid estimators: Augmented Inverse Probability Weighting
  7. Methodology
  8. Hybrid Augmented Inverse Probability Weighting
  9. Asymptotic validity and efficiency
  10. Step-by-step recipe with Large Language Models
  11. Experiments
  12. H-AIPW offers improved statistical precision
  13. Estimators and metrics
  14. Results
  15. Image treatments and contamination
  16. ...and 30 more sections

Key Result

Proposition 1

Let $\tilde{\mathcal{D}}$ be an auxiliary sample, independent of $\mathcal{D}$. Let $\widehat{h}$ be the outcome regression model trained on $\tilde{\mathcal{D}}$, and let $h^\dagger$ be a square-integrable limit such that for $a\in\{0,1\}$, where $\mathbb P^\ast$ denotes the joint law of $(\mathcal{D},\mathcal{\tilde{D}})$. Then, it follows that $\widehat{\theta}_{\textsc{aipw}}(\widehat{h})$ is

Figures (6)

  • Figure 1: Our estimator achieves the same statistical precision as the standard estimator with up to 20% fewer samples. Each study is subsampled to $n = 75$. We plot here the percentage reduction in the sample size needed to match the confidence interval width of the standard estimator using ours.
  • Figure 2: Illustration of the Hybrid Augmented Inverse Probability Weighting ($\textsc{H-Aipw}$) estimator. For each unit $i$ we observe covariates $X_i$, treatment $A_i$ and outcome $Y_i$; $(X_i,A_i,Y_i)_{i=1}^n$ forms the experimental data. An outcome regression model $\widehat{h}$ fitted to this sample yields the standard $\textsc{Aipw}$ estimate $\widehat{\theta}(\widehat{h})$. In addition, foundation models trained on external data provide the candidate outcome regression models $f_1,\ldots, f_k$, which result in $k$ competing $\textsc{Aipw}$ estimates $\widehat{\theta}(f_1), \ldots,\widehat{\theta}(f_k)$. By integrating the outcome regression models trained on a large external sample, rather than fitting a single model on the small experimental sample, $\textsc{H-Aipw}$ can reduce the variance of the average treatment effect estimate.
  • Figure 3: Examples of a system and user prompts used to generate synthetic responses for fahey2023principled.
  • Figure 4: Impact of model scale and inference-time compute on the performance of $\textsc{H-Aipw}$ in the study by fahey2023principled. (Left) Model scale: \ref{['fig:scaling_laws']} shows the relationship between the estimate of the $\textsc{H-Aipw}$ variance (average on $R=10k$ repetitions, sample size $n=50$) and mean squared error (MSE) for LLMs of varying sizes ($10$ prompts at inference time). (Right) Inference-time compute: \ref{['fig:test_compute']} shows the impact on the MSE of increasing the number of prompts at inference time and averaging the resulting predictions.
  • Figure 5: Impact of increasing the number of models in $\textsc{H-Aipw}$ on precision and validity in the study by fahey2023principled. Models are sequentially incorporated based on their mean squared error (MSE), starting with LLaMA 3 70B (lightest red, $k=1$) and ending with Gemma 2 27B (darkest red, $k=7$), following \ref{['fig:scaling_laws']}. The left panel shows the empirical variance, while the right panel shows empirical coverage. The standard $\textsc{Aipw}$ estimator is included for reference. Each experiment is averaged over $R=10k$ repetitions, with significance level set to $\alpha=0.05$.
  • ...and 1 more figures

Theorems & Definitions (4)

  • Proposition 1: Asymptotic behavior of Aipw
  • Theorem 2: Asymptotic behavior of H-Aipw
  • Lemma A.1
  • proof : Proof of Lemma \ref{['lem:variance_bound']}