Scalable Quasi-Bayesian Inference for Instrumental Variable Regression
Year: 2021

Abstract:
Recent years have witnessed an upsurge of interest in employing flexible machine learning models for instrumental variable (IV) regression, but the development of uncertainty quantification methodology is still lacking. In this work we present a scalable quasi-Bayesian procedure for IV regression, building upon the recently developed kernelized IV models. Contrary to Bayesian modeling for IV, our approach does not require additional assumptions on the data generating process, and leads to a scalable approximate inference algorithm with time cost comparable to the corresponding point estimation methods. Our algorithm can be further extended to work with neural network models. We analyze the theoretical properties of the proposed quasi-posterior, and demonstrate through empirical evaluation the competitive performance of our method.

heading: 1 Introduction
text: 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) cannot 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.

heading: 2 Notations and Setup
text: 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.

heading: 3 Quasi-Bayesian Analysis of Dual IV
text: Introduce the notations Sx : H → Rn, Sxf := (f(x1), . . . , f(xn)), so that Ĉzx = 1nS ∗ zSx, Ĉzz = 1 nS ∗ zSz . Define λ := nλ̄, ν = nν̄. Now we can re-express (3) in an equivalent form as:
L̄(f) := n λ L(f) = 1 2 (f(X)− Y )>(λ−1L)(f(X)− Y ) + 1 2 ‖f‖2H, (5)
where L := 1nSzĈ −1 zz,ν̄S ∗ z is a linear map from Rn to Rn, and thus can be identified with an n × n matrix. Since the first term above is equivalent to the log density of the multivariate normal distribution N (Y | f(X), λL−1), we can view (5) as the objective of a kernel ridge regression problem, which has a data-dependent noise covariance λL−1.3 The connection between kernel ridge regression and Gaussian process regression [28] thus motivates the use of the quasi-posterior
2With an abuse of notation, k refers to the reproducing kernel of the corresponding RKHS (kx forH or kz for I) whenever the denotation is clear.
3Here we assume the invertibility of L for brevity. Alternatively, observe that (5) defines a linear inverse problem with the finite-dimensional observation operator f 7→ √ Lf(X) and noise variance λ−1I , and we can follow [27, Chapter 6] to derive the same quasi-posterior.
Π(df | D(n)), defined through the following Radon-Nikodym derivative w.r.t. the standard GP prior Π = GP(0, kx):
dΠ(· | D(n)) dΠ
(f) ∝ exp ( − 1
2 (f(X)−Y )>(λ−1L)(f(X)−Y )
) = exp ( −n λ d2n(Ênf − b̂) ) . (6)
Note that contrary to standard Bayesian modeling, we do not assume Y ∼ N (f(X), λL−1) is part of the true data generating process, nor does the theoretical analysis below rely on it. Instead, the quasi-posterior (6) should be interpreted as a Gibbs distribution which trades off between the properly scaled evidence λ−1nd2n(Ênf − b̂), which characterizes the estimated violation of the GMM constraint (1), and our prior belief Π(df). This trade-off is most clear from the well-known variational characterization of the Gibbs distribution [29],
Π(· | D(n)) = arg min Ψ Ef∼Ψ[λ−1nd2n(Ênf − b̂)] + KL(Ψ ‖Π). (7)
Nonetheless, the fictitious data generating process f ∼ GP(0, k), Y ∼ N (f(X), λL−1) is useful for deriving the quasi-posterior, since its conditional distribution pfic(df | Y ) coincides with (6) [27, Chapter 6]. In the probability space of this fictitious data generating process, for any finite set of test inputs x∗, we have
pfic(f(x∗), Y ) ∼ N ( 0, [ K∗∗ K∗x Kx∗ Kxx + λL −1 ]) ,
where K(·) denote the corresponding Gram matrices with subscript ∗ denoting x∗ and x denoting X (so, e.g., K∗x := k(x∗, X)). Thus by the Gaussian conditioning formula, we have
Π(f(x∗) | D(n)) = pfic(f(x∗) | Y ) = N (m,S), (8) where m := K∗x(λI + LKxx)−1LY, (9)
S := K∗∗ −K∗xL(λI +KxxL)−1Kx∗, (10) L = Kzz(Kzz + νI) −1. (11)
In the above Kzz := k(Z,Z) denotes the Gram matrix, and (11) follows from the definition L = n−1Sz(n −1S∗zSz + ν̄I) −1S∗z , the Woodbury identity, and the observation that SzS ∗ z = Kzz .
Theoretical Analysis For the quasi-posterior Π(· | D(n)) to be a useful measure of uncertainty, it needs to satisfy the following informal criteria: as n→∞, we expect
(C1) Π(· | D(n)) will exclude incorrect solutions; (C2) In cases of non-identification, Π(· | D(n)) will not exclude any valid solution in the model.
In the following, we formalize these criteria, and demonstrate that with appropriately chosen hyperparameters, the quasi-posterior satisfies both criteria. We will work with the following assumptions: Assumption 3.1. The restriction of the conditional expectation operator f 7→ E(f(x) | z = ·) on H, denoted as E, has its image contained in I. E : H → I is bounded. Assumption 3.1 is intuitive, and is also slightly more general than some previous work [20, 26] that require the conditional expectation operator to be bounded on the entire hypothesis space, which typically corresponds to the “sample space” of the GP prior in our setting, and is much larger than H (see Appendix A). Nonetheless, we note that Hypothesis 4 in Singh et al. [7] may be more general, although they impose extra smoothness assumptions on E. Assumption 3.2. The true structural function f†(x) is such that, there exists a sequence {τm : m ∈ N} satisfying τm → 0, and
mτ2m ≥ inf f†m∈H:‖f†−f†m‖≤τm ‖f†m‖2H − log Π({f : ‖f‖ ≤ τm}), (12)
where ‖ · ‖ denotes the sup norm. Assumption 3.2 requires that f† can be well approximated in H; this is more general than previous work [e.g., 20, 7] that require f† ∈ H, but is typical in the GP literature [30, 31]. The sequence {τm} is determined by the complexity of H and its ability for approximating f†, and is usually the optimal posterior contraction rate for Gaussian process regression on the unconfounded dataset {(xi, f†(xi) + i) : i ∈ [m]} [31].
Assumption 3.3. X and Z are Polish spaces. The kernels kx, kz are continuous, supx∈X k(x, x) + supz∈Z k(z, z) ≤ κ2, and Mercer’s representations [32] of kx and kz exist. The random variable y − f†(x) is sub-exponential. Assumption 3.3 imposes technical conditions frequently assumed in literature [e.g., 33]. The requirements on the kernels can be satisfied by e.g. continuous kernels on compact subsets of Rd. Given the assumptions above, we characterize (C1) by showing that as n→∞, the posterior places vanishing mass on the region of functions that violates the GMM constraints (1). Concretely, Theorem 3.1 (Proof in Appendix B). There exist a constantM > 0 depending on the data distribution and the kernels of choice, and a sufficiently slowly growing sequence {γn} → ∞ (e.g., we can always have γn ≤ log log n) such that when taking λ̄ = n−1/2, ν̄ = min{τ2d√ne, n −1/2γn}, it has
ED(n)Π({f : E2(f(x)− y | z) > Mτ2d√ne} | D (n))→ 0. (13)
Remark 3.1. The above result should primarily be interpreted as a consistency result, although it also provides a posterior contraction rate [34] in the order of τd√ne in the semi-norm f 7→ ‖Ef‖L2(P (dz)), where {τm} is defined in Assumption 3.2. As an example, suppose the regularity of f† is similar to the Matérn-1/2 RKHS; then for suitable kernels we have τd√ne = O(n−1/8), see Remark A.3.
The rate is suboptimal, and can be immediately improved if we impose further regularity assumption on kz: if we assume the critical radius of the local Rademacher complexity [35] of the unit-norm ball I1 is δn τn n− b 2(b+1) , a rate of O(n− b−1 2(b+1) ) for the semi-norm can be established following similar arguments. This is still worse than [26], which provides the rate O(max{τn, δn}) for what corresponds to our posterior mean estimator. Nonetheless, their result is also generally suboptimal, because the choice of δn τn is actually suboptimal when the IV regression problem is ill-posed. To date, minimax optimal rates have only been established under additional assumptions on the relations between the conditional expectation operator and the model for f . Our extended version establishes such results, for both this semi-norm and more intuitive norms such as L2(P (dx)).
Remark 3.2. The use of an increasing λ = nλ̄ is a technical artifact due to our minimal assumptions on kz . It is common to impose extra regularization in nonparametric IV, due to the need to estimate E from data [19], and, in the quasi-Bayesian setting, also due to the additional technical challenges [18]. However, this is not necessary for our method if additional assumptions are imposed, as we show in the extended version. In practice, our hyperparameter selection procedure does not produce λ with significant growth.
The following proposition characterizes (C2): Proposition 3.1 (Proof in Appendix B). The scaled log quasi-likelihood estimate `n(f) = d2n(Ênf − b̂) satisfies
Π ({ ∀f : ∀δ > 0, lim
n→∞ PD(n)( ∣∣`n(f)− Ez∼P (z)(E(f(x)− y | z)2)∣∣ > δ) = 0}) = 1. In words, for Π-almost every f , the log quasi-likelihood estimate scaled by n−1 converges to Ez∼P (z)(E(f(x) − y | z)2), which characterizes the violation of (1). While the restriction to the probability-1 subset can be concerning in the nonparametric setting [34], the proof only depends on f satisfying similar approximability conditions to f†; see (45). Remark 3.3. To understand why the proposition characterizes (C2), suppose there are multiple f that satisfy (1). Then all of them will eventually have significantly higher log likelihood than functions violating (1): the difference is Θ(λ−1n). If these functions have a similar level of regularity in the sense that their concentration functions, which is the right-hand side of (12), have the same asymptotics as → 0, Borell’s inequality [e.g., 34, Proposition 11.17] implies that the posterior mass of small ‖ · ‖-norm balls around them will also have the same asymptotics.
Finally, we provide the following (over-)simplified example, which compares the behavior of the quasi-posterior and bootstrap in the context of (C2): Example 3.1. Suppose z is completely non-informative so that E(f(x) | z) ≡ Ef(x) for all f ; and suppose the estimated conditional expectation Ên is sufficiently accurate, so that we replace it
with E.4 In this case bootstrap on (3) will always return the point estimator f ≡ 0 due to the nonzero regularization on ‖f‖H, while the quasi-posterior behaves like the prior, correctly reflecting the complete lack of evidence in data.
While this example is oversimplified, and in practice the estimation error of E plays an important role, it is known that bootstrap uncertainty estimates can be problematic given weak IVs [36–38]. We also observe similar failures for bootstrap in the experiments (see Section 6.1).

