Machine Learning for Variance Reduction in Online Experiments
Year: 2022

Abstract:
We consider the problem of variance reduction in randomized controlled trials, through the use of covariates correlated with the outcome but independent of the treatment. We propose a machine learning regression-adjusted treatment effect estimator, which we call MLRATE. MLRATE uses machine learning predictors of the outcome to reduce estimator variance. It employs cross-fitting to avoid overfitting biases, and we prove consistency and asymptotic normality under general conditions. MLRATE is robust to poor predictions from the machine learning step: if the predictions are uncorrelated with the outcomes, the estimator performs asymptotically no worse than the standard difference-in-means estimator, while if predictions are highly correlated with outcomes, the efficiency gains are large. In A/A tests, for a set of 48 outcome metrics commonly monitored in Facebook experiments the estimator has over 70% lower variance than the simple differencein-means estimator, and about 19% lower variance than the common univariate procedure which adjusts only for pre-experiment values of the outcome.

heading: 1 Introduction
text: While sample sizes are typically larger for online experiments than traditional field experiments, the desired minimum detectable effect sizes may be small, and the outcome variables of interest may be heavy-tailed. Even with quite large samples, statistical power may be low. Variance reduction methods play a key role in these settings, allowing for precise inferences with less data [12, 32, 38].
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
One common technique involves "adjusting" the simple difference-in-means estimator to account for covariate imbalances between the test and control groups [12, 22], with the magnitude of the adjustment depending on both the magnitude of those imbalances, and the correlation between the covariates and the outcome of interest. If covariates are highly correlated with the outcome, then the treatment effect estimator’s variance will decrease substantially.
Performing this adjustment procedure with the pre-experiment values of the outcome variable itself as the covariate can greatly reduce confidence interval (CI) width, if the outcome exhibits high autocorrelation. A natural question is how to adjust for multiple covariates, which may have a complicated nonlinear relationship with the outcome variable. Using many covariates in a machine learning (ML) model, it may be possible to develop a proxy highly correlated with the outcome variable and hence generate further variance reduction gains. This raises both statistical and scalability issues, however, as it is unclear how traditional justifications for linear regression adjustment with a fixed number of covariates translate to the case of general, potentially very complex ML methods, and it may not be scalable to generate new predictions every time an experiment’s results are queried.
This paper makes three contributions. First, we propose an easy-to-implement and practical estimator which can take full advantage of ML methods to perform regression adjustment across a potentially large number of covariates, and derive its asymptotic properties. We name the procedure MLRATE, for “machine learning regression-adjusted treatment effect estimator”. MLRATE uses cross-fitting (e.g. [3, 9, 25]), which simplifies the asymptotic analysis and guarantees that the “naive” CIs which do not correct for the ML estimation step are asymptotically valid. We also ensure robustness of the estimator to poor quality predictions from the ML stage, by including those ML predictions as a covariate in a subsequent linear regression step. Our approach is agnostic or model-free in two key respects—we do not assume that the ML model converges to the truth, and in common with [22], in the subsequent linear regression step we do not assume that the true conditional mean is linear. Second, we demonstrate that the method works well for online experiments in practice. Across a variety of metrics, the estimator reduces variance in A/A tests by around 19% on average relative to regression adjustment for pre-experiment outcomes only. Some metrics see variance reduction of 50% or more. Variance reduction of this magnitude can amount to the difference between experimentation being infeasibly noisy and being practically useful. Third, we sketch how the computational considerations involved in implementing MLRATE at scale can be surmounted.

heading: 2 Outcome prediction for variance reduction

heading: 2.1 Setup & motivation
text: The data consist of a vector of covariates X , an outcome variable Y , and a binary treatment indicator T . The treatment is assigned randomly and independently of the covariates. For observations i = 1, 2, . . . , N , the vector (Yi, Xi, Ti) is drawn iid from a distribution P . To motivate our main estimator and illustrate some of the central ideas, first consider a “difference-in-difference”-style estimator, where we train an ML model g(X) predicting Y from X , and then compute the difference between the test and control group averages of Y − g(X). If we treat the estimated ML model g as non-random and ignore its dependence on the sample, the resulting estimator has the same expectation as the usual difference-in-means estimator where we compute the difference between the test and control group averages of Y . This is because g(X) and T are independent, and hence E[Y − g(X) | T = 1]− E[Y − g(X) | T = 0] = E[Y | T = 1]− E[Y | T = 0]. Furthermore, if g(X) is a good predictor of Y , then V ar(Y ) will exceed V ar(Y − g(X)), and the difference-indifference estimator based on averages of Y − g(X) will be lower variance than the difference-inmeans estimator based on averages of Y .
MLRATE differs in two main respects from the heuristic argument above. First, instead of directly subtracting the ML predictions g(X) from the outcome Y , we include them as a regressor in a subsequent linear regression step. This guarantees robustness of the estimator to poor, even asymptotically inconsistent predictions: regardless of how bad the outcome predictions from the ML step are, MLRATE has an asymptotic variance no larger than the difference-in-means estimator. Second, we use cross-fitting to estimate the predictive models, so that the predictions for every observation are generated by a model trained only on other observations. This allows us to control the randomness in the ML function ignored in the argument above. We derive the asymptotic distribution of this regression-adjusted estimator, and show that the usual, “naive” CIs for the average treatment
effect (ATE), which ignore the randomness generated by estimating the predictive models, are in fact asymptotically valid. Thus asymptotically the ML step can only increase precision, and introduces no extra complications in computing CIs.

