0.0.1 Introduction↩︎

Machine learning, statistics, and causal inference rely on risk minimization problems of the form \[\begin{align} \min_{\theta \in \Theta} \: \big[ L_0(\theta) := {\mathbb{E}}_{Z \sim \mathbb{P}}\left[\ell_0(\theta; Z)\right] \big], \label{eq:erm} \end{align}\tag{1}\] where \(\Theta \subseteq\mathbb{R}^d\) is a parameter space, \(Z\) is a \(\mathcal{Z}\)-valued random variable, and \(\ell_0: \Theta \times \mathcal{Z}\rightarrow \mathbb{R}\) is a loss function. The quantity \(\ell_0(\theta; z)\) describes the performance of a model parametrized by \(\theta \in \Theta\) on a test example \(z \in \mathcal{Z}\). Given only an oracle that provides a stochastic gradient estimate of the objective 1 , practitioners are able to train models ranging from linear functions on tabular data to billion-parameter neural networks on vision and language data.

The success of stochastic gradient descent (SGD) algorithms [1]–[4] has motivated an abundance of work on their theoretical properties under various algorithmic and risk conditions, such as class separability [5], random reshuffling [6], decomposable objectives [7], [8], quantization noise [9], and noise dominance [10]. This success has been fueled by machine learning and AI software libraries such as JAX, PyTorch, TensorFlow, and others, which offer a wide range of SGD variants, as long as a loss function can be clearly specified. The gradient is then evaluated automatically on a mini-batch of datapoints and used for stochastic updates.

Though powerful, this recipe takes one thing for granted: that the learner can always compute the risk (or an unbiased estimate thereof). Indeed, many complex learning problems rely on a risk function that is only partially specified up to a class \[\begin{align} \mathcal{L}:= \left\{L(\,\cdot\,, g): g \in \mathcal{G}\right\},\label{eq:loss95class} \end{align}\tag{2}\] where \(\mathcal{G}\) is a possibly infinite-dimensional set and \(L: \Theta \times \mathcal{G}\rightarrow \mathbb{R}\) is a function of both the target parameter \(\theta \in \Theta\) and an unknown nuisance parameter \(g \in \mathcal{G}\).

This framework originates from semiparametric inference [11]–[14], wherein the risk is a Kullback-Leibler (KL) divergence and \(g\) provides information about the true data-generating distribution \(\mathbb{P}\), but is not of primary scientific interest. While the partially specified loss framework from 2 originates from this specific literature, it is not limited to semiparametric inference problems, and connects to many areas of interest. For instance, profile likelihood methods [15]–[17] reduce 2 by applying a pointwise minimum over \(g \in \mathcal{G}\) to then construct a problem that can be solved in \(\theta \in \Theta\). Another example arises in applications with distribution shifts, for which \(g\) represents an unknown test data distribution that may differ from the one from which the training data were drawn. Distributionally robust optimization [18]–[20] reduces 2 by instead taking a pointwise maximum over \(\mathcal{G}\) and solving the resulting saddle-point problem. While the pointwise minimum and maximum are natural reductions, it is often the case that the application admits a “true” \(g_0 \in \mathcal{G}\) and the loss class is reduced by first estimating \(g_0\) with auxiliary data to produce some \(\hat{g}\), which we refer to as double/debiased machine learning, or DML, following [21]. The problem 1 is then thought to be derived via \(L_0(\theta) \equiv L(\theta, g_0)\) in this case (see examples in 0.0.2). This is the focus of this paper.

Despite the prominence of SGD and DML individually, the convergence guarantees of SGD for recovering the risk minimizer with a misspecified nuisance parameter remain unknown. Indeed, after producing \(\hat{g}\), the user typically solves a (full batch) empirical risk minimization problem, i.e. minimizing a sample average approximation of \(L(\cdot, \hat{g})\). In this paper, we aim to fill this gap by proving convergence guarantees on the sequence \((\theta^{(n)})_{n\geq 1}\) generated by updates of the form \[\begin{align} \theta^{(n)} = \theta^{(n-1)} - \eta S(\theta^{(n-1)}, \hat{g};Z_n),\label{eq:sgd} \end{align}\tag{3}\] where \(\eta > 0\) is a learning rate, \(\mathcal{D}_n := (Z_i)_{i=1}^n\) is a stream of independent data drawn from \(\mathbb{P}\), \(\hat{g}\) is a nuisance parameter estimate, and \(S: \Theta \times \mathcal{G}\times \mathcal{Z}\rightarrow \mathbb{R}^d\) is a stochastic gradient oracle satisfying of \(\mathbb{E}_{Z \sim \mathbb{P}}[S(\theta, g; Z)] = \nabla_\theta L(\theta, g)\) for all \((\theta, g) \in \Theta \times \mathcal{G}\). In particular, when \(\mathcal{G}\) lies within a Banach space equipped with norm \(\lVert \cdot \rVert_\mathcal{G}\), we wish to compare \(\theta^{(n)}\) to \[\begin{align} \label{minimization32problem} \theta_\star = \mathop{\mathrm{arg\,min}}_{\theta \in \Theta} L(\theta, g_0), \end{align}\tag{4}\] given conditions on the degree of misspecification \(\lVert \hat{g}- g_0 \rVert_{\mathcal{G}}\) and (approximate) Neyman orthogonality of the risk \(L\) [22].

Intuitively, Neyman orthogonal classes of objectives are instances of 2 whose curvature with respect to \(\theta\) is insensitive to the choice of \(g\) (see 0.0.2 for the formal description). When Neyman orthogonality is satisfied, the double machine learning framework is also known as orthogonal statistical learning (OSL) [23], [24]. Other than the obvious computational considerations, we argue that the SGD perspective in this paper also sheds light on methodological opportunities in DML/OSL. Indeed, while loss functions are typically specified by the chosen architecture, Neyman orthogonality is often achieved by specialized analytic calculations on the part of the user. While this property is generally seen as a second-order property of the loss, it can also be viewed as a first-order property of the gradient oracle \(S\). As we detail in 0.0.3, it may be easier, and more aligned with the spirit of modern machine learning, to craft Neyman orthogonal gradient oracles instead of losses.

0.0.1.0.0.1 Contributions.

We prove the first theoretical convergence guarantees for SGD under an unknown nuisance model. We find that \(\theta^{(n)}\) converges linearly to a neighborhood of \(\theta_\star\)—the optimum in the well-specified case—with a radius that has a fourth-power (resp. squared) dependence on \(\Vert \hat{g} - g_0 \Vert_\mathcal{G}\) when Neyman orthogonality is (resp. is not) satisfied. Our analysis can also apply to two-stream settings in which the nuisance parameter is learned online alongside the target. We further analyze a new algorithm, called orthogonalized SGD (OSGD), wherein the gradient oracle of a possibly non-orthogonal loss can be iteratively made orthogonal using an “approximately orthogonalized” gradient oracle, which is based on a separate estimation procedure. This algorithm enjoys a convergence guarantee that interpolates between the \(\Vert \hat{g} - g_0 \Vert_\mathcal{G}^4\) (nuisance insensitive) and \(\Vert \hat{g} - g_0 \Vert_\mathcal{G}^2\) (nuisance sensitive) regimes depending on the quality of the orthogonalizing operator.

We provide an introduction to the OSL/DML setting in 0.0.2. Our SGD and OSGD procedures are analyzed in 0.0.3. Related/future work and concluding remarks are given in 0.0.4. All proofs and numerical illustrations can be found in the Appendix.

0.0.2 Orthogonal Statistical Learning↩︎

We first introduce various examples of risk functions in the form of 2 , then formally introduce Neyman orthogonality and its implications. As is common in learning settings, the risk will be in the form of an expectation, \[\begin{align} L(\theta, g) = {\mathbb{E}}_{Z \sim \mathbb{P}}\left[\ell(\theta, g; Z)\right], \end{align}\] where \(\ell: \Theta \times \mathcal{G}\times \mathcal{Z}\rightarrow \mathbb{R}\) is an instance-level loss function. Various assumptions used in the analysis in 0.0.3 (e.g. convexity) may be placed on either the loss \(\ell\) or the risk \(L\). In each example, we provide the structure of the data point \(Z\), the set \(\mathcal{G}\), and the loss \(\ell\), and the true \(g_0 \in \mathcal{G}\) to fully specify the problem. Here, we interpret “true” to mean that \(g_0\) is a parameter of the data-generating distribution (e.g. a propensity score in causal inference), or that \(g_0\) satisfies a cost-minimizing or utility-maximizing criterion (as in the profile likelihood or saddle-point examples from 0.0.1).

The next example concerns a quantity widely studied in causal inference [26].

Finally, we maintain the data structure from the previous example, but consider a loss corresponding to a different target parameter according to [28].

While the choices of the loss function in Examples 2 and 3 might look opaque to readers outside of causal inference and statistics, they are carefully designed to be Neyman orthogonal. To motivate its definition, notice that, invariably, \(g_0\) is unknown to the user. In DML, the user may produce or access some \(\hat{g}\in \mathcal{G}\), which is an estimate of \(g_0\) based on independent training data other than the stream \((Z_i)_{i=1}^n\) used to produce \(\theta^{(n)}\). It is of clear interest how stochastic optimization algorithms (and their resulting minimizers) behave in light of the misspecification of \(g_0\), and what precise theoretical conditions govern this behavior. Moreover, as we demonstrate in 0.0.3, these same conditions can be used to analyze procedures for which the user may access additional data to progressively improve the estimate \(\hat{g}\) and learn \(\theta_\star\) simultaneously. We now formally introduce Neyman orthogonality, and by extension, the orthogonal statistical learning (OSL) variant of DML.

0.0.2.0.0.1 Neyman Orthogonality.

For a definition that accounts for a possibly infinite-dimensional function class \(\mathcal{G}\), we introduce the directional derivative, or equivalently, the derivative operator.

We denote by \(S_\theta(\theta, g;z) = \nabla_\theta \ell(\theta, g;z)\) the gradient of the loss function w.r.t. the target parameter \(\theta \in \Theta\). Borrowing terminology from statistics, we call this the score, whether \(\ell\) is based on a likelihood or not.¹ This constitutes one particular example of a stochastic gradient oracle \(S\) used in 3 . Overloading notation, the population gradient oracle is defined as \(S_\theta(\theta, g) = \mathbb{E}_{Z\sim\mathbb{P}}[\nabla_\theta \ell(\theta, g;Z)]\).