heading: 4 Scalable Approximate Inference via a Randomized Prior Trick
text: We now turn to approximate inference with parametric models such as random feature expansion or wide NNs. Scalable inference for the IV quasi-posterior appears difficult, since for any f , computing the quasi-likelihood involves computing Ênf , which in turn requires either inverting an n × n Gram matrix, judging from (5) and (11), or solving an optimization problem specific to f from (4). Nonetheless, we show that it is possible, by extending the “randomized prior” trick for Gaussian process regression [22] to work with (quasi-)likelihoods with an optimization formulation as in (4).
Our algorithm works with random feature models. A random feature model for kz approximates kz(z, z ′) ≈ k̃z,m(z, z′) := 1mφz,m(z) >φz,m(z
′), where φz,m takes value in Rm. Then the map ϕ 7→ 1√
m ϕ>φz,m(·) =: g(·;ϕ) parameterizes an approximate RKHS H̃; and for all c > 0, the
random function g(·;ϕ), where ϕ ∼ N (0, cI), is distributed as GP(0, ck̃z,m). The notations related to kx are similar and thus omitted. Now we can state the objective function: Proposition 4.1 (Proof in Appendix C.1). Let φ0 ∼ N (0, λν−1I), θ0 ∼ N (0, I), ỹi ∼ N (yi, λ). Then the optima θ∗ of
min θ∈Rm max φ∈Rm n∑ i=1 ( (f(xi; θ)− ỹi)g(zi;φ)− g(zi;φ) 2 2 ) − ν 2 ‖φ− φ0‖22 + λ 2 ‖θ − θ0‖22 (14)
parameterizes a random function which follows the quasi-posterior distribution (6), where the kernels are replaced by the random feature approximations.
Given the above proposition, we can sample from the random feature-approximated quasi-posterior by solving (14) with stochastic gradient descent-ascent; the approximation errors will be analyzed in the following. The objective (14) is closely related to (4); as we show in Appendix C.1.1, it is equivalent to
min f∈H̃ max g∈Ĩ n∑ i=1 ( (f(xi)− ỹi)g(zi)− g(zi) 2 2 ) − ν 2 ‖g − g0‖2Ĩ + λ 2 ‖f − f0‖2H̃, (15)
which differs from (4) only in the regularizers: instead of regularizing the norm of f and g, (15) encourages the functions to stay close to randomly sampled anchors [39]. Alternatively, we can view (15) as perturbing the point estimator (4), so that it has a covariance matching that of the quasi-posterior. A similar relation is also observed in the original randomized prior trick, which transforms GP regression to the optimization problem minf∈H̃ ∑n i=1(f(xi)− ỹi)2 + λ‖f − f0‖22. In both cases, the resultant algorithm for approximate inference has the same time complexity as ensemble training for point estimation.
While the algorithm can be directly applied to neural network models as in [22], we follow [40] and modify the objective, to account for the difference between the neural tangent kernel (NTK) [41] of a wide neural network architecture and the NNGP kernel of the corresponding infinite-width Bayesian neural network [42–44]. Concretely, we modify (14) as
min θ max φ n∑ i=1 ( (f̃θ(xi)− ỹi)g̃φ(zi)− g̃φ(zi) 2 2 ) − ν 2 ‖φ− φ0‖22 + λ 2 ‖θ − θ0‖22, (16)
where g̃φ(z) := g(z;φ)− g(z;φ0) + g̃0(z), g̃0(z) := √ λ
ν
〈 φ̄0, ∂g
∂φ
∣∣ φ=φ0 (z) 〉 ,
4This setting could be realistic in the sample splitting setup considered in [7]. Note that as long as we have finite samples for the estimation of f (the “second stage”), we still need to have λ̄ > 0.
and φ0 denotes the initial value of φ, and φ̄0 ∼ N (0, I) is a set of randomly initialized NN parameters independent of φ0; and f̃θ is defined similarly.
We only give a formal justification for the modification, under the assumption5 that the NNs remain in the kernel regime throughout training, so that g(z;φ)−g(z;φ0) = 〈φ−φ0, ∂g(z)∂φ |φ0〉2 [47]. Thus for the purpose of analyzing g(·;φ) − g(·;φ0), we can view g as a random feature model with the parameterization φ 7→ 〈φ, ∂g(z)∂φ |φ0〉2. Thus by the argument in Appendix C.1.1, we can show that the weight regularizer ‖φ − φ0‖2 is equivalent to ‖g(·;φ) − g(·;φ0)‖Ĩ = ‖g̃φ − g̃0‖Ĩ , where Ĩ is determined by the NTK kg,ntk(z, z′) := 〈∂g(z)∂φ |φ0 , ∂g(z′) ∂φ |φ0〉2. Similar arguments can be made for f̃θ and f̃0. Consequently, (16) is equivalent to an instance of (15) with H̃, Ĩ defined by the NTKs. Implementation details for the algorithm, including hyperparameter selection, are discussed in Appendix D.
Convergence analysis We now complete the analysis of the inference algorithm, by showing that for any fixed set of test points x∗, SGDA can approximate the marginal distribution Π(f(x∗) | D(n)) arbitrarily well given a sufficient computational budget. This implies that the approximate posterior is good enough for prediction purposes.
We place several mild assumptions on the random feature model, listed in Appendix C.2; they are satisfied by common approximations such as the random Fourier features [48]. The SGDA algorithm is described in detail in Appendix C.4. Under this setup, we have Proposition 4.2 (Proof in Appendix C.5). Fix D(n) and λ, ν > 0. Then there exist a sequence of choices of m and SGDA step-size schemes, such that for any l ∈ N, we have
supx∗∈X l max(‖µ̂m − µm‖2, ‖Ŝm − Sm‖F ) p→ 0.

