Scalable Quasi-Bayesian Inference for Instrumental Variable Regression

1 Introduction . Instrumental variable ( IV ) regression is a standard approach for estimating causal effect from confounded observational data . In the presence of confounding , any regression method estimating E ( y | x ) can not recover the causal relation f† between the outcome y and the treatment x , since the residual u = y − f† ( x ) is correlated with x due to the unobserved confounders . IV regression enables identification of the causal effect through the introduction of instruments , variables z that are known to influence y only through x. IV regression is widely used in economics [ 1 ] , epidemiology [ 2 ] and clinical research [ 3 ] , but modeling nonlinear effect in IV regression can be challenging . Recent years have seen great development in adopting modern machine learning models for IV regression [ 4–8 ] . However , there is still a lack of uncertainty quantification measures for these flexible IV models . Uncertainty quantification is especially important for IV analysis , since unlike in standard supervised learning scenarios , we do not have ( unconfounded ) validation data , from which we could deduce the error pattern of the estimated model . Moreover , the instrument of choice may be weak , meaning that it only provides limited information for x ; in such cases point estimators suffer from high variance [ 9 ] . This problem is exacerbated in the nonparametric setting , where IV estimation is usually an ill-posed inverse problem , where instruments provide vanishing information for the higher-order nonlinear effects [ 10 ] . The IV setting brings unique challenges for uncertainty quantification . For example , while it is natural to consider a Bayesian approach , specification of the likelihood requires knowledge of the entire data generating process , which is typically unavailable in IV regression . Consequently , most , if not all , existing work on Bayesian IV [ 11–15 ] assumes the following data generating process : ∗ZW and YZ contribute equally . ‡JZ is the corresponding author . An extended version is available at https : //arxiv.org/abs/2106.08750v2 . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . x = g ( z ) + u1 , y = f ( x ) + u2 , where ( u1 , u2 ) are correlated and independent of z . These methods then conduct posterior inference on g , f as well as the generative model for the unobserved confounders ( u1 , u2 ) . However , the additive error in x is an unnecessary assumption for most point estimation procedures , and is difficult to check in high dimensions . Furthermore , the need to model the generating process of ( u1 , u2 ) introduces an extra risk of model misspecification , and Bayesian inference on the generative model is computationally expensive , especially on complex high-dimensional datasets . None of these issues present if only point estimation is needed . For the above reasons , it is appealing to turn to an alternative quasi-Bayesian approach [ 16–18 ] . Quasi-Bayesian analysis views IV estimation as a generalized method-of-moments ( GMM ) procedure , and defines the quasi-posterior as a Gibbs distribution constructed from a chosen prior and violation of the moment constraints . It does not require full knowledge of the data generating process , and thus does not suffer from the aforementioned drawbacks of Bayesian inference . However , computation of the quasi-posterior is non-trivial , as its density contains a conditional expectation term E ( y − f ( x ) | z ) , which itself needs to be estimated from data . This makes both approximate inference and theoretical analysis difficult . So far , quasi-Bayesian analysis for nonparametric IV is only developed on classical models such as wavelet basis [ 17 , 18 ] or Nadaraya-Watson smoothing [ 19 ] , while adoption of flexible machine learning models such as kernel machines or neural networks remains an open challenge . Moreover , numerical study has been limited in previous work , so little is known about the empirical performance of the quasi-Bayesian approach . In this work , we present a novel quasi-Bayesian procedure for IV regression , building upon the recent development in kernelized IV models [ 7 , 20 ] . We employ a Gaussian process ( GP ) prior and construct a quasi-likelihood using a kernel conditional expectation estimator . We establish theoretical properties of the resultant quasi-posterior , proving its consistency and showing that it may quantify instrument strength . Furthermore , inspired from the minimax formulation of IV estimation [ 5 , 8 , 20 , 21 ] , we design a principled approximate inference algorithm using random feature expansion and a novel adaptation of the “ randomized prior trick ” [ 22 ] . The algorithm has the form of stochastic gradient descent-ascent and is thus scalable and easy to implement . It can also be adapted to work with flexible neural network models , in which case its behavior can be formally justified by analyzing the neural networks in the kernel regime . Empirical evaluation shows that the proposed method produces informative uncertainty estimates , scales to high-dimensional and complex nonlinear problems , and is especially advantageous when the instrument is weak . The rest of the paper is organized as follows : In Section 2 we set up the problem . We then derive the quasi-posterior and analyze its theoretical properties in Section 3 , and present the approximate inference algorithm in Section 4 . Section 5 reviews related work , and Section 6 presents numerical studies . Finally , we discuss conclusion and future work in Section 7 . 2 Notations and Setup . Notations We use boldface ( x , y , z ) to represent random variables on the space X × Y × Z , regular font ( x , y , z ) to denote deterministic values . [ n ] denotes the set { 1 , 2 , · · · , n } . { ( xi , yi , zi ) : i ∈ [ n ] } indicates the training data . We use the notations X : = ( x1 , . . . , xn ) ∈ Xn , f ( X ) : = ( f ( x1 ) , . . . , f ( xn ) ) ; likewise for Y , Z . For finite-dimensional vectors θ , θ′ ∈ Rm , we use ‖θ‖2 , 〈θ , θ′〉2 to denote the Euclidean norm and inner product , respectively . For any operator A : H1 → H2 between Hilbert spaces H1 and H2 , we denote its adjoint by A∗ : H2 → H1 . When H1 = H2 and λ ∈ R , we use the notation Aλ : = A+ λI for simplicity . IV regression Denote the treatment and response variables as x , y , the instrument as z , and the true structural function of interest as f† . Consider the data generating process y = f† ( x ) + u , where the unobserved u satisfies E ( u | z ) = 0 , but may correlate with x . Then f† satisfies E ( y − f† ( x ) | z ) = 0 a.s. [ P ( dz ) ] , ( 1 ) where P denotes the data distribution . This conditional moment restriction ( CMR ) formulation is the standard definition in literature [ e.g. , 23 , 18 ] , and is used in the recent work on machine learning models for IV . It connects to GMM as ( 1 ) can be viewed as a continuum of generalized moment constraints . Note that ( 1 ) does not place any structural constraint on the conditional distribution p ( x | z ) , such as additive noise ; hence , it does not require full knowledge of the data generating process . Also , as discussed in , e.g. , Hartford et al . [ 4 ] , the setup can also be extended to incorporate observed confounders v , by including v in both x and z. Kernelized IV and a dual view Let H , I be suitably chosen function spaces on X , Z , respectively . The generalized moment constraint ( 1 ) motivates the use of the following objective min f∈H L ( f ) : = d2n ( Ênf − b̂ ) + λ̄Ω ( f ) , ( 2 ) where Ên : H → I is an empirical approximation to the conditional expectation operator E : f 7→ E ( f ( x ) | z = · ) , b̂ is an estimator of E ( y | z = · ) , { dn } is a sequence of suitable ( semi- ) norm on I , and Ω : H → R is a regularization term . When H , I are reproducing kernel Hilbert spaces ( RKHS ) with corresponding kernels kx , kz , it is natural to set Ω ( f ) = 12‖f‖ 2 H , and define b̂ as the result of kernel ridge regression on y with respect to z . In this case Ên can be defined with the empirical kernel conditional expectation operator : let Czx = E ( k ( x , · ) ⊗ k ( z , · ) ) , Czz = E ( k ( z , · ) ⊗ k ( z , · ) ) . Assuming E maps all f ∈ H to Ef ∈ I , we have [ 24 ] CzzEf = Czxf , ∀f ∈ H. This motivates the use of Ên : = Ĉ−1zz , ν̄Ĉzx where Ĉzz : = 1 n ∑n i=1 k ( zi , · ) ⊗ k ( zi , · ) , Ĉzx is defined similarly , and ν̄ is a regularization hyperparameter.2 The choice of dn is flexible ; the dual IV formulation uses d2n ( g ) : = 1 2n ∑n j=1 g ( zj ) 2 + ν̄2‖g‖ 2 I = 1 2 〈g , Ĉzz , ν̄g〉I . Introducing the evaluation operator Sz : I → Rn , Szg : = ( g ( z1 ) , . . . , g ( zn ) ) , we have b̂ = Ĉ−1zz , ν̄ S∗zY n , and L ( f ) = 1 2 ∥∥∥∥Ĉ−1/2zz , ν̄ ( Ĉzxf − S∗zYn ) ∥∥∥∥2 I + λ̄ 2 ‖f‖2H ( 3 ) = max g∈I 1 n n∑ j=1 ( ( f ( xj ) − yj ) g ( zj ) − g ( zj ) 2 2 ) − ν̄ 2 ‖g‖2I + λ̄ 2 ‖f‖2H , ( 4 ) where ( 4 ) holds because of the Fenchel duality 12u 2 = supv∈R ( uv − 12v 2 ) and the equality E ( u ( x , y ) g ( z ) ) = E ( E ( u ( x , y ) | z ) g ( z ) ) ; see [ 25 , 8 , 21 , 26 ] . The dual formulation ( 4 ) circumvents the need to compute Ên directly , an operator between the typically infinite-dimensional spaces H and I , and leads to a scalable estimation procedure based on stochastic gradient descent-ascent ( SGDA ) . It can also be generalized to work with deep neural networks ( DNNs ) instead of kernel machines , by replacing the RKHS regularizer with a suitable regularizer for DNNs , although theoretical analysis for the resulted algorithm requires separate effort [ 8 , 26 ] . As we shall see , ( 4 ) will also enable the construction of a scalable approximate inference algorithm .

Summary: The paper proposes a scalable quasi-Bayesian inference for instrumental variable regression. Specifically, the paper builds a quasi-posterior on the recent quadratic kernel IV loss function. For scalability, the paper extends the randomized prior trick based on the recent dual formulation of the quadratic loss. The approach has a theoretical guarantee that incorrect solutions are excluded while valid kept. Moreover, the convergence analysis for the randomized prior trick with the stochastic gradient descent-ascent is present. Experiments validate the approach.