heading: 2.2 Related work
text: Our work is closely related to the large literature on semiparametric statistics and econometrics, in which low-dimensional parameters of interest are estimated in the presence of high-dimensional nuisance parameters [6, 29, 23, 33, 21, 37]. A common approach in this literature appeals to Donsker conditions and empirical process theory to control the randomness generated by the estimation error in the nuisance function [1, 36, 35, 18]. This approach is less appealing in this context, as it would greatly restrict the kind of ML methods that could be used for the prediction step [9], and hence the variance reduction attainable. The idea of instead using sample-splitting in semiparametric problems— estimating nuisance parameters on one subset of data and evaluating them on another—dates back at least to [6], with subsequent contributions by [30, 19, 7], among others. More recent applications of this idea, also referred to as “cross-fitting”, include [9, 3, 4, 40]. This paper is especially similar in spirit to “double machine-learning” [9], which combines sample-splitting with the use of Neymanorthogonal scores, which have the property of being insensitive to small errors in estimating the nuisance function. Although the results of [9] do not directly carry over to our setting, we use similar arguments to establish our results. A second strand of related literature concerns “agnostic” regression adjustment, which delivers consistent estimates of the average treatment estimate even when the regression model is misspecified [39, 14, 22, 16, 11]. The procedure described in [11] is particularly relevant, as it shares the same structure of first estimating nonlinear models and then calibrating them in a linear regression step, although in contrast to their work we use sample-splitting to allow for a very general class of nonlinear ML models. A third strand of the literature considers improving estimator precision in the context of large-scale online experiments [8, 12, 10, 38].
Relative to these literatures, our contribution is to describe an estimator which i) delivers substantial variance reduction when good ML predictors of the outcome variable are available, and ii) performs well in the presence of poor-quality predictions, even allowing for predictive models which never converge to the truth with infinite data. In particular, MLRATE is guaranteed to never perform worse, asymptotically, than the difference-in-means estimator, even with arbitrarily poor predictions. In contrast to low-dimensional regression adjustment, we give formal statistical guarantees on inference even when complex ML models are used to predict outcomes; in contrast to double-ML, as applied to randomized experiments, our proposed estimator need not be semiparametrically efficient, but allows for inconsistent estimates of the nuisance parameters. Finally, this methodology is practical and computationally efficient enough to be deployed at large scale, and we show with Facebook data that MLRATE can deliver substantial additional variance reduction beyond the existing state-of-the-art commonly used in practice, of linear regression adjustment for pre-experiment covariates [12, 38].

