The exact contents of citations.db main_text.text for this paper — one flattened LaTeX string, title through conclusion, appendix excluded, unmodified except for removing email addresses. This is what our citation measures are computed over.
52,726 characters
Fast Instrument Learning with Faster Rates
\maketitle
\begin{abstract}
We investigate nonlinear instrumental variable (IV) regression given high-dimensional instruments. We propose a simple algorithm which combines kernelized IV methods and an arbitrary, adaptive regression algorithm, accessed as a black box. Our algorithm enjoys faster-rate convergence and adapts to the dimensionality of informative latent features, while avoiding an expensive minimax optimization procedure, which has been necessary to establish similar guarantees. It further brings the benefit of flexible machine learning models to quasi-Bayesian uncertainty quantification, likelihood-based model selection, and model averaging. Simulation studies demonstrate the competitive performance of our method.
\end{abstract}
\doparttoc
\faketableofcontents
\section{Introduction}\label{sec:intro}
\todo{disable todonotes and the extra margin}
Instrumental variable (IV) analysis is widely used for causal inference \citep{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):
\begin{equation}\label{eq:npiv}
\mathbb{E}(\by-f_0(\bx)\mid \bz)=0\quad a.s.~[P(dz)]
\end{equation}
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 \citep[see, e.g.,][for examples]{chen_estimation_2012,liao_provably_2020}.
Starting from \cite{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 \citep{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 \citep{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 \citep{lin_regularization_2015}, and may also
arise from the general desire to use structured data
as instrument.
Following previous work in the parametric setting \citep{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 \citep{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
\eqref{eq:npiv}.
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 \citep[e.g.,][]{hartford2017deep}, or justified with crude {\em slow-rate} analyses, which
establish convergence rates that saturate at $\Omega(n^{-1/4})$ \citep[e.g.,][]{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 \citep{liao_provably_2020,bennett2019deep,dikkala_minimax_2020,muandet_dual_2020}. \cite{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 \citep[e.g.,][]{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 \citep{neyshabur2014search} and PAC-Bayesian analyses \citep{alquier2021user}.
From a practical perspective, minimax optimization is computationally expensive, yet it cannot be avoided in the framework of \cite{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 \citep{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{app:exo-algo}.}
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 \citep{dikkala_minimax_2020,muandet_dual_2020,singh_kernel_2020}, which have closed-form solutions and can be efficiently approximated (e.g.,~with Nystr\"om \cite{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 \citep{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 \cite{singh2020machine,chen_mostly_2021}, which consider a linear outcome model with fixed dimensionality, and use ML methods to learn the {\em optimal instruments} \citep{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 \cite{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{rmk:iv-alternative}).
Our analysis also connects to the multi-task learning literature, to which we make technical contributions. Section~\ref{sec:related-work} discusses related work in detail.
The remaining of the paper is organized as follows: Section~\ref{sec:setup} reviews background knowledge.
Section~\ref{sec:iv-to-klearn}
introduces the instrument learning problem, and reduces it to a general kernel learning problem. We solve the latter problem in Section~\ref{sec:main}, and return to IV
in Section~\ref{sec:iv} with our main results.
We review related work in Section~\ref{sec:related-work}, and present numerical experiments in Section~\ref{sec:experiments}.
\section{Notations and Setup}\label{sec:setup}
\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'' \citep{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 \eqref{eq:npiv} \citep{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 \eqref{eq:npiv} as a linear inverse problem:
\begin{equation}\label{eq:npiv-inv-problem}
E f_0 = g_0,
\end{equation}
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 \citep{angrist2008mostly},
we refer to the modeling of $f_0$ as the \emph{second stage}, and that of
$E$ -- or equivalently, that of $Ef$ for all $f$ in a hypothesis space -- as the \emph{first stage}.
The following assumption describes the setup in full detail:
\begin{assumption}[NPIV]\label{ass:npiv} {\em(i)}
The data variables $\bx,\by,\bz$ satisfy \eqref{eq:npiv}, and $\by$ is bounded by $B$. {\em(ii)}
We observe two sets of i.i.d.~samples, with matching (marginal) distributions: $\cD_{s2}^{(n_2)} := \{(z_i,x_i,y_i): i\in [n_2]\}$,
$\cD_{s1}^{(n_1)} := \{(\tilde z_i,\tilde x_i): i\in [n_1]\}$.
\end{assumption}
\vspace{-0.1em}
We impose {\em (ii) } since such additional samples are sometimes available, as discussed in e.g.~\cite{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.
\begin{assumption}[second stage RKHS]\label{ass:s2}
{\em(i)} $\cX$ is a bounded subset of $\mathbb{R}^{d_x}$; the reproducing kernel $k_x$ of $\cH$ is bounded and continuous.
{\em(ii)} The integral operator $T_x: f\mapsto \int k_x(x,\cdot)f(x)P(dx)$ has eigenvalues
$\lambda_i(T_x) \lesssim i^{-(b+1)}$, for some $b>0$.
{\em(iii)} {\em One of the following} holds true:
\begin{enumerate}[label={\em(iii)}.\alph*,leftmargin=*]
\item\label{it:kernel-scheme}\!({\underline{``kernel scheme''}}): $f_0\in\cH$.
\item\label{it:gp-scheme}\!({\underline{``GP scheme''}}): $b>1$; and for all $n$, $\exists f^\dagger_{n}\in\cH$ s.t.~$
\|f^\dagger_{n}-f_0\|_\cH\lesssim n^{\frac{1/2}{b+1}}, ~
\|f^\dagger_{n}-f_0\|_2\lesssim n^{-\frac{b/2}{b+1}}.
$
\end{enumerate}
\end{assumption}
\vspace{-0.1em}
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{it:kernel-scheme}} is intuitive. {\ref{it:gp-scheme}} is standard in the posterior contraction literature \citep{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 \cite[Theorem 2.1]{van_der_vaart_rates_2008}. This is different from \ref{it:kernel-scheme} because when $\cH$ is infinite dimensional, almost all GP samples fall out of $\cH$ \cite{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 \cite{wang2021quasibayesian}.
The RKHS assumption has been employed in a thread of recent work \citep{singh_kernel_2020,muandet_dual_2020,zhang_maximum_2020,wang2021quasibayesian}, and generalizes the sieve method in literature \citep{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{ex:matern}-\ref{ex:f0}).
Appendix~\ref{app:exo-algo} 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 \citep{horowitz_applied_2011}. We now quantify the degree of ill-posedness:
\begin{assumption}\label{ass:mildly-ill-posed}
The operator $E$ is compact, with singular values
$s_i(E)\asymp i^{-p}$, where $p>0$.
\end{assumption}
\vspace{-0.1em}
Such \emph{mildly ill-posed} settings \cite{cavalier2008nonparametric} match our polynomial eigendecay assumption for the kernel.
In the \emph{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 \citep{knapik_bayesian_2011,knapik2016bayes}.
\section{From Instrument to Kernel Learning}\label{sec:iv-to-klearn}
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{sec:main} 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{ex:informative-latents}).
Let us first consider the GP scheme (Assumption \ref{ass:s2} \ref{it:gp-scheme})
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:
\begin{lemma}[proof in Appendix~\ref{app:proof-sec-iv-to-kl}]\label{lem:iv-to-kl-1}
Denote by $[g]_\sim$ the $L_2$ equivalence class of $g$.\footnote{Recall the $L_2$ space is not a function space, and consists of equivalence classes of functions.
Note that for readability, we may occasionally ignore this distinction in the main text, and use (a version of) $E$ to also denote the corresponding map between function spaces. All such denotations can be made unambiguous (Remark~\ref{rmk:regular-cond-exp}), and all null set ambiguities in this section can be removed under mild additional assumptions (Lemma~\ref{lem:iv-to-kl-further}).}
Under Assumptions~\ref{ass:npiv}, \ref{ass:s2}, there exists a kernel $k_z$, with integral operator $T_z = E T_x E^\top$,
s.t.~for $f\sim\mathcal{GP}(0,k_x),g\sim\mathcal{GP}(0, k_z)$, $[g]_\sim$ has the same distribution as $E[f]_\sim$.
\end{lemma}
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 \cite{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) }.
\begin{lemma}[proof in Appendix~\ref{app:proof-sec-iv-to-kl}]\label{lem:optimal-I}
Let $\cI$ be the RKHS defined by $k_z$.
Under Assumptions~\ref{ass:npiv}, \ref{ass:s2}, {\em(i)}
for any $f\in\cH$, there exists $g\in\cI$ s.t.~$[g]_\sim = E[f]_\sim$; {\em(ii)}
for any $g\in\cI$, there exists $f\in\cH$ satisfying the above.
In both cases, we have $\|f\|_\cH = \|g\|_\cI$.
\end{lemma}
We now demonstrate that $\cI$ only depends on the informative latent features.
\begin{example}[informative latent structure]\label{ex:informative-latents}
Let $\Phi: \cZ\to\bar\cZ$ be a {\em feature extractor} that maps the observed instruments $\bz$ to latent features $\bar\bz:=\Phi(\bz)$, s.t.~$\mathbb{E}(f(\bx)\mid\bz)=\mathbb{E}(f(\bx)\mid\Phi(\bz))$ for all $L_2$-integrable $f$.
Then
we can apply Lemma~\ref{lem:iv-to-kl-1}, with $E$ replaced by $\bar E: f\mapsto \mathbb{E}(f(\bx)\mid \bar\bz)\in L_2(P(d\bar z))$, leading to a latent-space
RKHS $\bar\cI$ with kernel $\bar k_z$. $\bar\cI$ induces the input-space RKHS
$$
\cI := \{g = \bar g\circ \Phi: \bar g\in\bar\cI\}, ~~\|\bar g\circ\Phi\|_\cI := \|\bar g\|_{\bar\cI}; ~~ k_z(z,z') = \bar k_z(\Phi(z),\Phi(z')).
$$
The above $k_z$ satisfies Lemma~\ref{lem:iv-to-kl-1}-\ref{lem:optimal-I}.\footnote{There may be multiple kernels satisfying Lemma~\ref{lem:iv-to-kl-1}, but they are equivalent up to null sets (Claim~\ref{claim:ex-fs}); the ambiguity can be removed under mild assumptions (Lemma~\ref{lem:iv-to-kl-further}).}
Observe that $\cI$ perfectly approximates $\{E f:f\in\cH\}$, but its complexity only depends on $\bar\cI$.
In particular,
$k_z$ has the same Mercer eigenvalues as $\bar k_z$ (Claim~\ref{claim:ex-fs}),
the decay of which is a standard complexity measure \citep[e.g.,][Ch.~7]{steinwart2008support}.
\end{example}
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{sec:main}, 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{lem:iv-to-kl-1}, $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 \citep{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.
\vspace{-0.3em}
\section{Black-Box Kernel Learning}\label{sec:main}
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.
\begin{assumption}[RKHS]\label{ass:s1-rkhs}
There exist a continuous function $\Phi:\cZ\to \bar\cZ$, and a reproducing kernel $\bar k_z$ over $\bar\cZ$, s.t.
{\em(i)} the random variable $\bar \bz=\Phi(\bz)$ is supported on a bounded subset of $\mathbb{R}^{d_l}$; $\bar k_z$ is bounded.
{\em(ii)} The eigenvalues of the integral operator $T_{\bar z}: \bar g\mapsto \int \bar k_z(\bar z, \cdot) g(\bar z) P(d\bar z)$
satisfy $\lambda_i(T_{\bar z}) \lesssim i^{-({\bar b}+1)}$, for some ${\bar b}>0$.
{\em(iii)} $\bar g\sim\mathcal{GP}(0,\bar k_z)$ have finite sup norm with probability 1.
\end{assumption}
The above assumption applies to a latent-space kernel $\bar k_z$.
As shown in Example~\ref{ex:informative-latents}, $\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{app:technical}. Appendix~\ref{app:regularity} discusses its applicability in the IV setting, where $\cI$ is defined as in Section~\ref{sec:iv-to-klearn}. 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{ass:qb-1}, ${\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:
\begin{assumption}[regression oracle]\label{ass:s1-oracle}
Let $\cD^{(n_1)} := \{(\tilde z_i, g(\tilde z_i)+e_i)\}$ be $n_1$ iid replications of the rvs $(\bz,g(\bz)+\be)$,
s.t.~$\mathbb{E}(\be\mid \bz)=0$ and $g(\bz)+\be$ has a $1$-subgaussian distribution. Then
the oracle returns estimator $\hat g_{u,n_1}$ s.t.
$
\mathbb{E}_{g\sim\mathcal{GP}(0,k_z)}\mathbb{E}_{\cD^{(n_1)}} \|\hat g_{u,n_1} - g\|_2^2 \le \xi_{n_1}^2,
$
for some $\xi_{n_1}\to 0.$
\end{assumption}
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{lem:borell-tis}.
To provide some intuition on adaptivity,
we instantiate the assumption with the DNN model in \cite{schmidt-hieber_nonparametric_2020}, and compare the resulted $\xi_n$ with fixed-form kernels:
\begin{example}[adaptivity of DNN oracles]\label{ex:dnn-oracles}
Let $\cZ\subset \mathbb{R}^{d_z}$,
$\Phi: \cZ \to \bar\cZ$ be $\beta_1$-H\"older regular,
$\bar\cI$ be a Mat\'ern-$\beta_2$ RKHS, and $\beta_1,\beta_2 \ge 1$. Let the regression oracle return a $\epsilon_{opt}^2$-approximate empirical risk minimizer for the
model in \cite{schmidt-hieber_nonparametric_2020}. Then for any $\epsilon>0$, it holds that (see Appendix~\ref{app:derivation-examples} for derivations)
\begin{equation}\label{eq:xi-n-dnn}
\xi_n =\tilde \cO\Big(n^{-\frac{\beta_1}{2\beta_1+d_z}} + n^{-\frac{\beta_2-\epsilon}{2\beta_2+d_l}} + \epsilon_{opt}\Big) =:
\tilde \cO\Big(\epsilon_{fea,n} + n^{-\frac{\beta_2-\epsilon}{2\beta_2+d_l}} + \epsilon_{opt}\Big)
\end{equation}
In the above, $\epsilon_{fea,n}$ characterizes the hardness of feature learning, i.e., learning $\Phi$. The second term characterizes that of kernelized regression given the optimal features: it matches the optimal regression rate if we {\em had} full knowledge about $\Phi$, or equivalently, $\cI$, and would be attainable by kernel ridge regression (KRR) using $\cI$.
As long as $\epsilon_{opt}$ is small, \eqref{eq:xi-n-dnn} will match the minimax rate up to logarithms. When $\nicefrac{\beta_1}{d_z} < \nicefrac{\beta_2}{d_l}$, the minimax rate is
$
\epsilon_{fea,n} \gg n^{-\beta_2/(2\beta_2+d_l)},
$
meaning that the hardness of feature learning cannot be overlooked.
Otherwise, the rate $\xi_n$ nearly matches the
rate given full knowledge of the unknown $\cI$, up to the infinitesimal $\epsilon>0$;
this is realistic when, e.g., $d_z\gg d_l$ and $\Phi$ is linear ($\beta_1=\infty$).
We are interested in the high-dimensional regime where $d_z \gg d_l$. In this case,
fixed-form Mat\'ern or RBF kernels could only attain the rate of $\cO\big(n^{-\frac{\min\{\beta_1,\beta_2\}}{2\min\{\beta_1,\beta_2\}+d_z}}\big)$, which
can always be much worse than \eqref{eq:xi-n-dnn}, {\em regardless of the hardness of feature learning}. This comparison suggests
that fixed-form kernels cannot adapt to the latent feature structure
to avoid the curse of dimensionality.\footnote{\cite{schmidt-hieber_nonparametric_2020} establishes formal lower bounds.
Also, for small $\beta_2$, we can replace $d_z$ with a manifold dimensionality of $\cZ$, but it can still be much larger than $d_l$.}
\end{example}
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
\vspace{-0.1em}
\begin{equation}\label{eq:approx-rkhs}
\tilde\cI := \big\{g(z) =\theta^\top\hat G_{n_1}(z)\text{ for some }\theta\in\mathbb{R}^m\big\},
~~\text{with norm }\|g\|_{\tilde\cI} := \sqrt{m}\|\theta\|_2.
\end{equation}
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}$.
\vspace{-0.3em}
\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:
\begin{theorem}[proof in Appendix~\ref{app:proof-prop-approx}]\label{prop:approx}
Under Assumptions~\ref{ass:s1-rkhs}, \ref{ass:s1-oracle},
there exists a universal constant $c_r>0$, and an event $E_{n_1}$ determined by $g^{(1\ldots m)}$ and $\cD^{(n_1)}$ with $\mathbb{P}_{\cD^{(n_1)}} E_{n_1}\to 1$, on which
for any $g^*\in L_2(P(dz))$, there exists $\tilde g^*\in\tilde\cI$ s.t.
\begin{align}
\|\tilde g^*\|_{\tilde\cI} &\le c_r\|\mathrm{Proj}_{m'}{g^*}\|_\cI, \label{eq:approx-norm-bound}\\
\|\tilde g^* - g^*\|_2 &\le c_r\|\mathrm{Proj}_{m'}{g^*}\|_{\cI}(\xi_{n_1} + m^{-({\bar b}+1)/2}) \sqrt{\log n_1}
+ \|g^* - \mathrm{Proj}_{m'}{g^*}\|_2,\label{eq:approx-bound}
\end{align}
where $m'=[m/2]$, and $\mathrm{Proj}_{m'}$ denotes the projection onto the top $m'$ Mercer eigenfunctions of $k_z$.
\end{theorem}
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})$ \cite{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{sec:iv}, to analyze IV estimation in the kernel scheme (Assumption~\ref{ass:s2}~\ref{it:kernel-scheme}).
\begin{proposition}[proof in Appendix~\ref{app:estimation}]\label{prop:est}
Let $\tilde\cI,\cD^{(n_1)}$ be defined as above, and $\delta_{n_2}$ be the critical radius of the local Rademacher complexity of the norm ball $\tilde\cI_1$ \citep[Ch.~14]{wainwright2019high}.
On the event defined in Theorem~\ref{prop:approx}, we have
$
\delta_{n_2} = \tilde O(
n_2^{-({\bar b}+1)/2({\bar b}+2)} + m^{-({\bar b}+1)/2} +\xi_{n_1}
).
$
\end{proposition}
IV estimation in the GP scheme is more delicate, and requires additional analysis of $\tilde\cI$, which is deferred to App.~\ref{app:gp-fixed-design-regr}.
Before we proceed, however, we illustrate the results on a simple regression task:
\begin{example}[Example~\ref{ex:dnn-oracles}, cont'd]\label{ex:dnn-oracles-cont}
Let $\Phi,\cI$ be defined as before, and $\hat g^{(j)}_{n_1}$ be estimated by the DNN oracle. Suppose $\epsilon_{opt}$ is not greater than the other terms. Then
\begin{enumerate}[nosep,leftmargin=*,label=\roman*.]
\item
Let $m = \lceil n_1^{{\bar b}/({\bar b}+1)^2}\rceil$. On the event in Theorem~\ref{prop:approx},
for any $g^*\in\cI$, there exists $\tilde g^*\in\tilde\cI$ s.t.
$\|\tilde g^*\|_{\tilde\cI} \le c_r \|g^*\|_\cI,
\|\tilde g^*-g^*\|_2
= \tilde O(\|g^*\|_\cI \xi_{n_1})$.
\item Let $g^*\sim\mathcal{GP}(0,k_z)$.
A refined analysis, based on Corollary~\ref{thm:ml}, shows that when $m =\lceil n_1^{1/({\bar b}+1)}\rceil$, there exists $\tilde g$ s.t.
$
\|\tilde g\|_{\tilde\cI} \lesssim n_1^{1/2({\bar b}+1)},
\mathbb{E}_{g^*\sim\mathcal{GP}(0,k_z)} \|\tilde g-g^*\|_2 = \tilde \cO(\xi_{n_1}).
$
\end{enumerate}
(See Appendix~\ref{app:derivation-example-cont} for derivations, and another high-probability bound in the GP scheme.)
Let $\hat g^*_n$ be the truncated OLS estimate using $\tilde\cI$, on a dataset $\{(z_i,g^*(z_i)+e_i):i\in [n_2]\}$ where $\mathbb{E}(e_i\mid z_i)=0,\mathrm{Var}(e_i)\le 1,\|g^*\|_\infty\le B$. Then
$\mathbb{E}\|\hat g^*_n-g^*\|_2=\tilde\cO(\|\tilde g^*-g^*\|_2+B\sqrt{m/n_2})$ \citep[Thm.~11.3]{gyorfi2002distribution}. When $n_1=n_2$, the latter term is $\ll \xi_{n_2}$, and
case (ii) above always matches the DNN rate.
Case (i) matches the DNN rate when feature learning becomes harder ($\xi_{n_1}^2\gtrsim n_1^{-{\bar b}/({\bar b}+1)}$); otherwise the rate may be slightly inferior, but still approaches $n_1^{-1/2}$ as the regularity ${\bar b}$ improves.
\end{example}
As discussed in Example~\ref{ex:dnn-oracles}, 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.
\section{Results for IV Regression}\label{sec:iv}
We shall use the approximate first stage $\tilde\cI$ for IV regression, by plugging $\tilde\cI$ and $\cH$ to the kernelized estimators in \cite{dikkala_minimax_2020,wang2021quasibayesian}; see Algorithm~\ref{alg:main}. We analyze the resulted estimators in this section, while deferring implementation details, including hyperparameter selection, to
Appendix~\ref{app:algo}.
\begin{algorithm}[h]
\caption{Kernelized IV with learned instruments.}\label{alg:main}
\begin{algorithmic}[1]
\REQUIRE $\cD_{s1}^{(n_1)},\cD_{s2}^{(n_2)}$; regression algorithm $\mathsf{Regress}$; second-stage kernel $k_x$; $m\in\mathbb{N}$
\FOR{$j \gets 1$ to $m$}
\STATE Sample
$f^{(j)}\sim \mathcal{GP}(0, k_x)$
\STATE
$\hat g^{(j)}_{u,n_1}\gets \mathsf{Regress}(\{(\tilde z_i, f^{(j)}(\tilde x_i)): i\in [n_1]\})$
\ENDFOR~{\color{gray}\COMMENT{the $m$ invocations of $\mathsf{Regress}$ may be replaced with a single vector-valued regression}}
\STATE Define
$
\tilde k_z(z,z') :=
\frac{1}{m} \sum_{j=1}^{m} \hat g^{(j)}_{n_1}(z)\hat g^{(j)}_{n_1}(z'),
$
where $\hat g^{(j)}_{n_1} := \min\{\hat g^{(j)}_{u,n_1}(\cdot), C\log m\}$.
\STATE {\bf return} $\mathsf{KernelizedIV}(\cD_{s2}^{(n_2)},\tilde k_z,k_x)$~{\color{gray}\COMMENT{See \eqref{eq:minimax-estimator} below, or Appendix~\ref{app:algo} for the closed-form solution}}
\end{algorithmic}
\end{algorithm}
Both \cite{dikkala_minimax_2020} and the posterior mean estimator of \cite{wang2021quasibayesian} have the form
\begin{equation}\label{eq:minimax-estimator}
\arg\min_{f\in\cH}\ell_{n_2}(f)+\mu\|f\|_\cH^2 := \arg\min_{f\in\cH}\max_{g\in\tilde\cI} \frac{1}{n_2}\sum_{i=1}^{n_2}
(y_i-f(x_i)- \kappa g(z_i))g(z_i) - \lambda \|g\|_{\tilde\cI}^2 + \mu\|f\|_\cH^2.
\end{equation}
Their difference lies in the regularization scaling, which arises from the different assumptions about $f_0$ and $\cH$ (Assumption~\ref{ass:s2}).
Thus, we analyze the resulted two estimators separately, in Section~\ref{sec:dikkala-regime} and Section~\ref{sec:qb-regime} below.
In the setting of \cite{wang2021quasibayesian} we are also able to justify the use of likelihood-based model comparison and (quasi-)Bayesian model averaging (BMA).
\subsection{Estimation in the Kernel Scheme} \label{sec:dikkala-regime}
\cite{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{sec:main}.
\begin{proposition}[proof in Appendix~\ref{app:dikkala-proof}]\label{prop:dikkala-regime}
Assume Assumptions~\ref{ass:npiv}, \ref{ass:s2} (kernel scheme), \ref{ass:s1-rkhs} and \ref{ass:s1-oracle}.
Let $\tilde\cI$ be defined by $\tilde k_z$ in Algorithm~\ref{alg:main}, and
$\hat f_{n_2}$ be defined by \eqref{eq:minimax-estimator}, with $\kappa,\lambda,\nu$ set as in Appendix E.1.
On the event defined in Theorem~\ref{prop:approx}, we have
\begin{equation}\label{eq:iv-rate-dikkala-setting}
\!\|E(\hat f_{n_2} - f_0)\|_2 = \tilde \cO_p\big(\big(\xi_{n_1}+n_2^{-\frac{b+1}{2(b+2)}}\big)(1+\|f_0\|_\cH^2)\big).\!\!
\end{equation}
\end{proposition}
Let us compare the result with \cite{dikkala_minimax_2020} in the setting of Example~\ref{ex:dnn-oracles}. Suppose $n_1=n_2$. \cite{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{ex:dnn-oracles}, 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{rmk:xi-optimality-rkhs}). Thus, the two rates are equivalent if $\epsilon_{fea}$ is sufficiently large, meaning that
the difficulty of feature learning cannot be ignored. Otherwise, \cite{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 \cite{dikkala_minimax_2020} with DNNs requires solving a minimax problem similar to \eqref{eq:minimax-estimator}, while for our learned kernel \eqref{eq:minimax-estimator} can be evaluated in closed form.
We can also compare \eqref{eq:iv-rate-dikkala-setting} with kernelized IV using a fixed-form first stage. In the above setting, its best rate is also provided by \cite{dikkala_minimax_2020}, and is dominated by the kernel regression error in Ex.~\ref{ex:dnn-oracles} 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{ex:informative-latents}), at a rate that is also adaptive to the informative latent structure (Ex.~\ref{ex:dnn-oracles}).
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.
\subsection{Quasi-Bayesian Estimation and Uncertainty Quantification}\label{sec:qb-regime}
Quasi-Bayesian analysis enables efficient uncertainty quantification for NPIV, without introducing extra risks of model misspecification \citep{kato_quasi-bayesian_2013,chernozhukov_mcmc_2003}.
\cite{wang2021quasibayesian} studies a quasi-Bayesian posterior constructed from \eqref{eq:minimax-estimator} 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{app:algo}).
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, \cite{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{app:qb-ass} 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$ \citep{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.
}
\begin{theorem}[posterior contraction; proof in App.~\ref{app:qb-proof}]\label{prop:qb-regime}
Assume Asms.~\ref{ass:npiv}, \ref{ass:s2} (GP scheme), \ref{ass:mildly-ill-posed}, \ref{ass:s1-rkhs}, \ref{ass:s1-oracle}, \ref{ass:emb-general}, \ref{ass:oracle-sup-norm-err}, \ref{ass:qb-1}, \ref{ass:qb-n}.
Let $n_1$ be s.t.~$\xi_{n_1}^2\log n_2+n_1^{-(b+2p)/(b+2p+1)}\lesssim n_2^{-1}$, and $m\asymp n_1^{1/(b+2p+1)}$.
Let $\Pi_{n_1}(\cdot\mid\cD_{s2}^{(n_2)})$ be defined in \eqref{eq:quasi-posterior} in appendix.
Then, with $\cD_{s1}^{(n_1)}$-probability $\to 1$, we have
\begin{align*}
\mathbb{E}_{\cD_{s2}^{(n_2)}}\Pi_{n_1}(\{f:\|f-f_0\|_2\ge M\bar\epsilon_{n_2}\}\mid\cD_{s2}^{(n_2)}) &\to 0, \\
\mathbb{E}_{\cD_{s2}^{(n_2)}}\Pi_{n_1}(\{f:\|E(f-f_0)\|_2\ge M\bar\delta_{n_2}\}\mid\cD_{s2}^{(n_2)}) &\to 0,
\end{align*}
where $\bar\delta_{n_2} =
\tilde \cO(n_2^{-(b+2p)/2(b+2p+1)})$,
$\bar\epsilon_{n_2}=\tilde \cO(n_2^{-b/2(b+2p+1)})$.
\end{theorem}
Theorem~\ref{prop:qb-regime} immediately implies Theorems 5, 6 in \cite{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{prop:qb-regime} 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:
\begin{corollary}[proof in Appendix~\ref{app:proof-qb-mlh}]\label{lem:qb-mlh}
In the setting of Theorem~\ref{prop:qb-regime}, for some $C>0$ we have
\begin{align*}
\mathbb{P}_{\cD_{s2}^{(n_2)}}\big(
C^{-1}n_2^{\frac{1}{b+2p+1}}\log^{-\frac{6}{b}}n_2
&\le -\log \Pi_{n_1}(\cD_{s2}^{(n_2)})
\le Cn_2^{\frac{1}{b+2p+1}} \log^2 n_2\big) \to 1.
\end{align*}
\end{corollary}
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{defn:power-spaces}), 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 (\citep{fischer2020sobolev}; Ex.~\ref{ex:matern-gp}). We can verify that such $\cH^\gamma$ fulfills the assumptions about $\cH$. Thus, provided the other assumptions continue to hold,
Corollary~\ref{lem:qb-mlh} 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$ \citep{knapik2016bayes,szabo2015frequentist,jia2018posterior}, and
the $\gamma>1$ case requires additional assumptions \citep{szabo2015frequentist}.\footnote{We do not cover it here for brevity, noting that it is well-understood in inverse problem settings \citep{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 \citep{bennett2019deep,muandet_dual_2020,singh_kernel_2020} based on the estimated violation of \eqref{eq:npiv}.
When the approximation cannot be guaranteed,
it could be preferable to stick to the prior knowledge and fix a conservative choice for $\cH$.
\section{Related Work}\label{sec:related-work}
\paragraph{Multi-Task Learning} Our
Example~\ref{ex:dnn-oracles-cont} 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 \cite{tripuraneni_theory_2020,du_few-shot_2021}, a line of recent work establishes similar results.
Most related is \cite[Sec.~5]{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, \cite{du_few-shot_2021} allows for non-iid training tasks. \cite[Sec.~6]{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.
\vspace{-0.2em}
\paragraph{Causal Statistics}
The double machine learning framework \citep{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: \cite[p.~8]{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 \citep{chernozhukov_doubledebiased_2018,syrgkanis2019machine,singh2020generalized,jung2021double}, which are orthogonal to our goal.
It has long been known \citep{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{sec:iv-to-klearn} can be viewed as an infinite-dimensional generalization of this observation.\footnote{As noted in \cite{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 \cite{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 \cite{singh2020machine} for a review.
As we move to nonparametric models, we focus on estimation which becomes much more challenging,
in the spirit of \cite{foster2019orthogonal}.
Still, we have provided qualitative characterization for uncertainty estimates in Section~\ref{sec:qb-regime}.
For the use of ML for nonlinear IV,
\cite{hartford2017deep} studied a heuristic application of NNs.
We discussed the minimax formulation in introduction.
\cite{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 \citep{wei2019regularization,ghorbani2019limitations}.
\cite{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.
\cite{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{sec:qb-regime}, \cite{chen2021adaptive} prove the validity of bootstrap-based selection for the sieve estimator \citep{newey2003instrumental}.
\cite{zhang_maximum_2020,hsu2019bayesian} investigate the use of marginal likelihood for two different kernel-based IV estimators:
\cite{hsu2019bayesian} establish a crude $-\nicefrac{1}{4}\log n$ upper bound for the log marginal likelihood,
and \cite{zhang_maximum_2020} connect it to the empirical leave-one-out validation error.
Neither result fully justifies model selection as our Corollary~\ref{lem:qb-mlh}.
For kernelized IV models,
\cite{muandet_dual_2020,singh_kernel_2020,zhang_instrument_2021} proposed
validation statistics for comparing {\em a finite number of} first stage models.
\begin{figure*}[bt]
\centering
\includegraphics[width=0.98\linewidth,clip,trim={0.25cm 0cm 0.35cm 0.25cm}]{figs/main.pdf}
\caption{Predictive performance: test MSE vs sample size $n_1=n_2$ for all method, and $D\in\{2,100\}$. Full results are in App.~\ref{sec:exp-main}.
}\label{fig:predictive-main}
\end{figure*}
\section{Simulation Study}\label{sec:experiments}
Our main simulation setup is adapted from \cite{bennett2019deep,dikkala_minimax_2020}; Appendix~\ref{app:exp-exo} presents additional experiment on the demand dataset \cite{hartford2017deep,xu_learning_2020}. In \cite{bennett2019deep,dikkala_minimax_2020}, the observed $\bz,\bx,\by$ are generated by
\begin{align*}
\bar \bz &\sim \mathrm{Unif}[-3, 3]^{\lfloor\frac{D}{2}\rfloor}, ~ \bz=h(\bar\bz), ~\bu \sim\cN(0, 1), ~
\bx := \bar \bz_1 + \bu + \be_x, ~
\by := ({f_0(\bx) + \bu + \be_y - \mu})/\sigma,
\end{align*}
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 \cite{lecun1998mnist} or CIFAR-10 \cite{krizhevsky2009learning} image with matching label; the MNIST setting also appeared in previous work.
\begin{table}[bt]\centering\small
\caption{Runtime results for all methods in the predictive experiment, for $N=2500,D=100$.}\label{tbl:runtime}
\begin{tabular}{ccccc}\toprule
Method & AGMM-Tree & AGMM-NN & AGMM-RBF & Proposed \\
\midrule
Runtime / s & $1374\pm 418$ & $303\pm 16$ & $6.7\pm 0.1$ & $25.9\pm 5.6$ \\
\bottomrule
\end{tabular}
\end{table}
\begin{table}[tb]\centering\small
\caption{
Test MSE, radius and estimated coverage rate of the $90\%$ $L_2$ credible ball (CB), and the average coverage of pointwise $90\%$ credible interval (CI), for $f_0\sim\mathcal{GP},D=100$.
For the CB coverage rate estimate, we report its $95\%$ Wilson score interval \citep{wilson27}.
Full results are in App.~\ref{app:exp-uq}.
}\label{tbl:uq-gprand-main}
\begin{tabular}{cccccc}\toprule
Method &
$n_1=n_2$ & Test MSE & 90\% CB.~Rad. & 90\% CB.~Cvg. & 90\% CI.~Cvg. \\
\midrule
\multirow{3}{*}{Proposed}
& $500$ & $.097$ {\tiny $\pm .065$} & $.201$ {\tiny $\pm .025$} & $.923$ {\tiny $[.888, .948]$} & $.915$ {\tiny $\pm .123$}\\% 0.9233 300
& $2500$ & $.035$ {\tiny $\pm .024$} & $.074$ {\tiny $\pm .008$} & $.917$ {\tiny $[.880, .943]$} & $.908$ {\tiny $\pm .127$}\\% 0.9167 300
& $5000$ & $.024$ {\tiny $\pm .016$} & $.049$ {\tiny $\pm .004$} & $.920$ {\tiny $[.884, .946]$} & $.905$ {\tiny $\pm .134$}\\% 0.9200 300
\midrule \multirow{3}{*}{RBF}
& $500$ & $.431$ {\tiny $\pm .192$} & $.240$ {\tiny $\pm .036$} & $.187$ {\tiny $[.147, .235]$} & $.640$ {\tiny $\pm .191$}\\% 0.1867 300
& $2500$ & $.176$ {\tiny $\pm .089$} & $.175$ {\tiny $\pm .023$} & $.517$ {\tiny $[.460, .573]$} & $.822$ {\tiny $\pm .136$}\\% 0.5167 300
& $5000$ & $.126$ {\tiny $\pm .072$} & $.156$ {\tiny $\pm .019$} & $.660$ {\tiny $[.605, .711]$} & $.855$ {\tiny $\pm .143$}\\% 0.6600 300
\bottomrule
\end{tabular}
\end{table}
We consider two choices for $f_0$: {\bf (i)} a widely used collection of functions (e.g., $\sin,\mathrm{abs}$) in \cite{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 \eqref{eq:s2-log-qlh}.
We set $N_1=N_2\in\{500,2500,5000\}$. We defer
setup details and full results to Appendix~\ref{app:simulations}, and summarize the findings below:
\paragraph{Hyperparameter Selection (App.~\ref{sec:exp-s1-randgp})} 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 \eqref{eq:s1-val-stats} 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{sec:exp-main})}
For $h$ defined as in (i-ii), we compare our algorithm with \cite[AGMM]{dikkala_minimax_2020}, instantiated with kernel, tree and NN models. As shown in \cite{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{fig:predictive-main}: 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.
\cref{tbl:runtime}
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 \cite{zhang_maximum_2020}, which report the best results in the MNIST setting. Our method outperforms both baselines.
\paragraph{Uncertainty Quantification (App.~\ref{app:exp-uq})} Table~\ref{tbl:uq-gprand-main} 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 \cite{bennett2019deep}, we experiment with BMA over a grid of RBF kernels, and present visualizations in Appendix~\ref{app:exp-uq}. 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{app:exp-exo})} We evaluate the extended algorithm in \cref{app:exo-algo} on the demand dataset \citep{hartford2017deep}, which is a widely used simulation design with high-dimensional exogenous covariates. As shown in the appendix, our extended algorithm has competitive performance.
\section*{Acknowledgements}
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.
\bibliographystyle{IEEEtranN}
\bibliography{bib}
\newpage