In Definition 2, we allow \(\Theta' \times \mathcal{G}' \subseteq \Theta \times \mathcal{G}\) to be a proper subset, which not only provides a weaker condition, but also accounts for localization-style arguments. Moreover, since \(\mathrm{D}_\theta L(\theta_\star, g_0)[\theta - \theta_\star] = {\langle S_\theta(\theta_\star, g_0), \theta - \theta_\star \rangle}\), if the population risk \(L\) satisfies 5 , then 6 holds for any target parameter class \(\Theta' \subseteq \mathbb{R}^d\). As mentioned above, the risk functions in Examples 1–3 are all Neyman orthogonal at their respective value of \((\theta_\star, g_0)\). In the next section, we will discuss a procedure to make a non-orthogonal gradient oracle “approximately” orthogonal. We illustrate the intended outcome intuitively in 1 by comparing a generic loss and its orthogonalized counterpart.

0.0.3 Stochastic Gradient Optimization↩︎

In this section, we propose two stochastic gradient algorithms, which rely on different choices of the stochastic gradient oracle \(S\) used in 3 . The first is the familiar stochastic gradient oracle that provides a sample estimate of the gradient \(\nabla L(\cdot, \hat{g})\) for a fixed estimate \(\hat{g}\). The second employs an approximately orthogonalized gradient oracle, or OSGD oracle, to achieve a notion of approximate Neyman orthogonality (in a manner we make precise in this section). We analyze the first algorithm under both non-orthogonal and orthogonal settings, achieving an illustrative breakdown of “nuisance sensitive” and “nuisance insensitive” regimes regarding the theoretical convergence guarantee. For the OSGD algorithm, we prove a convergence guarantee that interpolates between the two regimes, depending on the quality of the oracle. We then give context to our results with respect to several lines of complementary work.

0.0.3.0.0.1 Notation and Assumptions.

For readers’ convenience, a table of all the notation we introduce throughout the paper is collected in 1.0.1. We maintain the prototypical bias/variance conditions on the stochastic gradient oracle \(S\), that is, an unbiasedness condition and a second-moment growth condition (see 1(d)). To dispel confusion, note that by “unbiased”, we mean specifically that \({\color{royalazure}{\mathbb{E}}_{Z \sim \mathbb{P}}\left[S(\theta, g; Z)\right] = \nabla_\theta L(\theta, g)}\) for all \((\theta, g)\), as opposed to the “bias” of replacing \(\hat{g}\) with \(g_0\), a terminology sometimes used in DML/OSL. Our analysis will rely partly on the initial distance \(r\) of the nuisance estimate \(\hat{g}\) to \(g_0\) in \(\mathcal{G}\), which defines the ball \[\begin{align} \mathcal{G}_{r}\left(g_0\right)= \left\{g \in \mathcal{G}: \Vert g - g_0\Vert_{\mathcal{G}} \leq r\right\}. \end{align}\] Various assumptions on the risk will be required to hold only locally, that is, within \(\mathcal{G}_{r}\left(g_0\right)\) as opposed to the entire linear space \(\mathcal{G}\). Thus, the assumptions become weaker as the estimate \(\hat{g}\) improves.

1 does not require Neyman orthogonality at \((\theta_\star, g_0)\). Instead, 1(a) is a standard differentiability condition. 1(b) and (c) implies that \(\theta_\star\) is a unique global minimizer. 1(d) generalizes the uniformly bounded second moment condition in stochastic optimization (e.g. [30]) by adding a quadratic form in \(\theta\), which allows us to consider an unbounded feasible set \(\Theta\) and more loss classes. Finally, 1(d) and (e) can be satisfied when the Hessian of the population risk is a bounded operator. Usually, \(K_1\), \(\kappa_1\), and \(\alpha_1\) would depend on the initial nuisance estimation distance \(r\). We compute the constants appearing in the assumption for our chosen examples in 1.0.2. We proceed to the main results regarding the convergence of SGD and OSGD, respectively.

0.0.3.0.0.2 Stochastic Gradient Algorithm.

Here, we use the standard single-sample stochastic gradient estimate \(S = S_\theta\) in 3 . This leads to the update \[\begin{align} \label{baseline32SGD32procedure} \theta^{(n)} = \theta^{(n-1)} - \eta S_\theta(\theta^{(n-1)}, \hat{g}; Z_n), \quad \theta^{(0)} \in \Theta. \end{align}\tag{7}\] While the SGD procedure can be easily extended to using a batch of unbiased gradient estimates, we keep our single-observation construction to highlight the most important aspects of the analysis. In order to achieve quantitative guarantees in the Neyman orthogonal setting, which essentially removes certain second-order terms that include \(\theta\) and \(g\), we will consider the following higher-order condition in some cases.

When satisfied, 2 results in the nuisance insensitivity alluded to at the beginning of this section. Notice that Neyman orthogonality is not necessary to construct a stochastic optimizer, and it is still possible to obtain a nuisance sensitive rate under only 1. We demonstrate this in 1.

Note that the assumption that the iterates remain in \(\Theta\) is satisfied in common cases. It is satisfied trivially for the first two examples in 0.0.2 because \(\Theta = \mathbb{R}^d\). Another case is when the loss decomposes into the sum of a \(G\)-Lipschitz continuous component and the \(\ell_2\)-norm regularizer, i.e. \(\ell(\theta, g; z) = h(\theta, g; z) + \frac{\mu}{2}\lVert \theta \rVert_2^2\). In this case, the iterates and the optimum remain in the ball \(\left\{\theta: \lVert \theta \rVert_2 \leq G/\mu\right\}\) (see, e.g., [31]), so 2 can be restricted to this compact set.

1 states that SGD converges linearly to a ball around \(\theta_\star\) with a radius that depends on the bias (due to the replacement of \(g_0\) with \(\hat{g}\)) and the variance due to gradient noise. Moreover, the variance component decays proportionally to the learning rate \(\eta\). Under 2, the bias component can have a significantly more favorable scaling with the error in the nuisance estimate \(\Vert\hat{g} - g_0\Vert_\mathcal{G}\)—specifically, \(\Vert\hat{g} - g_0\Vert_\mathcal{G}^4\) instead of \(\Vert\hat{g} - g_0\Vert_\mathcal{G}^2\). A similar breakdown into two regimes of the bias scaling occurs in the works of both [24] and [32] under 2 (called “slow rate” and “fast rate” there). Importantly, their bounds are based on an exact, offline empirical risk minimization procedure for a fixed training set, i.e. they provide excess risk bounds on the quantity \(L(\hat{\theta}_n, g_0) - L(\theta_\star, g_0)\), where \[\begin{align} \hat{\theta}_n = \mathop{\mathrm{arg\,min}}_{\theta \in \Theta} \frac{1}{n}\sum_{i=1}^n \ell(\theta, \hat{g}; Z_i). \end{align}\] In contrast, 1 accounts for both the expected distance to the optimum and the interplay between the bias incurred by \(\hat{g}\) and the progress bound on each step of the algorithm. In particular, even when using a constant learning rate, the bias does not accrue on each iterate and is in fact constant in \(n\).

0.0.3.0.0.3 Orthogonalized Stochastic Gradient Algorithm.

Given the marked improvement in the rate of decay of the bias term when an orthogonal loss is used, it is clearly beneficial to do so when possible. We now describe how the user can induce orthogonality by adjusting the stochastic gradient oracle through a secondary estimation problem. The construction of Neyman orthogonal losses has historically been motivated as a means to build efficient (or minimum asymptotic variance) estimators in semiparametric theory and statistical learning [14], [24], [33], [34]. However, the actual implementation of this technique occurs analytically at the stage of defining the loss, on a case-by-case basis. To our knowledge, orthogonalizing an existing oracle via computational methods is a relatively new area (see [35] and references therein). Our approach is inspired by the construction reviewed in [21]. While orthogonalization is usually described using the language of differential/information geometry, we offer an intuitive explanation based on least-squares regression.

While our construction holds in general dimensions, let us first consider the illustrative case when \(\mathcal{G}= \mathbb{R}^k\). At the true parameters \((\theta_\star, g_0)\), consider the problem of finding the best predictor of the \(\mathbb{R}^d\)-valued target variable \(S_\theta(\theta_\star, g_0; Z) = \nabla_\theta \ell(\theta_\star, g_0; Z)\) given the \(\mathbb{R}^k\)-valued predictor \(\nabla_g \ell(\theta_\star, g_0; Z)\) variable, specified by the program \[\begin{align} \Gamma_0 = \mathop{\mathrm{arg\,min}}_{\Gamma\in \mathbb{R}^{d \times k}} {\mathbb{E}}_{\mathbb{P}}\left[\lVert S_\theta(\theta_\star, g_0; Z) - \Gamma\nabla_g \ell(\theta_\star, g_0; Z) \rVert_2^2\right]. \label{eq:regression} \end{align}\tag{8}\] In the special case that \(\ell(\theta, g; z) = -\log p_{\theta,g}(z)\) for a density \(p_{\theta, g}\) on \(\mathcal{Z}\) that governs the random variable \(Z\), the projection direction solving 8 can be shown to satisfy \(\Gamma_0 = H_{\theta g}^\top H_{g g}^{-1}\), where \(H_{\theta g} = \nabla_gS_\theta(\theta_\star, g_0) \in \mathbb{R}^{k \times d}\) is the transposed Jacobian and \(H_{gg} = \nabla_g^2 L(\theta_\star, g_0) \in \mathbb{R}^{k \times k}\) is the Hessian. The prediction \(\Gamma_0 \nabla_g \ell(\theta_\star, g_0; Z)\) accounts for all correlative information between \(\nabla_\theta \ell(\theta_\star, g_0; Z)\) (the gradient w.r.t. \(\theta\)) and \(\nabla_g \ell(\theta_\star, g_0; Z)\) (the gradient w.r.t. \(g\)). It stands to reason that as \(\theta \rightarrow \theta_\star\), the random vector \[\begin{align} S(\theta, g_0; Z) &:= S_\theta(\theta, g_0; Z) - \Gamma_0 \nabla_g \ell(\theta, g_0; Z) =S_\theta(\theta, g_0; Z) - H_{\theta g}^\top H_{g g}^{-1} \nabla_g \ell(\theta, g_0; Z) \label{eq:no95score95finite} \end{align}\tag{9}\] would be less sensitive to perturbations of \(g_0\), as the component of \(S_\theta(\theta, g_0; Z)\) that is predictable through changes in \(g_0\) is subtracted out. Furthermore, if we are aware that the expectation of \(S\) is made zero at \(\theta_\star\), then a stochastic gradient scheme based on 9 could conceivably achieve a nuisance insensitive rate guarantee in lieu of 1. From the variance reduction viewpoint, the correction term in 9 subtracts the regression of the \(\theta\) gradient of the loss on the \(g\) “gradient” of the loss. By the law of total variance, the variance of the gradient reduces and improves the trajectory of stochastic optimization. Furthermore, this variational description 9 hints at how such an operator can be computed algorithmically, instead of the historical semiparametric inference approach of deriving the operator via calculation by hand.

Supported by this illustration, we define a generalization that will provide an alternative stochastic gradient oracle to use for optimization purposes. Without assuming that \(\ell\) is a negative log-likelihood, we generalize the formulas for \(\nabla_g \ell(\theta, g; z) \in \mathbb{R}^k\), \(H_{\theta g} \in \mathbb{R}^{k \times d}\) and \(H_{g g} \in \mathbb{R}^{k \times k}\) for when \(\mathcal{G}\equiv (\mathcal{G}, {\langle \cdot, \cdot \rangle}_\mathcal{G})\) is an infinite-dimensional Hilbert space. Under regularity conditions on the directional derivatives of \(L\), we have that \(\nabla_g \ell(\theta, g; z) \in \mathcal{G}\) for all \(z \in \mathcal{Z}\), \(H_{\theta g} = (H_{\theta g}^{\scriptscriptstyle(1)}, \ldots, H_{\theta g}^{\scriptscriptstyle(d)}) \in \mathcal{G}^d\), and \(H_{gg}: \mathcal{G}\rightarrow \mathcal{G}\) is a bounded and self-adjoint operator. The formal details of their construction are contained in 1.0.4. Just as in 9 , we may consider the operator \(\Gamma_0: \mathcal{G}\rightarrow \mathbb{R}^d\), defined element-wise by \([\Gamma_0 g]_j = {\langle H_{\theta g}^{\scriptscriptstyle(j)},H_{gg}^{-1} g \rangle}_\mathcal{G}\), where the invertibility of \(H_{gg}\) is satisfied by our assumptions preceding 2. As shown in 8 , the orthogonalizing \(\Gamma_0\) is defined by both the true nuisance \(g_0\) and the target \(\theta_\star\), where \(g_0\) can usually be learned as some conditional expectation and \(\theta_\star\) can be learned by our proposed methods. We then construct the central object of the upcoming 2: the Neyman orthogonalized (NO) gradient oracle \[\begin{align} \label{no32score32of32infinite32dim} S_{\text{no}}(\theta, g; z) = S_\theta(\theta, g; z) - \Gamma_0 \nabla_g \ell(\theta, g; z). \end{align}\tag{10}\]

We refer readers to 11 for the proof. In this context, we refer to the operator \(\Gamma_0\) as the “orthogonalizing operator”. As a natural sanity check, we note that for a risk function that is already Neyman orthogonal at \((\theta_\star, g_0)\), the NO score \(S_{\text{no}}\) is exactly equal to score function \(S_\theta\) itself since \(\Gamma_0 = 0\). To construct \(S_{\text{no}}\) for the non-orthogonal loss, we provide the following example in Partially linear model where the corresponding derivations of \(\Gamma_0\) and \(S_{\text{no}}\) are included in 1.0.2.1.2.

Motivated by the advantage of a Neyman orthogonal score, we now construct our OSGD algorithm using an estimated the NO score \(S_{\text{no}}\). While \(\Gamma_0\) (like \(g_0\)) is unknown to the user in general, using an arbitrary estimate \(\hat{\Gamma}\), we can define the estimated NO score \(\hat{S}_{\text{no}}\) oracle via \[\begin{align} \hat{S}_{\text{no}}(\theta, g; z) = S_\theta(\theta, g; z) - \hat{\Gamma}\nabla_g \ell(\theta, g; z).\label{eq:noscorehar} \end{align}\tag{11}\] Usually, one can obtain such an estimate \(\hat{\Gamma}\) using the same data stream of \(\hat{g}\); we discuss possible strategies in 1.0.6.3. Finally, using \(\hat{S}_{\text{no}}\) as the stochastic gradient oracle \(S\) in 3 , we derive the OSGD update \[\begin{align} \theta^{(n)} = \theta^{(n-1)} - \eta \hat{S}_{\text{no}}(\theta^{(n-1)}, \hat{g}; Z_n), \quad \theta^{(0)} \in \Theta.\label{eq:dsgd} \end{align}\tag{12}\]

To measure the quality of \(\hat{\Gamma}\) in our analysis, we use the Frobenius norm \(\lVert \Gamma \rVert^2_{\mathrm{Fro}} = \sum_{j=1}^d \lVert \Gamma^{(j)} \rVert^2_{\mathrm{op}}\) where \(\Gamma: \mathcal{G}\rightarrow \mathbb{R}^d, \;\Gamma^{(j)}: g \mapsto [\Gamma g]_j\) and \(\lVert \cdot \rVert_{\mathrm{op}}\) denotes the usual operator norm for linear functionals. As an example, by the uniqueness of Riesz representations, \(\lVert \Gamma_0 \rVert^2_{\mathrm{Fro}} = \sum_{j=1}^d \lVert H_{gg}^{-1} H_{\theta g}^{\scriptscriptstyle(j)} \rVert_\mathcal{G}^2\).

3(a) is necessary for the convergence of the OSGD to \(\theta_\star\). When \(S_\theta\) is Neyman orthogonal at \((\theta_\star, g_0)\), \(\Gamma_0 = 0\) is accessible and thus, \(S_{\text{no}}= S_\theta\). When \(S_\theta\) is non-orthogonal, 3(a) can be satisfied whenever \(\nabla_g L(\theta_\star, g_0) = 0\), implying that \((\theta_\star, g_0)\) is a local minimizer of \(L(\theta, g)\). 3(b) is related to the Schur complement of the population Hessian. Thus, the hypothetical objective relating to \(S_{\text{no}}\) inherits its strong convexity from that of the population risk \(L\) w.r.t. \((\theta, g) \in \Theta \times \mathcal{G}_r(g_0)\) when \(\mathcal{G}\) is finite-dimensional; see [36]. 3(c) and (d) are exactly analogous to 1(d) and (e), while 3(e) is analogous to 2.

With 3, we now establish the convergence result of OSGD.

Compared with 1, Theorem 2 shows that OSGD can outperform the nuisance sensitive rate through the correction term \(\Vert\hat{g} - g_0\Vert_\mathcal{G}^2 \cdot \Vert\hat{\Gamma}- \Gamma_0\Vert_{\mathrm{Fro}}^2\), and can align with the nuisance insensitive rate when \(\Vert\hat{\Gamma}- \Gamma_0\Vert_{\mathrm{Fro}}\) is of the order \(\mathcal{O}(\Vert\hat{g} - g_0\Vert_\mathcal{G})\). Depending on different assumptions, there is another approach that can be used to simplify the statement of 2—see 1.0.5 for details.

0.0.3.0.0.4 Interleaving Target and Nuisance Estimation.

The results seen thus far have considered the estimate \(\hat{g}\) to be a fixed element of \(\mathcal{G}\), and included terms that depend on the discrepancy \(\lVert \hat{g}- g_0 \rVert_\mathcal{G}\). Part of the convenience of these results is that if \(\hat{g}\equiv \hat{g}^{(m)}\) is the result of a learning procedure with \(m\) independent data points, then statistical bounds on \(\lVert \hat{g}^{(m)} - g_0 \rVert_\mathcal{G}\) (either in expectation or high probability, depending on the situation) can be plugged in to quantify the bias. While the results naturally account for batch learning procedures, they are also amenable to analyzing staggered procedures in which two data sources are queried to estimate \(\theta_\star\) and \(g_0\), respectively. To our knowledge, this is the first theoretical analysis of such an orthogonal statistical learning problem.

To be precise, suppose that we update the nuisance estimator for \(m\) times, leading to the sequence \(\hat{g}^{(1)}, \dots, \hat{g}^{(m)}\) on a stream of \(\mathcal{W}\)-valued data \(W_1, \ldots, W_m\), sampled i.i.d. from a probability measure \(\mathbb{Q}\). We define \(\theta^{(0, n)} = \theta^{(0)} \in \Theta\), and for the update of \(\hat{g}^{(i)}\) for \(1 \leq i \leq m\), we define \(\theta^{(i, 0)} = \theta^{(i-1, n)}\) and produce the sequence \(\theta^{(i, 1)}, \ldots, \theta^{(i, n)}\) using \(n\) steps of the SGD update 7 initialized at \(\theta^{(i, 0)}\). Consider, for example, the case in which \(\mathcal{G}\) is a reproducing kernel Hilbert space (RKHS) with kernel \(k(\cdot,\cdot)\). With the assumption that the eigenvalues \((\lambda_j)_{j\geq1}\) of covariance operator \(\mathbb{E}_{\mathbb{Q}}[k(W,\cdot) \otimes k(W,\cdot)]\) decay polynomially at order \(j^{-\alpha}\), the nonparametric stochastic approximation of [37] satisfies \({\mathbb{E}}_{\mathbb{Q}^m}\left[\Vert\hat{g}^{(m)} - g_0\Vert_{\mathcal{G}}^2\right] = \mathcal{O}(m^{-(2\alpha - 1)/(2\alpha)})\). This leads to the following nuisance sensitive rate for a non-Neyman orthogonal loss, by 3: \[\begin{align} \mathbb{E}_{\mathbb{P}^{mn}\otimes \mathbb{Q}^m}[\Vert\theta^{(m,n)} - \theta_\star\Vert_2^2] = \mathcal{O}\left(\left(1-\mu\eta/2\right)^{mn} + m^{-\frac{2\alpha - 1}{2\alpha}} + n^{-1} + \eta \right). \end{align}\] As another example, suppose that in addition, we can estimate \(\hat{\Gamma}\equiv \hat{\Gamma}^{(m)}\) using the stochastic gradient method of [37] and using the same data stream \((W_1, \dots, W_m)\). If there are high probability bounds for \(\Vert\hat{g}^{(m)} - g_0\Vert_{\mathcal{G}}^2\) and \(\lVert \hat{\Gamma}^{(m)} - \Gamma_0 \rVert_{\mathrm{Fro}}^2\) of the same order as \(\mathcal{O}(m^{-(2\alpha - 1)/(2\alpha)})\) and \(\Vert\theta^{(m,n)} - \theta_\star\Vert_2^2\) decays as described in 2, then we have in 4 that \(\Vert\theta^{(m,n)} - \theta_\star\Vert_2^2 = \mathcal{O}_p\left(\left(1-\mu\eta/2\right)^{mn} + m^{-(2\alpha - 1)/\alpha} + n^{-1} + \eta\right)\) where the \(O_p(m^{-(2\alpha - 1)/\alpha})\) nuisance bias term decays quadratically faster than the one for a non-Neyman orthogonal loss. We refer the reader to 1.0.6.3 for details on the analysis with multiple data streams.

0.0.4 Discussion↩︎

From an optimization perspective, it is helpful to know how our convergence bounds perform in the idealized case of a known nuisance, which is equivalent to 1 . In this case, 1 gives the convergence rate \(\mathbb{E}_{\mathcal{D}_n \sim \mathbb{P}^n}[\Vert \theta^{(n)} - \theta_\star\Vert_2^2] = \mathcal{O}(\left(1 - \mu\eta/2\right)^n + \eta)\), which aligns with the non-asymptotic SGD convergence rates, in mean-square error [38] and in high-probability [30]. While our result requires a smaller learning rate \(\eta < \mu/2(M\mu+ 2\kappa_1)\) compared with the requirement \(\eta < 1/(2M)\) of [30], this is entirely due to the our more general bounded moment assumption (see 1(d)), as opposed to a uniform boundedness assumption over all \(\Theta \times \mathcal{G}_{r}\left(g_0\right)\). In addition, when the uniform moment bound holds true, \(\kappa_1\) becomes zero, and our learning rate requirement returns to \(\eta < 1/(2M)\).

One out-of-the-box approach for analyzing SGD performance would be to use existing results for biased SGD methods. Here, the “bias” refers to the fact that \({\mathbb{E}}_{Z \sim \mathbb{P}}\left[S(\theta, \hat{g}; Z)\right] \neq \nabla_\theta L(\theta, g_0)\) in general, and the convergence radius depends on the average value of \(\lVert {\mathbb{E}}_{Z \sim \mathbb{P}}\left[S(\theta^{(n)}, \hat{g}; Z)\right] - \nabla_\theta L(\theta^{(n)}, g_0) \rVert_2^2\). Results along this line of work result in a radius that may not scale with \(\eta\) (see [39], for instance). While this form of bias may be related to \(\lVert \hat{g}- g_0 \rVert_\mathcal{G}\) under Lipschitzness conditions on the oracle, it is unclear how to effectively incorporate Neyman orthogonality into these general-purpose approaches.

In the general case of an unknown nuisance, [24], [34] consider learning methods based on analytically crafted Neyman orthogonal risk functions in various scenarios. For regression functionals, the statistical learning technique of [40] can ensure that the bias term \(\lVert \hat{g}- g_0 \rVert_\mathcal{G}^2\) is asymptotically negligible for large samples, in the sense that classical statistical confidence sets for \(\theta_\star\) are asymptotically valid. Our work is complementary to these, providing non-asymptotic guarantees for stochastic optimization. Our estimated OSGD oracle \(\hat{S}_{\text{no}}\) further offers a flexible solution to deal with general risk functions. When deriving an orthogonalized risk by hand is difficult or impossible, our proposed construction still applies, and 2 demonstrates that the convergence using \(\hat{S}_{\text{no}}\) would be improved to a rate of \[\begin{align} \mathcal{O}\Big(\left(1 - \frac{\mu_{\text{no}}\eta}{2}\right)^n + \underbrace{\Vert\hat{g} - g_0\Vert_\mathcal{G}^4 + \Vert\hat{g} - g_0\Vert_\mathcal{G}^2 \cdot \lVert \hat{\Gamma}- \Gamma_0 \rVert_{\mathrm{Fro}}^2}_{\text{improvement over \Vert\hat{g} - g_0\Vert_\mathcal{G}^2}} + \eta\Big). \end{align}\] When we have the true orthogonalizing \(\Gamma_0\), the improved rate recovers the nuisance insensitive one from 1. Besides, when \(\hat{g}\) converges but \(\Vert\hat{\Gamma}- \Gamma_0\Vert_\mathcal{G}^2 = \mathcal{O}_p(1)\), the improved rate resembles the nuisance sensitive rate of 1, plus a \(\mathcal{O}(\eta)\) bias term. Thus, the quality of the estimated orthogonalizing operator governs how the optimization interpolates between these two rates.

Having understood the performance of SGD when using an estimated orthogonalizing operator, a central question is how to compute or approximate such an operator. In 0.0.3, we used a correction based on least-squares regression as an intuitive example, reminiscent of score swindle methods for variance-reduced estimation [41]. However, solving a gradient-to-gradient regression problem is by no means the only way to derive an orthogonalizing operator; [35] recently demonstrated, from a computational viewpoint, that an orthogonalizing operator can be derived using algorithmic/reverse mode differentiation in many cases. We emphasize that none of the aforementioned works analyzes the convergence of stochastic gradient under nuisance. Our work addresses this important problem. Moreover, error guarantees on these (and other) orthogonalized gradient oracles can potentially be plugged in into our stochastic gradient method analysis for global statistical guarantees.

0.0.4.0.0.1 Acknowledgements.

The authors would like to thank L. Liu for valuable comments and suggestions. This work was supported by NSF DMS-2023166, CCF-2019844, DMS-2134012, DMS-2210216, NIH, and IARPA 2022-22072200003. Part of this work was performed while R. Mehta and Z. Harchaoui were visiting the Simons Institute for the Theory of Computing, and A. Luedtke was visiting the Institute of Statistical Mathematics.

1 Appendix↩︎

1.0.1 Notation↩︎

max width=1.0

1.0.2 Detailed Examples↩︎

In this section, we describe in detail how the three examples in 0.0.2 from the main text satisfy 1 and 2. We first talk about the partially linear model (PLM) in 1.0.2.1, and then introduce the conditional averaged treatment effect (CATE) based on the potential outcomes framework in 1.0.2.1. Under the same framework, finally we talk about the conditional relative risk (CRR) in 1.0.2.3. In addition, we also study a non-orthogonal loss usually used for PLM in 1.0.2.1.2 and an unrestricted loss function for CATE in 1.0.2.2.1. The constants for all examples are concluded in [tab:examples] and proofs of lemmas in this section are provided in 1.0.2.4.

max width=0.9

1.0.2.1 Partially Linear Model

1.0.2.1.1 Orthogonal Loss

We revisit Example 1 from the main text where we consider the target of interest as a solution of a partially linear model. Let \(Z = (X, Y, W)\), where \(X\) is an \(\mathbb{R}^d\)-valued input, \(Y\) is a real-valued outcome, and \(W\) is a \(\mathcal{W}\)-valued control or confounder. The space \(\mathcal{G}\) is a nonparametric class containing functions of the form \[\begin{align} g = (g_Y, g_X): \mathcal{W}\rightarrow \mathbb{R}\times \mathbb{R}^d. \end{align}\] Following the construction of [25], this \(g\) is supplied to the loss \[\begin{align} \ell(\theta, g; z) = \frac{1}{2}[y - g_Y(w) - {\langle \theta, x - g_X(w) \rangle}]^2. \end{align}\]

To ensure \(\theta_\star\) can be interpreted via the projection of \(\mathbb{E}_{\mathbb{P}}[Y|X,W]\) onto partially linear additive functions, the true nuisance is given by \(g_0 = (g_{0, X}, g_{0, Y})\), where \[\begin{align} g_{0, Y}(w) := {\mathbb{E}}_{\mathbb{P}}\left[Y \mid W = w\right] \text{ and } g_{0, X}(w) := {\mathbb{E}}_{\mathbb{P}}\left[X \mid W = w\right]. \end{align}\] We define the residual \(\epsilon\) at \((\theta_\star, g_0)\) as \[\begin{align} \epsilon = Y - g_{0,Y}(W) - {\langle \theta_\star, X - g_{0,X}(W) \rangle}. \end{align}\]

The proof of 2 is provided in 1.0.2.4.1.

1.0.2.1.2 Non-orthogonal Loss

Suppose that the outcome \(Y\) is generated under the partially linear model: \[\begin{align} Y = {\langle \theta_0, X \rangle} + g_0(W) + \epsilon, \end{align}\] where \(\theta_0 \in \mathbb{R}^d\) is the true parameter, \(g_0: \mathcal{W}\mapsto \mathbb{R}\) is the true nuisance function and \({\mathbb{E}}_{\mathbb{P}}\left[\epsilon \mid X, W\right] = 0\). The space \(\mathcal{G}\) is a nonparametric class containing functions of the form \[\begin{align} g: \mathcal{W}\mapsto \mathbb{R}. \end{align}\] We can also consider the following non-orthogonal squared loss function: \[\begin{align} \label{ex232loss} \ell(\theta, g; z) = \frac{1}{2}[y - g(w) - {\langle \theta, x \rangle}]^2. \end{align}\tag{13}\] We define the residual \(\epsilon\) at \((\theta_\star, g_0)\) as \[\begin{align} \epsilon = Y - g_{0}(W) - {\langle \theta_\star, X \rangle}. \end{align}\]

The proof of 3 is provided in 1.0.2.4.2.

1.0.2.1.2.1 Orthogonalization.

We can perform our orthogonalization method to obtain the Neyman orthogonal gradient oracle for this non-orthogonal loss. For any \(h_1, h_2 \in \mathcal{G}\), we define the inner product of \(\mathcal{G}\) as \[\begin{align} \label{inner32product} {\langle h_1, h_2 \rangle}_\mathcal{G}= {\mathbb{E}}_{\mathbb{P}}\left[h_1(W)h_2(W)\right]. \end{align}\tag{14}\]

For any \((\theta, g, z) \in \Theta \times G \times \mathcal{Z}\) By 1 the derivative of non-orthogonal loss 13 along the direction of \(h_1\) is given by \[\begin{align} \label{eq32Dg1} \mathrm{D}_g \ell(\theta, g; z)[h_1] = \frac{\mathrm{d}}{\mathrm{d}t}\left(\frac{1}{2}[y - (g + t h_1)(w) - {\langle \theta, x \rangle}]^2\right) = -(y - g(w) - {\langle \theta, x \rangle})h_1(w). \end{align}\tag{15}\] Do derivative on \(\mathrm{D}_g \ell(\theta, g; z)[h_1]\) along the direction of \(h_2\) and we have \[\begin{align} \label{eq32Dg2} \mathrm{D}_g^2 \ell(\theta, g; z)[h_1, h_2] = \frac{\mathrm{d}}{\mathrm{d}t}\left(-(y - (g+t h_2)(w) - {\langle \theta, x \rangle})h_1(w)\right) = h_1(w)h_2(w), \end{align}\tag{16}\] which implies \[\begin{align} \mathrm{D}_g^2 L(\theta_\star, g_0)[h_1, h_2] = {\mathbb{E}}_{\mathbb{P}}\left[ \mathrm{D}_g^2 \ell(\theta_\star, g_0; Z)[h_1, h_2]\right] = {\mathbb{E}}_{\mathbb{P}}\left[h_1(W)h_2(W)\right]. \end{align}\] By the definition in 66 , we have \(H_{gg} = \mathbf{I}\) the identity operator. In addition, do derivative on the score 93 along the direction of \(h \in \mathcal{G}\) and we have \[\begin{align} \mathrm{D}_g S_\theta(\theta, g; z)[h] = \frac{\mathrm{d}}{\mathrm{d}t}\left(-(y - (g+th)(w) - {\langle \theta, x \rangle})x\right) = h(w)x, \end{align}\] which implies that \[\begin{align} \mathrm{D}_g S_\theta(\theta, g)[h] = {\mathbb{E}}_{\mathbb{P}}\left[ S_\theta(\theta, g; Z)[h]\right] = {\mathbb{E}}_{\mathbb{P}}\left[h(W){\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right]\right]. \end{align}\] By the definition in 65 , we have \(H_{\theta g} = {\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right]\). Thus, by 67 we have \[\begin{align} \label{sim32gamma0} \Gamma_0: g \mapsto {\langle {\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right], g \rangle}_\mathcal{G}= {\mathbb{E}}_{\mathbb{P}}\left[{\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right]g(W)\right]. \end{align}\tag{17}\] Thus, the Neyman orthogonalized gradient oracle defined in 68 is given by \[\begin{align} S_{\text{no}}(\theta, g; z) &= S_\theta(\theta, g; z) - \mathrm{D}_g \ell(\theta, g; z)[{\mathbb{E}}_{\mathbb{P}}\left[X\mid W=w\right]]\notag \\ &= -(y - g(w) - {\langle \theta, x \rangle})(x - {\mathbb{E}}_{\mathbb{P}}\left[X\mid W=w\right]).\label{ex232noscore} \end{align}\tag{18}\]

The proof of 4 is provided in 1.0.2.4.3.

1.0.2.2 Conditional Averaged Treatment Effect

We now introduce examples in causal inference which are established based on the potential outcomes framework. The potential outcomes framework [42] has been widely used in causal inference. Let \(Z = (W, X, Y) \in \left\{0,1\right\}\times \mathbb{R}^d \times \mathbb{R}\) under some distribution \(\mathbb{P}\). We posit the existence of potential outcomes \(Y(1),Y(0) \in \mathbb{R}\). The conditional averaged treatment effect (CATE) is then defined as \[\begin{align} \tau_0(x) = \mathbb{E}_{\mathbb{P}}[Y(1) - Y(0) \mid X=x]. \end{align}\] To identify \(\tau_0(x)\) and the following causal assumptions are required:

Under 4, \(\tau_0\) can be identified by observed data since \[\begin{align} \tau_0(x) &= \mathbb{E}_{\mathbb{P}}[Y(1) - Y(0) \mid X=x]\\ &=\mathbb{E}_{\mathbb{P}}[Y(1) \mid W=1, X=x] - \mathbb{E}_{\mathbb{P}}[Y(0) \mid W=0, X=x]\\ &=\mathbb{E}_{\mathbb{P}}[Y \mid W=1, X=x] - \mathbb{E}_{\mathbb{P}}[Y \mid W=0, X=x]. \end{align}\]

1.0.2.2.1 Unrestricted Nuisance

We observe \(Z = (X, Y, W)\), where \(W\) is a binary treatment assignment. The functions in \(\mathcal{G}\) are of the form \[\begin{align} g = (g^{\text{out}}, g^{\text{prop}}): \mathbb{R}^d \rightarrow \mathbb{R}\times \mathbb{R}, \end{align}\] and are evaluated (see [43]) at the loss \[\begin{align} \ell\left(\theta, g; z\right) = \frac{1}{2}\left(y - g^{\text{out}}(x) - \left(w - g^{\text{prop}}(x)\right){\langle \theta, x \rangle}\right)^2. \end{align}\] For \(g_0 = (g^{\text{out}}_0, g^{\text{prop}}_0)\) nuisance functions \(g^{\text{out}}_0\) and \(g^{\text{prop}}_0\) represent the outcome regression and the propensity score, respectively: \[\begin{align} g^{\text{out}}_0(x) := \mathbb{E}_{\mathbb{P}}[Y \mid X = x] \text{ and } g^{\text{prop}}_0(x) := \mathbb{E}_{\mathbb{P}}[W \mid X = x]. \end{align}\] We define the residual \(\epsilon\) under the true model as \[\begin{align} \epsilon = Y - g_0^{\text{out}}(X) - \left(W- g_0^{\text{prop}}(X)\right)\tau_0(X). \end{align}\]

The proof of 5 is provided in 1.0.2.4.4.

1.0.2.2.2 Restricted Nuisance

We observe \(Z = (X, Y, W)\), where \(W\) is a binary treatment assignment. Here we restrict the propensity model as \(g^{\text{prop}}: \mathbb{R}^d \mapsto (0,1)\). The functions in \(\mathcal{G}\) are of the form \[\begin{align} g = (g^{\scriptscriptstyle(0)}, g^{\scriptscriptstyle(1)}, g^{\text{prop}}): \mathbb{R}^d \rightarrow \mathbb{R}\times \mathbb{R}\times (0, 1), \end{align}\] and are evaluated (see [27]) at the loss \[\begin{align} \ell(\theta, g; z) = \frac{1}{2}\left(g^{\scriptscriptstyle(1)}(x) - g^{\scriptscriptstyle(0)}(x) + \frac{w-g^{\text{prop}}(x)}{g^{\text{prop}}(x)(1-g^{\text{prop}}(x))}(y - g^{\scriptscriptstyle(w)}(x)) - {\langle \theta, x \rangle}\right)^2. \end{align}\] This loss also appears in [24]. For \(g_0 = (g^{\scriptscriptstyle(0)}_0, g^{\scriptscriptstyle(1)}_0, g^{\text{prop}}_0)\) nuisance functions \(g^{\scriptscriptstyle(0)}_0\) and \(g^{\scriptscriptstyle(1)}_0\) represent the outcome regressions \[\begin{align} g^{\scriptscriptstyle(0)}_0(x) := \mathbb{E}_{\mathbb{P}}[Y \mid W=1, X = x] \text{ and } g^{\scriptscriptstyle(1)}_0(x) := \mathbb{E}_{\mathbb{P}}[Y \mid W=0, X = x]. \end{align}\] We define the residual \(\epsilon\) as \[\begin{align} \epsilon = \frac{W-g_0^{\text{prop}}(X)}{g_0^{\text{prop}}(X)(1-g_0^{\text{prop}}(X))}(Y - g_0^{\scriptscriptstyle(W)}(X)). \end{align}\]

The proof of 6 is provided in 1.0.2.4.5.

1.0.2.3 Conditional Relative Risk

We retain all components of the previous example, changing only the loss and assuming that the outcome \(Y\) is binary/non-negative. First, consider the “label” function \[\begin{align} \mu_g^{\scriptscriptstyle(s)}(z) = g^{\scriptscriptstyle(s)}(x) + \frac{\mathbb{1}(w = s)}{sg^{\text{prop}}(x) + (1-s) (1-g^{\text{prop}}(x))}(y - g^{\scriptscriptstyle(s)}(x)), \end{align}\] where \(\mathbb{1}(\cdot)\) denotes the indicator function, and the log-linear predictor \(p_\theta(x) = e^{{\langle \theta, x \rangle}}/(1 + e^{{\langle \theta, x \rangle}})\). Following Example 2 in [28], we then employ the cross entropy-type loss function \[\begin{align} \ell(\theta, g; z) = - \left[\mu_g^{\scriptscriptstyle(1)}(z) \log p_\theta(x) + \mu_g^{\scriptscriptstyle(0)}(z) \log(1 - p_\theta(x))\right]. \end{align}\]

The proof of 7 is provided in 1.0.2.4.6.

1.0.2.4 Proofs

1.0.2.4.1 Proof of Lemma 2

1.0.2.4.2 Proof of Lemma 3

1.0.2.4.3 Proof of Lemma 4

1.0.2.4.4 Proof of Lemma 5

1.0.2.4.5 Proof of Lemma 6

1.0.2.4.6 Proof of Lemma 7

1.0.3 Convergence Proofs for Stochastic Gradient↩︎

This section is dedicated to demonstrate the SGD convergence in 1 from 0.0.3 of the main text using 1 and 2. We first give an overview of the problem settings and the expected results with the proof outline in 1.0.3.1. We then provide all the technical lemmas needed for 1 in 1.0.3.2, and finally prove our first main result in 1.0.3.3.

1.0.3.1 Overview

In this section, we demonstrate the convergence of SGD for a risk minimization problem with nuisance: \[\begin{align} \theta_\star = \mathop{\mathrm{arg\,min}}_{\theta \in \Theta} L(\theta, g_0), \end{align}\] where \(g_0 \in \mathcal{G}\) is the true nuisance, \(L(\theta, g) = {\mathbb{E}}_{Z\sim \mathbb{P}}\left[\ell(\theta,g; Z)\right]\), and \(\ell\) is a prespecified loss function. We consider the stochastic gradient method for learning \(\theta_\star\) when \(g_0\) is unknown but an estimate \(\hat{g}\) is accessible. Define \(\mathcal{D}_n = (Z_1, \ldots, Z_n)\), sampled from the product measure \(\mathbb{P}^n\). Recall the SGD \(\theta^{(n)}\) defined as \[\begin{align} \theta^{(n)} = \theta^{(n-1)} - \eta S_\theta(\theta^{(n-1)}, \hat{g}; Z_n). \end{align}\] Throughout the section, we take the following notations for simplicity: \[\begin{align} \delta^{(n)} &= \theta^{(n)} - \theta_\star,\\ S^{(n)} &= S_\theta(\theta^{(n-1)}, \hat{g}; Z_n),\\ v^{(n)} &= S^{(n)} - S_\theta(\theta^{(n-1)}, \hat{g}), \end{align}\] where \(S_\theta(\theta,g;z) = \nabla_\theta \ell(\theta,g;z)\) is the gradient and \(S_\theta(\theta,g) = {\mathbb{E}}_{Z \sim \mathbb{P}}\left[S_\theta(\theta,g;Z)\right]\) is the population gradient. We are interested in the mean squared error using an estimated nuisance \(\hat{g}\), and our results show that for non-Neyman orthogonal loss \(L\), the error \(\delta^{(n)}\) satisfies \[\begin{align} \label{sgd32overview32non-ortho} \mathbb{E}_{\mathcal{D}_{n} \sim \mathbb{P}^n}[\Vert\delta^{(n)}\Vert_2^2] \lesssim \left(1-\frac{\mu\eta}{2}\right)^n\Vert\delta^{(0)}\Vert_2^2+ \Vert\hat{g} - g_0\Vert_\mathcal{G}^2 + \eta, \end{align}\tag{56}\] where the nuisance estimator \(\hat{g}\) would lead to a bias of order \(\mathcal{O}(\Vert\hat{g} - g_0\Vert_\mathcal{G}^2)\) for the SGD convergence. If \(L\) is Neyman orthogonal, this bias introduced by the nuisance estimation would be further removed, resulting in the following convergence \[\begin{align} \label{sgd32overview32ortho} \mathbb{E}_{\mathcal{D}_{n} \sim\mathbb{P}^n}[\Vert\delta^{(n)}\Vert_2^2] \lesssim& \left(1-\frac{\mu\eta}{2}\right)^n\Vert\delta^{(0)}\Vert_2^2 + \Vert\hat{g} - g_0\Vert_\mathcal{G}^4 + \eta. \end{align}\tag{57}\]

1.0.3.1.0.1 Proof Outline.

The proofs for both results 56 and 57 follow the following four steps:

1. Upper bound the excess risk \(L(\theta^{(n)}, \hat{g}) - L(\theta_\star, \hat{g})\) in terms of the SGD improvement.

2. Lower bound \(L(\theta^{(n)}, \hat{g}) - L(\theta_\star, \hat{g})\) using strong convexity and Neyman orthogonality.

3. Derive a recursive formula of \(\mathbb{E}_{\mathcal{D}_n\sim\mathbb{P}^n}[\lVert \delta^{(n)} \rVert_2^2]\) from these bounds.

4. Perform the recursion and obtain the final result.

Follow these steps above, we provide technical lemma in 1.0.3.2, and then prove our first main result 1 in 1.0.3.3.

1.0.3.2 Technical Lemma

1.0.3.3 Proof of Theorem 1

1.0.4 Orthogonalization with respect to Nuisance↩︎

In this section, we establish our orthogonalization method for the possibly infinite-dimensional nuisance introduced in 0.0.3 of the main text. We demonstrate how we construct the orthogonalizing operator in 1.0.4.1, and provide all the technical lemmas in 1.0.4.2.

1.0.4.1 Orthogonalization via Riesz Representation

We consider \(\mathcal{G}\equiv (\mathcal{G}, {\langle \cdot, \cdot \rangle}_\mathcal{G})\) as a possibly infinite-dimensional Hilbert space. Recall the derivative operator \(\mathrm{D}_g\) defined as for any \(h \in \mathcal{G}\), \[\begin{align} \mathrm{D}_g \ell(\theta, g; z)[h] = \frac{\mathrm{d}}{\mathrm{d}t} \ell(\theta, g + th; z)\mid_{t = 0}. \end{align}\] This derivative operator is also known as the Gateaux derivative. We posit the usual assumption [44] that the derivative operator \(\mathrm{D}_g \ell(\theta, g; z)\) is linear and continuous in \(\mathcal{G}\) for any \((\theta, g, z) \in \Theta \times \mathcal{G}_{r}\left(g_0\right)\times \mathcal{Z}\). We also assume regularity conditions such that \(\mathrm{D}_g\mathrm{D}_\theta \ell\) is continuous and \(\mathrm{D}_g\mathrm{D}_\theta \ell= \mathrm{D}_\theta\mathrm{D}_g\ell\) at any \((\theta, g, z)\).

Since \(\mathrm{D}_g \ell(\theta, g; z)\) is linear and continuous, by the Riesz representation theorem, there uniquely exists some \(\nabla_g \ell(\theta, g; z) \in \mathcal{G}\) such that for any \(g \in \mathcal{G}\), \[\begin{align} \label{eq32Riesz32Dg} \mathrm{D}_g \ell(\theta, g; z)[g] = {\langle \nabla_g \ell(\theta, g; z), g \rangle}_{\mathcal{G}}. \end{align}\tag{64}\] 12 shows that the operator \(\mathrm{D}_gS_\theta(\theta_\star, g_0)\) is linear and continuous. By Riesz representation theorem, we can define \(H_{\theta g} = (H_{\theta g}^{(1)}, \dots, H_{\theta g}^{(d)}) \in \mathcal{G}^d\) such that for all \(g \in \mathcal{G}\), \[\begin{align} \label{def32H32theta32g} \mathrm{D}_g S_\theta^{(j)}(\theta_\star, g_0)[g] = {\langle H_{\theta g}^{(j)}, g \rangle}_{\mathcal{G}} \text{ for any j=1,\dots, n.} \end{align}\tag{65}\] The Hessian operator \(H_{gg}: \mathcal{G}\mapsto \mathcal{G}\) is defined as for any \(g_1, g_2 \in \mathcal{G}\), \[\begin{align} \label{def32H32gg} \mathrm{D}_g^2 L(\theta_\star, g_0)[g_1,g_2] = {\langle H_{gg}g_1, g_2 \rangle}_{\mathcal{G}}. \end{align}\tag{66}\] We will show in 13 that \(H_{gg}\) uniquely exists and is an self-adjoint and bounded linear operator when \(\mathrm{D}_g^2 L(\theta_\star, g_0)\) is bounded and symmetric bilinear. Assuming that \(H_{gg}\) is invertible, we define the orthogonalizing operator as \[\begin{align} \label{appx32gamma} \Gamma_0: \mathcal{G}\mapsto \mathbb{R}^d, \quad [\Gamma_0 g]_j = {\langle H_{\theta g}^{(j)}, H_{gg}^{-1} g \rangle}_\mathcal{G}, \forall g \in \mathcal{G}. \end{align}\tag{67}\] We now construct the Neyman orthogonalized (NO) gradient oracle \[\begin{align} S_{\text{no}}(\theta, g;z) = S_\theta(\theta, g; z)- \Gamma_0 \nabla_g\ell(\theta, g;z). \end{align}\] In addition, \(\Gamma_0 \nabla_g\ell(\theta, g;z)\) can be written as the derivative in the sense that for each \(j = 1,\dots,d\), \[\begin{align} [\Gamma_0 \nabla_g\ell(\theta, g;z)]_j &= {\langle H_{\theta g}^{(j)}, H_{gg}^{-1} \nabla_g\ell(\theta, g;z) \rangle}_\mathcal{G}\\ &={\langle H_{gg}^{-1}H_{\theta g}^{(j)}, \nabla_g\ell(\theta, g;z) \rangle}_\mathcal{G}= \mathrm{D}_g\ell(\theta, g;z)[H_{gg}^{-1}H_{\theta g}^{(j)}]. \end{align}\] That is, the NO gradient oracle can be easily obtain by \[\begin{align} \label{def32no32score32appx} S_{\text{no}}(\theta, g;z) = S_\theta(\theta, g; z)- \mathrm{D}_g\ell(\theta, g;z)[H_{gg}^{-1}H_{\theta g}]. \end{align}\tag{68}\] The following Lemma shows that \(S_{\text{no}}(\theta, g;z)\) is Neyman orthogonal at \((\theta_\star, g_0)\).

1.0.4.2 Technical Lemma

1.0.5 Convergence Proofs for Orthogonalized Stochastic Gradient↩︎

This section is dedicated to demonstrate the OSGD convergence in 2 from 1.0.3 of the main text. In 1.0.5.1, we give an overview of the OSGD settings, the additional assumptions, and the expected results with the proof outline. We then provide all the technical lemmas in 1.0.5.2, and finally prove 2 in 1.0.5.3.

1.0.5.1 Overview

Following the same problem settings in 1.0.3, we consider the orthogonalized SGD (OSGD) using the estimated NO score \(\hat{S}_{\text{no}}\) oracle defined as \[\begin{align} \hat{S}_{\text{no}}(\theta, g; z) = S_\theta(\theta, g; z) - \hat{\Gamma}\nabla_g \ell(\theta, g; z), \end{align}\] where \(\hat{\Gamma}\) is an estimator for the orthogonalizing operator defined in 67 . Specifically, we consider all continuous linear \(\hat{\Gamma}: \mathcal{G}\mapsto \mathbb{R}^d\) for estimating the orthogonalizing operator \(\Gamma_0\). By Riesz representation theorem, there exists some \(\hat{\gamma}^{(j)} \in \mathcal{G}\) for \(j=1,\dots,d,\) such that \[\begin{align} [\hat{\Gamma} g]_j = {\langle \hat{\gamma}^{(j)}, g \rangle}_\mathcal{G}\text{ for all } g \in \mathcal{G}. \end{align}\] For the orthogonalizing operator \(\Gamma_0\), we define \(\gamma_0^{(j)} = H_{gg}^{-1}H_{\theta g}^{(j)}, j = 1,\dots, d,\) such that \[\begin{align} [\Gamma_0 g]_j = {\langle \gamma_0^{(j)}, g \rangle}_\mathcal{G}\text{ for all } g \in \mathcal{G}. \end{align}\] We focus on the OSGD defined below: \[\begin{align} \label{estimated32OSGD32appendix} \theta^{(n)} = \theta^{(n-1)} - \eta \hat{S}_{\text{no}}(\theta^{(n-1)}, \hat{g};Z_n), \quad \theta^{(0)} \in \Theta. \end{align}\tag{69}\] Throughout the section, we take the following notations for simplicity: \[\begin{align} \delta_{\text{no}}^{(n)} &= \theta^{(n)} - \theta_\star,\\ \hat{S}_{\text{no}}^{(n)} &= \hat{S}_{\text{no}}(\theta^{(n-1)}, \hat{g}; Z_n),\\ v_{\text{no}}^{(n)} &= \hat{S}_{\text{no}}^{(n)} - \hat{S}_{\text{no}}(\theta^{(n-1)}, \hat{g}). \end{align}\]

We aim to show that the error \(\delta_{\text{no}}^{(n)}\) satisfies \[\begin{align} \label{appdx32osgd32rate} \mathbb{E}_{\mathcal{D}_n \sim \mathbb{P}^n}[\Vert\delta_{\text{no}}^{(n)}\Vert_2^2] &\lesssim \left(1 - \frac{\mu_{\text{no}}\eta}{2}\right)^n\Vert \delta_{\text{no}}^{(0)}\Vert_2^2+ \Vert\hat{g} - g_0\Vert_\mathcal{G}^4 + \Vert\hat{g} - g_0\Vert_\mathcal{G}^2 \cdot \lVert \hat{\Gamma}- \Gamma_0 \rVert_{\mathrm{Fro}}^2 + \eta. \end{align}\tag{70}\]

1.0.5.1.0.1 Proof Outline.

The proof the 70 follows the following four steps:

1. Upper bound \(\lVert I - \eta \nabla_\theta \hat{S}_{\text{no}}(\theta, g) \rVert_2\) w.r.t. the operator estimation error \(\Vert \hat{\Gamma}- \Gamma_0\Vert_\mathrm{Fro}\).

2. Upper bound \(\lVert \hat{S}_{\text{no}}(\theta_\star, \hat{g}) \rVert_2\) using Neyman orthogonality and the first order optimality.

3. Derive a recursive formula of \(\mathbb{E}_{\mathcal{D}_n\sim\mathbb{P}^n}[\lVert \delta_{\text{no}}^{(n)} \rVert_2^2]\) from these bounds.

4. Perform the recursion and obtain the final result.

Follow these steps above, we provide technical lemma in 1.0.5.2, and then prove our second main result 2 in 1.0.5.3.

Alternatively, the intuition of step 1 also suggests that we should focus on \(\hat{\Gamma}\) that lies in the neighborhood of \(\Gamma_0\) such that \(\Vert \hat{\Gamma}- \Gamma_0\Vert_\mathrm{Fro}\leq R\) for a small \(R>0\). Then, instead of assuming 3(b), we can directly assume that for all \(\theta \in \Theta\) and \(g \in \mathcal{G}_{r}\left(g_0\right)\), \(\lVert \nabla_\theta \hat{S}_{\text{no}}(\theta, g) \rVert_2 \leq M_{\text{no}}\) and \[\begin{align} \lambda_{\min}(\nabla_\theta \hat{S}_{\text{no}}(\theta, g) + \nabla_\theta \hat{S}_{\text{no}}(\theta, g)^\top) \geq 2\mu_{\text{no}}. \end{align}\] With this assumption, one can still show the same OSGD convergence rate by the identical proof while the constraint of the learning rate \(\eta\) will be simplified.

1.0.5.2 Technical Lemma

1.0.5.3 Proof of Theorem 2

1.0.6 Discussion↩︎

This section provides details on comparisons and remarks following the statements of the main results in the main text, and details on the discussions summarized in 0.0.4 from the main text. In 1.0.6.1, we compare our results to generic state-of-the-art results on biased SGD, that is SGD with errors in the stochastic gradients. In 1.0.6.2, we discuss different orthogonalization method in orthogonal statistical learning. In 1.0.6.3, we discuss how to interleave the target and nuisance estimation. In 1.0.6.4, we describe the connection between our orthogonalized gradient to variance reduction method in the Monte Carlo estimation literature.

1.0.6.1 Comparison to Biased SGD

There are several ways to think about the bias induced by using an imperfect estimate \(\hat{g}\) as opposed to the true nuisance \(g_0 \in \mathcal{G}\). For the sake of discussion, we will define \(L(\cdot, g_0)\) and \(L(\cdot, \hat{g})\) as the “original objective” and “shifted objective”, respectively. Accordingly, we will call \(\theta_\star\) the “original minimizer” and denote by \[\begin{align} \hat{\theta}_\star = \mathop{\mathrm{arg\,min}}_{\theta \in \Theta} L(\theta, \hat{g}). \end{align}\] the “shifted minimizer”. The bias can be measured in terms of (i) the error \(\lVert \hat{\theta}_\star - \theta_\star \rVert_2^2\) between the original and shifted minimizers, (ii) the uniform error \(\sup_{\theta}\left| L(\theta, g_0) - L(\theta, \hat{g}) \right|\) between the original and shifted objectives, and (iii) some summary of the gradient bias \[\begin{align} {\mathbb{E}}_{Z_t \sim \mathbb{P}}\left[S(\theta^{(t-1)}, \hat{g}; Z_t)\right] - \nabla_\theta L(\theta^{(t)}, g_0),\label{eq:a:discussion:gradient95bias} \end{align}\tag{72}\] of the oracle \(S\) (a vector-valued quantity) for step \(t = 1, \ldots, n\) of the algorithm. Whether one appeals to (i) or (ii) depends on whether the convergence guarantees are stated in terms of iterate convergence or function value convergence; because we analyze convergence of iterates, our discussion will cover (i) and (iii).

On (i), one applies the decomposition \[\begin{align} \lVert \theta^{(n)} - \theta_\star \rVert_2^2 \leq 2\lVert \theta^{(n)} - \hat{\theta}_\star \rVert_2^2 + 2\lVert \hat{\theta}_\star - \theta_\star \rVert_2^2, \end{align}\] and plugs an analysis of unbiased SGD from the current literature for the \(\lVert \theta^{(n)} - \hat{\theta}_\star \rVert_2^2\) term. The purpose of this substitution is to check how our theoretical results align with the known results on unbiased SGD.

[38] show that for constant learning rate \(\eta = \mathcal{O}(\mu/M^{2})\), the iterate \(\theta^{(n)}\) satisfies \[\begin{align} \label{eq32moulines} \mathbb{E}_{\mathcal{D}_{n} \sim \mathbb{P}^n}[\Vert\theta^{(n)} - \theta_\star\Vert_2^2] \lesssim& \exp\left(-\frac{\mu\eta n}{2}\right)\Vert\theta^{(0)} - \theta_\star\Vert_2^2 + \frac{K_1 \eta}{\mu} + \lVert \hat{\theta}_\star - \theta_\star \rVert_2^2. \end{align}\tag{73}\] [30] demonstrate that with the learning rate \(\eta = \mathcal{O}(1/M)\), the iterates would satisfy the following bound: \[\begin{align} \label{eq32cutler} \mathbb{E}_{\mathcal{D}_{n} \sim \mathbb{P}^n}[\Vert\theta^{(n)} - \theta_\star\Vert_2^2] \lesssim& \left(1 - \frac{\mu\eta}{2}\right)^n\Vert\theta^{(0)} - \theta_\star\Vert_2^2 + \frac{K_1 \eta}{\mu} + \lVert \hat{\theta}_\star - \theta_\star \rVert_2^2 \end{align}\tag{74}\] In addition, [30] provide the high probability bound of \(\theta^{(n)}\) that for \(\eta = \mathcal{O}(1/M)\), with probability at least \(1-\delta\), \[\begin{align} \label{eq32cutler32high32prob} \Vert\theta^{(n)} - \theta_\star\Vert_2^2 \lesssim& \left(1 - \frac{\mu\eta}{2}\right)^n\Vert\theta^{(0)} - \theta_\star\Vert_2^2 + \frac{K_1 \eta}{\mu} \log\left(\frac{e}{\delta}\right) + \lVert \hat{\theta}_\star - \theta_\star \rVert_2^2 \end{align}\tag{75}\] All of these bounds essentially agree, as we may apply \((1 - \mu\eta/2) \leq \exp(-\mu\eta/2)\). In comparison to our 1, our bias term is stated directly in terms of the nuisance error \(\lVert \hat{g}- g_0 \rVert_\mathcal{G}^2\). This can be viewed as a refinement of the less transparent bias measurement \(\lVert \hat{\theta}_\star - \theta_\star \rVert_2^2\). Moreover, although 73 –75 are of the same order as our results in 1 when the true nuisance \(g_0\) is available, all of the three bounds above require \(\kappa_1 = 0\) in 1(d). In this case, provide 74 and 75 use a learning rate of the order \(\mathcal{O}(1/M)\) (whereas the learning rate of 73 encounters an additional condition number \(M/\mu\)). Our learning rate recovers \(\mathcal{O}(1/M)\) when \(\kappa_1 = 0\), and adapts via the setting \(\eta = \mathcal{O}(\mu/(M\mu+ \kappa_1))\) when \(\kappa_1 > 0\). Finally, the high probability bound 75 requires a stronger assumption in the sense that \(S_\theta(\theta, g; Z) - S_\theta(\theta, g)\) is sub-Gaussian with uniform parameter \(K_1/2\) for all \(\theta \in \Theta\) and \(g \in \mathcal{G}_{r}\left(g_0\right)\).

Returning to the bias in the stochastic gradient oracle 72 , this case is handled quite generally in [39]. Their “ABC assumption” considers constants \(A, B, C, b, c \geq 0\) such that the inequalities \[\begin{align} {\langle \nabla L(\theta, g_0), {\mathbb{E}}_{Z \sim \mathbb{P}}\left[S_\theta(\theta, \hat{g}; Z)\right] \rangle} &\geq b \lVert \nabla L(\theta, g_0) \rVert_2^2 + c \tag{76}\\ \mathbb{E}_{Z \sim \mathbb{P}}\lVert S_\theta(\theta, \hat{g}; Z) \rVert_2^2 &\leq 2A \left(L(\theta, g_0) - L(\theta_\star, g_0)\right) + B \lVert \nabla L(\theta, g_0) \rVert_2^2 + C \tag{77}\\ A + M\left(B + 1 - 2b\right) &< \frac{\mu}{2} \tag{78} \end{align}\] hold for all \(\theta \in \mathbb{R}^d\) (where the expectations are conditional on any randomness in \(\hat{g}\)).² The bias is really captured in the first of the three inequalities, whereas the third inequality places conditions on the parameters of the problem that are not in the hands of the algorithm user. By strong convexity, our 1(d) satisfies 77 with \(A = \kappa_1/\mu\), \(B = 0\), \(C = K_1\). The resulting convergence guarantee [39] gives \[\begin{align} \mathbb{E}_{\mathcal{D}_{n} \sim \mathbb{P}^n}[\Vert\theta^{(n)} - \theta_\star\Vert_2^2] &\lesssim \left(1 - \frac{\mu\eta}{2}\right)^n\Vert\theta^{(0)} - \theta_\star\Vert_2^2 + \frac{\eta K_1}{\mu} + \frac{K_1 + c}{\mu^2}. \end{align}\] for a learning rate set as \[\begin{align} \eta \leq \min\left\{\frac{2}{\mu}, \frac{\mu- 2(\kappa_1/\mu+ M(1-2b))}{\kappa_1}\right\}. \end{align}\] We would like to check how our theoretical results compare with the above application of generic results on biased SGD. In our setting, the condition 78 reads \(\kappa_1/\mu+ M(1 - 2b) < \mu/2\). We assume neither this nor 76 , and in our 1, replace the \(\frac{K_1 + c}{\mu^2}\) term with a bias term that depends directly on the nuisance error \(\lVert \hat{g}- g_0 \rVert_\mathcal{G}^2\), either in the nuisance sensitive regime, or in the nuisance insensitive regime.

1.0.6.2 Discussion of Full-sample Orthogonal Statistical Learning and Related Methods

1.0.6.2.0.1 Comparison of Orthogonalizing Operators.

Constructing orthogonal losses or scores has been widely studied in semiparametric inference, hypothesis testing, and machine learning. In semiparametric statistics, such constructions often rely on the efficient influence function, which characterizes the asymptotic efficient estimation bound; see [13], [33], [14], [45]. In hypothesis testing, orthogonal scores were used by [22], [46] and [47] to guarantee the local unbiasedness of specific tests based on the likelihood with finite-dimensional nuisance. In machine learning, the construction of orthogonal scores was latter extended to non-likelihood losses in [48] and [32], which aligns with our construction limited to the finite-dimensional nuisance case. Recent work of orthogonalization in machine learning with infinite-dimensional nuisance relies on the approach named concentrating-out [21], [49]. However, although all these constructions produce Neyman orthogonal losses or scores, none of them consider the stochastic design. Our work is complementary to these, providing non-asymptotic guarantees for stochastic optimization.

Although these constructions might lead to different orthogonal scores, they can be the same at both the target and the true nuisance. Specifically, when \(\ell\) is the negative log-likelihood and \(\mathcal{G}= \mathbb{R}^k\), the concentrating-out approach and our NO gradient oracle \(S_{\text{no}}\) both produce the efficient score in the semiparametric theory literature; see [49], [14], and [33]. This identity can happen for infinite-dimensional nuisances as well. As an example, consider the partial linear model from 1.0.2.1.2, where the non-orthogonal loss is defined as \[\begin{align} \ell(\theta, g; z) = \frac{1}{2}(y - g(w) - {\langle \theta, x \rangle})^2. \end{align}\] [21] showed that the concentrating-out approach would produce an orthogonal score under the concentrated-out nuisance \(\varphi_0(\theta) = Z \mapsto \mathbb{E}_{\mathbb{P}}[Y - {\langle \theta, X \rangle} \mid W]\) as \[\begin{align} S(\theta_\star, \varphi_0(\theta_\star); Z) &= -(X - {\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right])(Y - {\mathbb{E}}_{\mathbb{P}}\left[Y\mid W\right] - {\langle \theta_\star, X - {\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right] \rangle}). \end{align}\] On the other hand, it is easy to verify that \(H_{gg} = \mathbf{I}\), \(H_{\theta g} = \mathbb{E}_{\mathbb{P}}[X \mid W]\), and \(\Gamma_0 \nabla_g \ell(\theta, g; z) = \mathrm{D}_g \ell(\theta,g;z)[H_{gg}^{-1}H_{\theta g}]\), which implies that our orthogonal gradient oracle \(S_{\text{no}}\) in 10 has the same form under the target \(\theta_\star\) and true nuisance \(g_0(W) = \mathbb{E}_{\mathbb{P}}[Y - {\langle \theta_\star, X \rangle} \mid W]\): \[\begin{align} S_{\text{no}}(\theta_\star, g_0; Z) &= -(X - {\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right])(Y - {\mathbb{E}}_{\mathbb{P}}\left[Y\mid W\right] - {\langle \theta_\star, X - {\mathbb{E}}_{\mathbb{P}}\left[X\mid W\right] \rangle}). \end{align}\]

1.0.6.2.0.2 Comparison with Debiased Machine Learning.

In machine learning, debiasing typically refers to reducing the impact of model selection error on the parameter or quantity of interest. In particular, mitigating the bias introduced by nuisance estimation is known as debiased machine learning (DML), which has been recently studied by [50]–[53]. Some of the calculations used by DML estimators have been shown to be amenable to computerization, simplifying their construction [35], [54], [55]. Statistical learning methods that use debiasing are also called orthogonal statistical learning (OSL) and have been studied in [23], [24], [32]. While the earlier studies focus on the empirical risk minimization, our paper provide a stochastic approximation method in DML/OSL and establish the convergence rate of the debiased estimation.

To strengthen the debiasing effect, one possible approach is to consider the higher-order Neyman orthogonality. If the loss function satisfies the \(k\)-th order orthogonality at \((\theta_\star, g_0)\), [23] show that we only need the nuisance estimator to converge at rate \(\mathcal{O}_p(n^{-\frac{1}{2(k+1)}})\) to have the nuisance effect in the order of \(\mathcal{O}_p(n^{-1})\), which aligns with the nuisance insensitive rate in 1, where \(k=1\) and the nuisance effect \(\Vert\hat{g} - g_0\Vert_{\mathcal{G}}^4 = \mathcal{O}_p(n^{-1})\) when \(\Vert\hat{g} - g_0\Vert_{\mathcal{G}} = \mathcal{O}_p(n^{-1/4})\). Similar improvements in sensitivity to nuisance estimation rates have been developed previously using higher-order influence functions [56], [57].

For a range of problems, debiasing methods often lead to cross-product estimations consisting of two nuisance estimators [58]. Such remainders frequently result from orthogonalization procedures used in missing data problems and causal inference problems [59]–[61]. [52] consider cases where \(Z = (W,Y)\), \(g_0(W)= {\mathbb{E}}_{\mathbb{P}}\left[Y \mid W\right]\), and the target can be written as the averaged moment of the form \[\begin{align} \theta_\star = \mathbb{E}_{\mathbb{P}}[m(g_0; W)], \end{align}\] where \(\mathbb{E}[m(g; W)]: \mathcal{G}\times \mathcal{W}\mapsto \mathbb{R}^d\) is a continuous linear functional of \(g: \mathcal{W}\mapsto \mathbb{R}\). By Riesz representation theorem, there uniquely exists a Riesz representer (RR) \(g_{0}^{\text{RR}} \in \mathcal{G}\) such that \(\mathbb{E}_{\mathbb{P}}[m(g; W)] = \mathbb{E}_{\mathbb{P}}[g_{0}^{\text{RR}}(W) g(W)]\). Then the debiased score for estimating \(\theta_\star\) is defined as \[\begin{align} S(\theta, g;Z) = m(g;W) + g_{0}^{\text{RR}}(W)(Y - g(W)) - \theta. \end{align}\] The debiasing effect on the nuisance turns out depending on the cross-product \(\Vert\hat{g}^{\text{RR}} - g_{0}^{\text{RR}}\Vert_{\mathcal{G}}\cdot\Vert\hat{g} - g_0\Vert_{\mathcal{G}}\). Specifically, [52] requires the cross product to be in the order \(\mathcal{O}_p(n^{-1/2})\) to construct \(\mathcal{O}_p(n^{-1})\) consistent target estimator. This aligns with the cross-product \(\Vert\hat{g} - g_0\Vert_{\mathcal{G}}\cdot \lVert \hat{\Gamma}- \Gamma_0 \rVert_\mathrm{Fro}\) in 2 where the same requirement needs to be satisfied to obtain a \(\mathcal{O}_p(n^{-1/2})\) consistent estimator. However, 2 also has a second, non-cross-product remainder \(\Vert\hat{g} - g_0\Vert_\mathcal{G}^4\) that will only be small if \(\hat{g}\) approximates \(g_0\), making it so that our consistency guarantee is robust to misspecification of \(\hat{\Gamma}\), but not to misspecification of \(\hat{g}\).

1.0.6.3 Discussion of Interleaving Target and Nuisance Estimation

To propose the interleaving approach, we consider the case where we learn the nuisance from the \(\mathcal{W}\)-valued data \(W = (U,V)\) from a probability measure \(\mathbb{Q}\). We assume that the true nuisance \(g_0\) satisfies \(g_0: \mathcal{U} \mapsto \mathbb{R}\) and is the minimizer of the mean squared error over \(\mathcal{G}\): \[\begin{align} g_0 = \mathop{\mathrm{arg\,min}}_{g \in \mathcal{G}}{\mathbb{E}}_{\mathbb{Q}}\left[(g(U) - V)^2\right]. \end{align}\]

Suppose that we observe another data stream \(W_1, \ldots, W_m\) sampled i.i.d. from \(\mathbb{Q}\), and that \(\mathcal{S}_m = \{W_1, \ldots, W_m\}\) is independent of the parameter stream \(\mathcal{D}_n\). We define the sigma algebra \(\mathcal{H}_{m} = \sigma(\mathcal{S}_m), m \geq 1\) as the nuisance filtration and the sigma algebra \(\mathcal{F}_{m,t} = \sigma(\mathcal{S}_m \cup \mathcal{D}_{(m-1)n+t}), 0\leq t \leq n\) as the parameter filtration. We assume that there are two stochastic processes \(\hat{g}^{(m)}, m \geq 1\) adapted to \(\mathcal{H}_{m}\) and \(\theta^{(m,t)}, 0\leq t \leq n\) adapted to \(\mathcal{F}_{m,t}\), to which we refer as the nuisance estimator and the parameter estimator, respectively. Intuitively, this means that the nuisance estimator \(\hat{g}^{(m)}\) can be updated now instead of being the fixed \(\hat{g}\), and the parameter estimator \(\theta^{(m,t)}\) can be updated \(n\) times between every two nuisance updates. Specifically, we use SGD as the parameter estimator. We define \(\theta^{(0, n)} = \theta^{(0)} \in \Theta\) and \(\theta^{(i, 0)} = \theta^{(i-1, n)}\) for \(1 \leq i \leq m\), and produce the sequence \(\theta^{(i, 1)}, \ldots, \theta^{(i, n)}\) using \(n\) steps of the SGD update 7 initialized at \(\theta^{(i, 0)}\).

1.0.6.3.0.1 Under Non-orthogonality.

Consider the case that \(\mathcal{G}\) is a reproducing kernel Hilbert space (RKHS) with kernel \(k(\cdot,\cdot)\). To obtain a sequence of nuisance estimator \({\hat{g}^{(m)}}\) on \(\mathcal{H}_{m}\), one possible approach is adopt the non-parametric stochastic approximation. With the assumption that the eigenvalues \((\lambda_j)_{j\geq1}\) of covariance operator \(\mathbb{E}_{\mathbb{Q}}[k(W,\cdot) \otimes k(W,\cdot)]\) decay polynomially at order \(j^{-\alpha}\), [37] suggests that the non-parametric stochastic approximation \(\hat{g}^{(m)}\) satisfies for some \(C>0\), \[\begin{align} \label{eq:32dieuleveut} \xi_m := {\mathbb{E}}_{\mathcal{S}_m \sim \mathbb{Q}^m}\left[\Vert\hat{g}^{(m)} - g_0\Vert_{\mathcal{G}}^2\right] \leq Cm^{-\frac{2\alpha - 1}{2\alpha}}. \end{align}\tag{79}\]

This leads to the following nuisance sensitive rate for non-Neyman orthogonal losses.

The proof is provided in 1.0.6.5. 3 demonstrates that interleaving the target and nuisance estimation allows \(\eta \asymp n^{-1}\) since the nuisance update iterations guarantees the shrinking of the term \(\left(1-\mu\eta/2\right)^{mn}\) in this case. This is an improvement to 1 where \(\eta\) should satisfy \((\eta n)^{-1} = o(1)\) to ensure \((1-\mu\eta/2)^n\) shrinking to zero.

1.0.6.3.0.2 Under Orthogonalized SGD.

To establish a similar probability bound for OSGD, we assume that the orthogonalizing operator \(\Gamma_0\) can be written as the minimizer of the following program: \[\begin{align} \Gamma_0 = \mathop{\mathrm{arg\,min}}_{\Gamma\in \mathcal{G}_*^d} {\mathbb{E}}_{\mathbb{P}}\left[\lVert S_\theta(\theta_\star, g_0; Z) - \Gamma\nabla_g \ell(\theta_\star, g_0; Z) \rVert_2^2\right], \end{align}\] where \(\mathcal{G}_*\) is the dual space of \(\mathcal{G}\). When \(d\) is fixed, we assume that \(\Gamma_0\) can be estimated (coordinate-wisely) from the data stream \(\mathcal{S}_m\) using the stochastic approximation of [37], which leads to a sequence of operator estimators \(\hat{\Gamma}^{(m)}, m\geq 1\). For any \(s>0\), we define the following events for \(i=0,1,\dots, m\), \[\begin{align} \mathcal{A}_{i}(s) = \left\{\Vert\hat{g}^{(i)} - g_0\Vert_{\mathcal{G}}^2 \leq Cs^{-1}i^{-\frac{2\alpha - 1}{2\alpha}}\right\} \text{ and } \mathcal{B}_{i}(s) = \left\{\lVert \hat{\Gamma}^{(i)} - \Gamma_0 \rVert_\mathrm{Fro}^2 \leq Cs^{-1}i^{-\frac{2\alpha - 1}{2\alpha}}\right\}. \end{align}\] We assume that for some constant \(c\geq 1\) the nuisance estimator \(\hat{g}^{(i)}\) satisfies \[\begin{align} \label{markov321} {\mathbb{E}}_{\mathcal{S}_i}\left[\Vert\hat{g}^{(i)} - g_0\Vert_{\mathcal{G}}^{2c} \mid \mathcal{A}_{i-1}(s^{1/c}),\dots,\mathcal{A}_{1}(s^{1/c})\right] \leq C^ci^{-\frac{(2\alpha - 1)c}{2\alpha}}. \end{align}\tag{80}\] Additionally, we assume that \(\hat{\Gamma}^{(i)}\) decays in the same rate such that \[\begin{align} \label{markov322} {\mathbb{E}}_{\mathcal{S}_i }\left[\lVert \hat{\Gamma}^{(i)} - \Gamma_0 \rVert_\mathrm{Fro}^{2c}\mid \mathcal{B}_{i-1}(s^{1/c}),\dots,\mathcal{B}_{1}(s^{1/c})\right] \leq C^ci^{-\frac{(2\alpha - 1)c}{2\alpha}}. \end{align}\tag{81}\] With all the assumptions above, it is possible to provide a convergence bound of \(\lVert \theta^{(m,n)}-\theta_\star \rVert_2^2\) in probability. The following proposition shows that estimations from \(\mathcal{S}_m\) using OSGD contribute to a nuisance insensitive rate of \(\mathcal{O}(m^{-\frac{2\alpha-1}{\alpha}})\), compared to the nuisance sensitive rate \(\mathcal{O}(m^{-\frac{2\alpha-1}{2\alpha}})\) in 3 for non-Neyman orthogonal losses.

We refer the reader to 1.0.6.6 for the proof.

1.0.6.4 Interpretation as Control Variate for Variance Reduction

The regression equation 8 , which provides an alternate characterization of the orthogonalized stochastic gradient oracle in the case of negative log-likelihood losses, yields an interesting connection to the Monte Carlo estimation literature. Variance reduction techniques (or “swindles”) are used in problems such as estimating the mean or variance of a statistic via Monte Carlo simulation. Consider a probability space \((\Omega, \mathscr{F}, \mathbb{P})\) with expectation denoted by \(\mathbb{E}\) and an unknown vector-valued target \(v \in \mathbb{R}^d\). We have \(\hat{v}: \Omega \rightarrow \mathbb{R}^d\), where we interpret \(\hat{v}\) as a (not necessarily unbiased) sample estimate of \(v\). Several variance reduction techniques fall into the category of control variates [62], where a random variable \(\hat{u}: \Omega \rightarrow \mathbb{R}^k\) with known expectations \(u = \mathbb{E}[\hat{u}]\) and a matrix \(\Gamma\in \mathbb{R}^{d \times k}\) are used in the variance-reduced estimator \[\begin{align} \tilde{v} = \hat{v} - \Gamma(\hat{u} - u). \end{align}\] A mean squared error decomposition yields the identity \[\begin{align} \mathbb{E}\lVert \tilde{v} - v \rVert_2^2 &= \mathbb{E}\lVert \hat{v} - v \rVert_2^2 - 2\mathbb{E}{\langle \hat{v} - v, \Gamma(\hat{u} - u) \rangle} + \mathbb{E}\lVert \Gamma(\hat{u} - u) \rVert_2^2\\ &= \mathbb{E}\lVert \hat{v} - v \rVert_2^2 - 2\mathbb{E}{\langle \hat{v} - v, \Gamma(\hat{u} - u) \rangle} + o(\lVert \Gamma \rVert_{\mathrm{op}}), \end{align}\] indicating that for sufficiently “small” \(\Gamma\), \(\tilde{v}\) provides an improved estimator if \(\hat{v} - v\) and \(\hat{u} - u\) have high (multiple) correlation. While in the Monte Carlo literature, \(\hat{u}\) and \(\Gamma\) can be chosen optimally provided knowledge of the underlying data-generating mechanism, as \(\Gamma\) can be interpreted as the regression function of \(\hat{v} - v\) on \(\hat{u} - u\).³ Outside of Monte Carlo simulation, this procedure can be applied more widely if the user chooses \(\hat{u}\) and \(\hat{\Gamma}\) based on intuition or limiting arguments.

In the stochastic optimization setting, \(v\) represents the true gradient of the objective at a particular parameter, while \(\hat{v}\) represents a stochastic gradient estimate from an oracle. Variance reduction techniques have previously been applied in an incremental setting, in which a fixed data set of size \(n\) is provided at initialization, and the algorithm may only make multiple passes through this same data set [63]. Note that this differs from our fully stochastic setting, in which we receive a fresh sample \(Z_t\) on each iterate \(t = 1, \ldots, n\). For negative log-likelihood losses, our orthogonalized oracle can be viewed in a similar light to control variate-based variance reduction methods (although in an infinite-dimensional setting). To summarize, we have from 8 that \[\begin{align} v &= S_\theta(\theta_\star, g_0)\\ \hat{v} &= S_\theta(\theta_\star, \hat{g})\\ u &= 0 &\text{ (by \Cref{asm:noscore})}\\ \hat{u} &= \nabla_g \ell(\theta_\star, g_0; Z)\\ \tilde{v} &= S_\theta(\theta_\star, g_0; Z) - \Gamma_0 \nabla_g \ell(\theta_\star, g_0; Z),\\ \intertext{using the idealized parameters. Using the approximations for \theta \neq \theta_\star, we have} v &= S_\theta(\theta, g_0)\\ \hat{v} &= S_\theta(\theta, \hat{g})\\ u &\approx 0 &\text{ (for \theta \approx \theta_\star)}\\ \hat{u} &= \nabla_g \ell(\theta, \hat{g}; Z)\\ \tilde{v} &= S_\theta(\theta, \hat{g}; Z) - \hat{\Gamma}\nabla_g \ell(\theta, \hat{g}; Z). \end{align}\] In fact, using the derivative of the log likelihood in a control variate procedure has been explored in the simulation literature as early as [41], as the correlation between a statistic and the score function has tight connections to the Cramér-Rao lower variance bound and exponential families. We emphasize, however, that our method does not require the loss to be of negative log-likelihood form nor any specific distributional knowledge to be applied.

1.0.6.5 Proof of Proposition 3

1.0.6.6 Proof of Proposition 4

1.0.7 Numerical Experiments↩︎

This section provides numerical experiments of the proposed stochastic methods in this paper. In 1.0.7.1, we design a numerical experiment to illustrate our orthogonalization method. In 1.0.7.2, we design simulations based on a partially linear model. In 1.0.7.4, we conduct a real data analysis with synthetic outcome to evaluate the performance of our methods. Code for reproduction can be found at <https://fachengyu.github.io/>.

1.0.7.1 Numerical Illustration

In this section, we design a numerical experiment to illustrate how our orthogonalization method effects the target estimation as shown in 1 from the main text.

1.0.7.1.0.1 Settings.

Consider \(\Theta\in R\) and \(\mathcal{G}=\mathbb{R}\). Let \(L(\theta, g)\) be a real-valued risk function defined as \[\begin{align} L(\theta, g) := L(u) = \frac{1}{2}{\langle u, Au \rangle} + \lambda \sin^2({\langle u, Bu \rangle}), \end{align}\] where \(u = (\theta, g)^\top \in \mathbb{R}^2\), \(\lambda = 0.02\) is the regularization parameter, and \[\begin{align} A = \begin{pmatrix}8 & 3\\ 3 & 2\end{pmatrix} \succ 0 \text{ and } B = \begin{pmatrix}2 & -1\\ -1 & 1.5\end{pmatrix}\succ 0. \end{align}\] It is easy to see that \((0,0)\) is the global minimizer of \(L\) since \(L(\theta, g) \geq 0\). Let \(q(u) = {\langle u, Bu \rangle}\). The gradient w.r.t. \(u\) is \[\begin{align} \nabla_u L(u) &= Au + 4 \lambda\sin(q(u))\cos(q(u))Bu\\ &=(A + 2\lambda\sin(2q(u))B)u. \end{align}\] Since \(A + 2\lambda\sin(2q(u))B \succcurlyeq A - 0.04 B \succ 0\), it is clear that \((0,0)\) is the only stationary point, implying that \((0,0)\) is the only minimizer of \(L\). Furthermore, we can obtain the Hessian w.r.t. \(u\) as \[\begin{align} \label{eq32colab32hessian} \nabla_u^2 L(u) = A + 2\lambda\sin(2q(u))B + 8\lambda \cos(2q(u))Bu(Bu)^\top, \end{align}\tag{84}\] which implies that \(L(\cdot,g)\) is not convex in \(\mathbb{R}\) given any \(g \in \mathcal{G}_{r}\left(g_0\right)\). However, when \(\Theta\) is a small neighborhood around zero, it is still possible to have \(L(\cdot,g)\) strongly convex for in \(\Theta\) given any \(g \in \mathcal{G}_{r}\left(g_0\right)\).

1.0.7.1.0.2 Orthogonalization.

To orthogonalize \(L\), we first derive the orthogonal gradient oracle using 10 , and then integral the oracle w.r.t. \(\theta\) to obtain the orthogonalized loss \(L_{\text{no}}\).

Let \(H\) be the Hessian at \((0,0)\). By 84 , we know that \(H = A\), implying \(H_{\theta g} = A_{12}\) and \(H_{gg} = A_{22}\). Since the gradient w.r.t. \(\theta\) satisfies \[\begin{align} \nabla_\theta L(\theta, g) &= [1, 0](A + 2\lambda\sin(2q(u))B )u \\ &= (A_{11} + 2\lambda\sin(2q(\theta, g))B_{11})\theta + (A_{12} + 2\lambda\sin(2q(\theta, g))B_{12})g, \end{align}\] and the gradient w.r.t. \(g\) satisfies \[\begin{align} \nabla_g L(\theta, g) &= [0, 1](A + 2\lambda\sin(2q(u))B )u \\ &= (A_{21} + 2\lambda\sin(2q(\theta, g))B_{21})\theta + (A_{22} + 2\lambda\sin(2q(u))B_{22})g, \end{align}\] follow the construction of 10 and we obtain the orthogonal gradient oracle as \[\begin{align} S_{\text{no}}(\theta ,g) &= \nabla_\theta L(\theta, g) - H_{\theta g} H_{gg}^{-1} \nabla_g L(\theta, g) \\ &=(a + 2b\lambda\sin(2q(\theta, g)))\theta + 2c\lambda\sin(2q(\theta, g))g, \end{align}\] where \(a = A_{11} - A_{12} A_{22}^{-1}A_{21}\), \(b = B_{11} - A_{12} A_{22}^{-1}B_{21}\), and \(c = B_{12} - A_{12} A_{22}^{-1}B_{22}\). Finally, we can integral \(S_{\text{no}}(\theta ,g)\) w.r.t. \(\theta\) and recover the orthognalized loss \(L_{\text{no}}\) as \[\begin{align} L_{\text{no}}(\theta, g) = \int_{0}^\theta S_{\text{no}}(s ,g) \mathrm{d}s. \end{align}\]

1.0.7.1.0.3 Numerical Computation.

Usually, \(S_{\text{no}}(s ,g)\) contains a form of integral, which needs to be numerically computed. For the example introduced above, we can simplify \(L_{\text{no}}(\theta, g)\) to stabilize the numerical computation. Note that \(\nabla_\theta \sin^2(q(\theta, g)) = \sin(2q(\theta, g))(B_{11}\theta + B_{12}g)\). Then \[\begin{align} 2b\lambda \int_0^\theta \sin(2q(s, g))s ds &= \frac{2b\lambda}{B_{11}}\left(\int_0^\theta \sin(2q(s, g))(B_{11}s + B_{12}g) \mathrm{d}s - B_{12} g\int_0^\theta \sin(2q(s, g)) \mathrm{d}s\right)\\ &= \frac{2b\lambda}{B_{11}}\left(\sin^2(q(\theta, g)) - B_{12} g\int_0^\theta \sin(2q(s, g)) \mathrm{d}s\right). \end{align}\] It follows that the orthogonalized loss \(L_{\text{no}}\) admits the following form \[\begin{align} L_{\text{no}}(\theta, g) &=\frac{a}{2}\theta^2 + \frac{2b\lambda}{B_{11}}\sin^2(q(\theta, g)) + 2\left(c - \frac{B_{12}}{B_{11}}b\right)\lambda g\int_0^\theta \sin(2q(s, g)) \mathrm{d}s, \end{align}\] which implies that only the integral of \(\sin(2q(s, g))\) w.r.t. \(s\) needs to be computed.

1.0.7.2 Simulations

1.0.7.2.1 Data Generating Process

To demonstrate 1 and 2, we revisit the partially linear model and the corresponding orthogonal and non-orthogonal losses in 1.0.2.1. Specifically, \((X,W,Y) \in \mathbb{R}^d \times \mathbb{R}^d \times \mathbb{R}\) satisfies the following partially linear model where the nonlinear function is determined by the distribution of \((W, U) \in \mathbb{R}^d \times \mathbb{R}\): \[\begin{align} Y &= {\langle \theta_0, X \rangle} + \alpha_0(W) + \epsilon \tag{85},\\ U &= \alpha_0(W) + \xi \tag{86}, \end{align}\] where \(\theta_0 \in \mathbb{R}^d\) is the true parameter, \(\alpha_0: \mathcal{W}\mapsto \mathbb{R}\) is the true nonlinear function, \(\epsilon\) and \(\xi\) are independent noises that satisfy \(\mathbb{E}[\epsilon \mid X, W] = 0\) and \(\mathbb{E}[\xi \mid W] = 0\). It is clear that \(\alpha_0(W) = \mathbb{E}[U \mid W]\). In our simulations, we choose \(d = 2\) and \(\theta_0 = [-0.5 ~~1]^\top\).

To get samples for simulations, we first generate \((X,W)\) under the Gaussian model \[\begin{align} \begin{bmatrix} X\\ W \end{bmatrix} = \mathcal{N}\left(\begin{bmatrix} \boldsymbol{\mu}_X \\ \boldsymbol{\mu}_W \end{bmatrix}, \begin{bmatrix} (1 + \delta)\mathbf{I}_2 &\lambda \mathbf{I}_2\\ \lambda \mathbf{I}_2 & (1 + \delta)\mathbf{I}_2 \end{bmatrix}\right), \end{align}\] where \(\boldsymbol{\mu}_X = [1 ~~1]^\top\), \(\boldsymbol{\mu}_W = [2 ~~2]^\top\), \(\mathbf{I}_2 \in \mathbb{R}^{2 \times 2}\) is the identity matrix, \(\lambda \in [0,1]\) is used to control the correlation between \(X\) and \(W\), and \(\delta = 0.05\) is used to prevent the degeneration of the covariance matrix. For simplicity, we define the nonlinear function \(\alpha_0\) as \[\begin{align} \alpha_0(w) = 0.5\times \cos\left(\frac{w_1 + w_2}{2}\right) + 0.5\times \sin\left(\frac{w_1 + w_2}{2}\right), \end{align}\] where \(w = [w_1~~ w_2]^\top \in \mathbb{R}^2\). We then generate \(Y\) and \(U\) using independent Gaussian noises \(\epsilon \sim \mathcal{N}(0,1)\) and \(\xi \sim \mathcal{N}(0,1)\) based on 85 and 86 , respectively.

1.0.7.2.2 Stochastic Gradient Oracles

To estimate the true parameter \(\theta_0\) using stochastic gradients, we need to design a correspond loss whose minimizer \(\theta_\star\) is equal to \(\theta_0\). Based on 1.0.2.1, there are two types of loss, the orthogonal loss and the non-orthogonal loss, available for this goal. We will derive the stochastic gradient oracle for these two losses and further derive the orthogonalized gradient oracle for the non-orthogonal loss.

1.0.7.2.2.1 Orthogonal loss.

Recall the orthogonal loss in 1.0.2.1.1: \[\begin{align} \label{sim32ortho32loss} \ell(\theta, g; z) = \frac{1}{2}[y - g_Y(w) - {\langle \theta, x - g_X(w) \rangle}]^2, \end{align}\tag{87}\] where \(g = (g_Y, g_X): \mathcal{W}\rightarrow \mathbb{R}\times \mathbb{R}^d\) and the norm \(\Vert\cdot\Vert_{\mathcal{G}}\) is defined in 20 . The true nuisance for this loss is \(g_0 = (g_{0, X}, g_{0, Y})\), where \(g_{0, Y}(w) := {\mathbb{E}}_{\mathbb{P}}\left[Y \mid W = w\right] \text{ and } g_{0, X}(w) := {\mathbb{E}}_{\mathbb{P}}\left[X \mid W = w\right].\) In fact, the explicit expression for \(g_0\) can be easily obtained as \[\begin{align} g_{0, Y}(w) &= {\langle \theta_0, g_{0, X}(w) \rangle} + \alpha_0(w),\tag{88}\\ g_{0, X}(w) &= \mu_X + \frac{\lambda}{1.05}(w - \mu_W). \tag{89} \end{align}\] From 85 and 88 , it is clear that \[\begin{align} Y - g_{0, Y}(W) = {\langle \theta_0, X - g_{0, X}(W) \rangle} + \epsilon, \end{align}\] which implies that \(\theta_\star = \theta_0\) by 2. The stochastic gradient oracle for the orthogonal loss 87 is then defined as \[\begin{align} \label{sim32ortho32score} S_\theta(\theta, g; z) = -(y - g_Y(w) - {\langle \theta, x - g_X(w) \rangle})(x - g_X(w)). \end{align}\tag{90}\]

1.0.7.2.2.2 Non-orthogonal loss.

We also provide the non-orthogonal loss in 1.0.2.1.2 as \[\begin{align} \label{sim32nonortho32loss} \ell(\theta, g; z) = \frac{1}{2}[y - g(w) - {\langle \theta, x \rangle}]^2, \end{align}\tag{91}\] where \(g: \mathcal{W}\mapsto \mathbb{R}\) and the norm \(\Vert\cdot\Vert_{\mathcal{G}}\) is now defined in 28 . The true nuisance for this non-orthogonal loss satisfies \[\begin{align} g_0(w) = \alpha_0(w) = {\mathbb{E}}_{}\left[U \mid W = w\right]. \label{sim32g0} \end{align}\tag{92}\] By 3, we have \(\theta_\star = \theta_0\). The stochastic gradient oracle for the orthogonal loss 91 is then defined as \[\begin{align} \label{sim32nonortho32score} S_\theta(\theta, g; z) = -(y - g(w) - {\langle \theta, x \rangle})x. \end{align}\tag{93}\]

1.0.7.2.2.3 Orthogonalized gradient oracle.

Since we perform orthogonalization on the non-orthogonal loss, we have \(\theta_\star = \theta_0\) being the same target parameter. By 18 in 1.0.2.1.2, the Neyman orthogonalized gradient oracle for this non-orthogonal loss 91 is given by \[\begin{align} \label{sim32noscore} S_{\text{no}}(\theta, g; z) = -(y - g(w) - {\langle \theta, x \rangle})(x - {\mathbb{E}}_{}\left[X\mid W=w\right]). \end{align}\tag{94}\]

1.0.7.3 Estimation Methods

Throughout the experiments, we estimate the nuisances and the orthogonalizing operator using full-batch data and stream data, respectively.

1.0.7.3.0.1 Nuisance estimation.

Note that the true nuisances for the orthogonal loss and the non-orthogonal loss are conditional expectation given \(W\). To conduct nonparametric regression, we use random Fourier feature (RFF) [64] using the kernel \(w \mapsto \exp\left(-\gamma \cdot\lVert w \rVert_2^2\right)\) to generate a randomized feature map for \(W\).

The nuisance estimation procedure for obtaining \(\hat{g}^{(m)}\) using full batch data can be summarized as

1. Fit RFF sampler with \(20\) components using \(m\) i.i.d. samples from \(P_{W\mid \lambda}\).

2. Fit Ridge regressions where the regularization parameter is set to be \(0.01/m\). Specifically,

   • For the orthogonal loss, fit two Ridge regressions using \(m\) i.i.d. samples from the joint distribution \(P_{X,W,Y\mid \lambda}\) and the fitted RFF sampler to coordinate-wisely estimate \(\mathbb{E}[X \mid W]\). With the same data, fit one Ridge regression using the fitted RFF sampler to estimate \(\mathbb{E}[Y \mid W]\).

   • For the non-orthogonal loss, fit one Ridge regression using \(m\) i.i.d. samples from the joint distribution \(P_{X,W,Y\mid \lambda}\) and the fitted RFF sampler to estimate \(\mathbb{E}[U \mid W]\).

To estimate nuisances using stream data, instead of fit a Ridge regression each time, we perform SGD for the Ridge regression loss. The procedure can be summarized as

1. Initialize RFF sampler with \(20\) components using \(n_0\) i.i.d. samples \((W_i)_{i=1}^{n_0}\) from \(P_{W\mid \lambda}\).

2. Perform SGD update once observing a mini-batch of i.i.d. samples from the joint distribution \(P_{X,W,Y\mid \lambda}\) with size \(n_g\). Specifically,

   • For the orthogonal loss, perform two SGD with the Ridge loss for \(m\) iterations to estimate \(\mathbb{E}[X\mid W]\) coordinate-wisely. With the same data perform another SGD with the Ridge loss for \(m\) iterations to estimate \(\mathbb{E}[Y\mid W]\).

   • For the non-orthogonal loss, perform one SGD with the Ridge loss for \(m\) iterations to estimate \(\mathbb{E}[U\mid W]\).

1.0.7.3.0.2 Orthogonalizing operator estimation.

To approximate the orthogonalizing operator \(\Gamma_0\), it suffices to estimate \({\mathbb{E}}_{}\left[X \mid W\right]\) by 17 . To that end, we use the same method as the nuisance estimation. The orthogonalizing operator estimation procedure for obtaining \(\hat{\Gamma}^{(k)}\) can be summarized as

1. Fit RFF sampler with \(50\) components using \(k\) i.i.d. samples \((W'_i)_{i=1}^{k}\) from \(P_{W\mid \lambda}\).

2. Fit two Ridge regressions with the regularization parameter being \(0.01/k\) using the fitted RFF sampler and another \(k\) i.i.d. samples \((X'_i, W'_i)_{i=k}^{2k}\) to coordinate-wisely estimate \(\mathbb{E}[X \mid W]\).

1.0.7.3.0.3 Target estimation.

After the estimation of nuisances and orthogonalizing operator, we perform stochastic gradient descent (SGD) to estimate \(\theta_\star\) using each of the three stochastic gradient oracles in 90 , 93 , and 94 on \(n\) i.i.d. samples drawn from the joint distribution \(P_{X,W,U,Y}\). The learning rates of all the three SGDs are fixed during the training.

1.0.7.3.1 Simulation Results

1.0.7.3.1.1 Setup.

For each nuisance estimation setting, we study three types of estimation methods for learning \(\theta_0\) established in this paper: (1) (orthogonal loss) obtain nuisance estimator \(\hat{g}^{(m)} = (\hat{g}_{Y}^{(m)}, \hat{g}_{X}^{(m)})\) of 88 and 89 and then perform SGD to obtain \(\theta^{(n)}\) using the gradient oracle 90 after plugging in \(\hat{g}^{(m)}\); (2) (non-orthogonal loss) obtain the nuisance estimator \(\hat{g}^{(m)} = \hat{\alpha}^{(m)}\) of 92 and then perform SGD to obtain \(\theta^{(n)}\) using the gradient oracle 93 after plugging in \(\hat{g}^{(m)}\); (3) (OSGD) obtain the nuisance estimator \(\hat{g}^{(m)}\) of 92 and the orthogonalizing operator estimator \(\hat{\Gamma}^{(k)}\) of 17 , and then perform SGD to obtain \(\theta^{(n)}\) using the gradient oracle 94 after plugging in \(\hat{g}^{(m)}\) and \(\hat{\Gamma}^{(k)}\). Each method is independently repeated 20 times. For nuisance estimated using stream data, we allow the procedure repeated by plugging in updated nuisance estimators and an updated operator estimator, where the nuisances get updated for 2000 iterations after every 2000 target SGD iterations.

1.0.7.3.1.2 Evaluation.

We evaluate the performance of nuisance estimators using the corresponding norms defined in 20 and 28 . Specifically, for method (1), we evaluate the nuisance estimation by \[\begin{align} \label{sim32gnorm32for32ortho} \Vert\hat{g}^{(m)} - g_0\Vert_{\mathcal{G}} = \max\left\{{\mathbb{E}}_{}\left[\lVert \hat{g}_{X}^{(m)}(W) - g_{0,X}(W) \rVert_2^4\right]^{\frac{1}{4}}, {\mathbb{E}}_{}\left[(\hat{g}_{Y}^{(m)}(W) - g_{0,Y}(W))^4\right]^{\frac{1}{4}}\right\}. \end{align}\tag{95}\] For method (2) and (3), we evaluate the nuisance estimation by \[\begin{align} \label{eq:32sim32gnorm32non-ortho} \Vert\hat{g}^{(m)} - g_0\Vert_{\mathcal{G}} = {\mathbb{E}}_{}\left[\lVert \hat{\alpha}^{(m)}(W) - \alpha_{0}(W) \rVert_2^2\right]^{\frac{1}{2}}. \end{align}\tag{96}\] We evaluate \(\hat{\Gamma}^{(k)}: g \mapsto \mathbb{E}[\hat{g}_X^{(k)}(W) g(W)]\) in method (3) using the Frobenius norm \(\Vert \hat{\Gamma}^{(k)} - \Gamma_0\Vert_\mathrm{Fro}\), which is defined as \[\begin{align} \label{sim32operator32norm} \Vert \hat{\Gamma}^{(k)} - \Gamma_0\Vert_\mathrm{Fro}= {\mathbb{E}}_{}\left[\lVert \hat{g}_X^{(k)}(W) - \hat{g}_{0,X}(W) \rVert_2^2\right]^{\frac{1}{2}}. \end{align}\tag{97}\] Finally, we evaluate the target estimation using two kinds of criterion: (a) the relative error \(\frac{\lVert \theta^{(n)} - \theta_0 \rVert_2}{\lVert \theta_0 \rVert_2}\), and (b) the risk \(L(\theta^{(n)}, g_0) - L(\theta_\star, g_0)\) where \(L(\theta, g) = \mathbb{E}[\ell(\theta, g; Z)]\). For method (1), \(\ell(\theta, g; z)\) is the orthogonal loss defined in 87 while for method (2) and (3), \(\ell(\theta, g; z)\) is the non-orthogonal loss defined in 91 .

1.0.7.3.1.3 Results using nuisances fitted on full-batch data.

We first estimate the target using prefitted nuisances and operator. The estimation errors of nuisances and the operator fitted on full-batch data are shown in 2, where all estimation converges when the sample size \(m\) increases and less samples are usually required to obtain the same error level when \(\lambda\) increases.

The performances of SGDs using prefitted nuisances and stochastic gradient oracles 90 , 93 , and 94 are shown in 3, 4, and 5, respectively. These figures suggest that when \(\lambda\) increases, i.e., the correlation between \(X\) and \(W\) increases, usually more iterations are required to have SGD converged due to the difficulty of separating the effect of \(X\) from \(W\). In addition, a well prefitted nuisance estimator would largely reduce the SGD estimation error, which aligns with 1. This improvement would be more obvious as \(\lambda\) increases. 5 also shows that either using a well estimated nuisance or a well estimated orthogonalizing operator can improve the OSGD performance, and that OSGD using both well prefitted nuisance and operator would perform nearly the same as OSGD using the true nuisance and the true operator.

1.0.7.3.1.4 Results using nuisances fitted on stream data.

We then study the interleaving the nuisance and target estimations discussed in 1.0.6.3. Here, Both the nuisance and the operator are learned using the same data stream and the results are shown in 6. Compared with 2, nuisances estimated using stream data usually has larger error and need more iterations to converge due to mini-batch, learning rate, and other tuning parameters.

The performances of SGDs by interleaving nuisance and target updates with stochastic gradient oracles 90 , 93 , and 94 are shown in 7, 8, and 9, respectively. For all the three stochastic gradients, when \(\lambda\) increases, the relative errors of the target SGD always get larger and their convergence rates become slower. There are obvious errors for SGDs using gradient oracles 90 and 93 in 7 and 8 since nuisances are not well estimated. However, OSGD performs perfectly as shown in 9, which verifies the analysis of 2 that using an estimated orthogonalizing operator would reduce the bias from nuisance estimation.

1.0.7.4 Real Data Analysis

We consider the Diabetes 130-Hospitals Dataset [65] as the real dataset example. We use six of these features as covariates, which are summarized in [tab:real32data]. We take the binary feature “change” as the input \(X \in \{0, 1\}\) and take the rest five features as the control \(W \in \mathbb{R}^5\).

max width=0.9

1.0.7.4.1 Synthetic outcomes

To evaluate the performance of our proposed methods, we use the synthetic outcome instead of a real outcome to examine the performance of our proposed methods. Using the synthetic outcome is common in causal inference; see [43][Sec. 4.1]. In this real data analysis, we generate outcome according to the following partially linear model: \[\begin{align} Y &= \tilde{\theta}_0 \cdot X + \tilde{\alpha}_0(W) + 0.5\times \epsilon \tag{98},\\ U &= \tilde{\alpha}_0(W) + 0.5 \times \xi \tag{99}, \end{align}\] where \(\tilde{\theta}_0 = -1\), \(\epsilon \sim \mathcal{N}(0,1)\) and \(\xi \sim \mathcal{N}(0,1)\) are independent noises, and \(\tilde{\alpha}_0: \mathbb{R}^5 \mapsto \mathbb{R}\) satisfies that for \(w = (w_1, \dots, w_5)\), \[\begin{align} \tilde{\alpha}_0(w) = 0.5\times \cos\left(5^{-1}\sum_{i=1}^5 w_i\right) + 0.5\times \sin\left(5^{-1}\sum_{i=1}^5 w_i\right). \end{align}\] Similar to 1.0.7.2.2, we have \(\theta_\star = \tilde{\theta}_0\) in this case.

1.0.7.4.2 Real Data Results

1.0.7.4.2.1 Setup.

We consider the same three stochastic gradient oracles as 1.0.7.2.2 and the same two nuisance estimation methods as 1.0.7.3 except that we use logistic regression on full batch data and SGD of the logistic loss on stream data for estimating \(\mathbb{E}[X\mid W]\). The setup of SGD using prefitted nuisances for this real data analysis is the same as 1.0.7.3.1. For nuisance estimated using stream data, we update nuisances for 100 iterations after every 500 target SGD iterations.

1.0.7.4.2.2 Evaluation.

Since the true nuisances \(\mathbb{E}[X\mid W]\) and \(\mathbb{E}[Y \mid W]\) are unknown, we evaluate the performance of nuisance estimation \(\hat{g}^{(m)} = (\hat{g}_{Y}^{(m)}, \hat{g}_{X}^{(m)})\) for the orthogonal loss by the mean squared error: \[\begin{align} \label{eq32mse321} \text{MSE}_1(\hat{g}^{(m)}) = \max\left\{{\mathbb{E}}_{\mathbb{P}}\left[(\hat{g}_{Y}^{(m)}(W) - Y)^2\right], {\mathbb{E}}_{\mathbb{P}}\left[(\hat{g}_{X}^{(m)}(W) - X)^2\right]\right\}. \end{align}\tag{100}\] We adopt the nuisance estimation error \(\Vert\hat{g}^{(n)} - g_0\Vert_{\mathcal{G}}\) defined in 96 as the nuisance evaluation for non-orthogonal loss due to the synthetic outcome, where now \(g_0 = \tilde{\alpha}_0\). For the operator estimation \(\hat{\Gamma}^{(m)}: g \mapsto \mathbb{E}[\hat{g}_X^{(m)}(W) g(W)]\), evaluate its performance by the mean squared error: \[\begin{align} \label{eq32mse322} \text{MSE}_2(\hat{\Gamma}^{(m)}) = {\mathbb{E}}_{\mathbb{P}}\left[(\hat{g}_X^{(m)}(W) - X)^2\right]. \end{align}\tag{101}\]

1.0.7.4.2.3 Results using nuisances fitted on full-batch data.

We first estimate the target using prefitted nuisances and operator. The estimation errors of nuisances and the operator using full-batch real data are shown in 10, which suggests that the estimation of \(\tilde{\alpha}_0\) converges to zero due to our design while there exists obvious bias for estimating the nuisance \((g_{0,X}, g_{0,Y})\) and the orthogonalizing operator \(\Gamma_0\) possibly due to model misspecification.

The performances of SGDs using prefitted nuisances and stochastic gradient oracles 90 , 93 , and 94 are shown in 11. Overall, the relative error and the risk are small when well estimated nuisances are used. In addition, both relative errors and risks become smaller when we use more samples to estimate nuisances for the orthogonal loss and the non-orthogonal loss.

1.0.7.4.2.4 Results using nuisances fitted on stream data.

We then estimate the target by interleaving the nuisance and target updates. Here, Both the nuisance and the operator are learned using the same data stream and the results are shown in 12. Compared with 10, nuisances estimated using stream data converges similarly.

The performances of SGDs by interleaving nuisance and target updates with stochastic gradient oracles 90 , 93 , and 94 are shown in 13. The figure on the left in 13 shows that the target estimation has small relative error using the estimated nuisance sequence \(\{\hat{g}^{(m)}:m \geq 1\}\). The figure in the middle suggests that there is still some bias for the target estimation while this bias is negligible. The figure on the right shows the performance of OSGD, where the relative error of OSGD using the estimated nuisance sequence is similar to OSGD using the true nuisance, which aligns with 2.

1.0.8 Extension to Variants of SGD↩︎

In this section, we discuss strategies for analyzing other variants of SGD under nuisances. In 1.0.8.1, we discuss the relationship between SGD with momentum and averaged SGD and provide a convergence analysis example of the averaged SGD. In 1.0.8.2, we discuss Adam as a generalization of SGD with momentum and the difficulties to analyze the convergence rate of Adam.

1.0.8.1 SGD with Momentum and Averaged SGD

For the gradient oracle sequence \(S^{(n)}\), SGD with momentum following the description of [66] can be expressed as \[\begin{align} m^{(n+1)} = \beta_n m^{(n)} + S^{(n)} \text{ and } \bar{\theta}^{(n+1)} = \bar{\theta}^{(n)} - \alpha_n m^{(n)}, \end{align}\] where \(\bar{\theta}^{(n)}\) is the SGD estimation sequence, \(m^{(n)}\) is the momentum sequence, and \((\alpha_n)_{n \geq 0}\) and \((\beta_n)_{n \geq 0}\) can be any positive sequence. The following example shows that the averaged SGD is a special case of SGD with momentum.

Example 5 demonstrates that the convergence rate of SGD with momentum can be analyzed in the same way as averaged SGD. While it is not the focus of this paper, we provide a convergence result of the averaged SGD based on the analysis of [67].

Before we prove 5, recall the example of non-orthogonal loss for the partially linear model in 1.0.2.1.2, where \(Z = (X, W, Y) \sim \mathbb{P}\) satisfies \[\begin{align} Y = {\langle \theta_0, X \rangle} + g_0(W) + \epsilon, \quad {\mathbb{E}}_{\mathbb{P}}\left[\epsilon \mid X, W\right] = 0. \end{align}\] The target parameter \(\theta_\star = \mathop{\mathrm{arg\,min}}_{\theta \in \Theta} {\mathbb{E}}_{\mathbb{P}}\left[\ell(\theta, g; Z)\right]\) where \(\ell\) is the non-orthogonal loss defined as \[\begin{align} \ell(\theta, g; z) = \frac{1}{2}[y - g(w) - {\langle \theta, x \rangle}]^2, \end{align}\] By 3, we have \(\theta_\star = \theta_0\). The stochastic gradient oracle for this non-orthogonal loss is \[\begin{align} S_\theta(\theta, g; z) = -X(y - g(w) - {\langle \theta, x \rangle}), \end{align}\] and the SGD iteration is defined by \(\theta^{(0)} \in \Theta\) and \[\begin{align} \label{appx32asgd:32eq32sgd} \theta^{(n)} = \theta^{(n-1)} - \eta S_\theta(\theta, \hat{g}; Z_{n-1}) = \theta^{(n-1)} + \eta X_{n-1}(Y_{n-1} - \hat{g}(W_{n-1}) - {\langle \theta, X_{n-1} \rangle}), \end{align}\tag{103}\] where \(\hat{g} \in \mathcal{G}_{r}\left(g_0\right)\) is any nuisance estimator independent of \(\{Z_i\}_{i=1}^n\). Note that 103 can be written as \[\begin{align} \theta^{(n)} - \theta_\star &= \left(I - \eta X_nX_n^\top\right)\theta^{(n-1)} + \eta X_n\left(Y_n - \hat{g}\left(W_n\right)\right) - \theta_\star\\ &=\left(I - \eta X_nX_n^\top\right)(\theta^{(n-1)}-\theta_\star) + \eta X_n\left(Y_n - \hat{g}\left(W_n\right)-X_n^\top \theta_\star\right)\\ &=\left(I - \eta X_nX_n^\top\right)(\theta^{(n-1)}-\theta_\star) + \eta X_n\epsilon_n - \eta X_n\left(\hat{g}\left(W_n\right) - g_0\left(W_n\right)\right). \end{align}\] Let \(\beta^{(n)} = \theta^{(n)} - \theta_\star\), \(r_n = \hat{g}\left(W_n\right) - g_0\left(W_n\right)\), and \[\begin{align} M_{k,j} = \left(\prod_{i=k+1}^j \left(I - \eta X_iX_i^\top\right)\right)^\top \in \mathbb{R}^{d\times d}. \end{align}\] By recursion, we have \[\begin{align} \beta^{(n)} &= \left(I - \eta X_nX_n^\top\right)\beta^{(n-1)} + \eta X_n\epsilon_n - \eta X_nr_n\\ &= M_{0,n}\beta_0 + \eta \sum_{k=1}^n M_{k,n} X_k\epsilon_k - \eta \sum_{k=1}^n M_{k,n} X_k r_k. \end{align}\] Let \(\bar \beta^{(n)} = \bar \theta^{(n)} - \theta_\star = \left(n+1\right)^{-1}\sum_{j=0}^n \beta_j\), we have \[\begin{align} \bar \beta^{(n)} &= \frac{1}{n+1}\sum_{j=0}^nM_{0,j}\beta_0 + \frac{\eta}{n+1}\sum_{k=1}^n\left(\sum_{j=k}^n M_{k,j}\right)\left(X_k\left(\epsilon_k - r_k\right) + \mathbb{E}\left[X_kr_k\right]\right) \\ &\quad - \frac{\eta}{n+1}\sum_{k=1}^n\left(\sum_{j=k}^n M_{k,j}\right)\mathbb{E}\left[X_kr_k\right]. \end{align}\] In the above formula, first two terms are usually interpreted as the bias term and the variance term under the true nuisance, respectively according to [67], and the last term can be viewed as the error term caused by the nuisance estimation.

To analyze the bias term and the variance term, we adopt the notations of [67] for matrices and operators. First, Define \(H = \mathbb{E}_{\mathbb{P}}[XX^\top]\). Let \(H_L\) (resp. \(H_R\)) be the matrix operator representing left multiplication (resp. right multiplication) by \(H\), and \(T\) be the linear operator such that for any square matrix \(M \in \mathbb{R}^{d \times d}\), \(TM = HM + MH - \eta {\mathbb{E}}_{\mathbb{P}}\left[(X^\top M X)XX^\top\right]\). Let \(\rho = \max\{\Vert I - \eta H \Vert_{\text{op}}, \Vert I - \eta T \Vert_{\text{op}}\}\) where the operator norm \(\Vert \cdot\Vert_{\text{op}}\) is defined as the largest singular value. Finally, let \(\eta_{\max}\) be the same as in 5. With definitions above, the asymptotic covariances of the bias and the variance follow directly from Theorems 1 and 2 of [67][Appendix 3].

In fact, the convergence rate of averaged SGD depends on \(\text{tr}\left(B_nB_n^\top\right)\) and \(\text{tr}\left(V_nV_n^\top\right)\). When \(\rho<1\), 18 demonstrate that \(\text{tr}\left(B_nB_n^\top\right)\) is of order \(n^{-2}\), while 19 shows that \(\text{tr}\left(V_nV_n^\top\right)\) is of order \(n^{-1}\), which is reasonable due to the randomness of the noise \(\epsilon_k - r_k\).

For the error term, we can analyze it in a similar way to the bias term. Let \(\Delta = \mathbb{E}_{\mathbb{P}}\left[X_nr_n\right]\) and \[\begin{align} E_n = \frac{\eta}{n+1}\sum_{k=1}^n\left(\sum_{j=k}^n M_{k,j}\right)\Delta. \end{align}\] Note that by Jensen’s inequality, \[\begin{align} {\mathbb{E}}_{\mathbb{P}}\left[\lVert E_n \rVert_2^2\right] &\leq {\mathbb{E}}_{\mathbb{P}}\left[\frac{\eta^2}{n}\sum_{k=1}^n\left\Vert\left(\sum_{j=k}^n M_{k,j}\right)\Delta\right\Vert_2^2 \right]\\ &= \frac{\eta^2}{n}\sum_{k=1}^n\text{tr}\left({\mathbb{E}}_{\mathbb{P}}\left[\left(\sum_{j=k}^n M_{k,j} \Delta\right)\left(\sum_{j=k}^n M_{k,j} \Delta\right)^\top\right]\right). \end{align}\] It is clear to see that the asymptotic covariance of \(\Delta_{k,n} := \sum_{j=k}^n M_{k,j} \Delta\) is of the same form as the bias term in 18. Let \(G_0 = {\mathbb{E}}_{\mathbb{P}}\left[UU^\top\right]\) and we have \[\begin{align} {\mathbb{E}}_{\mathbb{P}}\left[\Delta_{k,n} \Delta_{k,n}^\top\right] = \frac{1}{\eta^2}\left(H_L^{-1}+H_R^{-1}-\eta I\right)\left(T^{-1} G_0\right)+O\left((n-k)\rho^{n-k}\lVert G_0 \rVert_F\right). \end{align}\] Thus, the trace of the covariance summation over \(k=1, \dots, n\) satisfies \[\begin{align} \label{eq32asgd32error32term} \frac{\eta^2}{n}\sum_{k=1}^n\text{tr}\left({\mathbb{E}}_{\mathbb{P}}\left[\Delta_{k,n} \Delta_{k,n}^\top\right]\right) = \text{tr}\left((H_L^{-1}+H_R^{-1}-\eta I)(T^{-1} G_0)\right) + O\left(\frac{\eta^2}{n}\sum_{k=0}^{n-1} k\rho^{k}\lVert G_0 \rVert_F\right). \end{align}\tag{104}\]

Gathering the bias term, the variance term, and the error term, we are now ready to proof 5.

1.0.8.2 Adam

The primary updates for Adam under nuisance estimate \(\hat{g}\) are given by the following recursive equations. Below, we let \(i \in \left\{1, \ldots, d\right\}\) denote a particular dimension of the finite-dimensional parameter of interest. Following the description of [68], for the gradient oracle sequence \(S^{(n)}\) the Adam generates the target estimator \(\tilde{\theta}^{(n)}\) as below: \[\begin{align} m^{(n)}_i &= \beta_1 m^{(n-1)}_i + S_{i}^{(n)}\tag{105}\\ v^{(n)}_i &= \beta_2 v^{(n-1)}_i + \left(S_{i}^{(n)}\right)^2\tag{106}\\ \tilde{\theta}^{(n)}_i &= \tilde{\theta}^{(n-1)}_i - \eta \frac{m^{(n)}_i}{\sqrt{\epsilon + v^{(n)}_i}}, \tag{107} \end{align}\] where \(\beta_2 \in (0, 1]\), \(\beta_1 \in [0, \beta_2]\) are the momentum and variance parameters, \(m^{(n)}, v^{(n)} \in \mathbb{R}^d\) are the momentum and variance sequences, and \(\epsilon > 0\) is a numerical stability parameter. Adam differs from the SGD with momentum by adding a variance sequence \(v^{(n)}\). When \(v^{(n)}\) is chosen to be a constant, then 105 and 107 would reduce to the special case of SGD with momentum where \(\beta_n\) and \(\alpha_n\) are constant.

The analysis of Adam is often done in the case of smooth non-convex optimization, in which it is shown that the gradient of the objective tends to zero [68], [69]. Specifically, [68] consider a momentum-free Adam (\(\beta_1 = 0\)) to analyze the essential ingredients that differ from momentum: the variance pre-conditioning and element-wise updates, which suggests that under the true nuisance, i.e., \(S^{(n)} = S_\theta(\theta^{(n)}, g_0; Z_n)\) for an i.i.d. sample \(\left\{Z_i\right\}_{i=1}^n \sim \mathbb{P}^n\), the convergence result of Adam satisfies \[\begin{align} {\mathbb{E}}_{\mathbb{P}^n}\left[\lVert S_\theta(\tilde{\theta}^{(n)}, g_0) \rVert_2^2\right] \lesssim \frac{1}{\sqrt{n}} \left(1+ \log\left(\frac{1}{\epsilon}\right)\right). \end{align}\] Note that this result is not comparable to our convergence criterion (in terms of iterations), which differs non-trivially from a stationarity or function value analysis. While the convergence of Adam without nuisance has been studied in the literature, it still remains unclear that how to analysis Adam under an estimated nuisance \(\hat{g}\) and what should be the nuisance effect on the gradient norm criterion.

References↩︎

---

1. This notion of (Fisher) score differs from the “score” used in score-based generative modeling [29]. If \(\ell\) is based on a log-likelihood, then \(S_\theta\) is the gradient w.r.t. the parameter \(\theta \in \Theta\), not the input \(z \in \mathcal{Z}\).↩︎

2. The third inequality is actually \(A + M\left(B + 1 - 2b\right) < \mu\), but the constant \((1/2)\) to make the resulting bound more comparable, in that their bound can only improve over ours for the stronger inequality.↩︎

3. In the Monte Carlo settings, it often holds that \(d = k\) and \(\Gamma= \alpha I\) for some constant \(\alpha \in \mathbb{R}\). Then, \(\mathbb{E}{\langle \hat{v} - v, \Gamma(\hat{u} - u) \rangle}\) can be replaced by \(\alpha \operatorname{\boldsymbol{T}r}(\operatorname{\mathbb{C}ov}(\hat{v}, \hat{u}))\) and \(o(\lVert \Gamma \rVert_{\mathrm{op}})\) can be replaced by \(o(\alpha)\).↩︎