ML-assisted Randomization Tests for Detecting Treatment Effects in A/B Experiments

TL;DR

This paper introduces ML-assisted randomization tests for A/B experiments that leverage machine learning models to detect complex treatment effects, including heterogeneity and interference, with theoretical and empirical validation.

Contribution

It presents a novel method combining ML models with randomization tests to detect complex treatment effects in experiments, ensuring finite-sample validity.

Findings

Effective detection of heterogeneous treatment effects

Robustness to interference in experimental data

Theoretical guarantees and empirical validation

Abstract

Experimentation is widely utilized for causal inference and data-driven decision-making across disciplines. In an A/B experiment, for example, an online business randomizes two different treatments (e.g., website designs) to their customers and then aims to infer which treatment is better. In this paper, we construct randomization tests for complex treatment effects, including heterogeneity and interference. A key feature of our approach is the use of flexible machine learning (ML) models, where the test statistic is defined as the difference between the cross-validation errors from two ML models, one including the treatment variable and the other without it. This approach combines the predictive power of modern ML tools with the finite-sample validity of randomization procedures, enabling a robust and efficient way to detect complex treatment effects in experimental settings. We demonstrate this combined benefit both theoretically and empirically through applied examples.