heading: 2.3 Estimation and inference with MLRATE
text: The linear regression-adjusted estimator of the ATE is the OLS estimate of α1 in the regression
Yi = α0 + α1Ti + α2Xi + α3Ti(Xi −X) + i, (1)
where X is the average of Xi over all i. The covariates Xi may be multivariate, but are of fixed dimension that does not grow with the sample size. The analysis in [22] establishes that the OLS estimator for α̂1 is a consistent and asymptotically normal estimator of the ATEE[Y | T = 1]−E[Y | T = 0], and the robust, Huber-White standard errors are asymptotically valid. In contrast to this setting, we wish to capture complex interactions and nonlinearities in the relationship between the outcome and covariates, and allow for a vector of covariates with dimension potentially increasing
with the sample size. To this end, we propose the following procedure. We assume throughout that N is evenly divisible by K, to simplify notation.
Algorithm 1: Estimation and inference with MLRATE Input: Data (Yi, Xi, Ti)Ni=1 split uniformly at random into K equal-sized splits. Ik := index set of the k-th split and Ick := {1, 2, . . . , N} \ Ik, ∀k.M, a supervised learning algorithm. Result: ML regression-adjusted ATE estimator α̂1, asymptotic variance σ̂2. for k ← 1 to K do
Generate the function ĝk predicting Yi given Xi, by applyingM to the sample (Yi, Xi)i∈Ick . end Compute g = 1N ∑ i ĝk(i)(Xi), where k(i) := the split index containing observation i; Compute α̂1 as the OLS estimator for α1 in Yi = α0 + α1Ti + α2ĝk(i)(Xi) + α3Ti ( ĝk(i)(Xi)− g ) + εi; Compute σ̂2 according to (11).
Section 2.4 proves the statistical validity of this estimation and inference procedure. Note that if the cross-fitted, random functions {ĝk}Kk=1 were replaced by a single, fixed function g, MLRATE reduces to the standard linear regression-adjusted estimator. Instead the relation between the covariatesXi and the outcome Yi is itself estimated from the data. Intuitively this should help with variance reduction, as the estimated proxy ĝk(i)(Xi) may be highly correlated with Yi, but with the challenge that the dependence of the ĝk’s on the data complicates the analysis of the statistical properties of the treatment effect estimator α̂1. Our main technical result assuages this concern, showing that the asymptotic distribution of α̂1 is not impacted and thus it is a consistent, asymptotically normal estimator of the ATE. CIs with level 100(1− a) percent are given in the usual way by α̂1 ± Φ−1(1− a/2)σ̂/ √ N , where Φ is the CDF of the standard normal distribution.
Remark 1. The chief purpose of cross-fitting is to avoid bias from overfitting. With sufficiently flexible ML models, in-sample predictions would be close to the outcomes Yi. The linear regression step would then amount to adjusting for the outcome variable itself, which is correlated with the treatment, and this may introduce severe attenuation bias into estimates of the treatment effect. By generating predictions only on out-of-sample data, we ensure the adjustment covariate is independent of the treatment.
Remark 2. In online experiments, only a subset of users are typically assigned to any given experiment. To maximize training data and minimize compute costs, an equally valid variation on the above is to perform the cross-fitting ML step once, using data from all users, whereas the linear regression step must occur separately for every experiment of interest, using only the users in that experiment.
Remark 3. Alternatively, one may estimate a single ML model entirely on pre-experiment data. For an experiment starting at time t, we may train a model predicting time t − 1 outcomes from time t − 2 covariates, and then use that model to predict time t outcomes from time t − 1 covariates. Those model predictions can then be treated as any other covariate, as they are entirely a function of pre-experiment data, and the results of [22] apply. Although simpler, this approach suffers from the drawbacks that it requires some history of the outcome metric to exist even pre-experiment, and that the predictive model may perform worse if the relationship between covariates and outcomes changes over time.
Remark 4. The choice of K does not affect the asymptotic distribution of the estimator, although it may matter in finite samples. As [9] note, in cross-fitting applications involving estimating highdimensional nuisance functions with small samples, larger values of K (e.g. K = 4 or 5) may perform better. Much larger values of K may be unattractive, however, given diminishing returns in model performance and the extra compute cost. In the simulations and empirical examples in Section 3 we show that good performance is achievable even with the low computation choice of K = 2.
We now sketch the main technical result. Beyond standard regularity conditions, the main assumption is that for each split k, the estimated functions ĝk converge to some g0 in the sense that ∫ [ĝk(X)− g0(X)] 4dP →p 0. This condition is quite weak in two aspects: On the one hand, it only requires convergence of the ĝk to g0, and not convergence at a particular rate. Such consistency results are available for many common ML algorithms, including random forests (see [2] and references therein), gradient boosted decision trees [5], deep feedforward neural nets [13], and regularized linear regression in some asymptotic regimes [20]. On the other hand, we allow the ML modelling step
to be misspecified and inconsistent: there is no requirement that g0(X) = E[Y | X], although a poorly-specified ML model may limit the variance reduction obtained. Allowing for inconsistent estimators is an especially important advantage in the presence of high-dimensional covariates, as in such settings there is no general guarantee that ML estimators will be consistent if the number of covariates grows faster than log(N), due to the curse of dimensionality [31].
To make explicit the dependence on the ĝk’s, we denote MLRATE by α̂1({ĝk}Kk=1). We denote by α̂1(g0) the linear regression adjustment estimator as in (1) where we adjust for the covariate g0(Xi). This latter estimator is infeasible as g0 is unknown, but we prove in Theorem 1 below that the two estimators are asymptotically equivalent, i.e. √ N [α̂1({ĝk}Kk=1)− α̂1(g0)]→p 0. Deriving the asymptotic distribution of α̂1(g0) is straightforward, and from this equivalence we conclude that α̂1({ĝk}Kk=1) shares the same asymptotic distribution.

