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