heading: In the above, µ̂m, Ŝm denote the mean and covariance of the approximate marginal posterior for f(x∗), µ, S correspond to the true posterior, ‖·‖F denotes the Frobenius norm, and the convergence in probability is defined with respect to the sampling of random feature basis.

heading: 5 Related Work
text: Quasi-Bayesian analysis for GMM estimation problems was first developed in [49, 50, 16], which provided theoretical analysis for parametric models. The use of the quasi-posterior is motivated from the maximum entropy principle, based on which similar ideas have been developed in the machine learning literature [51–53]. Alternatively, it can also be obtained as a non-informative limit of certain Bayesian posteriors [54]. [17, 18] first analyzed the use of nonparametric models for quasi-Bayesian IV; no numerical study was presented. Closer to our work is [19] which constructed a quasi-posterior using Nadaraya-Watson smoothing as Ên. While similar in form to our method, the smoothing approach relies on stronger assumptions, especially for high-dimensional data; see [7, Appendix A.2.1] which compared the corresponding point estimators. Their approach is also harder to scale to large datasets.
Our quasi-Bayesian procedure builds upon the kernelized IV models [7, 20] and the dual formulation of IV regression [6, 20, 21, 26]. The formulation (3)-(4) is from [20, 26]. [7] proposes a similar method. Although it was motivated differently as a kernelized two-stage least squares (2SLS) estimator, its objective is asymptotically equivalent to (3); however, the slight difference in regularization prevents the use of estimation procedures similar to (4), as we discuss in Appendix C.1.4. An alternative kernel-based estimator is proposed in [55, 56]; in particular, [55] also includes a Gaussian process construction, but for the different purpose of computing a leave-one-out validation statistics. Other recent work on ML models for IV include [4, 57, 58]. It remains interesting future work to develop scalable quasi-Bayesian procedures for these methods, although the mean estimator derived from our quasi-posterior implementation has competitive performance.
[59] studies semi-parametric IV estimation using an exponentially tilted empirical likelihood (ETEL) estimator. ETEL also connect to Bayesian methods. However, similar to the quasi-Bayesian
5The same linearization assumption has been employed in [40, 45]. We refer readers to [46, 21] for analysis of the linearization error in various settings similar to ours.
approach, the need to estimate the conditional moment restrictions complicates the understanding of its frequentist behavior, as well as the design of scalable inference algorithms; in their case, it is especially unclear if the empirical likelihood can be computed in sublinear time. Other approaches for uncertainty quantification include Bayesian inference and bootstrap. We have discussed previous works on Bayesian IV and their limitations in Section 1. Bootstrap is typically justified in the asymptotic regime, which does not cover many scenarios where uncertainty is most needed; this is different from the quasi-Bayesian approach which can always be justified through (7). Moreover, standard bootstrap inference on 2SLS is known to be unreliable when the instrument strength is weak [36–38]; while remedies exist (e.g., [38]), they heavily rely on the linearity and additive noise assumptions, and are thus difficult to generalize to the nonlinear setting we are interested in. As the kernelized IV methods generalize 2SLS, we expect similar issues to exist in our setting.