heading: 2.4 The asymptotic behavior of MLRATE
text: Define the covariate vector as a function of an arbitrary (possibly random) function g, Z(g) = (1, T, g(X), T g(X))>, and define Zi(g) = (1, Ti, g(Xi), Tig(Xi))> for i = 1, . . . , N . We adopt the notation Pg = ∫ gdP .1 In what follows, matrix norms refer to the operator norm, and λmin(M) denotes the minimum eigenvalue of the symmetric matrix M . Recall that given the assumption of equally-sized splits, N = Kn. All proofs are in the Appendix. Assumption 1. i) p ∈ (0, 1). ii) For all k = 1, 2, . . . ,K, the estimated functions ĝk belong to a vector space of functions G with probability one, with G satisfying supg∈G P [|g|4+δ] <∞ for some δ > 0. iii) P [Y 4] <∞. iv) For each k = 1, 2, . . . ,K, ĝk converges to some function g0 ∈ G in the sense that ∫ [ĝk(X)− g0(X)]4dP →p 0. v) infg∈G V ar(g(X)) > 0.
Condition i) is a standard assumption in randomized controlled trials, while conditions ii) and iii) are standard boundedness requirements. Condition iv) is the convergence assumption discussed above. Condition v) can be motivated with reference to the scientific question at hand: restricting attention only to adjustment functions which exhibit nontrivial variation with respect to the value of the covariate is unlikely to hurt the amount of variance reduction achieved.
The following proposition ensures that the inverse of P [Z(g)Z(g)>] exists for all g.
Proposition 1. Given Assumption 1, infg∈G λmin(P [Z(g)Z(g)>]) > 0.
Define
β̂({ĝk}Kk=1) =
[ 1
N ∑ k ∑ i∈Ik Zi(ĝk)Zi(ĝk) >
]−1 [ 1
N ∑ k ∑ i∈Ik Zi(ĝk)Yi
] , (2)
and
β({ĝk}Kk=1) =
[ 1
K ∑ k P [Z(ĝk)Z(ĝk) >]
]−1 [ 1
K ∑ k P [Z(ĝk)Y ]
] . (3)
These are the sample and population OLS coefficients, from the regression of Yi on Zi(ĝk(i)). We also define the corresponding quantities for the limiting function g0,
β̂(g0) =
[ 1
N ∑ i Zi(g0)Zi(g0) >
]−1 [ 1
N ∑ i Zi(g0)Yi
] , (4)
and
β(g0) = [ P [Z(g0)Z(g0) >] ]−1 P [Z(g0)Y ]. (5)
The key intermediate step in deriving the asymptotic distribution of MLRATE is the following result, which states that the distribution of β̂({ĝk}Kk=1), centered around the random variable β({ĝk}Kk=1), is asymptotically equivalent to that of β̂(g0), centered around β(g0).
1If the input function ĝ is random, the quantity P ĝ is also a random variable. If it is a deterministic function g, then Pg is the same as the expectation E[g(X)].
Proposition 2. Under Assumption 1, √ N ∥∥∥[β̂({ĝk}Kk=1)− β({ĝk}Kk=1)]− [β̂(g0)− β(g0)]∥∥∥→p 0.
Having established Proposition 2, we turn to the limiting distribution of MLRATE. This is not quite immediate: the estimator is defined as the coefficient on Ti in the regression of Yi on a constant, Ti, ĝk(i)(Xi), and Ti(ĝk(i) − g). By contrast, Proposition 2 concerns the regression of Yi on a constant, Ti, ĝk(i)(Xi), and Tiĝk(i). To conclude the argument we write MLRATE in terms of the coefficients from the latter regression, and apply Proposition 2. MLRATE, α̂1({ĝk}Kk=1), can be written as
α̂1({ĝk}Kk=1) = β̂1({ĝk}Kk=1) + β̂3({ĝk}Kk=1) 1
N ∑ i ĝk(i)(Xi). (6)
We also define α̂1(g0), the corresponding quantity using the unknown g0:
α̂1(g0) = β̂1(g0) + β̂3(g0) 1
N ∑ i g0(Xi). (7)
The ATE α1 satisfies
α1 = β1(g0) + β3(g0)Pg0 = β1({ĝk}Kk=1) + β3({ĝk}Kk=1) 1
K K∑ k=1 P ĝk. (8)
Here βi({ĝk}Kk=1) denotes the i-th entry of β({ĝk}Kk=1); βi(g0) and β̂i({ĝk}Kk=1) are similarly defined. To see why (8) hold, note that α1 = β1(g) + β3(g)Pg holds for any function g: this is essentially a restatement of the observation that regardless of the particular covariate we adjust for, the regression-adjusted estimator will still be consistent for the ATE [39, 34]. This argument resembles the idea of Neyman orthogonality ([26, 9]), where the estimate of the parameter of interest is not heavily influenced by an undesirable estimate of the nuisance function.
We now state our main theorem, which asserts that the randomness from the ML function fitting step in MLRATE does not affect its asymptotic distribution. Theorem 1. Under Assumption 1,
√ N [ α̂1({ĝk}Kk=1)− α̂1(g0) ] →p 0.
Consequently √ N [ α̂1({ĝk}Kk=1)− α1 ] and √ N [α̂1(g0)− α1] are asymptotically equivalent.
Given Theorem 1, the problem of finding the asymptotic distribution of MLRATE reduces to finding the asymptotic distribution of α̂1(g0). The latter, summarized in the following proposition, can be established by standard asymptotic arguments, and is already known in the literature.2 Define p = E(T ), σ2g = V ar(g0(Xi)), σ 2 YC
= V ar(Yi | Ti = 0), and σ2YT = V ar(Yi | Ti = 1). For notational convenience, below we use β0,i to denote βi(g0) for each i. Proposition 3. If E(g0(X)2) < ∞, E(Y 2) < ∞, and 0 < p < 1, then √ N [α̂1(g0)− α1] N (0, σ2), where
(9)σ2 = σ2YC 1− p + σ2YT p − σ2g p(1− p) [β0,2p+ (β0,2 + β0,3)(1− p)]2 .
Putting together the previous results, we arrive at the asymptotic distribution for MLRATE, which is asymptotically normal and centered around the ATE α1.
Corollary 1. Under Assumption 1, √ N [ α̂1({ĝk}Kk=1)− α1 ] N (0, σ2), where
(10)σ2 = σ2YC 1− p + σ2YT p − σ2g p(1− p)
[β0,2p+ (β0,2 + β0,3)(1− p)]2 . 2See, for example, equation (10) in [39]. Note that there is a small typo in that display: it should
read Σ2 = 11−δσ (0) 22 + 1 δ σ (1) 22 − 1δ(1−δ)σ11 { (1 − δ)σ(1)12 + δσ (0) 12 }2 instead of Σ2 = 11−δσ (0) 22 + 1 δ σ (1) 22 −
1 δ(1−δ)σ11 { (1 − δ)σ(1)12 + δσ (0) 22 }2 .
It follows directly from this corollary that the asymptotic variance of MLRATE is smaller than variance of the simple difference-in-means estimator by the amount
σ2g p(1− p) [β0,2p+ (β0,2 + β0,3)(1− p)]2 ≥ 0.
Thus ML regression adjustment, like ordinary linear regression adjustment [39, 22], cannot reduce asymptotic precision. For some intuition about the determinants of variance reduction, consider the special case where β0,3 = 0 (i.e. the slope of the best-fitting linear relationship between Y and g0(X) does not vary from test to control groups), and σ2YC = σ 2 YT
. The unadjusted, difference-in-means estimator has asymptotic variance σ2YC/(1 − p) + σ 2 YT /p. The relative efficiency of the adjusted estimator, σ2/[σ2YC/(1 − p) + σ 2 YT /p], equals 1 − Corr(Y, g0(X))2. If Corr(Y, g0(X)) = 0.5,
regression adjustment shrinks CIs by 1 − √
1− 0.52 = 13.4%; with a correlation of 0.8, they are 40% smaller.
The following proposition shows that the sample analog of (10) is a consistent estimator of the asymptotic variance, and it can thus be used to construct asymptotically valid CIs. Proposition 4. Let σ̂2 be the sample analog of σ2, that is,
(11) σ̂2 = V̂ ar(Yi | Ti = 0) 1− p̂ + V̂ ar(Yi | Ti = 1) p̂
− V̂ ar(ĝk(i)(Xi))
p̂(1− p̂)
[ β̂2({ĝk}Kk=1)p̂+ ( β̂2({ĝk}Kk=1) + β̂3({ĝk}Kk=1) ) (1− p̂) ]2 ,
where p̂ = ∑ i Ti/N . Under Assumption 1, σ̂ 2 →p σ2.

