# Stratification Trees for Adaptive Randomization in Randomized Controlled Trials

arXiv: https://arxiv.org/abs/1806.05127 · HTML: https://mlsanity.com/paper/1806.05127/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 13 Jun 2018.

## Built on by (most cited first, 14 of 14)

- [Active Adaptive Experimental Design for Treatment Effect Estimation with Covariate Choices](https://mlsanity.com/paper/2403.03589/) (2024)
- [The Role of Contextual Information in Best Arm Identification](https://mlsanity.com/paper/2106.14077/) (2021)
- [Prediction-Guided Active Experiments](https://mlsanity.com/paper/2411.12036/) (2024)
- [Optimal Best Arm Identification in Two-Armed Bandits with a Fixed Budget under a Small Gap](https://mlsanity.com/paper/2201.04469/) (2022)
- [Best Arm Identification with Contextual Information under a Small Gap](https://mlsanity.com/paper/2209.07330/) (2022)
- [Asymptotically Optimal Fixed-Budget Best Arm Identification with Variance-Dependent Bounds](https://mlsanity.com/paper/2302.02988/) (2023)
- [Worst-Case Optimal Multi-Armed Gaussian Best Arm Identification with a Fixed Budget](https://mlsanity.com/paper/2310.19788/) (2023)
- [Locally Optimal Fixed-Budget Best Arm Identification in Two-Armed Gaussian Bandits with Unknown Variances](https://mlsanity.com/paper/2312.12741/) (2023)
- [Adaptive Experimental Design for Policy Learning](https://mlsanity.com/paper/2401.03756/) (2024)
- [Causal-EPIG: A Prediction-Oriented Active Learning Framework for CATE Estimation](https://mlsanity.com/paper/2509.21866/) (2025)
- [Inference for Batched Adaptive Experiments](https://mlsanity.com/paper/2512.10156/) (2025)
- [Observationally Informed Adaptive Causal Experimental Design](https://mlsanity.com/paper/2603.03785/) (2026)
- [Covariance-Driven Regression Trees: Reducing Overfitting in CART](https://mlsanity.com/paper/2601.07281/) (2026)
- [Optimal Multi-way Decision Trees for Stratified Sampling in Online Controlled Experiments](https://mlsanity.com/paper/2609.23308/) (2026)

## Abstract (arXiv)

This paper proposes an adaptive randomization procedure for two-stage randomized controlled trials. The method uses data from a first-wave experiment in order to determine how to stratify in a second wave of the experiment, where the objective is to minimize the variance of an estimator for the average treatment effect (ATE). We consider selection from a class of stratified randomization procedures which we call stratification trees: these are procedures whose strata can be represented as decision trees, with differing treatment assignment probabilities across strata. By using the first wave to estimate a stratification tree, we simultaneously select which covariates to use for stratification, how to stratify over these covariates, as well as the assignment probabilities within these strata. Our main result shows that using this randomization procedure with an appropriate estimator results in an asymptotic variance which is minimal in the class of stratification trees. Moreover, the results we present are able to accommodate a large class of assignment mechanisms within strata, including stratified block randomization. In a simulation study, we find that our method, paired with an appropriate cross-validation procedure ,can improve on ad-hoc choices of stratification. We conclude by applying our method to the study in Karlan and Wood (2017), where we estimate stratification trees using the first wave of their experiment.