heading: 6 Experiments
text: Code to reproduce the experiments is available at https://github.com/meta-inf/qbdiv.

heading: 6.1 1D Simulation
text: We first experiment on a variety of 1D synthetic datasets, constructed by modifying the setup in [5] to incorporate a nonlinear first stage, in a way similar to [7, 60]:
z := sigmoid(w), x := sigmoid ( αw + (1− α)u′√ α2 + (1− α)2 ) , yi ∼ N (f0(2x− 1) + 2u, 0.1),
where (u, u′) are normal random variables with unit variance and a correlation of 0.5, w ∼ N (0, 1) is independent of (u, u′), α is a parameter controlling the instrument strength, and f0 is constructed from the sine, step, abs or a linear function. We choose N ∈ {200, 1000} and α ∈ {0.05, 0.5}. Our baselines include BayesIV [15], a state-of-the-art Bayesian model based on B-splines and Dirichlet process mixture; we also include bootstrap on 2SLS with ridge regularization, either applied directly to the input features (Linear), on their polynomial expansion (Poly), or on the same kernelized models (KIV) as ours.6 Hyperparameter for the kernelized IV methods are selected by cross validation based on the observable first-stage and second-stage losses as in previous work [7, 20]; see Appendix D.1. For kernels we choose the RBF and Matérn kernels, although results for Matérn kernels are deferred to appendix for brevity. See Appendix D.3 for the detailed setup.
Normalized MSE and coverage rate of 95% credible intervals (CI) on the sine datasets are plotted in Figure 1. We report results on 20 independently generated datasets. As we can see, quasi-Bayesian inference provides the most reliable uncertainty estimates, and is especially advantageous in the weak IV setting (α = 0.05). While its CI can be conservative, we note that it is still informative, and properly reflects the sample size and instrument strength; see Appendix D.4 for visualizations. These results connect to previous work showing (in a different setting) that the radius of credible set produced by a GP posterior can have the correct order of magnitude [61].
We report the average run time in Table 1. We can see that in addition to the improved flexibility, our method is also more scalable compared with the standard Bayesian baseline, as it avoids the costly modeling of the joint residual distribution.
6We do not compare with [19] since their source code is unavailable. Note that their smoothing-based method does not scale to large datasets, and there is no numerical study on the credible interval in [19].
Full results on all datasets, visualizations and additional experiments are deferred to Appendix D.4. As a summary, (i) on the abs and linear datasets where all modeling assumptions hold, the results are qualitatively similar to the sin dataset. Moreover, the over-smoothed RBF kernel appears to have similar coverage comparing with the optimal kernel, and follow a similar contraction rate, as the previous work on GP regression [62] suggests. (ii) On the step dataset which violates Assumption 3.2, the quasi-posterior still provides more coverage than the baselines. (iii) Uncertainty estimates produced by the approximate inference algorithm are similar to that from the exact quasi-posterior.

