RieszNet and ForestRiesz: Automatic Debiased Machine Learning with Neural Nets and Random Forests

√ n-consistent and asymp- totically normal estimation of the object of interest requires debiasing to reduce the effects of regularization and/or model selection on the object of interest . Debiasing is typically achieved by adding a correction term to the plug-in estimator of the functional , that is derived based on a functional-specific theoretical derivation of what is known as the influence function and which leads to properties such as double robustness and Neyman orthogonality . We instead implement an automatic debiasing procedure based on automatically learning the Riesz representation of the linear functional using Neural Nets and Random Forests . Our method solely requires value query oracle access to the linear functional . We propose a multi-tasking Neural Net debiasing method with stochastic gradient descent minimization of a combined Riesz representer and regression loss , while sharing representation layers for the two functions . We also propose a Random Forest method which learns a locally linear representation of the Riesz function . Even though our methodology applies to arbitrary functionals , we experimentally find that it beats state of the art performance of the prior neural net based estimator of Shi et al . ( 2019 ) for the case of the average treatment effect functional . We also evaluate our method on the more challenging problem of estimating average marginal effects with continuous treatments , using semi-synthetic data of gasoline price changes on gasoline demand . 1 INTRODUCTION . A large number of problems in causal inference , off-policy evaluation and optimization and interpretable machine learning can be viewed as estimating the average value of a moment function that depends on an unknown regression function : θ0 = E [ m ( W ; g0 ) ] , where W : = ( Y , Z ) and g0 ( Z ) : = E [ Y | Z ] . In most cases , Y will be the outcome of interest , and inputs Z = ( T , X ) will include a binary or continuous treatment T and covariates X . Prototypical examples include the estimation of average treatment effects , average policy effects , average derivatives and incremental policy effects . Example 1 ( Average treatment effect ) . Here Z = ( T , X ) where T is a binary treatment indicator , and X are covariates . The object of interest and associated moment function are : θ0 = E [ g0 ( 1 , X ) − g0 ( 0 , X ) ] , m ( W ; g ) = g ( 1 , X ) − g ( 0 , X ) . If potential outcomes are conditionally independent of treatment T given covariates X , then this object is the average treatment effect ( Rosenbaum & Rubin , 1983 ) .. Example 2 ( Average policy effect ) . In the context of offline policy evaluation and optimization , our goal is to optimize over a space of assignment policies π : X → { 0 , 1 } , when having access to observational data collected by some unknown treatment policy . The policy value can also be formulated as the average of a linear moment : θ0 = E [ π ( X ) ( g0 ( 1 , X ) − g0 ( 0 , X ) ) + g0 ( 0 , X ) ] , m ( W ; g ) = π ( X ) ( g ( 1 , X ) − g ( 0 , X ) ) + g ( 0 , X ) . A long-line of prior work has considered doubly-robust approaches to optimizing over a space of . candidate policies from observational data ( see e.g . Dudı́k et al. , 2011 ; Athey & Wager , 2021 ) . Example 3 ( Average marginal effect / derivative ) . Here Z = ( T , X ) , where T is a continuously distributed policy variable of interest , X are covariates , and : θ0 = E [ ∂g0 ( T , X ) ∂t ] , m ( W ; g ) = ∂g ( T , X ) ∂t . This object is essentially the average slope in the partial dependence plot frequently used in work on interpretable machine learning ( see e.g . Zhao & Hastie , 2021 ; Friedman , 2001 ; Molnar , 2020 ) . Example 4 ( Incremental policy effects ) . Here Z = ( T , X ) , where T is a continuously distributed policy variable of interest , X are covariates , and π : X → [ −1 , 1 ] is an incremental policy of infinitesimally increasing or descreasing the treatment from its baseline value ( see e.g . Athey & Wager , 2021 ) . The incremental value of such an infinitesimal policy change takes the form : θ0 = E [ π ( X ) ∂g0 ( T , X ) ∂t ] , m ( W ; g ) = π ( X ) ∂g ( T , X ) ∂t . Such incremental policy effects can also be useful within the context of policy gradient algorithms in deep reinforcement learning , so as to take gradient steps towards a better policy , and debiasing techniques have already been used in that context ( Grathwohl et al. , 2017 ) . Even though the non-parametric regression function is typically estimable only at slower than parametric rates , one can often achieve parametric rates for the average moment function . However , this is typically not achieved by simply pluging a non-parametric regression estimate into the moment formula and averaging , but requires debiasing approaches to reduce the effects of regularization when learning the non-parametric regression . Typical debiasing techniques are tailored to the moment of interest . In this work we present automatic debiasing techniques that use the representation power of neural nets and random forests and which only require oracle access to the moment of interest . Our resulting average moment estimators are typically consistent at parametric √ n rates and are asymptotically normal , allowing for the construction of confidence intervals with approximately nominal coverage . The latter is essential in social science applications and can also prove useful in policy learning applications , so as to quantify the uncertainty of different policies and implement automated policy optimization algorithms which require uncertainty bounds ( e.g . algorithms that use optimism in the face of uncertainty ) . Relative to previous works in the automatically debiased ML ( Auto-DML ) literature , the contribution of this paper is twofold . On the one hand , we provide the first practical implementation of Auto-DML using neural networks ( RieszNet ) and random forests ( ForestRiesz ) . As such , we complement the theoretical guarantees of Chernozhukov et al . ( 2021 ) for generic machine learners . On the other hand , we show that our methods perform better than existing benchmarks and that inference based on asymptotic confidence intervals obtains coverage close to nominal in two settings of great relevance in applied research : the average treatment effect of a binary treatment and the average marginal effect ( derivative ) of a continuous treatment . The rest of the paper is structured as follows . Section 2 provides some background on estimation of average moments of regression functions . In Section 2.1 we describe the form of the debiasing term , and in Section 2.2 we explain how it can be automatically estimated . Sections 3 and 4 introduce our proposed estimators : RieszNet and ForestRiesz , respectively . Finally , in Section 5 we present our experimental results . 2 ESTIMATION OF AVERAGE MOMENTS OF REGRESSION FUNCTIONS . 2.1 DEBIASING THE AVERAGE MOMENT . We focus on problems where there exists a function α0 ( Z ) with finite second moment such that : E [ m ( W ; g ) ] = E [ α0 ( Z ) g ( Z ) ] , for all g with E [ g ( Z ) 2 ] < ∞ . By the Riesz representation theorem , existence of such an α0 ( X ) is equivalent to E [ m ( W ; g ) ] being a mean-square continuous linear functional of g. We will refer to this α0 ( X ) as the Riesz representer ( RR ) . Existence of the RR is equivalent to the semiparametric variance bound for θ0 being finite , as in Newey ( 1994 ) and Hirshberg and Wager ( 2018 ) . This RR exists in each of Examples 1 to 4 . For instance , in Example 1 the RR is α0 ( Z ) = T/p0 ( X ) − ( 1 − T ) / ( 1 − p0 ( X ) ) where p0 ( X ) = Pr ( T = 1 | X ) is the propensity score and in Example 3 , integration by parts gives α0 ( Z ) = − ( ∂f0 ( T , X ) /∂t ) /f0 ( Z ) where f0 ( Z ) is the joint probability density function ( pdf ) of T and Z . In general , the RR involves ( unknown ) nonparametric functions of the data , like the propensity score or the density f0 ( Z ) and its derivative . The RR is a crucial object in the debiased ML literature , since it allows us to construct a debiasing term for the moment function m ( W ; g ) ( see Chernozhukov et al. , 2018a , for details ) . The debiasing term in this case takes the form α ( Z ) ( Y − g ( Z ) ) . To see that , consider the score m ( W ; g ) + α ( Z ) ( Y − g ( Z ) ) − θ . It satisfies the following property : E [ m ( W ; g ) + α ( Z ) ( Y − g ( Z ) ) − θ0 ] = −E [ ( α ( Z ) − α0 ( Z ) ) ( g ( Z ) − g0 ( Z ) ) ] . This is sometimes known as double robustness , since the score will be zero in expectation when either α ( Z ) = α0 ( Z ) or g ( Z ) = g0 ( Z ) are correctly specified . An estimator of θ0 can be constructed from this score and first-stage estimators ĝ and α̂ . Let En [ · ] denote the empirical expectation over a sample of size n , i.e . En [ Z ] = 1n ∑n i=1 Zi . We consider : θ̂ = En [ m ( W ; ĝ ) + α̂ ( Z ) ( Y − ĝ ( Z ) ) ] . ( 1 ) What double robustness means in practice is that the bias of this estimator will vanish at a rate equal to the product of the mean-square convergence rates of α̂ and ĝ . Therefore , in cases where the regression function g0 can be estimated very well , the rate requirements on α̂ will be less strict , and vice versa . More notably , whenever the product of the mean-square convergence rates of α̂ and ĝ is larger than √ n , the estimator θ̂ will be asymptotically normal at the parametric rate √ n , as proven formally in Theorem 4 of Chernozhukov et al . ( 2021 ) . The regression estimator ĝ and the RR estimator α̂ may use samples different than the i-th , which constitutes cross-fitting . Cross-fitting reduces bias from using the i-th sample in estimating α and g. Also ĝ and α̂ may use different samples , which constitutes double cross-fitting ( see e.g . Newey & Robins , 2018 , for the benefits of double cross-fitting in reducing the requirements on the quality of the regression and RR estimates ) .

Summary: This paper considers the problem of estimating an expectation over covariates of some functional of an unknown regression function. The authors propose two estimators, one based on neural networks and one based on random forests. In contrast to many (all?) previous estimators which were derived for specific functionals, these estimators are applicable to general functionals. They employ a (new?) debiasing technique. Moreover, they use a novel multi-task architecture based on the observation that it suffices to estimate the regression function as a function of the Riesz representer. The paper demonstrates improved accuracies compared to prior work and good coverage on semi-synthetic tasks derived from two data sets.