Machine Learning for Variance Reduction in Online Experiments

1 Introduction . 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 . 2 Outcome prediction for variance reduction . 2.1 Setup & motivation . 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 . 2.2 Related work . 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 ] . 2.3 Estimation and inference with MLRATE . 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 .

Summary: 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.