# Optimal Multi-way Decision Trees for Stratified Sampling in Online Controlled Experiments

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 7 of its 23 references that are in the map.
- First posted on arXiv 20 Sep 2026.

## Builds on (most cited first, 7 of 7)

- [Machine Learning for Variance Reduction in Online Experiments](https://mlsanity.com/paper/2106.07263/) (2021)
- [Stratification Trees for Adaptive Randomization in Randomized Controlled Trials](https://mlsanity.com/paper/1806.05127/) (2018)
- [Optimal binning: mathematical programming formulation](https://mlsanity.com/paper/2001.08025/) (2020)
- [Scalable Optimal Multiway-Split Decision Trees with Constraints](https://mlsanity.com/paper/2302.06812/) (2023)
- [Subset Selection for Stratified Sampling in Online Controlled Experiments](https://mlsanity.com/paper/2509.15576/) (2025)
- [Improving Sensitivity in A/B Tests: Integrating CUPED with Trimmed Mean Techniques](https://mlsanity.com/paper/2510.03468/) (2025)
- [Interpretable clustering via optimal multi-way decision trees](https://mlsanity.com/paper/2602.13586/) (2026)

## Abstract (arXiv)

Online controlled experiments, or A/B tests, are widely used to estimate causal effects on digital platforms. A central challenge is to improve experimental sensitivity, or statistical power, without increasing the experimental sample size. Stratified sampling is a classical variance reduction technique; however, its effectiveness depends critically on how the strata are constructed. We thus propose an optimization-based stratification framework for stratified sampling using optimal multi-way decision trees. Our method, called Optimal Multi-way Stratification Trees (OMST), formulates stratification as a path-selection problem over a feature graph. The selected paths define interpretable stratification rules and are optimized using an exact variance-minimizing binary optimization formulation under continuous proportional allocation and a Neyman-type optimal allocation. We incorporate supervised optimal binning to generate outcome-relevant candidate splits for numerical features. Furthermore, we introduce reduction procedures for redundant candidate paths and assignment constraints, substantially reducing the optimization problem size. Experiments on both a real-world and a simulated dataset demonstrate that OMST achieves comparable or superior variance reduction to existing methods while maintaining shallow and interpretable stratification trees.