heading: 6.2 Airline Demand
text: We now turn to the more challenging demand simulation first proposed by [4]. The dataset simulates a scenario where we need to predict the demand of airline tickets y, as a function of the price x, and two observed confounders: customer type s, and time of year t. The data generating process is
x := (z + 3)ψ(t) + 25 + u′, y := f0(x, t, s) + u, f0(x, t, s) := 100 + (10 + x) · s · ψ(t)− 2x where (u, u′) are standard normal variables with correlation ρ, z ∼ N (0, 1) is independent of (u, u′), and ψ is a nonlinear function whose shape is given in Figure 2. The variable s either varies across {0, . . . , 6} (the lower-dimensional setting), or is observed as an MNIST image representing the corresponding digit; the latter case represents the real-world scenario where only highdimensional surrogates of the true confounder is observed. We only report results for ρ = 0.5, noting that results using other choices of ρ have been similar. We use n ∈ {1000, 10000} for the lower-dimensional setting, and use n = 50000 for the image-based setting.
We compare our method with bootstrap on the same model, BayesIV, and bootstrap on linear or polynomial models. Performance of other point estimators on this dataset has been reported in [7, 20, 58], compared with which our method is generally competitive. We implement the dual IV model using both an RBF kernel and DNN models. See Appendix D.5 for details.
We report the test MSE and CI coverage for N = 1000 in Table 2, and visualize all approximate quasi-posteriors for the NN models in Figure 2. As we can see, when implemented with DNNs, our method produces uncertainty estimates with excellent coverage, which also correctly reflects the information available in the dataset: the CI is wider when N is smaller, or in the highdimensional experiment where estimation is harder. Bootstrap has a noticeably worse performance when N = 1000. Still, it performs well in the (arguably less interesting) large-sample setting, with a CI coverage similar to our method; see Appendix D.6. This is because on this dataset, the total instrument strength is stronger due to the presence of observed confounders, and the NN model is a good fit. Consequently, the asymptotic behavior of bootstrap can be observed when N is large.
Both methods have poorer performances when we switch to the RBF kernel, although the quasiposterior is still more reliable. We conjecture that both H and I are misspecified in this case, and
thus the credible intervals can only reflect the uncertainty within the prior model. These results show that NN models can be advantageous, which our inference algorithm supports.
The other baselines perform poorly due to their inflexibility; in particular, note that BayesIV uses additive models for both stages (e.g., f(x, t, s) = f1(x) + f2(t) + f3(s)) which do not approximate this data generating process well. Full results and visualizations are deferred to Appendix D.6.