heading: 3 Simulations & empirical results
text: We now validate MLRATE in practice, on both simulated data, and real Facebook user data. These two validation exercises serve complementary purposes: simulations allow us to verify that the CIs’ empirical coverage is indeed close to their nominal coverage for the data generating process of our choice, while the Facebook data gives an indication of the magnitude of variance reduction that can be expected in practice. All computation is done on an internal cluster, on a standard 64GB ram machine.
Our simulated data generating process hasN = 10,000 iid observations and 100 covariates distributed as Xi ∼ N (0, I100×100). The outcome variable is Yi = b(Xi) + Tiτ(Xi) + ui, where b(·) is the Friedman function b(Xi) = 10 sin(πXi1Xi2) + 20(Xi3 − 0.5)2 + 10Xi4 + 5Xi5 and the treatment effect function is τ(Xi) = Xi1 + log(1 + exp(Xi2)) [15, 27]. The treatment indicator is Ti ∼ Bernoulli(0.5), and the error term is ui ∼ N (0, 252). Treatment is independent of covariates and the error term, and the error term is independent of the covariates. This data generating process involves non-trivial complexity, with nonlinearities and interactions in the baseline outcome, many extraneous covariates that do not affect outcomes, and heterogeneous treatment effects correlated with some covariates. We find the ATE by Monte Carlo integration, and compute the average number of times the MLRATE CIs contain this ATE, over 10,000 simulation repetitions, as well as 95% CIs for this coverage percentage. Both in these simulations and the subsequent analysis of Facebook data, we choose gradient boosted regression trees (GBDT) and elastic net regression as two examples of ML prediction procedures in MLRATE, with scikit-learn’s implementation [28]. Moreover, we choose K = 2 splits for cross-fitting.
Table 1 shows the simulation results. “CI Coverage” displays the average coverage percentage rate over the 10,000 simulations, and the CI width for these estimated coverage rates. “Relative CI Width” displays the CI width for each method divided by the simple difference-in-means CI width (“Unadjusted”), averaged over the 10,000 simulations. Empirical coverage is close to the nominal coverage for all three estimators, with the CIs for empirical coverage including the nominal rate. Both the GBDT and elastic net versions of MLRATE demonstrate efficiency gains over the difference-in-means estimator. As might be expected given the highly nonlinear dependence of the outcomes on covariates, GBDT performs substantially better than the linear, elastic net model:
Table 1: CI coverage and variance reduction results of MLRATE-GBDT and MLRATE-Elastic Net on complex nonlinear simulated data. “CI Coverage” displays the average coverage percentage rate over 10,000 simulations, and the CI width for these estimated coverage rates. “Relative CI Width” displays the CI width for each method divided by the simple difference-in-means CI width (“Unadjusted”), averaged over the 10,000 simulations.
MLRATE-GBDT MLRATE-Elastic Net Unadjusted
CI Coverage (%) 95.18± 0.42 95.34± 0.41 94.88± 0.43 Relative CI Width 0.62 0.86 1.00
Elastic Net
GBDT
LinInteract
0.00 0.25 0.50 0.75 1.00
mean sd max min
Li I t
MLRATEGBDT
MLRATEElastic Net
(a) (b)
Figure 1: Variance reduction results on 48 real metrics used in online experiments run by Facebook. Confidence intervals (CI) are calculated by sampling ∼ 400, 000 observations for each metric. (a) Mean, standard deviation, maximum and minimum value of relative variance of LinInteract, MLRATE-GBDT and MLRATE-Elastic Net compared to the difference-in-means estimator among the metrics. (b) Distribution of CI width of MLRATE-GBDT and MLRATE-Elastic Net relative to LinInteract, by metric.
on average across simulations, the MLRATE-GBDT CIs are 62% the width of the unadjusted CIs, whereas the analogous figure for the elastic net CIs is 86%. Also of interest is the comparison to the semiparametric efficiency bound [24, 17], which can be calculated explicitly for this data generating process: despite the fact that MLRATE is agnostic, and does not assume consistency of the ML procedure employed, the MLRATE-GBDT CIs are only 11.3% wider than those implied by the semiparametric efficiency bound.
The variance reduction numbers above are of course dependent on the particular data generating process specified in the simulation. To get a better sense of the magnitudes of variance reduction one might expect in practice, we evaluate the estimator on 48 real metrics used in online experiments run by Facebook, capturing a broad range of the most commonly consulted user engagement and app performance measurements. We focus on A/A tests in this evaluation rather than A/B tests run in production. This is because the true effect is unknown in the latter, which makes it impossible to evaluate the coverage properties of the CI. Because treatments in online experiments are typically subtle and are unlikely to greatly change the relationship between outcomes and covariates, the magnitude of variance reduction will likely be very similar in A/B tests.
For each outcome metric, we select a random sample of approximately 400,000 users, and simulate an A/A test by assigning a treatment indicator for each user, drawn from a Bernoulli(0.5) distribution. The features used in the ML model vary for each metric and consist of the pre-experiment values of the metric, as well as the pre-experiment values of other metrics that have been grouped together as belonging to the same product area. There are between 20 and 100 other such metrics, with the exact number depending on the outcome metric in question. Outcome values are calculated as the sum of the daily values over a period of one week, and the features values are calculated as the sum of the daily values over the three weeks leading up to the experiment start date.
For each metric, we calculate variances and CI width for four estimators of the ATE: The differencein-means estimator; A univariate linear regression adjustment procedure of equation (1) where the only covariate Xi is the pre-experiment value of the outcome metric Yi (for simplicity, we denote it by ‘LinInteract’); And MLRATE-GBDT/Elastic Net with all available pre-experiment metrics used as features.
Figure 1a shows that LinInteract substantially outperforms the simple difference-in-means estimator, and MLRATE delivers additional gains still. Unlike in the simulated data generating process above, MLRATE-GBDT and MLRATE-Elastic Net perform similarly. The variance reduction relative to the difference-in-means estimator is 72 - 74% on average across metrics, and relative to LinInteract is 19%. The corresponding figures for reduction of the average CI width are 50 - 51%, and 11 - 12%, respectively. Alternatively, to achieve the same precision as the MLRATE-GBDT estimator, the difference-in-means estimator would require sample sizes on average 5.44 times as large on average across metrics and the univariate procedure would require sample sizes 1.56 times as large.
Figure 1b displays the metric-level distribution of CI widths relative to the univariate adjustment case. There is substantial heterogeneity in performance across metrics: for some, ML regression adjustment delivers only quite modest gains relative to univariate adjustment, while for others, it drastically shrinks CIs. This is natural given the variety of metrics in the analysis: some, especially binary or discrete outcomes, may benefit more from more sophisticated predictive modelling, whereas for others simple linear models may perform well. For some metrics, CIs are shrunk by half or more, which may be the difference between experimentation for those metrics being practical and not. As in the simulations, the coverage rates for ML regression adjusted CIs for these metrics are close to the nominal level. For the metric experiencing the largest variance reduction gains from MLRATE—where one might be the most concerned with coverage—we find an average coverage rate of 94.90% over 10,000 simulated A/A tests, where each simulated A/A test is carried out on a 10% subsample drawn at random with replacement from the initial user dataset.
We remark that we design our evaluation to give a realistic sense of the potential variance reduction gains that can be attained with minimal effort and common software implementations of standard ML algorithms. In fact, the supervised learning models we use in this analysis–GBDT and elastic net regression–are deliberately simple, and the training data sample sizes of around 400,000 observations are not especially large by the standards of online A/B tests. The input features to the models are not heavily preprocessed: they are typically raw logged metric values, as opposed to, say, embeddings generated by a prior ML layer. Moreover, as already mentioned in remark 4, we always choose the number of splits K = 2 instead of treating it as a hyperparameter and tuning for better performance. We expect that with more sophisticated supervised learning techniques (e.g. deep, recurrent neural networks with transfer learning across metrics), larger datasets, and better choice of K through cross-validation, the precision gains could be considerably greater still.3
In the simulations and the empirical study above, the dimension of the covariates is not large compared to the sample size. However, our algorithm applies equally to the high-dimensional regime. In many high-dimensional applications, Assumption 1 can be easily satisfied, and our theory fully extends to this case.

