Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
52,726 characters · 10 sections · 100 citation commands
Fast Instrument Learning with Faster Rates
\doparttoc \faketableofcontents
\todo{disable todonotes and the extra margin} Instrumental variable (IV) analysis is widely used for causal inference angrist2008mostly,brookhart2010instrumental,burgess2017review. Given confounded observational data, IV analysis identifies the causal effect through the use of {\em instruments}. Nonlinear IV regression is typically defined by the following conditional moment restrictions (CMRs):
where $f_0$ is the causal effect function of interest, and $\bx,\by,\bz$ denote the observed treatment, response and instrument, respectively. Similar CMR problems also appear in other applications of causal statistics and machine learning chen_estimation_2012,liao_provably_2020.
Starting from hartford2017deep, recent works have demonstrated great promise in applying flexible machine learning (ML) methods to IV regression. Modern ML methods are appealing due to their {\em adaptivity} to the {\em informative latent structure} in data bach2017breaking: they may adapt to the low dimensionality of the informative latent features even if the observed input is high-dimensional and its signal-to-noise ratio is low. Sample complexity gaps have been established in such settings, between deep models based on neural networks (NNs) or trees, and linear models such as fixed-form kernels wei2019regularization,ghorbani2019limitations,schmidt-hieber_nonparametric_2020. In IV regression, such adaptivity will be highly desirable when the observed instruments are high-dimensional, which is prevalent in applications such as genomics lin_regularization_2015, and may also arise from the general desire to use structured data as instrument. Following previous work in the parametric setting singh2020machine,chen_mostly_2021, we refer to this problem of learning informative latents in instruments as {\em instrument learning}. It generalizes the classical problem of instrument selection okui2011instrumental,carrasco2012regularization,belloni2012sparse.
Comparing with standard supervised learning, IV regression is more challenging, due to the need to estimate a conditional expectation operator which defines (ref). Consequently, establishing adaptivity guarantees becomes more difficult. While many recent works have demonstrated promising empirical results using deep models, they are often used as heuristics hartford2017deep, or justified with crude {\em slow-rate} analyses, which establish convergence rates that saturate at $\Omega(n^{-1/4})$ xu_learning_2020. This is in contrast to {\em faster rates} which approach $n^{-1/2}$ as the regularity of model improves. The only exception is a minimax formulation of IV estimation liao_provably_2020,bennett2019deep,dikkala_minimax_2020,muandet_dual_2020. dikkala_minimax_2020 establish faster rate convergence for this formulation, for models with local Rademacher complexity bounds. Though local Rademacher analysis covers many adaptive ML procedures schmidt-hieber_nonparametric_2020,syrgkanis_estimation_2020, it still does not fully explain the success of modern ML approaches, with prominent alternatives including implicit regularization neyshabur2014search and PAC-Bayesian analyses alquier2021user. From a practical perspective, minimax optimization is computationally expensive, yet it cannot be avoided in the framework of dikkala_minimax_2020, unless we instantiate their method with the less flexible kernel models. It requires additional hyperparameter tuning, which can be challenging in causal problems where validation is indirect and more difficult. It also prevents the use of such flexible models for {\em uncertainty quantification}, or {\em inference}, for which reliable methods have only been developed for linear nonparametric models kato_quasi-bayesian_2013,chen_optimal_2015,wang2021quasibayesian.
This work bridges the gap between sharp theoretical guarantees and robust, practical implementation. We assume our prior knowledge about the causal effect function $f_0$ is characterized by a reproducing kernel Hilbert space (RKHS) $\cH$, and focus on the flexibility of conditional expectation estimation. This is often possible, because the treatment variable is determined by the problem at hand, and thus has a fixed dimensionality.\footnote{Our method can also be applied when $\bx$ contains high-dimensional exogenous covariates; see Appendix (ref).} We then present a surprisingly simple algorithm, with faster rate guarantees which in many cases match the best known in literature. The algorithm defines the conditional expectation estimates using a learned kernel, the basis of which is defined by applying adaptive regression algorithms to random draws from a Gaussian process (GP) prior. Given this learned first-stage kernel and $\cH$, we can estimate $f_0$ using kernelized IV methods dikkala_minimax_2020,muandet_dual_2020,singh_kernel_2020, which have closed-form solutions and can be efficiently approximated (e.g., with Nystr\"om dikkala_minimax_2020). Our method allows easy hyperparameter tuning, and exhibits competitive performance in simulations. It accesses the regression algorithm as a black box, thus allowing for the use of any ML methods and benefits from their established theoretical guarantees. It also enables fast quasi-Bayesian uncertainty quantification kato_quasi-bayesian_2013,chernozhukov_mcmc_2003,florens_nonparametric_2012 with improved flexibility.
Our algorithm connects to many ideas in literature. Most notably, it can be viewed as an infinite-dimensional generalization of singh2020machine,chen_mostly_2021, which consider a linear outcome model with fixed dimensionality, and use ML methods to learn the {\em optimal instruments} kelejian1971two,chamberlain_asymptotic_1987. Our setting requires different analyses, for defining an infinite-dimensional estimation target and quantifying errors with its intrinsic complexity. Additionally, analysis of the resultant IV estimator is complicated by the ill-posedness of infinite dimensional IV models horowitz_ill-posed_2014. A more subtle distinction is in the choice of basis: while for finite-dimensional function spaces we can pick any set of basis (i.e., features) and apply the black-box regressor separately, in our case seemingly obvious choices of basis lead to inferior results (Appendix (ref)). Our analysis also connects to the multi-task learning literature, to which we make technical contributions. Section (ref) discusses related work in detail.
The remaining of the paper is organized as follows: Section (ref) reviews background knowledge. Section (ref) introduces the instrument learning problem, and reduces it to a general kernel learning problem. We solve the latter problem in Section (ref), and return to IV in Section (ref) with our main results. We review related work in Section (ref), and present numerical experiments in Section (ref).
\paragraph{Notations} Denote the joint data distribution as $P(dz\times dx\times dy)$, its marginal distributions as $P(dx),P(dz)$, etc., and their support as $\cX,\cZ$. For functions of observed variables (e.g., $\bx$ or $\bz$), $\|\cdot\|_2$ denotes the $L_2$ norm w.r.t. the respective marginal data distribution. $\|\cdot\|_\infty$ denotes the $L_\infty$ norm. We use the notation $[m]:=\{1,\ldots,m\}$. Boldface ($\bx,\by,\bz$) emphasizes the denotation of random variables. For any kernel $k$, $\mathcal{GP}(0,k)$ refers to the “standard Gaussian process” van2008reproducing with zero mean, and covariance defined by $k$. $\lesssim,\gtrsim,\asymp$ represent (in)equalities up to constants; the hidden constants will not depend on any sample size. $\tilde \cO(\cdot)$ denotes inequality up to logarithm factors.
\paragraph{Problem Setup} Nonparametric IV regression (NPIV) is formulated as (ref) newey2003instrumental,horowitz_applied_2011. Introduce the {\em conditional expectation operator} $E: L_2(P(dx))\to L_2(P(dz)), f\mapsto \mathbb{E}(f(\bx)\mid\bz=\cdot)$, and define $g_0(z) := \mathbb{E}(\by\mid \bz=z)$. We can then express (ref) as a linear inverse problem:
where we observe $g_0$ up to regression error. NPIV deviates from standard inverse problems in its need to estimate both $f_0$ and $E$. Following conventions in the two stage least square method angrist2008mostly, we refer to the modeling of $f_0$ as the second stage, and that of $E$ -- or equivalently, that of $Ef$ for all $f$ in a hypothesis space -- as the first stage. The following assumption describes the setup in full detail:
We impose {\em (ii) } since such additional samples are sometimes available, as discussed in e.g. singh_kernel_2020. If only $n$ samples from the joint distribution are available, we can set $n_1=n_2=n/2$.
In the main text, we assume that our prior knowledge about $f_0$ is fully characterized by an RKHS $\cH$, in the following sense.
In the above, {\em (i) } and {\em (ii) } are common technical assumptions: {\em (i) } ensures the existence of Mercer's representation, and {\em (ii) } is a complexity measure, with a larger value of $b$ indicating a smaller hypothesis space. {\em (iii) } requires $\cH$ is correctly specified for regression; its two cases cover the different assumptions in standard RKHS-based estimation and GP modeling. {(ref)} is intuitive. {(ref)} is standard in the posterior contraction literature van_der_vaart_rates_2008; it roughly requires $f_0$ to be (at least) as regular as “typical” samples from $\mathcal{GP}(0, k_x)$, in the sense of van_der_vaart_rates_2008. This is different from (ref) because when $\cH$ is infinite dimensional, almost all GP samples fall out of $\cH$ driscoll1973reproducing,van_der_vaart_information_2011. Our algorithm applies to both settings, but analysis of the GP scheme requires additional effort. It is useful as it allows for quasi-Bayesian uncertainty quantification using a $\mathcal{GP}(0,k_x)$ prior wang2021quasibayesian.
The RKHS assumption has been employed in a thread of recent work singh_kernel_2020,muandet_dual_2020,zhang_maximum_2020,wang2021quasibayesian, and generalizes the sieve method in literature newey2003instrumental,blundell_semi-nonparametric_2007,chen_estimation_2012.
It is most reasonable when $\bx$ has moderate dimensions; for example, when $f_0$ satisfies certain $L_2$-Sobolev regularity conditions, we can set $k_x$ to be a suitable Mat\'ern kernel (Example (ref)-(ref)). Appendix (ref) studies a more general setting, where $\bx$ and $\bz$ include additional, high-dimensional exogenous covariates.
Nonetheless, the assumption will be less reasonable when the treatment variable is high-dimensional and variable selection is needed for it.
NPIV is typically an ill-posed inverse problem horowitz_applied_2011. We now quantify the degree of ill-posedness:
Such mildly ill-posed settings cavalier2008nonparametric match our polynomial eigendecay assumption for the kernel. In the severely ill-posed setting where the decay of $s_i(E)$ is exponential, kernels with a similar eigendecay should be used. While the analyses of the two settings share many ideas, the Bayesian inverse problem literature typically restricts to the former for technical reasons knapik_bayesian_2011,knapik2016bayes.
As we assume $\cH$ is a correctly specified second-stage model, it remains to determine the first stage. In this section, we show that an ideal first stage model can be defined using another RKHS $\cI$, determined by $\cH$ and $E$. Although its kernel $k_z$ has an unknown form, we demonstrate that we can access noisy samples from $\mathcal{GP}(0, k_z)$, which, as Section (ref) below shows, enable efficient learning of $\cI$. This can be viewed as instrument learning, as $\cI$ will only depend on the informative features in $\bz$ (Example (ref)).
Let us first consider the GP scheme (Assumption (ref) (ref)) which roughly requires $f_0$ to be similar to typical samples from $\mathcal{GP}(0, k_x)$. From a Bayesian perspective, an ideal prior for $E f_0$ should match the distribution of $E f$, for $f\sim\mathcal{GP}(0, k_x)$. This distribution is “almost equivalent” to another GP:
Informally, the lemma shows that $\mathcal{GP}(0,k_z)$ matches the distribution of $E f$. It is thus intuitive that $k_z$ could be a good choice for the first stage. The following lemma further motivates its use in the kernel scheme: its {\em (i) } shows that $\cI$ fulfills the conditions in previous work dikkala_minimax_2020,singh_kernel_2020: the restriction of $E$ on $\cH$ has image contained in $\cI$, and is a bounded linear map to $\cI$. {\em (ii) } shows that $\cI$, as a set of functions, cannot be made smaller while maintaining {\em (i) }.
We now demonstrate that $\cI$ only depends on the informative latent features.
While $k_z$ has ideal properties, it cannot be used directly as it involves the unknown operator $E$. Instead, we need to construct an approximation from data. Our main insight is that {\em we can effectively draw noisy samples from $\mathcal{GP}(0,k_z)$}; as we develop in Section (ref), such samples enable the approximation of $k_z$. To see how the noisy samples are obtained, consider $f\sim\mathcal{GP}(0,k_x)$. By Lemma (ref), $g = E f$ is $L_2$-equivalent to clean samples from $\mathcal{GP}(0,k_z)$; and we have $f(\bx) = g(\bz) + (f(\bx)-(Ef)(\bz))$, where the latter term is unpredictable given $\bz$, and from this perspective can be viewed as noise. Thus, if we apply any regression algorithm to $f\sim\mathcal{GP}(0,k_x)$, with $\bz$ as input, we will recover a “denoised” sample from $\mathcal{GP}(0,k_z)$, up to regression errors.
In the informative latent feature setting, optimal regression error can only be achieved by methods that adapt to such structures wei2019regularization,ghorbani2019limitations,schmidt-hieber_nonparametric_2020. Approximating $\cI$ with such “denoised” samples can then be viewed as a knowledge distillation procedure, which results in a compact representation of the adaptive regression algorithm. This is particularly beneficial in the NPIV setting: as discussed in the introduction, using a learned kernel eliminates the need of minimax optimization in estimation, and allows the (indirect) use of adaptive methods for uncertainty quantification.
In this section, we address the problem of kernel learning given noisy GP samples. As our results apply to more general settings, we first state the assumptions with full generality.
The above assumption applies to a latent-space kernel $\bar k_z$. As shown in Example (ref), $\Phi$ and $\bar k_z$ induce an input-space kernel $k_z$, and RKHS $\cI$, which inherit the assumed regularity conditions. Our goal is to estimate $k_z$. This is harder than the estimation of $\bar k_z$, as it also involves $\Phi$.
All conditions for $\bar k_z$ are satisfied by Mat\'ern kernels with a sutiable order; see Appendix (ref). Appendix (ref) discusses its applicability in the IV setting, where $\cI$ is defined as in Section (ref). Briefly, {\em(ii)} always holds for ${\bar b}\ge \max\{b, 2p-1\}$, and if $\cH$ is further correctly specified in the sense of Assumption (ref), ${\bar b} = b+2p$. {\em(i)} and {\em(iii)} hold under mild technical assumptions.
We will “denoise” noisy $\mathcal{GP}(0,k_z)$ samples using a regression oracle, which is specified below:
In the IV setting, we have $g=Ef \sim\mathcal{GP}(0,k_z)$, and $\mathbb{E}(\be\mid\bz)= \mathbb{E}(f(\bx)-(Ef)(\bz)\mid\bz)=0$; the subgaussian condition is verified by Lemma (ref).
To provide some intuition on adaptivity, we instantiate the assumption with the DNN model in schmidt-hieber_nonparametric_2020, and compare the resulted $\xi_n$ with fixed-form kernels:
We now define the approximate RKHS. Let $\{g^{(j)}:j\in[m]\}$ be i.i.d. samples from the GP prior, and $\hat g^{(j)}_{u,n_1}$ be the respective estimate returned by the regression oracle, constructed from the shared dataset $ \cD^{(n_1)} = \{(\tilde z_i, g^{(j)}(\tilde z_i)+\epsilon^{(j)}_i): i\in [n_1], j\in [m]\} $ where $\epsilon^{(j)}_i$ are subgaussian, mean-zero noise. Let $\hat g^{(j)}_{n_1} := \min\{C_k \log n, \hat g^{(j)}_{u,n_1}(\cdot)\}$, where $C_k$ is a constant determined by $\cI$. Define $\hat G_n(z) := (\hat g^{(1)}_{n_1}(z),\ldots,\hat g^{(m)}_{n_1}(z))$. Our approximate RKHS is defined as
As $\tilde\cI$ is a finite-dimensional linear space, it is an RKHS. We can check that $\|g\|_\infty\le C_k\log n_1\|g\|_{\tilde\cI}$.
\paragraph{Theoretical Results} Under a given model, regression error is decomposed into approximation and estimation (i.e., generalization) errors. We first present the approximation error bound:
We will use $\tilde\cI$ to estimate functions on a separate dataset with $n_2$ samples. For a single regression task, the estimation error can be simply bounded as $\tilde\cO(\sqrt{m/n_2})$ gyorfi2002distribution. However, our analysis of IV estimation will require quantifying the intrinsic complexity of $\tilde\cI$, which will also allow the use of a larger $m$ in practice. The following proposition provides one such result; it will be used in Section (ref), to analyze IV estimation in the kernel scheme (Assumption (ref) (ref)).
IV estimation in the GP scheme is more delicate, and requires additional analysis of $\tilde\cI$, which is deferred to App. (ref). Before we proceed, however, we illustrate the results on a simple regression task:
As discussed in Example (ref), when $d_z>d_l$, the DNN rate can outperform fixed-form kernels by a large margin. The above example demonstrates a similar superiority of the learned kernel.
We shall use the approximate first stage $\tilde\cI$ for IV regression, by plugging $\tilde\cI$ and $\cH$ to the kernelized estimators in dikkala_minimax_2020,wang2021quasibayesian; see Algorithm (ref). We analyze the resulted estimators in this section, while deferring implementation details, including hyperparameter selection, to Appendix (ref).
Both dikkala_minimax_2020 and the posterior mean estimator of wang2021quasibayesian have the form
Their difference lies in the regularization scaling, which arises from the different assumptions about $f_0$ and $\cH$ (Assumption (ref)). Thus, we analyze the resulted two estimators separately, in Section (ref) and Section (ref) below. In the setting of wang2021quasibayesian we are also able to justify the use of likelihood-based model comparison and (quasi-)Bayesian model averaging (BMA).
dikkala_minimax_2020 establish faster rate convergence of the point estimator under simple assumptions. We now provide corresponding results using our learned $\tilde\cI$, by plugging in the results in Section (ref).
Let us compare the result with dikkala_minimax_2020 in the setting of Example (ref). Suppose $n_1=n_2$. dikkala_minimax_2020 establishes the rate of $\cO_p((\xi'_{n_2}+n_2^{-(b+1)/(b+2)})(1+\|f_0\|_\cH^2))$, where $\xi'_{n_2}$ is comparable with a first-stage regression rate established from local Rademacher analysis. For the DNN model in Example (ref), we have $\xi'_n=\tilde\Theta(\xi_n)=\tilde\cO(\epsilon_{fea,n} + n^{-{\bar b}/2({\bar b}+1)})$ in the general case,\footnote{With some abuse of notation, we also use $\tilde\cO$ to hide the infinitesimal deterioration of the polynomial order.} or $\xi_n' = \tilde\cO(\epsilon_{fea,n} + n^{-({\bar b}+1)/2({\bar b}+2)})$ assuming additional regularity (Remark (ref)). Thus, the two rates are equivalent if $\epsilon_{fea}$ is sufficiently large, meaning that the difficulty of feature learning cannot be ignored. Otherwise, dikkala_minimax_2020 may be better if ${\bar b} < b+1$; this is a somewhat narrow range, as ${\bar b}\ge \max\{2p-1,b\}$. Recall that our method is more appealing computationally: directly instantiating dikkala_minimax_2020 with DNNs requires solving a minimax problem similar to (ref), while for our learned kernel (ref) can be evaluated in closed form.
We can also compare (ref) with kernelized IV using a fixed-form first stage. In the above setting, its best rate is also provided by dikkala_minimax_2020, and is dominated by the kernel regression error in Ex. (ref) which, as we discussed, can be much worse than $\xi_n$. Our improved rate has been made possible by the fact that we are approximating a first-stage model with optimal adaptivity (Ex. (ref)), at a rate that is also adaptive to the informative latent structure (Ex. (ref)).
In summary, {\em our algorithm combines the best of both worlds}: it maintains the sharp guarantees of adaptive models, and the simplicity of kernel methods.
Quasi-Bayesian analysis enables efficient uncertainty quantification for NPIV, without introducing extra risks of model misspecification kato_quasi-bayesian_2013,chernozhukov_mcmc_2003. wang2021quasibayesian studies a quasi-Bayesian posterior constructed from (ref) and a $\mathcal{GP}(0,k_x)$ prior. It is defined through the Radon-Nikodym derivative w.r.t. the prior, $ \big(\nicefrac{d\Pi(\cdot\mid\cD_{s2}^{(n_2)})}{d\Pi}\big)(f) \propto e^{-n_2\ell_{n_2}(f)} $. For kernel first-stage models, the quasi-posterior can be evaluated in closed form (App. (ref)). For general models, however, it is entirely unclear if approximate inference can be possible, since for any parameter $f$, evaluation of $\ell_n(f)$ involves solving a separate optimization problem. Our kernel learning algorithm enables the (indirect) use of such models.
Analysis of (quasi-)Bayesian procedures is more challenging, partly because of the weaker regularization. Thus, wang2021quasibayesian introduced additional assumptions. Our analysis is further complicated by a different assumption on $\cI$, and approximation errors in $\tilde\cI$, which necessitate further assumptions. App. (ref) discusses these assumptions in detail. For simplicity, we state the result in a “rate-optimal” case:\footnote{ Classical NPIV lower bounds continue to hold given full knowledge of $E$ chen_rate_2007, so the rate $n_2^{-\nicefrac{b}{2(b+2p+1)}}$ is minimax optimal irrespective of $n_1$. In our setting, it is certainly desirable to improve the dependency on $n_1$, and our restriction is only employed to simplify proof. In simulations we find the choice of $n_1=n_2$ works well. }
Theorem (ref) immediately implies Theorems 5, 6 in wang2021quasibayesian for our $\tilde\cI$, with the extra logarithms, as their proofs do not involve the first stage. Those results establish Sobolev norm rates, and justify uncertainty quantification by lower bounding the magnitude of posterior spread.
In the nonparametric Bayes literature, contraction results like Theorem (ref) often lead to the justification of marginal likelihood-based model selection and averaging. This is also the case here. The key ingredient is the following marginal quasi-likelihood bound:
This result allows the comparison of a finite number of second-stage RKHSes. Of particular interest is the comparison between {\em power RKHSes} (Defn. (ref)), which often have intuitive interpretations: e.g., for a Mat\'ern RKHS $\cH$ and $\gamma\in (2/(b+1),1)$, the power RKHS $\cH^\gamma$ is equivalent to lower-order Mat\'ern RKHSes (fischer2020sobolev; Ex. (ref)). We can verify that such $\cH^\gamma$ fulfills the assumptions about $\cH$. Thus, provided the other assumptions continue to hold, Corollary (ref) will hold for all such $\cH^\gamma$, with $b+1$ replaced by $\gamma(b+1)$, showing the marginal likelihood has a different asymptotics. Consequently, it establishes asymptotically valid comparison between such models, and justifies the use of BMA.
Analysis of more general settings requires additional effort: NPIV is an inverse problem, and we anticipate the subtleties of model selection in nonparametric inverse problems. For example, analyses are usually restricted to the selection of $\gamma$ knapik2016bayes,szabo2015frequentist,jia2018posterior, and the $\gamma>1$ case requires additional assumptions szabo2015frequentist.\footnote{We do not cover it here for brevity, noting that it is well-understood in inverse problem settings knapik2016bayes,szabo2015frequentist.} In the IV setting, it should also be noted that valid model comparison requires a good approximation to $E|_\cH$, since otherwise the quasi-likelihood becomes less meaningful at any finite sample size. The same intuition applies to other model selection procedures bennett2019deep,muandet_dual_2020,singh_kernel_2020 based on the estimated violation of (ref). When the approximation cannot be guaranteed, it could be preferable to stick to the prior knowledge and fix a conservative choice for $\cH$.
\paragraph{Multi-Task Learning} Our Example (ref) can also be viewed as quantifying sample efficiency improvements in multi-task learning, if we view the GP prior draws as the labeling functions for a handful of diverse training tasks, which share the representation $\Phi$. This general idea is not new: starting from tripuraneni_theory_2020,du_few-shot_2021, a line of recent work establishes similar results. Most related is du_few-shot_2021, which assumes a fixed-dimensional linear model for $\bar g$, and an adaptive $\Phi$ with metric entropy bounds. We assume more general models for both components, and do not require different training tasks to have separate inputs. On the flip side, du_few-shot_2021 allows for non-iid training tasks. du_few-shot_2021 investigated infinite-dimensional $\bar g$, but established a slow rate. We are unaware of any work that established fast-rate convergence for infinite-dimensional top-level models, or used ML models as a black box. Both aspects may be interesting for multi-task learning, and are necessary for instrument learning.
\paragraph{Causal Statistics} The double machine learning framework chernozhukov_doubledebiased_2018 also uses black-box ML models to estimate certain nuisance parameters in the model. While the operator $E$ can be viewed as a nuisance parameter, the structure of the NPIV problem is quite different: foster2019orthogonal noted that it is very unclear if such a view can be helpful for NPIV estimation; consistent with their remarks, we have also been unable to cast our problem into the double ML framework. Note that double ML has been applied to semiparametric estimation and inference for IV chernozhukov_doubledebiased_2018,syrgkanis2019machine,singh2020generalized,jung2021double, which are orthogonal to our goal.
It has long been known kelejian1971two,chamberlain_asymptotic_1987 that under a linear outcome model $f_0(\bx)=\theta^\top\bx$, using $\mathbb{E}(\bx\mid \bz)$ as instrument leads to $\sqrt{n}$-consistent estimates. Our Section (ref) can be viewed as an infinite-dimensional generalization of this observation.\footnote{As noted in singh_kernel_2020, when $f_0\in\cH$ for some RKHS $\cH$, the first stage should model $\mathbb{E}(f(\bx)\mid\bz)$ for all $f\in\cH$, as opposed to merely modeling $\mathbb{E}(\bx\mid\bz)$. Note that singh_kernel_2020 did not study the optimal choice of the first stage. } Given high-dimensional instruments and a parametric outcome model, there is a large body of literature on efficient inference; see singh2020machine for a review. As we move to nonparametric models, we focus on estimation which becomes much more challenging, in the spirit of foster2019orthogonal. Still, we have provided qualitative characterization for uncertainty estimates in Section (ref).
For the use of ML for nonlinear IV, hartford2017deep studied a heuristic application of NNs. We discussed the minimax formulation in introduction. liao_provably_2020,wang2021quasibayesian justified the use of NNs with the respective neural tangent kernels (NTKs) which, like other fixed-form kernels, cannot adapt to the informative latent structure wei2019regularization,ghorbani2019limitations. zhang_maximum_2020,chen_efficient_2021 investigated the combination of an NN-based second stage and a linear first stage, which could be useful in complementary scenarios. xu_learning_2020 considered feature learning in both stages, but only established a slow rate; as the authors noted, it is also unclear if their algorithm reliably minimizes the empirical risk.
For model selection in the setting of Section (ref), chen2021adaptive prove the validity of bootstrap-based selection for the sieve estimator newey2003instrumental. zhang_maximum_2020,hsu2019bayesian investigate the use of marginal likelihood for two different kernel-based IV estimators: hsu2019bayesian establish a crude $-\nicefrac{1}{4}\log n$ upper bound for the log marginal likelihood, and zhang_maximum_2020 connect it to the empirical leave-one-out validation error. Neither result fully justifies model selection as our Corollary (ref). For kernelized IV models, muandet_dual_2020,singh_kernel_2020,zhang_instrument_2021 proposed validation statistics for comparing {\em a finite number of} first stage models.
Our main simulation setup is adapted from bennett2019deep,dikkala_minimax_2020; Appendix (ref) presents additional experiment on the demand dataset hartford2017deep,xu_learning_2020. In bennett2019deep,dikkala_minimax_2020, the observed $\bz,\bx,\by$ are generated by
where $\bu$ is the confounder, $\be_x,\be_y\sim \cN(0, 0.1^2)$ are independent noise, and the constants $\mu,\sigma$ standardize $\by$. We consider three choices for $h$: {\bf(i)} $D=2$, $h$ is the identity function; this recovers the setup in previous work, and quantifies the hardness of the NPIV problem given true instruments. {\bf(ii)} $\mathrm{dim}\:\bz=D\in\{40,100\}$, $h$ is a three-layer DNN; this simulates a feature learning scenario, and ensures the observation has a low signal-to-noise ratio ($O(1/D)$). {\bf (iii)} $h$ maps $\bar\bz_1$ to a MNIST lecun1998mnist or CIFAR-10 krizhevsky2009learning image with matching label; the MNIST setting also appeared in previous work.
We consider two choices for $f_0$: {\bf (i)} a widely used collection of functions (e.g., $\sin,\mathrm{abs}$) in bennett2019deep. {\bf (ii)} $f_0\sim\mathcal{GP}(0,k_x)$. (ii) ensures the correct specification of $\cH$ and allows us to focus on the first stage.
We use a DNN as the black-box learner, and a RBF kernel for $\cH$, with bandwidth determined by marginal likelihood (ref). We set $N_1=N_2\in\{500,2500,5000\}$. We defer setup details and full results to Appendix (ref), and summarize the findings below:
\paragraph{Hyperparameter Selection (App. (ref))} We first study hyperparameter selection in instrument learning. We set $f_0\sim\mathcal{GP}$, $D\in\{2,40,100\}$. We find our validation statistics (ref) always correlates with the counterfactual MSE $\|\hat f_n-f_0\|_2^2$, and that across a large hyperparameter space, trained DNNs always outperform first-stage models based on RBF kernels, or randomly initialized DNNs.
\paragraph{Predictive Performance (App. (ref))} For $h$ defined as in (i-ii), we compare our algorithm with dikkala_minimax_2020, instantiated with kernel, tree and NN models. As shown in dikkala_minimax_2020, the baselines have competitive performance on this setup; the latter two models also enjoy adaptivity guarantees. A representative subset of results are plotted in Fig. (ref): our method has stable performance as we move to high dimensions, demonstrating excellent adaptivity. In contrast, fixed-form kernels fail to identify the informative features. AGMM-tree and AGMM-NN also have deteriorated performance as $D$ increases, despite their theoretical guarantees, presumably due to the challenges in optimization.
(ref) reports the run time of all methods in this experiment. As we can see, our method is more efficient than both adaptive baselines.
For image-based $h$, we compare with AGMM-NN and zhang_maximum_2020, which report the best results in the MNIST setting. Our method outperforms both baselines.
\paragraph{Uncertainty Quantification (App. (ref))} Table (ref) presents a subset of results for $f_0\sim\mathcal{GP}$. Comparing with a fixed-form RBF first stage, our method produces sharper credible intervals, which also have better coverage. For $f_0$ specified as in bennett2019deep, we experiment with BMA over a grid of RBF kernels, and present visualizations in Appendix (ref). We find that when the model is more correctly specified, BMA produces conservative uncertainty estimates which are nonetheless informative. However, when all models are severely misspecified (e.g., when $f_0$ is a step function), we cannot expect model-based uncertainty estimates to have ideal coverage.
\paragraph{Exogenous Covariates (App. (ref))} We evaluate the extended algorithm in (ref) on the demand dataset hartford2017deep, which is a widely used simulation design with high-dimensional exogenous covariates. As shown in the appendix, our extended algorithm has competitive performance.
This work was supported by NSFC Projects (Nos. 62061136001, 62076145, 62076147, U19B2034, U1811461, U19A2081, 61972224), Beijing NSF Project (No. JQ19016), BNRist (BNR2022RC01006), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University. J.Z is also supported by the XPlorer Prize.