heading: 7 Conclusion
text: In this work we propose a scalable quasi-Bayesian procedure for IV regression. We analyze the theoretical properties of the proposed quasi-posterior, and derive a scalable algorithm for approximate inference. Empirical evaluations show that the proposed method scales to large and highdimensional datasets, and can be particularly advantageous when the instrument strength is weak.
Beyond IV regression, formulations like (1) also appear in various other problems in causal inference and statistics, as discussed in, e.g., [21]; our method can be readily applied to these problems. Future work includes extension to more general conditional moment restriction problems with nonlinear constraints [63, 17].

heading: Acknowledgement
text: We thank the anonymous reviewers for their helpful feedback and references. Z.W., Y.Z. and J.Z. were supported by NSFC Projects (Nos. 61620106010, 62061136001, 61621136008, 62076147, U19B2034, U19A2081, U1811461), Beijing NSF Project (No. JQ19016), Beijing Academy of Artificial Intelligence (BAAI), Tsinghua-Huawei Joint Research Program, a grant from Tsinghua Institute for Guo Qiang, Tiangong Institute for Intelligent Computing, and the NVIDIA NVAIL Program with GPU/DGX Acceleration.
Broader Impact
IV analysis is widely used in social science and clinical research, where it can be highly undesirable to have a systematic error or increased prediction variance on subpopulations which are underrepresented in the training dataset, or on which the instrument is locally weak. Uncertainty quantification is an important first step in detecting such issues, and the quasi-Bayesian approach we take can be preferable since it has a clear interpretation in a non-asymptotic setting, and when the instrument strength is weak. However, epistemic uncertainty estimates are only as good as the model allows, and it should not create a false sense of complete security. Caution must be taken when employing complex models in scenarios with potential fairness implications.

Review summaries:
The paper concerns the uncertain quantification for instrumental variable (IV) regression that it is an approach for estimating causal effect from confounded observational data and, proposes a scalable quasi-Bayesian procedure. The procedure was built upon the kernelized IV models (Muandet et al (2020)). The contribution is forming a quasi-posterior by employing a Gaussian process prior and constructing a kernel conditional expectation estimator using a randomized prior trick. Theoretical properties and simulation studies are provided. 

The paper studies a quasi-Bayesian approach to instrumental variable (IV) regression, based on using an empirical estimate of the ‘likelihood’ to obtain a Gibbs posterior. Such an approach has been studied in earlier literature due to the difficulty of incorporating moment conditions in IV via a likelihood. The paper pursues a Gaussian process approach and proves some theory, namely that the resulting posterior puts most of its mass on functions (almost) satisfying the moment constraint. The paper provides a computation approach based on the randomized prior trick and illustrates the method on various simulated datasets.

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.

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