heading: 4 Implementation
text: The key guiding principle for selecting features for the ML model is that we can use any variables independent of treatment assignment. Thus any variable extracted before the experiment start is eligible. This simple rule facilitates collaboration with engineering and data science partners familiar with forecasting: they can freely apply their domain expertise to engineer features and build predictive models for specific metrics, without concerns about statistical validity as long as the cross-fitting step in MLRATE is enforced.
The ML step in MLRATE means that the analyst can err on the side of inclusivity in deciding what features to use, as irrelevant features will tend to be omitted from the fitted model. In contrast to [12], for example, we are automatically learning the one ‘feature’ (ĝk(i)(Xi)) that has the best predictive
3In Figure 1b, one metric in the GBDT case has substantially larger variance than the univariate adjustment case, indicating that the default GBDT fit performs quite poorly on this sample. Larger sample sizes or more customized ML modeling will have the benefit of attenuating such anomalies.
power instead of restricting ourselves to a particular pre-experiment feature, thus allowing for greater overall variance reduction. Moreover, this method is highly scalable as the ML step does not need to be performed once per experiment. Once predictions have been generated for a given metric, they can be used to improve precision for all experiments starting after the period used for feature construction.
For real-world applications, the linear regression step in MLRATE, which ensures non-inferiority relative to the difference-in-means estimator, is an important safeguard. There may be no guarantee in practice that the predictive models produced by modeling teams will always be well-calibrated, and without the linear regression layer this non-inferiority guarantee need not hold.
Finally, we note that an additional “censoring” step may be useful when the metric has substantial mass close to zero, reducing computation cost without significantly affecting estimation accuracy. After training the models {ĝk}k=1,2,...,K , instead of regression adjustment using {ĝk(i)(Xi)}, define ĝτ (Xi) = T (ĝk(i)(Xi), τ) for some pre-determined threshold τ , where T is the hard-thresholding operator T (u, τ) = u1{u≥τ}. Then one can perform regression adjustment with ĝτ (Xi) in place of ĝk(i)(Xi), with the same statistical theory applying. Small values of τ will cause small efficiency losses, but can greatly reduce the computation cost on the linear regression when N is large.4

