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 authors introduce a scalable quasi-Bayesian method for nonparametric instrumental variable (IV) regression that provides uncertainty quantification for the IV models. They base their approach on kernelized IV models and define a quasi-posterior as the Radon-Nikodym derivative of a generalized method of moments (GMM) objective function with respect to the standard Gaussian process prior. They go on to show that the quasi-posterior exhibits satisfies certain posterior contraction criteria that lead to asymptotically sound results, and prove a conservative posterior contraction rate under a set of typical assumptions. They then show how to achieve scalable inference for the quasi-posterior using a previously proposed “randomized prior trick” for Gaussian process regression, and provide a proof for the convergence of the inference algorithm. Finally, the authors showcase the performance of their method in a series of simulated experiments of varying complexity, by comparing it against other approaches that provide uncertainty quantification (Bayesian, bootstrap).