heading: 5 Conclusion
text: MLRATE is a scalable methodology that allows ML algorithms to be used for variance reduction, while still giving formal statistical guarantees on consistency and CI coverage. Of particular practical importance is the methodology’s robustness to the ML algorithm used, both in the sense that the ML algorithm used need not be consistent for the truth, and in the sense that no matter how bad the ML predictions are, MLRATE has asymptotic variance no larger than the difference-in-means estimator. Our application to Facebook data demonstrates variance reduction gains using pre-experiment covariates and even simple predictive algorithms. We expect that more sophisticated predictive algorithms, and incorporating other covariates into this framework–for example, generating user covariates by synthetic-control inspired strategies that incorporate contemporaneous data on outcomes for individuals outside the experiment–could lead to more substantial efficiency gains still.

Review summaries:
This paper proposes a new approach for variance reduction in randomized controlled trials. The proposed method first predicts the outcome variable Y from a set of covariates X that are independent of the treatment assignment T, and measures the average treatment effect by regressing the difference between this prediction and the observed Y against the predicted $\hat{Y}$. The method avoids the potential pitfall of invalid confidence intervals by training the predictor for a given $Y_i$ using data from $j \neq i$ ("cross-fitting"), and avoids the potential pitfall of over-reliance on poor estimates of Y by including the predictions as part of a linear regression (instead of directly subtracting the predictions from the observed Y). The consistency and asymptotic normality of the estimator are proven, and experiments on both simulated and real data are presented.

The setting of this paper is that of a randomized controlled trial with data $(Y_i, X_i, T_i)$ for $i=1,...,N$. $Y_i$ is the response of interest, $T_i \in \\{0,1\\}$ is the treatment indicator and $X_i$ is a vector of covariates. The goal is to estimate and provide confidence intervals for the causal effect,  $$ \tau = E[Y_i \mid T_i = 1] - E[Y_i \mid T_i = 0].$$  The simplest estimator here is the difference in means estimator  $$ \hat{\tau}_{\text{diff. means}} = \text{average}(Y_i: T_i=1) - \text{average}(Y_i: T_i=0).$$  Confidence intervals can be constructed using a simple central limit theorem. One can do substantially better (shorter confidence intervals, better point estimator) by utilizing information in $X_i$ to reduce variance and there have been many approaches to do so. This paper proposes an approach that uses cross-fitting (splitting all $i$'s into $K$ folds) and arbitrary predictive models of $E[Y_i \mid X_i]$ (learned out-of-fold) and combines the results using a linear regression adjustment. The authors also show how to compute confidence intervals for the method and prove they are asymptotic valid even if inconsistent predictive models are used.  The asymptotic variance is never worse than the difference in means estimator. The authors also show substantial gains for the analysis of randomized controlled experiments at Facebook.

This paper proposes what the authors call a machine learning regression-adjusted treatment effect estimator to reduce the variance when estimating the average treatment effect in randomized controlled trials, by using covariates correlated with the outcome but independent of the treatment. The algorithm first trains a machine learning model to predict the outcome from the covariates, and then runs a regression of the outcome on the treatment, predicted model and their cross-product, giving the regression coefficient of the treatment as the estimator. The authors prove consistency and asymptotic normality for their estimator under general conditions. As a main strength, the proposed estimator offers variance reduction when the machine learning prediction is good, and at the same time is robust to poor predictions - in the sense that its asymptotic behavior is no worse than the baseline difference-in-means estimator even if the prediction is bad.  The authors' main contributions are: 1) From the theoretical perspective, they propose an estimator with lower variance which is robust and easy to implement, and prove the consistency and asymptotic normality. 2) From the practical perspective, their estimator performs better than the baseline difference-in-means estimator and the univariate linear regression adjustment procedure in Deng et al. (2013) when applying to 48 outcome metrics commonly monitored in Facebook.

This paper considers the analysis of randomized (online) experiments (A/B tests) through a linear regression that adjusts for pre-treatment covariates, which are independent of the treatment assignment. It is well-known that this type of estimate of the treatment effect can be more efficient than the difference-in-means estimator when the covariates are predictive of the outcome.   The authors propose to incorporate the covariates into the linear regression through a predictor of the outcome, which itself is estimated with data-adaptive, machine learning methods. Cross-fitting is used to control the bias introduced by the data-adaptive estimation. It is shown that the resulting estimator is asymptotically normal with a limiting variance that is no larger than that of the difference-in-mean estimator, under relatively weak conditions. In particular, the method does not require the consistency of the nuisance predictor function. This model-agnosticism is clearly an advantage when compared to influence-function-based methods (e.g., double ML), which typically require the convergence of the nuisance to the truth at a certain rate.   The method developed is shown to deliver correct coverage from simulations and A/A testing data, and to typically outperform competing methods.