Demystifying SGD with Doubly Stochastic Gradients
June 03, 2024
Abstract
Optimization objectives in the form of a sum of intractable expectations are rising in importance (e.g., diffusion models, variational autoencoders, and many more), a setting also known as “finite sum with infinite data.” For these problems, a popular strategy is to employ SGD with doubly stochastic gradients (doubly SGD): the expectations are estimated using the gradient estimator of each component, while the sum is estimated by subsampling over these estimators. Despite its popularity, little is known about the convergence properties of doubly SGD, except under strong assumptions such as bounded variance. In this work, we establish the convergence of doubly SGD with independent minibatching and random reshuffling under general conditions, which encompasses dependent component gradient estimators. In particular, for dependent estimators, our analysis allows fined-grained analysis of the effect correlations. As a result, under a per-iteration computational budget of \(b \times m\), where \(b\) is the minibatch size and \(m\) is the number of Monte Carlo samples, our analysis suggests where one should invest most of the budget in general. Furthermore, we prove that random reshuffling (RR) improves the complexity dependence on the subsampling noise.
1 Introduction↩︎
Stochastic gradient descent (SGD; [1]–[4]) is the de facto standard for solving large scale optimization problems of the form of finite sums such as \[\mathop{\mathrm{minimize}}_{\vx \in \mathcal{X} \subseteq \mathbb{R}^d}\; \left\{\; F\left(\vx\right) \;\triangleq\; {\textstyle\frac{1}{n} \sum_{i=1}^n} f_i\left(\vx\right) \;\right\}. \label{eq:objective}\tag{1}\] When \(n\) is large, SGD quickly converges to low-accuracy solutions by subsampling over components \(f_1, \ldots, f_n\). The properties of SGD on the finite sum class have received an immense amount of interest [5] as it includes empirical risk minimization (ERM; [6]).
Unfortunately, for an emerging large set of problems in machine learning, we may not have direct access to the components \(f_1, \ldots, f_n\). That is, each \(f_i\) may be defined as an intractable expectation, or an “infinite sum” \[f_i(\vx) = \mathbb{E}_{\rvveta \sim \varphi} f_i\left(\vx; \rvveta\right), \label{eq:expectation95subcomponent}\tag{2}\] where we only have access to the noise distribution \(\varphi\) and the integrand \(f_i\left(\vx; \rvveta\right)\), and \(\rvveta\) is a potentially continuous and unbounded source of stochasticity; a setting [7], [8] have previously called “finite sum with infinite data.” Such problems include the training of diffusion models [9]–[11], variational autoencoders [12], [13], solving ERM under differential privacy [14], [15], and also classical problems such as variational inference [16]–[18], and variants of empirical risk minimization [8], [19]–[22]. In contrast to the finite sum setting where SGD has traditionally been applied, our problem takes the form of \[\mathop{\mathrm{minimize}}_{\vx \in \mathcal{X} \subseteq \mathbb{R}^d}\; \left\{\; F\left(\vx\right) \;\triangleq\; { {\textstyle\frac{1}{n} \sum_{i=1}^n} \mathbb{E}_{\rvveta \sim \varphi} f_i\left(\vx; \rvveta\right) } \;\right\}.\] These optimization problems are typically solved using SGD with doubly stochastic gradients (doubly SGD; coined by [17], [19]), so-called because, in addition to subsampling over \(f_{i}\), stochastic estimates of each component \(f_i\) are used.
Previous studies have relied on strong assumptions to analyze doubly stochastic gradients. For instance, [7], [8], [23] have
assumed that the variance of each component estimator is bounded by a constant, which contradicts componentwise strong convexity [24] when \(\mathcal{X} = \mathbb{R}^d\),
or that the integrand \(\nabla f_i\left(\vx; \veta\right)\), is \(L\)-Lipschitz smooth “uniformly” over \(\veta\).
That is, for any fixed \(\veta\) and \(i\), \[{\left\lVert \nabla f_i\left(\vx; \veta\right) - \nabla f_i\left(\vy; \veta\right) \right\rVert} \leq L {\left\lVert \vx - \vy \right\rVert}_2^2\] holds for all \((\vx, \vy) \in \mathcal{X}^2\). Unfortunately, this only holds for additive noise and is otherwise unrealizable when \(\rvveta\) has an unbounded support. Therefore, analyses relying on uniform smoothness obscure a lot of interesting behavior. Meanwhile, weaker assumptions such as expected smoothness (ES; [25], [26]) have shown to be realizable even for complex gradient estimators [27], [28]. Therefore, a key question is how these ES-type assumptions propagate to doubly stochastic estimators. Among these, we focus on the expected residual (ER; [29]) condition.
Furthermore, in practice, certain applications of doubly SGD share the randomness \(\rvveta\) across the batch \(\rvB\). (See 2.2 for examples.) This introduces dependence between the gradient estimate for each component such that \(\nabla f_i\left(\vx; \rvveta\right) \not{\mathpalette{\independenT}{\perp}} \nabla f_j\left(\vx; \rvveta\right)\) for \(i, j \in \rvB\). Little is known about the effect of this practice apart from some empirical results [30]. For instance, when \(m\) Monte Carlo samples of \(\rvveta\) and a minibatch of size \(b\) are used, what is the trade-off between \(m\) and \(b\)? To answer this question, we provide a theoretical analysis of doubly SGD that encompasses dependent gradient estimators.
[thm:doubly95stochastic95variance]: For doubly stochastic estimators, we establish a general variance bound of the form of \[\mathcal{O}\,\Bigg(\, \frac{\textstyle \frac{1}{n} \sum^n_{i=1} \sigma^2_{i} }{m b} + \rho \frac{ {\textstyle\left(\frac{1}{n} \sum^n_{i=1} \sigma_{i} \right)}^2 }{m} + \frac{\tau^2}{b} \,\Bigg),\] where \(\sigma_{i}^2\) is the variance of the estimator of \(\nabla f_i\), \(\rho \in [0, 1]\) is the correlation between the estimators, and \(\tau^2\) is the variance of subsampling.
[thm:general95conditions] [thm:bounded95variance]: Using the general variance bound, we show that a doubly stochastic estimator subsampling over correlated estimators satisfying the ER condition and the bounded variance (BV; 2; bounded only on the solution set) condition equally satisfies the ER and BV conditions as well. This is sufficient to guarantee the convergence of doubly SGD on convex, quasar convex, and non-convex smooth objectives.
[thm:strongly95convex95reshuffling95sgd95convergence]: Under similar assumptions, we also prove the convergence of doubly SGD with random reshuffling (doubly SGD-RR), instead of independent subsampling, on a strongly convex objective with strongly convex components.
Should I invest in (increase) \(m\) or \(b\)? When dependent gradient estimators are used, increasing \(m\) or \(b\) does not have the same impact on the gradient variance as the subsampling strategy also affects the resulting correlation between the estimators. Through [thm:expected95variance95lemma], our analysis provides insight into this effect. In particular, we reveal that reducing subsampling variance also reduces Monte Carlo variances. Therefore, for a fixed budget \(m \times b\), increasing \(b\) should always be preferred over increasing \(m\).
Random Reshuffling Improves Complexity. Our analysis of doubly SGD-RR reveals that, for strongly convex objectives, random reshuffling improves the iteration complexity of doubly SGD from \(\mathcal{O}\left( \frac{1}{\epsilon} \sigma^2_{\mathrm{mc}} + \frac{1}{\epsilon} \sigma^2_{\mathrm{sub}} \right)\) to \(\mathcal{O}\left( \frac{1}{\epsilon} \sigma^2_{\mathrm{mc}} + \frac{1}{\sqrt{\epsilon}} \sigma_{\mathrm{sub}} \right)\). Furthermore, for dependent gradient estimators, doubly SGD-RR is “super-efficient”: for a batch taking \(\Theta(m b)\) samples to compute, it achieves a \(n/b\) tighter asymptotic sample complexity compared to full-batch SGD.
2 Preliminaries↩︎
2.0.0.1 Notation
We denote random variables (RVs) in serif (e.g., \(\rvx\), \(\rvvx\), \(\rvmX\), \(\rvB\)), vectors and matrices in bold (e.g., \(\vx\), \(\rvvx\), \(\mA\), \(\rvmA\)). For a vector \(\vx\), we denote the \(\ell_2\)-norm as \({\left\lVert \vx \right\rVert}_2 \triangleq \sqrt{\left\langle \vx, \vx \right\rangle} = \sqrt{\vx^{\top}\vx}\), where \(\left\langle \vx, \vx \right\rangle = \vx^{\top}\vx\) is the inner product. Lastly, \(\rvX \mathpalette{\independenT}{\perp}\rvY\) denotes independence of \(\rvX\) and \(\rvY\).
.3em
| Symb. | Description | Ref. |
|---|---|---|
| \(F\left(\vx\right)\) | Objective function | 1 |
| \(f_i\left(\vx\right)\) | \(i\)th component of \(F\) | 1 |
| \(\nabla f_{\rvB}\left(\vx\right)\) | Minibatch subsampling estimator of \(\nabla F\) | 4 |
| \(\rvB\) | Minibatch of component indices | 3 |
| \(\pi\) | Minibatch subsampling strategy | 3 |
| \(b_{\mathrm{eff}}\) | Effective sample size of \(\pi\) | 5 |
| \(\rvvg_i\left(\vx\right)\) | Unbiased stochastic estimator of \(\nabla f_i\) | 7 |
| \(\vg_i\left(\vx; \rvveta\right)\) | Integrand of estimator \(\rvvg_i\left(\vx\right)\) | 7 |
| \(\rvvg_{\rvB}\left(\vx\right)\) | Doubly stochastic estimator of \(\nabla F\) | 8 |
| \(\mathcal{L}_{\mathrm{sub}}\) | ER constant (1) of \(\pi\) | Assu. |
| \(\mathcal{L}_i\) | ER constant (1) of \(\rvvg_i\) | Assu. |
| \(\tau^2\) | BV constant (2) of \(\pi\) | Assu. |
| \(\sigma_i^2\) | BV constant (2) of \(\rvvg_i\) | Assu. |
2.1 Stochastic Gradient Descent on Finite-Sums↩︎
Stochastic gradient descent (SGD) is an optimization algorithm that repeats the steps \[\vx_{t+1} = \Pi_{\mathcal{X}}\left( \vx_t - \gamma_t \rvvg\left(\vx_t\right) \right),\] where, \(\Pi_{\mathcal{X}}\) is a projection operator onto \(\mathcal{X}\), \({(\gamma_t)}_{i=0}^{T-1}\) is some stepsize schedule, \(\rvvg\left(\vx\right)\) is an unbiased estimate of \(\nabla F\left(\vx\right)\).
2.1.0.1 Finite-Sum Problems.
When the objective can be represented as a “finite sum” it is typical to approximate the gradients of the objective as \[\begin{align} \nabla F\left(\vx\right) = \frac{1}{n} {\textstyle\sum^n_{i=1}} \nabla f_i\left(\vx\right) = \mathbb{E}_{\rvB \sim \pi}\left[ \frac{1}{b} {\textstyle\sum_{i \in \rvB}} \nabla f_i\left(\vx\right) \right], \label{eq:finite95sum} \end{align}\tag{3}\] where \(\rvB \sim \pi\) is an index set of cardinality \({\left| \rvB \right|} = b\), or “minibatch,” formed by subsampling over the datapoint indices \(\{1, \ldots, n\}\). More formally, we are approximating \(\nabla F\) using the (minibatch) subsampling estimator \[\nabla f_{\rvB}\left(\vx\right) \triangleq \frac{1}{b} \sum_{i \in \rvB} \nabla f_i\left(\vx\right), \label{eq:minibatch95gradient}\tag{4}\] where the performance of this estimator, or equivalently, of the subsampling strategy \(\pi\), can be quantified by its variance \[\mathrm{tr}\mathbb{V}\left[ \nabla f_{\rvB}\left(\vx\right) \right] = \frac{1}{b_{\mathrm{eff}}} \underbrace{\textstyle \frac{1}{n} \sum^n_{i=1} {\left\lVert \nabla f_i\left(\vx\right) - \nabla F\left(\vx\right) \right\rVert}_2^2, }_{\text{(unit) subsampling variance}} \label{eq:ess}\tag{5}\] where we say \(b_{\mathrm{eff}}\) is the “effective sample size” of \(\pi\). For instance, independent subsampling achieves \(b_{\mathrm{eff}} = b\), and sampling without replacement, also known as “\(b\)-nice sampling” [29], [31], [32], achieves \(b_{\mathrm{eff}} = \frac{(n-1) b}{n - b}\) ([thm:variancewithoutreplacement]).
2.2 Doubly Stochastic Gradients↩︎
For problems where the components are defined as intractable expectations as in 2 , we have to rely on an additional Monte Carlo approximation step such as \[\begin{align} {2} \nabla F\left(\vx\right) &= \frac{1}{n} \sum^n_{i=1} \nabla f_i\left(\vx\right) = \mathbb{E}_{\rvB \sim \pi}\left[ \frac{1}{b} \sum_{i \in \rvB} \mathbb{E}_{\rvveta \sim \varphi}\left[ \nabla f_i\left(\vx; \rvveta\right) \right] \right] \nonumber \\ &= \mathbb{E}_{\rvB \sim \pi,\;\; \rvveta_{j} \sim \varphi}\left[ \frac{1}{m b} \sum_{i \in \rvB} \sum_{j=1}^m \nabla f_i\left(\vx; \rvveta_{j}\right) \right], \label{eq:doubly95stochastic95full} \end{align}\tag{6}\] where \(\rvveta_{j} \sim \varphi\) are \(m\) independently and identically distributed (i.i.d.) Monte Carlo samples from \(\varphi\).
2.2.0.1 Doubly Stochastic Gradient
Consider an unbiased estimator of the component gradient \(\nabla f_i\) such that \[\mathbb{E} \rvvg_{i}\left(\vx\right) = \mathbb{E}_{\rvveta \sim \varphi} \vg_i\left(\vx; \rvveta\right) = \nabla f_i\left(\vx\right), \label{eq:monte95carlo95gradient}\tag{7}\] where \(\vg_i\left(\vx; \rvveta\right)\) is the measurable integrand. Using these, we can estimate \(\nabla F\) through the doubly stochastic gradient estimator \[\rvvg_{\rvB}\left(\vx\right) \triangleq \frac{1}{b} \sum_{i \in \rvB} \rvvg_{i}\left(\vx\right), \label{eq:doubly95stochastic95gradient}\tag{8}\] We separately define the integrand \(\vg\left(\vx; \veta\right)\) since, in practice, a variety of unbiased estimators of \(\nabla f_i\) can be obtained by appropriately defining the integrand \(\vg_i\). For example, one can form the \(m\)-sample “naive” Monte Carlo estimator by setting \[\vg_i\left(\vx; \rvveta\right) = {\textstyle\frac{1}{m} \sum^m_{j=1}} \nabla f_i\left(\vx; \rvveta_{j}\right),\] where \(\rvveta = [\rvveta_1, \ldots, \rvveta_m] \sim \varphi^{\otimes m}\).
2.2.0.2 Dependent Component Gradient Estimators.
Notice that, in 6 , the subcomponents in the batch share the Monte Carlo samples, which may occur in practice. This means \(\rvvg_i\) and \(\rvvg_j\) in the same batch are dependent and, in the worst case, positively correlated, which complicates the analysis. While it is possible to make the estimators independent by sampling \(m\) unique Monte Carlo samples for each component (\(m b\) Monte Carlo samples in total) as highlighted by [30], it is common to use dependent estimators for various practical reasons:
ERM with Randomized Smoothing: In the ERM context, recent works have studied the generalization benefits of randomly perturbing the model weights before computing the gradient [21], [22]. When subsampling is used, perturbing the weights independently for each datapoint is computationally inefficient. Therefore, the perturbation is shared across the batch, creating dependence.
Black-Box Variational inference [17], [18]: Here, each component can be decomposed as \[f_i\left(\vx; \rvveta\right) = \ell_i\left(\vx; \rvveta\right) + r\left(\vx; \rvveta\right),\] where \(\ell_i\) is the log likelihood and \(r\) is the log-density of the prior. By sharing \({(\rvveta_j)}_{j=1}^m\), \(r\) only needs to be evaluated \(m\) times. To create independent estimators, it needs to be evaluated \(m b\) times instead, but \(r\) can be expensive to compute.
Random feature kernel regression with doubly SGD [19]: The features are shared across the batch 1. This reduces the peak memory requirement from \(b m d_{\veta}\), where \(d_{\veta}\) is the size of the random features, to \(m d_{\veta}\).
One of the analysis goals of this work is to characterize the effect of dependence in the context of SGD.
2.3 Technical Assumptions on Gradient Estimators↩︎
To establish convergence of SGD, contemporary analyses use the “variance transfer” strategy [24]–[26], [29], [33]. That is, by assuming the gradient noise satisfies some condition resembling smoothness, it is possible to bound the gradient noise on some arbitrary point \(\vx\) by the gradient variance on the solution set \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\).
2.3.0.1 ER Condition.
In this work, we will use the expected residual (ER) condition by [34]:
Definition 1 (Expected Residual; ER). A gradient estimator \(\rvvg\) of \(F : \mathcal{X} \to \mathbb{R}\) is said to satisfy \(\mathrm{ER}\left(\mathcal{L}\right)\) if \[\mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right] \leq 2 \mathcal{L} \left( F(\vx) - F(\vx_*) \right),\] for some \(0 < \mathcal{L} < \infty\) and all \(\vx\in \mathcal{X}\) and all \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\).
When \(f\) is convex, a weaker form can be used: We will also consider the convex variant of the ER condition that uses the Bregman divergence defined as \[{\textstyle \mathrm{D}_{\phi}\left(\vy, \vx\right) \triangleq \phi\left(\vy\right) - \phi\left(\vx\right) - \left\langle \nabla \phi\left(\vx\right), \vy - \vx \right\rangle, }\] \(\forall(\vx, \vy) \in \mathcal{X}^2\), where \(\phi : \mathcal{X} \to \mathbb{R}\) is a convex function.
2.3.0.2 Why the ER condition?
A way to think about the ER condition is that it corresponds to the “variance form” equivalent of the expected smoothness (ES) condition by [26] defined as \[\begin{align} &\mathbb{E}{\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right\rVert}_2^2 \leq 2 \mathcal{L} \left( F\left(\vx\right) - F\left(\vx_*\right) \right), &\quad \text{(ES)} \end{align}\] but is slightly weaker, as shown by [34]. The main advantage of the ER condition is that, due to the properties of the variance, it composes more easily:
Let \(\rvvg\) satisfy \(\mathrm{ER}\left(\mathcal{L}\right)\). Then, the \(m\)-sample i.i.d. average of \(\rvvg\) satisfy \(\mathrm{ER}\left(\frac{\mathcal{L}}{m}\right)\).
2.3.0.3 BV Condition.
From the ER property, the gradient variance on any point \(\vx \in \mathcal{X}\) can be bounded by the variance on the solution set as long as the following holds:
Definition 2 (Bounded Gradient Variance). A gradient estimator \(\rvvg\) of \(F : \mathcal{X} \to \mathbb{R}\) satisfies \(\mathrm{BV}\left(\sigma^2\right)\) if \[\mathrm{tr} \mathbb{V}\left[ \rvvg\left(\vx_*\right) \right] \leq \sigma^2\] for some \(\sigma^2 < \infty\) and all \(\vx_* \in \mathop{\mathrm{arg\,max}}_{\vx \in \mathcal{X}} F\left(\vx\right)\).
2.4 Convergence Guarantees for SGD↩︎
2.4.0.1 Sufficiency of ER and BV.
From the ER and BV conditions, other popular conditions such as ES [26] and ABC [35] can be established with minimal additional assumptions. As a result, we retrieve the previous convergence results on SGD established for various objective function classes:
(Note: quasar convexity is strictly weaker than convexity [36]; PL: Polyak-Łojasiewicz.) (See also the comprehensive treatment by [37].) Therefore, ER and BV are sufficient conditions for SGD to converge on problem classes typically considered in SGD convergence analysis.
In this work, we will specifically focus on smooth and strongly convex objectives:
Assumption 1. There exists some \(\mu, L\) satisfying \(0 < \mu \leq L < \infty\) suc that the objective function \(F : \mathcal{X} \to \mathbb{R}\) is \(\mu\)-strongly convex and \(L\)-smooth as \[\begin{align} F\left(\vy\right) - F\left(\vx\right) &\geq \left\langle \nabla F\left(\vx\right), \vy - \vx \right\rangle + \frac{\mu}{2} {\left\lVert \vx - \vy \right\rVert}_2^2 \\ F\left(\vy\right) - F\left(\vx\right) &\leq \left\langle \nabla F\left(\vx\right), \vy - \vx \right\rangle + \frac{L}{2} {\left\lVert \vx - \vy \right\rVert}_2^2 \end{align}\] hold for all \((\vx, \vy) \in \mathcal{X}^2\).
Also, we will occasionally assume that \(F\) is comprised of a finite sum of convex and smooth components:
Assumption 2. The objective function \(F : \mathcal{X} \to \mathbb{R}\) is a finite sum as \(F = \frac{1}{n} \left(f_1 + \ldots + f_n\right)\), where each component is \(L_i\)-smooth and convex such that \[{\left\lVert \nabla f_i\left(\vx\right) - \nabla f_i\left(\vy\right) \right\rVert}_2^2 \leq 2 L_i \, \mathrm{D}_{f_i}\left(\vx, \vy\right)\] holds for all \((\vx, \vy) \in \mathcal{X}^2\).
Note that 2 alone already implies that \(F\) is convex and \(L_{\mathrm{max}}\)-smooth with \(L_{\mathrm{max}} = \max\left\{L_1, \ldots, L_n\right\}\).
2.4.0.2 Why focus on strongly convex functions?
We focus on strongly convex objectives as the effect of stochasticity is the most detrimental: in the deterministic setting, one only needs \(\mathcal{O}\left(\log\left(\frac{1}{\epsilon}\right)\right)\) iterations to achieve an \(\epsilon\)-accurate solution. But with SGD, one actually needs \(\mathcal{O}\left(\frac{1}{\epsilon}\right)\) iterations due to noise. As such, we can observe a clear contrast between the effect of optimization and noise in this setting.
With that said, for completeness, we provide full proof of convergence on strongly convex-smooth objectives:
lemma Let \(F:\mathcal{X} \to \mathbb{R}\) be \(L\)-smooth function. Then, the expected squared norm of a gradient estimator \(\rvvg\) satisfying both \(\mathrm{ER}\left(\mathcal{L}\right)\) and \(\mathrm{BV}\left(\sigma^2\right)\) is bounded as \[\mathbb{E} {\left\lVert \rvvg\left(\vx\right) \right\rVert}_2^2 \leq 4 \left(\mathcal{L} + L\right) \left(F\left(\vx\right) - F\left(\vx_*\right)\right) + 2 \sigma^2,\] for any \(\vx \in \mathcal{X}\) and \(\vx_* \in \mathop{\mathrm{arg\,max}}_{\vx \in \mathcal{X}} F\left(\vx\right)\).
The proof is a minor modification of Lemma 2.4 by [29] and Lemma 3.2 by [34].
By applying the bound \({(a + b)}^2 \leq 2 a^2 + 2 b^2\), we can “transfer” the variance on \(\vx\) to the variance of \(\vx_*\). That is, \[\begin{align} \mathbb{E} {\left\lVert \rvvg\left(\vx\right) \right\rVert}_2^2 &= \mathbb{E} {\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) + \rvvg\left(\vx_*\right) \right\rVert}_2^2 \\ &\;\leq 2 \, \underbrace{ \mathbb{E} {\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right\rVert}_2^2 }_{V_1} + 2 \, \underbrace{ \mathbb{E} {\left\lVert \rvvg\left(\vx_*\right) \right\rVert}_2^2 }_{V_2} \end{align}\]
The key is to bound \(V_1\). It is typical to do this using expected-smoothness-type assumptions such as the ER assumption. That is, \[\begin{align} V_1 &= \mathbb{E} {\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right\rVert}_2^2 \\ &= \mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right] + \left( \nabla F \left(\vx\right) - \nabla F\left(\vx_*\right) \right), \shortintertext{from the \(L\)-smoothness of \(F\),} &\leq \mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right] + 2 L \left(F\left(\vx\right) - F\left(\vx_*\right)\right), \shortintertext{and the ER condition,} &\leq 2 \mathcal{L} \left( F\left(\vx\right) - F\left(\vx_*\right) \right) + 2 L \left(F\left(\vx\right) - F\left(\vx_*\right)\right) \\ &= 2 \left(L + \mathcal{L}\right) \big( F\left(\vx\right) - F\left(\vx_*\right) \big). \end{align}\] Finally, \(V_2\) immediately follows from the BV condition as \[V_2 = \mathbb{E} {\left\lVert \rvvg\left(\vx_*\right) \right\rVert}_2^2 \leq \sigma^2.\]
lemma Let the objective function \(F\) satisfy 1 and the gradient estimator \(\rvvg\) be unbiased and satisfy both \(\mathrm{ER}\left(\mathcal{L}\right)\) and \(\mathrm{BV}\left(\sigma^2\right)\). Then, the last iterate of SGD guarantees \[\mathbb{E}\left[ {\left\lVert \vx_T - \vx^* \right\rVert}_2^2 \right] \leq {\left(1 - \mu \gamma\right)}^{T} {\left\lVert \vx_0 - \vx_* \right\rVert}_2^2 + \frac{2 \sigma^2}{\mu} \gamma\] where \(\vx_* = \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\) is the global optimum.
Firstly, we have \[\begin{align} &{\left\lVert \vx_{t+1} - \vx_* \right\rVert}_2^2 \\ &\;= {\left\lVert \Pi_{\mathcal{X}}\left( \vx_{t} - \gamma \rvvg\left(\vx_t\right) \right) - \Pi\left(\vx_*\right) \right\rVert}_2^2, \shortintertext{and since the projection onto a convex set under a Euclidean metric is non-expansive,} &\;\leq {\left\lVert \vx_{t} - \gamma \rvvg\left(\vx_t\right) - \vx_* \right\rVert}_2^2 \\ &\;= {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 - 2 \gamma \left\langle \rvvg\left(\vx_t\right) , \vx_{t} - \vx_* \right\rangle + \gamma^2 {\left\lVert \rvvg\left(\vx_t\right) \right\rVert}_2^2. \end{align}\] Denoting the \(\sigma\)-algebra formed by the randomness and the iterates up to the \(t\)th iteration as \(\mathcal{F}_t\) such that \({(\mathcal{F}_t)}_{t \geq 1}\) forms a filtration, the conditional expectation is \[\begin{align} &\mathbb{E}\left[ {\left\lVert \vx_{t+1} - \vx_* \right\rVert}_2^2 \mid \mathcal{F}_t \right] \\ &\;= {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 - 2 \gamma \left\langle \mathbb{E}\left[ \rvvg\left(\vx_t\right) \mid \mathcal{F}_t \right], \vx_{t} - \vx_* \right\rangle \\ &\qquad+ \gamma^2 \mathbb{E}\left[ {\left\lVert \rvvg\left(\vx_t\right) \right\rVert}_2^2 \mid \mathcal{F}_t \right]. \\ &\;= {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 - 2 \gamma \left\langle \nabla F \left(\vx_t\right), \vx_{t} - \vx_* \right\rangle \\ &\qquad+ \gamma^2 \mathbb{E}\left[ {\left\lVert \rvvg\left(\vx_t\right) \right\rVert}_2^2 \mid \mathcal{F}_t \right], \shortintertext{applying the \(\mu\)-strong convexity of \(F\),} &\;\leq {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 - 2 \gamma \left(F\left(\vx_t\right) - F\left(\vx_*\right) + \frac{\mu}{2} {\left\lVert \vx_t - \vx_* \right\rVert}_2^2\right) \\ &\qquad+ \gamma^2 \mathbb{E}\left[ {\left\lVert \rvvg\left(\vx_t\right) \right\rVert}_2^2 \mid \mathcal{F}_t \right] \\ &\;= \left(1 - \gamma \mu\right) {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 - 2 \gamma \left(F\left(\vx_t\right) - F\left(\vx_*\right)\right) \\ &\qquad+ \gamma^2 \mathbb{E}\left[ {\left\lVert \rvvg\left(\vx_t\right) \right\rVert}_2^2 \mid \mathcal{F}_t \right] \end{align}\]
From [thm:expected95residual95gradient95variance95bound], we have \[\begin{align} \mathbb{E}\left[ {\left\lVert \rvvg\left(\vx_t\right) \right\rVert}_2^2 \mid \mathcal{F}_t \right] \leq \left(4 \left(\mathcal{L}+L\right) \left(F\left(\vx_t\right) - F\left(\vx_*\right) \right) + 2 \sigma^2 \right). \end{align}\] Therefore, \[\begin{align} &\mathbb{E}\left[ {\left\lVert \vx_{t+1} - \vx_* \right\rVert}_2^2 \mid \mathcal{F}_t \right] \\ &\leq \left(1 - \gamma \mu\right) {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 - 2 \gamma \left(F\left(\vx_t\right) - F\left(\vx_*\right)\right) \\ &\quad+ \gamma^2 \left(4 \left(\mathcal{L}+L\right) \left(F\left(\vx_t\right) - F\left(\vx_*\right) \right) + 2 \sigma^2 \right) \\ &= \left(1 - \gamma \mu\right) {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 \\ &\quad- 2 \gamma \left(1 - 2 \gamma \left(\mathcal{L} + L\right)\right) \left(F\left(\vx_t\right) - F\left(\vx_*\right)\right) + 2 \gamma^2 \sigma^2, \shortintertext{and with a small-enough stepsize satisfying \(\gamma < \frac{1}{2 \left(\mathcal{L} + L\right)}\), we can guarantee a partial contraction as} &\leq \left(1 - \gamma \mu\right) {\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 + 2 \gamma^2 \sigma^2. \end{align}\] Note that the coefficient \(1 - \gamma \mu\) is guaranteed to be strictly smaller than 1 since \(\mu \leq L\), which means that we indeed have a partial contraction.
Now, taking full expectation, we have \[\mathbb{E}{\left\lVert \vx_{t+1} - \vx_* \right\rVert}_2^2 \leq \left(1 - \gamma \mu\right) \mathbb{E}{\left\lVert \vx_{t} - \vx_* \right\rVert}_2^2 + 2 \gamma^2 \sigma^2.\] Unrolling the recursion from \(0\) to \(T-1\), we have \[\begin{align} \mathbb{E}{\left\lVert \vx_{T} - \vx_* \right\rVert}_2^2 &\leq {\left(1 - \gamma \mu\right)}^T \mathbb{E}{\left\lVert \vx_{0} - \vx_* \right\rVert}_2^2 \\ &\qquad+ 2 \gamma^2 \sigma^2 \sum^{T-1}_{t=0} {\left(1 - \gamma \mu\right)}^{t}. \\ &\leq {\left(1 - \gamma \mu\right)}^T \mathbb{E}{\left\lVert \vx_{0} - \vx_* \right\rVert}_2^2 + \frac{2 \sigma^2}{\mu} \gamma. \end{align}\] where the last inequality follows from the asymptotic bound on geometric sums.
lemma Let the objective \(F\) satisfy 1 and the gradient estimator \(\rvvg\) satisfy \(\mathrm{ER}\left(\mathcal{L}\right)\) and \(\mathrm{BV}\left(\sigma^2\right)\). Then, the last iterate of SGD is \(\epsilon\)-close to the global optimum \(\vx_* = \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\) such that \(\mathbb{E}{\left\lVert \vx_T - \vx_* \right\rVert}_2^2 \leq \epsilon\) if \[T \geq 2 \max\left(\frac{\sigma^2}{\mu^2} \frac{1}{\epsilon}, \frac{\mathcal{L} + L}{\mu}\right) \log\left(2 {\left\lVert \vx_0 - \vx_* \right\rVert}_2^2 \frac{1}{\epsilon} \right)\] and the fixed stepsize \[\gamma = \min\left( \frac{\epsilon \mu}{2 \sigma^2}, \frac{1}{2 \left( \mathcal{L} + L \right)} \right).\]
We can apply [thm:geometric95complexity] to the result of [thm:strongly95convex95sgd95convergence] with the constants \[r_0 = {\left\lVert \vx_0 - \vx_* \right\rVert}_2^2, \quad B = \frac{2 \sigma^2}{\mu}, \;\text{and}\;\; C = 2 \left( \mathcal{L} + L \right).\] Then, we can guarantee an \(\epsilon\)-accurate solution with the stepsize \[\gamma = \min\left( \frac{\epsilon \mu}{2 \sigma^2}, \frac{1}{2 \left( \mathcal{L} + L \right)} \right)\] and a number of iterations of at least \[\begin{align} T &\geq \frac{1}{\mu} \max\left( \frac{2 \sigma^2}{\mu}, 2 \left(\mathcal{L} + L\right) \right) \log\left(2 {\left\lVert \vx_0 - \vx_* \right\rVert}_2^2 \frac{1}{\epsilon} \right) \\ &= 2 \max\left( \frac{\sigma^2}{\mu^2}, \frac{\mathcal{L} + L}{\mu} \right) \log\left(2 {\left\lVert \vx_0 - \vx_* \right\rVert}_2^2 \frac{1}{\epsilon} \right). \end{align}\]
Note that our complexity guarantee is only \(\mathcal{O}(\frac{1}{\epsilon} \log\left(\frac{1}{\epsilon}\right))\) due to the use of a fixed stepsize. It is also possible to establish a \(\mathcal{O}(\frac{1}{\epsilon})\) guarantee using decreasing stepsize schedules proposed by [29], [38]. In practice, these schedules are rarely used, and the resulting complexity guarantees are less clear than with fixed stepsizes. Therefore, we will stay on fixed stepsizes.
3 Main Results↩︎
3.1 Doubly Stochastic Gradients↩︎
r0.35 .3em
| §3.1.1 | §3.1.2 | |
|---|---|---|
| \(\rvvx_i\) | \(\leftrightarrow\) | \(\rvvg_i\) |
| \(\rvvx_{\rvB}\) | \(\leftrightarrow\) | \(\rvvg_{\rvB}\) |
| \(\bar{\vx}_i\) | \(\leftrightarrow\) | \(\nabla f_i\) |
| \(\bar{\vx}\) | \(\leftrightarrow\) | \(\nabla F\) |
First, while taming notational complexity, we will prove a general result that holds for combining unbiased but potentially correlated estimators through subsampling. All of the later results on SGD will fall out as special cases following the correspondence in 2.
3.1.1 General Variance Bound↩︎
3.1.1.1 Theoretical Setup.
Consider the problem of estimating the population mean \(\bar{\vx} = \frac{1}{n} \sum_{i=1}^n \bar{\vx}_i\) with a collection of RVs \(\rvvx_{1}, \ldots, \rvvx_{n}\), each an unbiased estimator of the component \(\bar{\vx}_i = \mathbb{E}\rvvx_{i}\). Then, any subsampled ensemble \[\rvvx_{\rvB} \triangleq \frac{1}{b} \sum_{i \in \rvB} \rvvx_i \quad\text{with}\quad \rvB \sim \pi,\] where \(\pi\) is an unbiased subsampling strategy with an effective sample size of \(b_{\mathrm{eff}}\), is also an unbiased estimator of \(\bar{\vx}\). The goal is to analyze how the variance of the component estimators \(\mathrm{tr}\mathbb{V}\rvvx_i\) for \(i = 1, \ldots, n\) and the variance of \(\pi\) affect the variance of \(\rvvx_{\rvB}\). The following condition characterizes the correlation between the component estimators:
Assumption 3. The component estimators \(\rvvx_1, \ldots, \rvvx_n\) have finite variance \(\mathrm{tr}\mathbb{V}\rvvx_i < \infty\) for all \(i = 1, \ldots, n\) and, there exists some \(\rho \in [0, 1]\) for all \(i \neq j\) such that \[\mathrm{tr}\, \mathrm{Cov}\left( \rvvx_i, \rvvx_j \right) \leq \rho \, \sqrt{\mathrm{tr}\mathbb{V}\rvvx_i} \sqrt{\mathrm{tr}\mathbb{V}\rvvx_j}.\]
3 always holds with \(\rho = 1\) as a basic consequence of the Cauchy-Schwarz inequality.
For a collection of mutually independent estimators \(\rvvx_1, \ldots, \rvvx_n\) such that \(\rvvx_i \mathpalette{\independenT}{\perp}\rvvx_j\) for all \(i \neq j\), 3 holds with \(\rho = 0\).
The equality in 3 holds with \(\rho = 0\) for independent estimators, while it holds with \(\rho = 1\) when they are perfectly positively correlated such that, for all \(i \neq j\), there exists some constant \(\alpha_{ij} \geq 0\) such that \(\rvvx_i = \alpha_{i,j} \rvvx_j\)
lemma Consider a finite population of \(n\) vector-variate random variables \(\vx_1, \ldots, \vx_n\). Then, the variance of the average of \(b\) samples chosen without replacement is \[\mathrm{tr}\mathbb{V}\left[ \frac{1}{b} \sum_{i=1}^b \rvvx_{\rvB_i} \right] = \frac{n - b}{b \left(n - 1\right)} \sigma^2,\] where \(\rvB = \{\rvB_1, \ldots, \rvB_b\}\) is the collection of random indices of the samples and \(\sigma^2\) is the variance of independently choosing a single sample.
From the variance of the sum of random variables, we have \[\begin{align} {2} \mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^b \rvvx_{\rvB_i} \right] &= \sum_{i=1}^b \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB_i} \right] + \sum_{i=1}^b \sum_{i \neq j}^b \mathrm{Cov}\left( \rvvx_{\rvB_i}, \rvvx_{\rvB_j} \right), \nonumber \shortintertext{and noticing that the covariance is independent of the index in the batch,} &= b \, \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB_i} \right] + b (b-1) C, \label{eq:sampling95wihout95replacement95eq1} \end{align}\tag{9}\] where \(C = \mathrm{Cov}\left(\rvvx_{\rvB_i}, \rvvx_{\rvB_j}\right)\). Using the fact that the variance is 0 for \(b = n\), we can solve for \(C\) such that \[C = - \frac{1}{n-1}\, \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB_i} \right],\] which is negative, and a negative correlation is always great. Plugging this expression to 9 , we have \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^b \rvvx_{\rvB_i} \right] &= b \, \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB_i} \right] - b (b-1) \frac{1}{n-1}\, \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB_i} \right], \\ &= b \left(1 - \frac{b-1}{n-1} \right) \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB_i} \right] \\ &= b\left(\frac{n - b}{n - 1}\right) \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB_i} \right]. \end{align}\] Dividing both sides by \(b^2\) yields the result.
lemma Let \(\rvvx_1, \ldots, \rvvx_n\) be vector-variate random variables. Then, the variance of the sum is upper-bounded as \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^n \rvvx_i \right] &\leq {\left( {\textstyle\sum^n_{i=1} \sqrt{ \mathrm{tr}\mathbb{V}\left[ \rvvx_i \right] } } \right)}^2 \tag{10} \\ &\leq n {\textstyle\sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]}. \tag{11} \end{align}\] The equality in 10 holds if and only if \(\rvvx_i\) and \(\rvvx_j\) are constant multiples such that there exists some \(\alpha_{ij} \geq 0\) such that \[\rvvx_i = \alpha_{ij} \rvvx_j\] for all \(i, j\).
The variance of a sum is \[\begin{align} {2} \mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^n \rvvx_i \right] = \sum_{i=1}^n \sum_{j=1}^n \mathrm{tr}\, \mathrm{Cov}\left( \rvvx_i, \rvvx_j \right). \end{align}\] From the Cauchy-Schwarz inequality for expectations, \[\begin{align} \mathrm{tr}\, \mathrm{Cov}\left( \rvvx_i, \rvvx_j \right) &= \mathbb{E}{\left(\rvvx_i - \mathbb{E}\rvvx_i\right)}^{\top} {\left(\rvvx_j - \mathbb{E}\rvvx_j\right)} \\ &\leq \mathbb{E}{\left\lVert \rvvx_i - \mathbb{E}\rvvx_i \right\rVert}_2 \mathbb{E}{\lVert\rvvx_j - \mathbb{E}\rvvx_j \rVert}_2 \\ &= \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_j \right]}. \end{align}\] This implies \[\begin{align} {2} \mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^n \rvvx_i \right] &= \sum_{i=1}^n \sum_{j=1}^n \mathrm{tr}\,\mathrm{Cov}\left( \rvvx_i, \rvvx_j \right) \nonumber \\ &\leq \sum_{i=1}^n \sum_{j=1}^n \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_j \right]} \nonumber \\ &= {\left( \sum_{i=1}^n \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]} \right)}^2. \label{eq:sumofvariance95eq3} \end{align}\tag{12}\] The equality statement comes from the property of the Cauchy-Schwarz inequality. Lastly, 11 follows from additionally applying Jensen’s inequality as \[\begin{align} {2} {\left( \sum_{i=1}^n \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]} \right)}^2 &= n^2 {\left( \frac{1}{n} \sum_{i=1}^n \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]} \right)}^2 \\ &\leq n^2 \frac{1}{n} \sum_{i=1}^n {\left( \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]} \right)}^2 \\ &= n \sum_{i=1}^n \mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]. \end{align}\] An equivalent proof strategy is to expand the quadratic in 12 and apply the arithmetic mean-geometric mean inequality to the cross terms.
lemma[Lemma A.2; [37]] For a recurrence relation given as \[r_{T} \leq {\left(1 - \gamma \mu\right)}^{T} r_{0} + B \gamma,\] for some constant \(0 < \gamma < 1/C\), \[r_T \leq \epsilon\] can be guaranteed by setting \[\begin{align} \gamma &= \min\left(\frac{\epsilon}{2 B}, \frac{1}{C} \right) \text{ and } \\ T &\geq \frac{1}{\mu} \max\left(2 B \, \frac{1}{\epsilon}, C\right) \log\left(2 \frac{r_0}{\epsilon}\right), \end{align}\] where \(\mu, B > 0\) and \(0 < C < \mu\) are some finite constants.
First, notice that the recurrence \[r_{T} \leq \underbrace{{\left(1 - \gamma \mu\right)}^{T} r_{0}}_{\text{bias}} + \underbrace{ B \gamma }_{\text{variance}} ,\] is a sum of monotonically increasing (variance) and decreasing (bias) terms with respect to \(\gamma\). Therefore, the bound is minimized when both terms are equal. This implies that \(r_t \leq \epsilon\) can be achieved by solving for \[{\left(1 - \gamma \mu\right)}^{T} r_{0} \leq \frac{\epsilon}{2} \quad\text{and}\quad B \gamma \leq \frac{\epsilon}{2}\] First, for the variance term, \[\begin{align} {3} B \gamma &\leq \frac{\epsilon}{2} \quad\Leftrightarrow\quad \gamma \leq \frac{\epsilon}{2 B}. \end{align}\]
For the bias term, as long as \(\gamma < \frac{1}{\mu}\), \[\begin{align} {3} & & {\left(1 - \gamma \mu\right)}^{T} r_{0} &\leq \frac{\epsilon}{2} \\ &\Leftrightarrow&\qquad T \log{\left(1 - \gamma \mu\right)} &\leq \log \frac{\epsilon}{2 r_{0}} \\ &\Leftrightarrow&\qquad T &\leq \frac{\log \frac{\epsilon}{2 r_{0}}}{\log{\left(1 - \gamma \mu\right)}} \\ &\Leftrightarrow&\qquad T &\geq \frac{\log \frac{2 r_{0}}{\epsilon}}{ \log{\left(1/\left(1 - \gamma \mu\right)\right)}} \end{align}\] Furthermore, using the bound \(\log 1/x \geq 1 - x\) for \(0 < x < 1\), we can achieve the guarantee with \[T \geq \frac{1}{\gamma \mu} \log\left( \frac{2 r_{0}}{\epsilon} \right).\] Therefore, \(1/\gamma\) determines the iteration complexity. Plugging in the minimum over the constraints on \(\gamma\) yields the iteration complexity.
lemma For a recurrence relation given as \[r_{T} \leq {\left(1 - \gamma \mu\right)}^{T} r_{0} + A \gamma^{2} + B \gamma,\] for some constant \(0 < \gamma < 1/C\), \[r_T \leq \epsilon\] can be guaranteed by setting \[\begin{align} \gamma &= \min\left(\frac{-B + \sqrt{B^2 + 2A\epsilon}}{2 A}, \frac{1}{C} \right) \text{ and } \\ T &\geq \frac{1}{\mu} \max\left( 2 B \, \frac{1}{\epsilon} + \sqrt{2 A}\,\frac{1}{\sqrt{\epsilon}}, C \right) \log\left(2 \frac{r_0}{\epsilon}\right), \end{align}\] where \(\mu, A, B > 0\) and \(0 < C < \mu\) are some finite constants.
This theorem is a generalization of Lemma A.2 by [37]. First, notice that the recurrence \[r_{T} \leq \underbrace{{\left(1 - \gamma \mu\right)}^{T} r_{0}}_{\text{bias}} + \underbrace{ A \gamma^{2} + B \gamma }_{\text{variance}} ,\] is a sum of monotonically increasing (variance) and decreasing (bias) terms with respect to \(\gamma\). Therefore, the bound is minimized when both terms are equal. This implies that \(r_t \leq \epsilon\) can be achieved by solving for \[{\left(1 - \gamma \mu\right)}^{T} r_{0} \leq \frac{\epsilon}{2} \quad\text{and}\quad A \gamma^{2} + B \gamma \leq \frac{\epsilon}{2}\] First, for the variance term, \[\begin{align} {3} & & A \gamma^{2} + B \gamma &\leq \frac{\epsilon}{2} \\ &\Leftrightarrow&\qquad A \gamma^{2} + B \gamma - \frac{\epsilon}{2} &\leq 0 \end{align}\] The solution to this equation is given by the positive solution of the quadratic equation as \[0 < \gamma \leq \frac{- B + \sqrt{ B^2 + 2 A \epsilon }}{2 A}.\]
For the bias term, as long as \(\gamma < \frac{1}{\mu}\), the solution is identical to [thm:geometric95complexity]. Therefore, \[T \geq \frac{1}{\gamma \mu} \log\left( \frac{2 r_{0}}{\epsilon} \right) \label{eq:reshuffle95recurrence95t95lowerbound}\tag{13}\] can guarantee the bias term to be smaller than \(\frac{\epsilon}{2}\), while \(1/\gamma\) determines the iteration complexity. Plugging in the minimum over the constraints on \(\gamma\), \[\gamma = \min\left( \frac{- B + \sqrt{ B^2 + 2 A \epsilon }}{2 A}, \frac{1}{C}\right) \label{eq:reshuffling95reccurrence95stepsize}\tag{14}\] yields the iteration complexity.
Now, since the quadratic formula is not very interpretable, let us simplify the expression for \(1/\gamma\) using the bound \[\frac{a}{2 \sqrt{b^2 + a}} \leq - b + \sqrt{ b^2 + a},\] which holds for any \(a,b > 0\) and is tight for \(\epsilon \to 0\). With our constants, this reads \[\begin{align} \frac{A \epsilon}{\sqrt{B^2 + 2 A \epsilon}} \leq - B + \sqrt{ B^2 + 2 A \epsilon }, \end{align}\] and therefore \[\begin{align} \frac{2 A}{- B + \sqrt{ B^2 + 2 A \epsilon }} &\leq \frac{2 \sqrt{B^2 + 2 A \epsilon}}{\epsilon} \\ &\leq \frac{2 B + \sqrt{2 A \epsilon}}{\epsilon} \\ &= 2 B \frac{1}{\epsilon} + \sqrt{2 A} \frac{1}{\sqrt{\epsilon}}. \end{align}\] Therefore, for the stepsize choice of 14 , \[\frac{1}{\gamma} \leq \min\left(2 B \frac{1}{\epsilon} + \sqrt{2 A} \frac{1}{\sqrt{\epsilon}}, \frac{1}{C}\right).\] Plugging this into 13 yields the statement.
lemma Let \(F : \mathcal{X} \to \mathbb{R}\) be a finite sum of convex functions as \(F = \frac{1}{n}\left(f_1 + \ldots + f_n\right)\), where \(f_i : \mathcal{X} \to \mathbb{R}\). Then, \[\frac{1}{n} \sum^n_{i=1} \mathrm{D}_{f_i}\left(\vx, \vx'\right) = \mathrm{D}_{F}\left(\vx, \vx'\right),\] for any \(\vx, \vx' \in \mathcal{X}\).
The result immediately follows from the definition of Bregman divergences as \[\begin{align} &\frac{1}{n} \sum^n_{i=1} \mathrm{D}_{f_i}\left(\vx, \vx'\right) \\ &\;= \frac{1}{n} \sum^n_{i=1} \left( f_i\left(\vx\right) - f_i\left(\vx'\right) - \left\langle \nabla f_i\left(\vx'\right), \vx - \vx' \right\rangle \right) \\ &\;= \left(\frac{1}{n} \sum^n_{i=1} f_i\left(\vx\right) \right) - \left(\frac{1}{n} \sum^n_{i=1} f_i\left(\vx'\right) \right) \\ &\qquad- \left\langle \frac{1}{n} \sum^n_{i=1} \nabla f_i\left(\vx'\right), \vx - \vx' \right\rangle \\ &\;= F\left(\vx\right) - F\left(\vx'\right) - \left\langle \nabla F\left(\vx'\right), \vx - \vx' \right\rangle \\ &\;= \mathrm{D}_{F}\left(\vx, \vx'\right). \end{align}\]
theorem Let the component estimators \(\rvvx_1, \ldots, \rvvx_n\) satisfy 3. Then, the variance of the doubly stochastic estimator \(\rvvx_{\rvB}\) is bounded as \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB} \right] \leq V_{\mathrm{com}} + V_{\mathrm{cor}} + V_{\mathrm{sub}}, \end{align}\] where \[\begin{align} V_{\mathrm{com}} &= {\textstyle \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) {\left(\frac{1}{n} \sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvx_i \right] \right)}, } \\ V_{\mathrm{cor}} &= {\textstyle \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\textstyle\left(\frac{1}{n} \sum^n_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \right]} \right)}^2, \text{ and} } \\ V_{\mathrm{sub}} &= {\textstyle \frac{1}{b_{\mathrm{eff}}} \frac{1}{n}\sum^{n}_{i=1} {\left\lVert \bar{\vx}_i - \bar{\vx} \right\rVert}_2^2. } \end{align}\] Equality holds when the equality in 3 holds.
We start from the law of total (co)variance, \[\mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB} \right] = \underbrace{ \mathbb{E}_{\pi}\left[ \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB} \mid \rvB \right] \right] }_{\text{Variance of ensemble}} + \underbrace{ \mathrm{tr}\mathbb{V}_{\pi}\left[ \mathbb{E}\left[ \rvvx_{\rvB} \mid \rvB \right] \right]. }_{\text{Variance of subsampling}}\] This splits the variance into the variance of the specific ensemble of \(\rvB\) and subsampling variance. The main challenge is to relate the variance of the ensemble of \(\rvB\) with the variance of the individual estimators in the sum \[\begin{align} \mathbb{E}_{\pi}\left[ \textstyle \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB} \mid \rvB \right] \right] &= \mathbb{E}_{\pi}\left[ \textstyle \mathrm{tr}\mathbb{V}\left[ \frac{1}{b} \sum_{i \in \rvB} \rvvx_{i} \right] \right]. \end{align}\] Since the individual estimators may not be independent, analyzing the variance of the sum can be tricky. However, the following lemma holds generally:
lemma[]expectedvariancelemma Let \(\rvvx_1, \ldots, \rvvx_b\) be a collection of vector-variate RVs dependent on some random variable \(\rvB\) satisfying 3. Then, the expected variance of the sum of \(\rvvx_1, \ldots, \rvvx_b\) conditioned on \(\rvB\) is bounded as \[\begin{align} \mathbb{E}\left[ {\textstyle \mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^b \rvvx_i \mid \rvB \right]} \right] \leq \rho \mathbb{V}\left[ \rvS \right] + \rho {\left(\mathbb{E}\rvS\right)}^2 + \left(1 - \rho\right) \mathbb{E}\left[ \rvV \right], \end{align}\] where \[{\textstyle \rvS = \sum^{b}_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \mid \rvB \right]} \;\;\text{and}\;\; \rvV = \sum^{b}_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvx_i \mid \rvB \right]. }\] Equality holds when the equality in 3 holds.
Here, \(\rvS\) is the sum of conditional standard deviations, while \(\rvV\) is the sum of conditional variances. Notice that the “variance of the variances” is playing a role: if we reduce the subsampling variance, then the variance of the ensemble, \(V_{\mathrm{com}}\), also decreases.
The rest of the proof, along with the proof of [thm:expected95variance95lemma], can be found in 7.3 page .
Starting from the law of total covariance, we have \[\begin{align} \mathbb{V}\left[ \rvvx_{\rvB} \right] &= \underbrace{ \mathbb{E}_{\rvB \sim \pi}\left[ \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB} \mid \rvB \right] \right] }_{\text{Ensemble Variance}} + \underbrace{ \mathrm{tr}\mathbb{V}_{\rvB \sim \pi}\left[ \mathbb{E}\left[ \rvvx_{\rvB} \mid \rvB \right] \right]. }_{\text{Subsampling Variance}} \label{eq:doubly95stochastic95variance95main} \end{align}\tag{15}\]
3.1.1.2 Ensemble Variance
Bounding the variance of each ensemble is key. From [thm:expected95variance95lemma], we have \[\begin{align} \mathbb{E}\left[ \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB} \mid \rvB \right] \right] &= \mathbb{E}\left[ \mathrm{tr}\mathbb{V}\left[ \frac{1}{b} \sum_{i \in \rvB} \rvvx_{i} \,\middle|\, \rvB \right] \right] \nonumber \\ &= \mathbb{E}\left[ \mathrm{tr}\mathbb{V}\left[ \sum_{i \in \rvB} \left( \frac{1}{b} \rvvx_{i} \right) \,\middle|\, \rvB \right] \right] \nonumber \\ &\leq \rho \mathbb{V}{\rvS} + \rho {\left(\mathbb{E}{\rvS}\right)}^2 + \left(1 - \rho\right) \mathbb{E}{\rvV} , \label{eq:doubly95stochastic95variance95eq1} \end{align}\tag{16}\] where \[\begin{align} \rvS &\triangleq \sum_{i \in \rvB} \sqrt{\mathrm{tr} \mathbb{V}\left[ \frac{1}{b} \rvvx_{i} \right]} = \frac{1}{b} \sum_{i \in \rvB} \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]}, \\ \rvV &\triangleq \sum_{i \in \rvB} \mathrm{tr} \mathbb{V}\left[ \frac{1}{b} \rvvx_{i} \right] = \frac{1}{b^2} \sum_{i \in \rvB} \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]. \end{align}\] In our context, \(\rvS\) is the batch average of the standard deviations, and \(\rvV\) is the batch average of the variance (scaled with a factor of \(1/b\)).
Notice that \(\rvS\) is an \(b\)-sample average of the standard deviations. Therefore, if \(\pi\) is an unbiased subsampling strategy, we retrieve the population average standard deviation as \[\begin{align} \mathbb{E}_{\rvB \sim \pi}\left[ \rvS \right] = \mathbb{E}_{\rvB \sim \pi}\left[ \frac{1}{b} \sum_{i \in \rvB} \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]} \right] = \frac{1}{n} \sum_{i = 1}^n \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]}. \label{eq:doubly95stochastic95variance95eq2} \end{align}\tag{17}\] Under a similar reasoning, the variance of the standard deviations follows as \[\begin{align} &\mathbb{V}_{\rvB \sim \pi}\left[ \rvS \right] \nonumber \\ &= \mathbb{V}_{\rvB \sim \pi}\left[ \frac{1}{b} \sum_{i \in \rvB} \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]} \right] \nonumber \\ &\;= \frac{1}{b_{\mathrm{eff}}} \mathbb{V}_{i \sim \mathrm{Uniform}\left\{1, \ldots, n\right\}}\left[ \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]} \right] \nonumber \\ &\;= \frac{1}{b_{\mathrm{eff}}} \left( \frac{1}{n} \sum_{i=1}^n \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right] - {\left( \frac{1}{n} \sum_{i=1}^n \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]} \right)}^2 \right), \label{eq:doubly95stochastic95variance95eq3} \end{align}\tag{18}\] where the last identity is the well-known formula for the variance: \(\mathbb{V}\rvX = \mathbb{E}\rvX^2 - {(\mathbb{E}\rvX)}^2\). Likewise, the average variance follows as \[\begin{align} \mathbb{E}_{\rvB \sim \pi} \rvV &= \frac{1}{b^2} \mathbb{E}_{\rvB \sim \pi}\left[ \sum_{i \in \rvB } \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right] \right] \nonumber \\ &= \frac{1}{b} \mathbb{E}_{\rvB \sim \pi}\left[ \frac{1}{b} \sum_{i \in \rvB } \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right] \right] \nonumber \\ &= \frac{1}{b} \left( \frac{1}{n} \sum_{i=1}^n \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right] \right) \label{eq:doubly95stochastic95variance95eq4} \end{align}\tag{19}\] Plugging 17 18 19 into 16 , we have \[\begin{align} &\mathbb{E}_{\rvB \sim \pi}\left[ \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB} \mid \rvB \right] \right] \nonumber \\ &\;\leq \rho \mathbb{V}{\rvS} + \rho {\left(\mathbb{E}{\rvS}\right)}^2 + \left(1 - \rho\right) \mathbb{E}{\rvV} \nonumber \\ &\;= \frac{\rho}{b_{\mathrm{eff}}} \left( \frac{1}{n} \sum_{i=1}^n \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right] - {\left( \frac{1}{n} \sum_{i=1}^n \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]} \right)}^2 \right) \nonumber \\ &\quad+ \rho {\left( \frac{1}{n} \sum_{i = 1}^n \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]} \right)}^2 \nonumber \\ &\quad+ \frac{1 - \rho}{b} {\left( \frac{1}{n} \sum_{i=1}^n \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right] \right)} \nonumber \\ &\;= \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \left( \frac{1}{n} \sum_{i=1}^n \mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right] \right) \nonumber \\ &\quad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\left( \frac{1}{n} \sum_{i=1}^n \sqrt{\mathrm{tr} \mathbb{V}\left[ \rvvx_{i} \right]} \right)}^2. \label{eq:doubly95stochastic95variance95ensemble95variance} \end{align}\tag{20}\]
3.1.1.3 Subsampling Variance
The subsampling noise is straightforward. For this, we will denote the minibatch subsampling estimator of the component means as \[\bar{\vx}_{\rvB} \triangleq \frac{1}{b} {\sum_{i \in \rvB}} \bar{\vx}_{i}.\] Since each component estimator \(\rvvx_i\) is unbiased, the expectation conditional on the minibatch \(\rvB\) is \[\begin{align} \mathbb{E}\left[ \bar{\vx}_{\rvB} \mid \rvB \right] = \frac{1}{b} {\sum_{i \in \rvB}} \bar{\vx}_{i}. \end{align}\] Therefore, \[\begin{align} \mathrm{tr}\mathbb{V}_{\rvB \sim \pi}\left[ \mathbb{E}\left[ \rvvx_{\rvB} \mid \rvB \right] \right] &= \mathrm{tr}\mathbb{V}_{\rvB \sim \pi}\left[ \bar{\vx}_{\rvB} \right] \nonumber \\ &= \frac{1}{b_{\mathrm{eff}}} \left(\frac{1}{n} {\sum^n_{i=1}} {\left\lVert \bar{\vx}_{i} - \bar{\vx} \right\rVert}_2^2 \right). \label{eq:doubly95stochastic95variance95subsampling95variance} \end{align}\tag{21}\]
Combining 20 21 into 15 yields the result. Notice that the only inequality we used is 16 , [thm:expected95variance95lemma], in which equality holds if the equality in 3 holds.
In [thm:doubly95stochastic95variance], \(V_{\mathrm{com}}\) is the contribution of the variance of the component estimators, while \(V_{\mathrm{cor}}\) is the contribution of the correlation between component estimators , and \(V_{\mathrm{sub}}\) is the subsampling variance.
3.1.1.4 Monte Carlo with Subampling Without Replacement.
[thm:doubly95stochastic95variance] is very general: it encompasses both the correlated and uncorrelated cases and matches the constants of all of the important special cases. We will demonstrate this in the following corollary along with variance reduction by Monte Carlo averaging of \(m\) i.i.d. samples. That is, we subsample over \(\rvvx_1^m, \ldots, \rvvx_n^m\), where each estimator is an \(m\)-sample Monte Carlo estimator: \[\rvvx_i^{m} \triangleq {\textstyle\frac{1}{m} \sum^m_{j=1}} \rvvx_i^{(j)},\] where \(\rvvx_i^{(1)}, \ldots, \rvvx_i^{(m)}\) are i.i.d replications with mean \(\bar{\vx}_i = \mathbb{E}\rvvx^{(j)}_i\). Then, the variance of the doubly stochastic estimator \(\rvvx_{\rvB}\) of the mean \(\bar{\vx} = \frac{1}{n} \sum^n_{i=1} \bar{\vx}_i\) defined as \[\rvvx_{\rvB}^m \triangleq {\textstyle\frac{1}{b} \sum_{i \in \rvB}} \rvvx_i^m \quad\text{with}\quad \rvB \sim \pi,\] can be bounded as follows:
Corollary 1. For each \(j=1, \ldots, m\), let \(\rvvx_1^{(j)}, \ldots, \rvvx_n^{(j)}\) satisfy 3. Then, the variance of the doubly stochastic estimator \(\rvvx_{\rvB}^m\) of the mean \(\bar{\vx} = \frac{1}{n} \sum^n_{i=1} \bar{\vx}_i\), where \(\pi\) is \(b\)-minibatch sampling without replacement, satisfy the following corollaries:
\(\rho=1\) and \(1 < b < n\): \[\begin{align}\mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB}^m \right] &\leq \frac{n-b}{(n-1) m b} {\textstyle\left(\frac{1}{n} \sum^n_{i=1} \sigma^2_{i}\right) } \\ &\;\;+ \frac{n \left(b-1\right)}{(n-1) m b} {\textstyle\left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}\right)}^2 + \frac{n-b}{(n-1) b} \tau^2 \end{align}\]
\(\rho = 1\) and \(b = 1\): \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB}^m \right] &\leq \frac{1}{m} {\textstyle\left(\frac{1}{n} \sum^n_{i=1} \sigma_i^2 \right)} + \tau^2 \end{align}\]
\(\rho=1\) and \(b = n\): \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB}^m \right] \leq \frac{1}{m} {\textstyle\left(\frac{1}{n} \sum^n_{i=1} \sigma_i \right)}^2 \end{align}\]
\(\sigma_{i} = 0\) for all \(i = 1, \ldots, n\): \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB}^m \right] \leq \frac{n-b}{(n-1) b} \tau^2, \end{align}\]
\(\rho = 0\): \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \rvvx_{\rvB}^m \right] \leq \frac{1}{m b} {\textstyle\left(\frac{1}{n} \sum^n_{i=1} \sigma_i^2 \right)} + \frac{n - b}{(n-1) b} \tau^2 \end{align}\]
where, for all \(i = 1, \ldots, n\) and any \(j = 1, \ldots, m\), \[\begin{align} {2} \sigma_i^2 &= \mathrm{tr}\mathbb{V}\,\rvvx_i^{(j)} &&\;\;\text{is invidual variance and} \\ \tau^2 &= {\textstyle \frac{1}{n}\sum^{n}_{i=1}} {\left\lVert \bar{\vx}_i - \bar{\vx} \right\rVert}_2^2 &&\;\;\text{is the subsampling variance.} \end{align}\]
Notice that, for case of \(\rho=1\), 1 (i), the term with \(\frac{1}{n} \sum^n_{i=1} \sigma_{i}^2\) is reduced in a rate of \(\mathcal{O}\left(\frac{1}{m b}\right)\). This means reducing the subsampling noise by increasing \(b\) also reduces the noise of estimating each component. Furthermore, the first term dominates the second term as \[{\textstyle {\left( \frac{1}{n} \sum^n_{i=1} \sigma_{i} \right)}^2 \leq \frac{1}{n} \sum^n_{i=1} \sigma_{i}^2, }\] as stated by Jensen’s inequality. Therefore, despite correlations, increasing \(b\) will have a more significant effect since it reduces both dominant terms \(\frac{1}{n}\sum_{i=1}^n\sigma_{i}^2\) and \(\tau^2\).
When independent estimators are used, 1 (v) shows that increasing \(b\) reduces the full variance in a \(\mathcal{O}(1/b)\) rate, but increasing \(m\) does not.
1 achieves all known endpoints in the context of SGD: For \(b=n\) (full batch), doubly SGD reduces to SGD with a Monte Carlo estimator, where there is no subsampling noise (no \(\tau^2\)). When the Monte Carlo noise is 0, then doubly SGD reduces to SGD with a subsampling estimator (no \(\sigma_{i}\)), retrieving the result of [29].
3.1.2 Gradient Variance Conditions for SGD↩︎
From [thm:doubly95stochastic95variance], we can establish the ER and BV conditions (2.3) of the doubly stochastic gradient estimators. Following the notation in 2.2, we will denote the doubly stochastic gradient estimator as \(\rvvg_{\rvB}\), which combines the estimators \(\rvvg_1, \ldots, \rvvg_n\) according to the subsampling strategy \(\rvB \sim \pi\), which achieves an effective sample size of \(b_{\mathrm{eff}}\). We will also use the corresponding minibatch subsampling estimator \(\nabla f_{\rvB}\) for the analysis.
Assumption 4. For all \(\vx \in \mathcal{X}\), the component gradient estimators \(\rvvg_1\left(\vx\right), \ldots, \rvvg_{n}\left(\vx\right)\) satisfy 3 with some \(\rho \in [0, 1]\).
Again, this assumption is always met with \(\rho = 1\) and holds with \(\rho = 0\) if the estimators are independent.
Assumption 5. The subsampling estimator \(\nabla f_{\rvB}\) satisfies the \(\mathrm{ER}\left(\mathcal{L}_{\mathrm{sub}}\right)\) condition in 1.
This is a classical assumption used to analyze SGD on finite sums and is automatically satisfied by 2. (See [thm:expected95residual95without95replacement] in 7.4.3 for a proof.)
Assumption 6. For all \(i = 1, \ldots, n\) and \(\vx \in \mathcal{X}\) and global minimizers \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx_*\right)\), the component gradient estimator \(\rvvg_i\) satisfies at least one of the following variants of the ER condition:
\(\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vy\right) \right] \leq 2 \mathcal{L}_i \mathrm{D}_{f_i}\left(\vx, \vy\right),\)
where \(f_i\) is convex.\(\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vy\right) \right] \leq 2 \mathcal{L}_i \left( f_i\left(\vx\right) - f_i\left(\vx_*\right) \right)\)
where \(f_i\left(\vx\right) \geq f_i\left(\vx_*\right)\).\(\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vy\right) \right] \leq 2 \mathcal{L}_i \left( F\left(\vx\right) - F\left(\vx_*\right) \right)\).
Each of these assumptions holds under different assumptions and problem setups. For instance, \(\mathrm{A}^{\mathrm{CVX}}\) holds only under componentwise convexity, while \(\mathrm{A}^{\mathrm{ITP}}\) requires majorization \(f_i\left(\vx\right) \geq f_i\left(\vx_*\right)\), which is essentially assuming “interpolation” [34], [39], [40] in the ERM context. Among these, (B) is the strongest since it directly relates the individual components \(f_1, \ldots, f_n\) with the full objective \(F\).
We now state our result establishing the ER condition:
theorem Let 5 4 6 hold. Then, we have:
If \((\mathrm{A}^{\mathrm{CVX}})\) or \((\mathrm{A}^{\mathrm{ITP}})\) hold, \(\rvvg_{\rvB}\) satisfies \(\mathrm{ER}\left(\mathcal{L}_{\rm{A}}\right)\).
If \(\rm{(B)}\) holds, \(\rvvg_{\rvB}\) satisfies \(\mathrm{ER}\left(\mathcal{L}_{\rm{B}}\right)\).
where \(\mathcal{L}_{\mathrm{max}} = \max\big\{ \mathcal{L}_{1}, \ldots, \mathcal{L}_{n} \big\}\), \[\begin{align} \mathcal{L}_{\rm{A}} &= {\textstyle \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \mathcal{L}_{\mathrm{max}} + \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\textstyle{\left( \frac{1}{n} \sum^{n}_{i=1} \mathcal{L}_i \right)}} + \frac{\mathcal{L}_{\mathrm{sub}}}{b_{\mathrm{eff}}} } \\ \mathcal{L}_{\rm{B}} &= {\textstyle \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \left( \frac{1}{n} \sum^{n}_{i=1} \mathcal{L}_{i} \right) } \\ &\qquad+ {\textstyle \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\left( \frac{1}{n} \sum^{n}_{i=1} \sqrt{\mathcal{L}_{i}} \right) }^2 + \frac{\mathcal{L}_{\mathrm{sub}}}{b_{\mathrm{eff}}}. } \end{align}\]
From [thm:doubly95stochastic95variance], we have \[\begin{align} &\mathrm{tr}\mathbb{V}\left[ \rvvg_{\rvB}\left(\vx\right) - \rvvg_{\rvB}\left(\vx_*\right) \right] \nonumber \\ &\;\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \left( \frac{1}{n} \sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right] \right) \nonumber \\ &\;\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right]} \right)}^2 \nonumber \\ &\;\qquad + \frac{1}{b_{\mathrm{eff}}} \mathrm{tr}\mathbb{V}\left[ \nabla f_{\rvB}\left(\vx\right) - \nabla F\left(\vx\right) \right], \nonumber \shortintertext{where \ref{assumption:subsampling95er} yields} &\;\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \underbrace{ \left( \frac{1}{n} \sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right] \right) }_{\triangleq \mathrm{T_{\mathrm{var}}}} \nonumber \\ &\;\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) \underbrace{ {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right]} \right)}^2 }_{\triangleq T_{\mathrm{cov}}} \nonumber \\ &\;\qquad + \frac{2 \mathcal{L}_{\mathrm{sub}} }{b_{\mathrm{eff}}} \left( F\left(\vx\right) - F\left(\vx_*\right) \right) \nonumber \\ &\;= \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) T_{\mathrm{var}} \nonumber + \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) T_{\mathrm{cov}} \\ &\qquad+ \frac{2 \mathcal{L}_{\mathrm{sub}} }{b_{\mathrm{eff}}} \left( F\left(\vx\right) - F\left(\vx_*\right) \right). \label{eq:master95equation95er} \end{align}\tag{22}\]
3.1.2.1 Proof of with \((\rm{A}^{\rm{CVX}})\)
Since 6 \((\rm{A}^{\rm{CVX}})\) requires \(f_1, \ldots, f_n\) to be convex, \(F\) is also convex. Therefore, we can use the identity in [thm:average95bregman] and \[\mathrm{D}_{F}\left(\vx, \vx_*\right) = F\left(\vx\right) - F\left(\vx_*\right).\] With that said, under \((\rm{A}^{\rm{CVX}})\), we have \[\begin{align} T_{\text{var}} &\leq \frac{1}{n} \sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right] \nonumber \\ &= \frac{1}{n} \sum^n_{i=1} 2 \mathcal{L}_i \mathrm{D}_{f_i}\left(\vx, \vx_*\right), \nonumber \shortintertext{applying \(\mathcal{L}_{\mathrm{max}} \geq \mathcal{L}_i\) for all \(i = 1, \ldots, n\),} &\leq 2 \mathcal{L}_{\mathrm{max}} \frac{1}{n} \sum^n_{i=1} \mathrm{D}_{f_i}\left(\vx, \vx_*\right) \nonumber \shortintertext{and \ref{thm:average95bregman},} &= 2 \mathcal{L}_{\mathrm{max}} \mathrm{D}_{F}\left(\vx, \vx_*\right). \label{eq:ercondition95item195tvar} \end{align}\tag{23}\]
For \(T_{\text{cov}}\), since \[\begin{align} T_{\text{cov}} &= {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right]} \right)}^2 \nonumber \shortintertext{ is monotonic w.r.t. the variance, we can apply \((\rm{A}^{\rm{CVX}})\) as } &\leq \frac{2}{n^2} {\left( \sum^n_{i=1} \sqrt{ \mathcal{L}_i \mathrm{D}_{f_i}\left(\vx, \vx_*\right) } \right)}^2. \nonumber \shortintertext{Now, applying the Cauchy-Schwarz inequality yields} &\leq \frac{2}{n^2} \left( \sum^n_{i=1} \mathcal{L}_i \right) \left( \sum^n_{i=1} \mathrm{D}_{f_i}\left(\vx, \vx_*\right) \right) \nonumber \\ &= 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \left( \frac{1}{n} \sum^n_{i=1} \mathrm{D}_{f_i}\left(\vx, \vx_*\right) \right) \nonumber \shortintertext{and by \ref{thm:average95bregman},} &= 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \mathrm{D}_{F}\left(\vx, \vx_*\right). \label{eq:ercondition95item195tcov} \end{align}\tag{24}\] Plugging 23 24 into 22 , we have \[\begin{align} &\mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right] \\ &\;\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) T_{\mathrm{var}} \nonumber + \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) T_{\mathrm{cov}} \\ &\qquad+ \frac{2 \mathcal{L}_{\mathrm{sub}} }{b_{\mathrm{eff}}} \left( F\left(\vx\right) - F\left(\vx_*\right) \right) \\ &\;\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) 2 \mathcal{L}_{\mathrm{max}} \mathrm{D}_{F}\left(\vx, \vx_*\right) \\ &\;\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \mathrm{D}_{F}\left(\vx, \vx_*\right) \\ &\;\qquad+ \frac{1}{b_{\mathrm{eff}}} 2 \mathcal{L}_{\mathrm{sub}} \mathrm{D}_{F}\left(\vx, \vx_*\right). \\ &= 2 \,\Bigg(\, \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \mathcal{L}_{\mathrm{max}} + \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \\ &\qquad\qquad+ \frac{1}{b_{\mathrm{eff}}} \mathcal{L}_{\mathrm{sub}} \,\Bigg)\, \mathrm{D}_{F}\left(\vx, \vx_*\right) \\ &= 2 \,\Bigg(\, \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \mathcal{L}_{\mathrm{max}} + \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \\ &\qquad\qquad+ \frac{1}{b_{\mathrm{eff}}} \mathcal{L}_{\mathrm{sub}} \,\Bigg)\, \left( F\left(\vx\right) - F\left(\vx_*\right) \right). \end{align}\]
3.1.2.2 Proof of with \((\rm{A}^{\rm{ITP}})\)
From 6 \((\rm{A}^{\rm{ITP}})\), we have \[\begin{align} T_{\text{var}} &= \frac{1}{n} \sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right] \nonumber \\ &\leq \frac{1}{n} \sum^n_{i=1} 2 \mathcal{L}_i \left(f_i\left(\vx\right) - f_i\left(\vx_*\right)\right), \nonumber \shortintertext{applying \(\mathcal{L}_{\mathrm{max}} \geq \mathcal{L}_i\) for all \(i = 1, \ldots, n\),} &\leq 2 \mathcal{L}_{\mathrm{max}} \frac{1}{n} \sum^n_{i=1} \left(f_i\left(\vx\right) - f_i\left(\vx_*\right)\right) \nonumber \\ &= 2 \mathcal{L}_{\mathrm{max}} \left( F\left(\vx\right) - F\left(\vx_*\right) \right). \label{eq:ercondition95item295tvar} \end{align}\tag{25}\]
Similarly, \[\begin{align} T_{\text{cov}} &= {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right]} \right)}^2, \nonumber \shortintertext{applying \((\rm{A}^{\rm{ITP}})\),} &\leq \frac{2}{n^2} {\left( \sum^n_{i=1} \sqrt{ \mathcal{L}_i \left( f_i\left(\vx\right) - f_i\left(\vx_*\right) \right) } \right)}^2, \nonumber \shortintertext{and applying the Cauchy-Schwarz inequality,} &\leq \frac{2}{n^2} \left( \sum^n_{i=1} \mathcal{L}_i \right) \left( \sum^n_{i=1} f_i\left(\vx\right) - f_i\left(\vx_*\right) \right) \nonumber \\ &= 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \left( \frac{1}{n} \sum^n_{i=1} f_i\left(\vx\right) - f_i\left(\vx_*\right) \right) \nonumber \\ &= 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \left( F\left(\vx\right) - F\left(\vx_*\right) \right). \label{eq:ercondition95item295tcov} \end{align}\tag{26}\]
Plugging 25 26 into 22 , we have \[\begin{align} &\mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right] \\ &\;\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) T_{\mathrm{var}} + \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) T_{\mathrm{cov}} \\ &\qquad+ \frac{2 \mathcal{L}_{\mathrm{sub}} }{b_{\mathrm{eff}}} \left( F\left(\vx\right) - F\left(\vx_*\right) \right) \\ &\;\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) 2 \mathcal{L}_{\mathrm{max}} \left( F\left(\vx\right) F\left(\vx_*\right) \right) \\ &\;\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \left(F\left(\vx\right) - F\left(\vx_*\right)\right) \\ &\;\qquad+ \frac{1}{b_{\mathrm{eff}}} 2 \mathcal{L}_{\mathrm{sub}} \left(F\left(\vx\right) - F\left(\vx_*\right)\right). \\ &= 2 \,\Bigg(\, \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \mathcal{L}_{\mathrm{max}} + \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \\ &\qquad\qquad+ \frac{1}{b_{\mathrm{eff}}} \mathcal{L}_{\mathrm{sub}} \,\Bigg)\, \left( F\left(\vx\right) - F\left(\vx_*\right)\right). \end{align}\]
3.1.2.3 Proof of
From 6 (B), we have \[\begin{align} T_{\text{var}} &= \frac{1}{n} \sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right] \nonumber \\ &\leq \frac{1}{n} \sum^n_{i=1} 2 \mathcal{L}_i \left(F\left(\vx\right) - F\left(\vx_*\right)\right) \nonumber \\ &= 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_{i} \right) \left(F\left(\vx\right) - F\left(\vx_*\right)\right). \label{eq:ercondition95item395tvar} \end{align}\tag{27}\]
And, \[\begin{align} T_{\text{cov}} &= {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx\right) - \rvvg_i\left(\vx_*\right) \right]} \right)}^2 \nonumber \\ &\leq \frac{2}{n^2} {\left( \sum^n_{i=1} \sqrt{ \mathcal{L}_i \left( F\left(\vx\right) - F\left(\vx_*\right) \right) } \right)}^2 \nonumber \\ &= 2 {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{ \mathcal{L}_i } \right)}^2 \left( F\left(\vx\right) - F\left(\vx_*\right) \right). \label{eq:ercondition95item395tcov} \end{align}\tag{28}\]
Plugging 28 27 into 22 , we have \[\begin{align} &\mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right] \\ &\;\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) 2 \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \left(F\left(\vx\right) - F\left(\vx_*\right) \right) \\ &\;\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {2 \left( \frac{1}{n} \sum^n_{i=1} \sqrt{ \mathcal{L}_i } \right)}^2 \left(F\left(\vx\right) - F\left(\vx_*\right) \right) \\ &\;\qquad+ \frac{1}{b_{\mathrm{eff}}} 2 \mathcal{L}_{\mathrm{sub}} \left(F\left(\vx\right) - F\left(\vx_*\right) \right), \\ &\;= 2 \, \Bigg( \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) \\ &\quad\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{ \mathcal{L}_i } \right)}^2 \\ &\quad\qquad+ \frac{1}{b_{\mathrm{eff}}} \mathcal{L}_{\mathrm{sub}} \Bigg) \, \left(F\left(\vx\right) - F\left(\vx_*\right) \right). \end{align}\]
Assuming the conditions in 6 hold with the same value of \(\mathcal{L}_i\), the inequality \[{\textstyle {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathcal{L}_{i} }\right)}^{2} \leq \frac{1}{n} \sum^n_{i=1} \mathcal{L}_{i} \leq \mathcal{L}_{\mathrm{max}}, }\] implies \(\mathcal{L}_{\rm{B}} \leq \mathcal{L}_{\rm{A}}\).
Meanwhile, The BV condition follows by assuming equivalent conditions on each component estimator:
Assumption 7. Variance is bounded for all \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\) such that the following hold:
\({\textstyle \frac{1}{n} \sum^n_{i=1} {\left\lVert \nabla f_i\left(\vx_*\right) \right\rVert}_2^2 \leq \tau^2 }\) for some \(\tau^2 < \infty\) and,
\(\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx_*\right) \right] \leq \sigma^2_{i}\) for some \(\sigma^2_{i} < \infty\), for all \(i = 1, \ldots, n\).
Based on these, [thm:doubly95stochastic95variance] immediately yields the result:
theorem Let 4 7 hold. Then, \(\rvvg_{\rvB}\) satisfies \(\mathrm{BV}\left(\sigma^2\right)\), where \[\begin{align} \sigma^2 &= \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \left( \textstyle\frac{1}{n} \sum^n_{i=1} \sigma_{i}^2 \right) \\ &\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\textstyle\left( \frac{1}{n} \sum^n_{i=1} \sigma_{i} \right)}^2 + \frac{\tau^2}{b_{\mathrm{eff}}}. \end{align}\] Equality in 2 holds if equality in 4 holds.
For any element of the solution set \(\vx_* = \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\), by [thm:doubly95stochastic95variance], we have \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \rvvg_{\rvB}\left(\vx_*\right) \right] &\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \left( \frac{1}{n} \sum^n_{i=1} \mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx_*\right) \right] \right) \\ &\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvg_i\left(\vx_*\right) \right]} \right)}^2 \\ &\qquad+ \frac{1}{b_{\mathrm{eff}}} \mathrm{tr}\mathbb{V}\left[ \nabla f_{\rvB}\left(\vx_*\right) \right]. \shortintertext{Applying \ref{assumption:bounded95variance95both}, we have} &\leq \left( \frac{\rho}{b_{\mathrm{eff}}} + \frac{1 - \rho}{b} \right) \left( \frac{1}{n} \sum^n_{i=1} \sigma_{i}^{2} \right) \\ &\qquad+ \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{ \sigma^2_{i} } \right)}^2 \\ &\qquad+ \frac{1}{b_{\mathrm{eff}}} \tau^2 \\ &= \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) \left( \frac{1}{n} \sum^n_{i=1} \sigma_{i}^{2} \right) \\ &\qquad+ \rho \left(1 - \frac{1}{b_{\mathrm{eff}}}\right) {\left( \frac{1}{n} \sum^n_{i=1} \sigma_{i} \right)}^2 \\ &\qquad+ \frac{1}{b_{\mathrm{eff}}} \tau^2. \end{align}\]
As discussed in 2.4, [thm:general95conditions] [thm:bounded95variance] are sufficient to guarantee convergence of doubly SGD. For completeness, let us state a specific result for \(\rho = 1\):
lemma Let the objective function \(F\) satisfy 2, \(\pi\) be sampling \(b\) samples without replacement, and all elements of the solution set \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\) be stationary points of \(F\). Then, the subsampling estimator \(\nabla f_{\rvB}\) satisfies the \(\mathrm{ER}\) condition as \[\begin{align} &\mathrm{tr}\mathbb{V}_{\rvB \sim \pi}\left[ \nabla f_{\rvB}\left(\vx\right) - \nabla f_{\rvB}\left(\vx_*\right) \right] \\ &\qquad\leq 2 \frac{n - b}{b \left(n - 1\right)} L_{\mathrm{max}} \left(F\left(\vx\right) - F\left(\vx_*\right)\right), \end{align}\] where \(L_{\mathrm{max}} = \max\left\{L_1, \ldots, L_n\right\}\).
Consider that, for any random vector \(\rvvx\), \[\mathrm{tr}\mathbb{V}\left[ \rvvx^2 \right] \leq \mathbb{E}{\left\lVert \rvvx \right\rVert}_2^2\] holds. Also, sampling without replacement achieves \(b_{\mathrm{eff}} = \frac{(n-1) b}{n-b}\). Therefore, we have \[\begin{align} &\mathrm{tr}\mathbb{V}_{\rvB \sim \pi}\left[ \nabla f_{\rvB}\left(\vx\right) - \nabla f_{\rvB}\left(\vx_*\right) \right] \\ &\;= \frac{n - b}{b \left(n - 1\right)} \mathrm{tr}\mathbb{V}\left[ \nabla f_{i}\left(\vx\right) - \nabla f_{i}\left(\vx_*\right) \right] \\ &\;\leq \frac{n - b}{b \left(n - 1\right)} \left( \frac{1}{n} \sum_{i=1}^n {\left\lVert \nabla f_{i}\left(\vx\right) - \nabla f_{i}\left(\vx_*\right) \right\rVert}_2^2 \right), \shortintertext{and from \ref{assumption:components},} &\;= \frac{n - b}{b \left(n - 1\right)} \left( \frac{1}{n} \sum_{i=1}^n 2 L_i \mathrm{D}_{f_i}\left(\vx, \vx_* \right) \right). \shortintertext{Using the bound \(L_{\mathrm{max}} \geq L_i\) for all \(i = 1, \ldots, n\),} &\;\leq 2 L_{\mathrm{max}} \frac{n - b}{b \left(n - 1\right)} \left( \frac{1}{n} \sum_{i=1}^n \mathrm{D}_{f_i}\left(\vx, \vx_* \right) \right), \shortintertext{applying \ref{thm:average95bregman},} &\;= 2 L_{\mathrm{max}} \frac{n - b}{b \left(n - 1\right)} \mathrm{D}_{F}\left(\vx, \vx_* \right), \shortintertext{and since \(\vx_*\) is a stationary point of \(F\),} &\;= 2 \frac{n - b}{b \left(n - 1\right)} L_{\mathrm{max}} \left(F\left(\vx\right) - F\left(\vx_*\right) \right). \end{align}\]
corollary Let the objective \(F\) satisfy 1 2, the global optimum \(\vx_* = \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\) be a stationary point of \(F\), the component gradient estimators \(\rvvg_1, \ldots, \rvvg_n\) satisfy 6 (B) and , and \(\pi\) be \(b\)-minibatch sampling without replacement. Then the last iterate of SGD with \(\rvvg_{\rvB}\) is \(\epsilon\)-close to \(\vx_*\) as \(\mathbb{E}{\left\lVert \vx_T - \vx_* \right\rVert}_2^2 \leq \epsilon\) if \[\begin{align} T \geq 2 \max\left(C_{\mathrm{var}} \frac{1}{\epsilon},\; C_{\mathrm{bias}}\right) \log\left(2 {\lVert \vx_0 - \vx_* \rVert}_2^2 \frac{1}{\epsilon} \right) \end{align}\] for some fixed stepsize where \[\begin{align} C_{\mathrm{var}} &= \frac{2}{b} \left( {\frac{1}{n} \sum^n_{i=1} \frac{\sigma_{i}^2}{\mu^2} }\right) + { 2\left( \frac{1}{n} \sum^n_{i=1} \frac{\sigma_{i}}{\mu} \right)}^2 + \frac{2}{b} \frac{\tau^2}{\mu^2}, \\ C_{\mathrm{bias}} &= \frac{2}{b} \left( {\frac{1}{n} \sum^n_{i=1} \frac{\mathcal{L}_i}{\mu} }\right) + {\textstyle 2 \left( \frac{1}{n} \sum^n_{i=1} \sqrt{\frac{\mathcal{L}_i}{\mu}} \right)}^2 + \frac{2}{b} \frac{L}{\mu}. \end{align}\]
From 2 and the assumption that \(\vx_*\) is a stationary point, [thm:expected95residual95without95replacement] establishes that \(\nabla f_{\rvB}\) satisfies the \(\mathrm{ER}\left(\mathcal{L}_{\mathrm{sub}}\right)\) holds with \[\mathcal{L}_{\mathrm{sub}} = \frac{n-b}{(n-1) b} L_{\mathrm{max}}.\] Therefore, 5 holds. Furthermore, since the component gradient estimators satisfy 6 (B) and 3 always hold with \(\rho = 1\), we can apply [thm:general95conditions] which estblishes that \(\rvvg_{\rvB}\) satisfies \(\mathrm{ER}\left(\mathcal{L}\right)\) with \[\begin{align} \mathcal{L} &= \frac{n-b}{(n-1) b} \left( \frac{1}{n} \sum^n_{i=1} \mathcal{L}_i \right) + \frac{n (b-1)}{(n-1) b} {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\mathcal{L}_i} \right)}^2 \\ &\quad+ \frac{n-b}{(n-1) b} L_{\mathrm{max}}. \end{align}\] Furthermore, under 7, [thm:bounded95variance] shows that \(\mathrm{BV}\left(\sigma^2\right)\) holds with \[\begin{align} \sigma^2 &= \frac{n-b}{(n-1) b} \left(\frac{1}{n} \sum^n_{i=1} \sigma^2_{i}\right) + \frac{n (b-1)}{(n-1) b} {\left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}\right)}^2 \\ &\quad+ \frac{n-b}{(n-1) b} \tau^2. \end{align}\] Since both \(\mathrm{ER}\left(\mathcal{L}\right)\) and \(\mathrm{BV}\left(\sigma^2\right)\) hold and \(F\) satisfies 1, we can now invoke [thm:strongly95convex95sgd95complexity], which guarantees that we can obtain an \(\epsilon\)-accurate solution after \[T \geq 2 \max\Bigg( \underbrace{\frac{\sigma^2}{\mu^2}}_{C_{\mathrm{var}}} \frac{1}{\epsilon}, \underbrace{\frac{\mathcal{L} + L}{\mu}}_{C_{\mathrm{bias}}} \Bigg) \log\left(2 {\left\lVert \vx_0 - \vx_* \right\rVert}_2^2 \frac{1}{\epsilon} \right)\] iterations and fixed stepsize of \[\gamma = \min\left( \frac{\epsilon \mu}{2 \sigma^2}, \frac{1}{2 \left(\mathcal{L} + L\right)} \right).\]
The constants in the lower bound on the number of required iterations can be made more precise as \[\begin{align} C_{\mathrm{var}} &= { \frac{n-b}{(n-1) b} \left( {\frac{1}{n} \sum^n_{i=1} \frac{\sigma_{i}^2}{\mu^2} }\right) } \\ &\quad+ { \frac{n (b-1)}{(n-1) b} {\left( \frac{1}{n} \sum^n_{i=1} \frac{\sigma_{i}}{\mu} \right)}^2 + \frac{n-b}{(n-1)b} \frac{\tau^2}{\mu^2} } \\ C_{\mathrm{bias}} &= { \frac{n-b}{(n-1) b} \left( {\frac{1}{n} \sum^n_{i=1} \frac{\mathcal{L}_i}{\mu} }\right) } \\ &\quad+ { \frac{n (b-1)}{(n-1) b} {\left( \frac{1}{n} \sum^n_{i=1} \sqrt{\frac{\mathcal{L}_i}{\mu}} \right)}^2 + \frac{n-b}{(n-1) b} \frac{L}{\mu}. } \end{align}\] Using the fact that \((n-b)/n \leq (n-1)/n \leq 2\) for all \(n \geq 2\) yields the simplified constants in the statement.
3.2 Random Reshuffling of Stochastic Gradients↩︎
We now move to our analysis of SGD with random reshuffling (SGD-RR). In the doubly stochastic setting, this corresponds to reshuffling over stochastic estimators instead of gradients, which we will denote as doubly SGD-RR. In practice, doubly SGD-RR is often observed to converge faster than doubly SGD, even when dependent estimators are used.
3.2.1 Algorithm↩︎
3.2.1.1 Doubly SGD-RR
The algorithm is stated as follows:
Reshuffle and partition the gradient estimators into minibatches of size \(b\) as \(\rvP = \{\rvP_1, \ldots, \rvP_p\}\), where \(p = n/b\) is the number of partitions or minibatches.
Perform gradient descent for \(i = 1, \ldots, p\) steps as \[\vx_{k}^{i+1} = \Pi_{\mathcal{X}}\left( \vx_k^{i} - \gamma \rvvg_{\rvP_i}\,(\vx_k^i) \right)\]
\(k \leftarrow k + 1\) and go back to step .
(We assume \(n\) is an integer multiple of \(b\) for clarity.) Here, \(i = 1, \ldots, p\) denotes the step within the epoch, \(k = 1, \ldots, K\) denotes the epoch number.
3.2.2 Proof Sketch↩︎
3.2.2.1 Why SGD-RR is Faster
A key aspect of random reshuffling in the finite sum setting (SGD-RR) is that it uses conditionally biased gradient estimates. Because of this, on strongly convex finite sums, [41] show that the Lyapunov function for random reshuffling is not the usual \({\lVert \vx_k^i - \vx_* \rVert}_2^2\), but some biased Lyapunov function \({\lVert \vx_k^i - \vx_*^k \rVert}_2^2\), where the reference point is \[{\textstyle \vx_*^i \triangleq \Pi_{\mathcal{X}}\left( \vx_* - \gamma \sum^{i-1}_{j=0} \nabla f_{\rvP_i}\left(\vx_*\right) \right). } \label{eq:lyapunov95bias}\tag{29}\] Under this definition, the convergence rate of SGD is not determined by the gradient variance anymore; it is determined by the squared error of the Lyapunov reference point, \({\lVert \vx_*^i - \vx_* \rVert}_2^2\). There are two key properties of this quantity:
The peak mean-squared error decreases at a rate of \(\gamma^2\) with respect to the stepsize \(\gamma\).
The squared error is 0 at the following two endpoints: beginning of the epoch and at the end of the epoch.
For some stepsize achieving a \(\mathcal{O}(1/T)\) rate on SGD, these two properties combined result in SGD-RR attaining a \(\mathcal{O}(1/T^2)\) rate at exactly the end of each epoch.
3.2.2.2 Is doubly SGD-RR as Fast as SGD-RR?
Unfortunately, doubly SGD-RR does not achieve the same rate as SGD-RR. Since stochastic gradients are used in addition to reshuffling, doubly SGD-RR deviates from the path that minimizes the biased Lyapunov function. Still, doubly SGD-RR does have provable benefits.
3.2.3 Complexity Analysis↩︎
We provide the general complexity guarantee for doubly SGD-RR on strongly convex objectives with \(\mu\)-strongly convex components and fully correlated component estimators (\(\rho=1\)):
lemma For any \(b\)-minibatch reshuffling strategy, the squared error of the reference point of the Lyapunov function (29 ) under reshuffling is bounded as \[\begin{align} \mathbb{E} {\lVert\vx_{*}^{i} - \vx_*\rVert}_2^2 \leq \frac{\gamma^2 n }{4 b^2} \, \tau^2 \end{align}\] for all \(i = 1, \ldots, p\), where \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\).
The proof is a generalization of [41], where we sample \(b\)-minibatches instead of single datapoints. Recall that \(\rvP\) denotes the (possibly random) partitioning of the \(n\) datapoints into \(b\)-minibatches \(\rvP_1, \ldots, \rvP_p\). From the definition of the squared error of the Lyapunov function in 29 , we have \[\begin{align} &\mathbb{E}\left[ {\lVert\vx_{*}^{i} - \vx_*\rVert}_2^2 \right] \\ &\;= \mathbb{E}\left[ {\left\lVert \Pi_{\mathcal{X}}\left( \vx_* - \sum_{k=0}^{i-1} \gamma \nabla f_{\rvP_i}\left(\vx_*\right) \right) - \Pi_{\mathcal{X}}\left(\vx_*\right) \right\rVert}_2^2 \right], \shortintertext{and since the projection onto a convex set under a Euclidean metric is non-expansive,} &\;\leq \mathbb{E}\left[ {\left\lVert \vx_* - \sum_{k=0}^{i-1} \gamma \nabla f_{\rvP_i}\left(\vx_*\right) - \vx_* \right\rVert}_2^2 \right] \\ &\;= \mathbb{E}\left[ {\left\lVert \sum_{k=0}^{i-1} \gamma \nabla f_{\rvP_i}\left(\vx_*\right) \right\rVert}_2^2 \right], \shortintertext{introducing a factor of \(i\) in and out of the squared norm,} &\;= \frac{i^2}{2} \, \mathbb{E}\left[ {\left\lVert \frac{1}{i} \sum_{k=0}^{i-1} \gamma \nabla f_{\rvP_i}\left(\vx_*\right) \right\rVert}_2^2 \right] \\ &\;= \frac{\gamma^2 i^2}{2} \, \mathbb{E}\left[ {\left\lVert \frac{1}{i} \sum_{k=0}^{i-1} \nabla f_{\rvP_i}\left(\vx_*\right) \right\rVert}_2^2 \right]. \end{align}\] Now notice that \(\frac{1}{i} \sum^{i-1}_{j=0} \nabla f_{\rvP_i} \left(\vx_*\right)\) is a sample average of \(i b\) samples drawn without replacement. Therefore, it is an unbiased estimate of \(\nabla F\left(\vx_*\right)\). This implies \[\begin{align} \mathbb{E}\left[ {\lVert\vx_{*}^{i} - \vx_*\rVert}_2^2 \right] &= \frac{\gamma^2 i^2}{2} \, \mathbb{E}\left[ {\left\lVert \frac{1}{i} \sum_{k=0}^{i-1} \nabla f_{\rvP_i}\left(\vx_*\right) \right\rVert}_2^2 \right] \\ &= \frac{\gamma^2 i^2}{2} \, \mathrm{tr}\mathbb{V}\left[ \frac{1}{i} \sum_{k=0}^{i-1} \nabla f_{\rvP_i}\left(\vx_*\right) \right], \shortintertext{and from \ref{thm:variancewithoutreplacement} with a sample size of \(ib\), } &= \frac{\gamma^2 i^2}{2} \, \frac{n - ib}{\left(n-1\right) ib} \frac{1}{n} \sum^n_{i=1} {\left\lVert \nabla f_{i}\left(\vx_*\right) \right\rVert}_2^2 \\ &= \frac{\gamma^2 i \left(\frac{n}{b} - i\right)}{2 \left(n-1\right)} \tau^2. \end{align}\] Notice that this is a quadratic with respect to \(i\), where the maximum is obtained by \(i = \frac{n}{2b}\). Then, \[\begin{align} \mathbb{E}\left[ {\lVert\vx_{*}^{i} - \vx_*\rVert}_2^2 \right] &\leq \frac{\gamma^2 {\left(\frac{n}{2b}\right)}^2 }{2 \left(n-1\right)} \tau^2 \\ &= \frac{\gamma^2 n^2 }{8 b^2 \left(n-1\right)} \tau^2, \shortintertext{and using the bound \(n/(n-1) \leq 2\) for all \(n \geq 2\),} &\leq \frac{\gamma^2 n }{4 b^2} \, \tau^2. \end{align}\]
lemma Let the objective function satisfy 2, \(B\) be any \(b\)-minibatch of indices such that \(B \subseteq \{1, \ldots, n\}\) and the component gradient estimators \(\rvvg_1, \ldots, \rvvg_n\) satisfy 6 (\(\rm{A}^{\rm{CVX}}\)). Then, \(\rvvg_{B}\) is convex-smooth in expectation such that \[\begin{align} \mathbb{E}_{\varphi}{\left\lVert \rvvg_{B}\left(\vx\right) - \rvvg_{B}\left(\vx_*\right) \right\rVert}_2^2 \leq 2 \left( \mathcal{L}_{\rm{max}} + L_{\mathrm{max}} \right) \mathrm{D}_{f_{B}} \left(\vx, \vx_*\right), \end{align}\] for any \(\vx \in \mathcal{X}\), where \[\begin{align} {3} \vx_* &= \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right), \\ \mathcal{L}_{\rm{max}} &= \max\left\{\mathcal{L}_{1}, \ldots, \mathcal{L}_{n} \right\}, \\ L_{\mathrm{max}} &= \max\left\{L_1, \ldots, L_n\right\}. \end{align}\]
Notice that, for this Lemma, we do not assume that the minibatch \(B\) is a random variable. Therefore, the only randomness is the stochasticity of the component gradient estimators \(\rvvg_1, \ldots, \rvvg_n\).
Now, from the property of the variance, we can decompose the expected squared norm as \[\begin{align} &\mathbb{E} {\lVert \rvvg_B\left(\vx\right) - \rvvg_B\left(\vx_*\right)\rVert}_2^2 \\ &\;= \underbrace{ \mathrm{tr}\mathbb{V}_{\varphi}\left[ \rvvg_B\left(\vx\right) - \rvvg_B\left(\vx_*\right) \right] }_{V_{\rm{com}}} + \underbrace{ {\lVert \nabla f_{B} \left(\vx\right) - \nabla f_{B}\left(\vx_*\right) \rVert}_2^2 }_{V_{\mathrm{sub}}}. \end{align}\]
First, the contribution of the variances of the component gradient estimators follows as \[\begin{align} V_{\rm{com}} &= \mathrm{tr}\mathbb{V}_{\varphi}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right] \nonumber \\ &= \mathrm{tr}\mathbb{V}_{\varphi}\left[ \frac{1}{b} \sum_{i \in B} \rvvg_{i}\left(\vx\right) - \rvvg_{i}\left(\vx_*\right) \right], \nonumber \shortintertext{applying \ref{eq:varianceofsumeq2} of \ref{thm:varianceofsum},} &\leq \frac{1}{b} \sum_{i \in B} \mathrm{tr}\mathbb{V}_{\varphi}\left[ \rvvg_{i}\left(\vx\right) - \rvvg_{i}\left(\vx_*\right) \right], \label{eq:loose95bound95for95reshuffling} \shortintertext{and then \ref{assumption:montecarlo95er} (\(\rm{A}^{\rm{CVX}}\)),} &\;\leq \frac{1}{b} \sum_{i \in B} 2 \mathcal{L}_{i} \, \mathrm{D}_{f_i}\left(\vx, \vx_*\right). \nonumber \shortintertext{Now, since \(\mathcal{L}_{\mathrm{max}} \geq \mathcal{L}_i\) for all \(i= 1, \ldots, n\),} &\;\leq 2 \mathcal{L}_{\rm{max}} \frac{1}{b} \sum_{i \in B} \mathrm{D}_{f_i}\left(\vx, \vx_*\right) \nonumber \\ &\;= 2 \mathcal{L}_{\rm{max}} \mathrm{D}_{f_{B}}\left(\vx, \vx_*\right). \nonumber \end{align}\tag{30}\]
On the other hand, the squared error of subsampling (it is not the variance since we do not take expectation over the batches) follows as \[\begin{align} V_{\mathrm{sub}} &= {\lVert \nabla f_{\rvB} \left(\vx\right) - \nabla f_{\rvB}\left(\vx_*\right) \rVert}_2^2 \\ &= {\left\lVert \frac{1}{b} \sum_{i \in B} \nabla f_{i} \left(\vx\right) - \nabla f_{i}\left(\vx_*\right) \right\rVert}_2^2, \shortintertext{by Jensen's inequality,} &\leq \frac{1}{b} \sum_{i \in B} {\lVert \nabla f_{i} \left(\vx\right) - \nabla f_{i}\left(\vx_*\right) \rVert}_2^2, \shortintertext{from \ref{assumption:components},} &\leq \frac{1}{b} \sum_{i \in B} 2 L_{i} \mathrm{D}_{f_i}\left(\vx, \vx_*\right) \shortintertext{and since \(L_{\mathrm{max}} \geq L_i\) for all \(i= 1, \ldots, n\),} &\leq 2 L_{\mathrm{max}} \frac{1}{b} \sum_{i \in B} \mathrm{D}_{f_i}\left(\vx, \vx_*\right) \\ &= 2 L_{\mathrm{max}} \mathrm{D}_{f_B}\left(\vx, \vx_*\right). \end{align}\] Combining the bound on \(V_{\rm{com}}\) and \(V_{\rm{sub}}\) immediately yields the result.
theorem Let the objective \(F\) satisfy 1 2, where, each component \(f_i\) is additionally \(\mu\)-strongly convex and 6 (\(\rm{A}^{\rm{CVX}}\)), hold. Then, the last iterate \(\vx_T\) of doubly SGD-RR with a stepsize satisfying \(\gamma < 1/\left(\mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}}\right)\) guarantees \[\begin{align} &\mathbb{E}{\lVert \vx_{K+1}^{0} - \vx_* \rVert}_2^2 \leq r^{K p} {\lVert \vx_1^{0} - \vx_* \rVert}_2^2 + C_{\mathrm{var}}^{\rm{sub}} \, \gamma^2 + C_{\mathrm{var}}^{\rm{com}} \, \gamma \end{align}\] where \(p = n/b\) is the number of epochs, \(\vx_* = \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\), \(r = 1 - \gamma \mu\) is the contraction coefficient, \[\begin{align} {3} C_{\mathrm{var}}^{\rm{com}} &= { \frac{4}{\mu b} {\left( \frac{1}{n} \sum^n_{i=1} \sigma^2_{i} \right)} + \frac{4}{\mu} \left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}\right)}^2, \text{ and} \\ C_{\mathrm{var}}^{\rm{sub}} &= \frac{1}{4} \frac{L_{\mathrm{max}}}{\mu} \frac{n}{b^2} \left( \frac{1}{n} \sum^n_{i=1} {\left\lVert \nabla f_i\left(\vx_*\right) \right\rVert}_2^2 \right). \end{align}\]
The key element of the analysis of random reshuffling is that the Lyapunov function that achieves a fast convergence is \({\lVert \vx_k^{i+1} - \vx^{i+1}_{*} \rVert}_2^2\) not \({\lVert \vx_{k}^{i+1} - \vx_* \rVert}_2^2\). This stems from the well-known fact that random reshuffling results in a conditionally biased gradient estimator.
Recall that \(\rvP\) denotes the partitioning of the \(n\) datapoints into \(b\)-minibatches \(\rvP_1, \ldots, \rvP_p\). As usual, we first expand the Lyapunov function as \[\begin{align} &{\lVert \vx^{i+1}_k - \vx_*^{i+1}\rVert}_2^2 \\ &\;= {\lVert \Pi_{\mathcal{X}}(\vx^{i}_k - \gamma \, \rvvg_{\rvP_i}(\vx^i_k)) - \Pi_{\mathcal{X}}(\vx_*^i - \gamma \, \nabla f_{\rvP_i}\left(\vx_*\right) ) \rVert}_2^2 \shortintertext{and since the projection onto a convex set under a Euclidean metric is non-expansive,} &\;\leq {\lVert (\vx^{i}_k - \gamma \, \rvvg_{\rvP_i}(\vx^i_k)) - (\vx_*^i - \gamma \, \nabla f_{\rvP_i}\left(\vx_*\right) ) \rVert}_2^2 \\ &\;= {\lVert \vx^{i}_k - \vx_*\rVert}_2^2 -2 \gamma \left\langle \vx^{i}_k - \vx_*^i \; , \rvvg_{\rvP_i}(\vx^i_k) - \nabla f_{\rvP_i} \left(\vx_*\right) \right\rangle \\ &\qquad+ \gamma^2 {\lVert \rvvg_{\rvP_i}(\vx^i_k) - \nabla f_{\rvP_i}\left(\vx_*\right)\rVert}_2^2. \end{align}\]
Taking expectation over the Monte Carlo noise conditional on the partitioning \(\rvP\), \[\begin{align} &\mathbb{E}_{\varphi}{\lVert \vx^{i+1}_k - \vx_*^{i+1}\rVert}_2^2 \\ &\;= {\lVert \vx^{i}_k - \vx_*^i\rVert}_2^2 -2 \gamma \left\langle \vx^{i}_k - \vx_*^i \; , \mathbb{E}_{\varphi}[\rvvg_{\rvP_i}(\vx^i_k)] - \nabla f_{\rvP_i} \left(\vx_*\right) \right\rangle \\ &\qquad+ \gamma^2 \mathbb{E}_{\varphi} {\lVert \rvvg_{\rvP_i}(\vx^i_k) - \nabla f_{\rvP_i} \left(\vx_*\right) \rVert} \\ &\;= {\lVert \vx^{i}_k - \vx_*^i \rVert}_2^2 -2 \gamma \left\langle \vx^{i}_k - \vx_*^i \; , \nabla f_{\rvP_i}(\vx_k^i) - \nabla f_{\rvP_i}\left(\vx_*\right) \right\rangle \\ &\qquad+ \gamma^2 \, \mathbb{E}_{\varphi} {\lVert \rvvg_{\rvP_i}(\vx^i_k) - \nabla f_{\rvP_i} \left(\vx_*\right) \rVert}_2^2. \end{align}\]
From the three-point identity, we can more precisely characterize the effect of the conditional bias such that \[\begin{align} &\left\langle \vx^{i}_k - \vx_*^i \; , \nabla f_{\rvP_i}(\vx_k^i) - \nabla f_{\rvP_i}\left(\vx_*\right) \right\rangle \\ &= \mathrm{D}_{f_{\rvP_i}}(\vx_*^i, \vx_k^i) + \mathrm{D}_{f_{\rvP_i}}(\vx_k^i, \vx_*) - \mathrm{D}_{f_{\rvP_i}}(\vx_*^i, \vx_*). \end{align}\]
For the gradient noise, \[\begin{align} &\mathbb{E}_{\varphi} {\lVert \rvvg_{\rvP_i}(\vx^i_k) - \nabla f_{\rvP_i} \left(\vx_*\right) \rVert}_2^2 \\ &\;= \mathbb{E}_{\varphi} {\lVert \rvvg_{\rvP_i}(\vx^i_k) - \rvvg_{\rvP_i} \left(\vx_*\right) + \rvvg_{\rvP_i}\left(\vx_*\right) - \nabla f_{\rvP_i} \left(\vx_*\right) \rVert}_2^2 \\ &\;\leq 2 \mathbb{E}_{\varphi} {\lVert \rvvg_{\rvP_i}(\vx^i_k) - \rvvg_{\rvP_i} \left(\vx_*\right) \rVert}_2^2 + 2 \mathbb{E}_{\varphi} {\lVert \rvvg_{\rvP_i}\left(\vx_*\right) - \nabla f_{\rvP_i} \left(\vx_*\right) \rVert}_2^2 \\ &\;= 2 \, \mathbb{E}_{\varphi} {\lVert \rvvg_{\rvP_i}(\vx^i_k) - \rvvg_{\rvP_i} \left(\vx_*\right) \rVert}_2^2 + 2 \, \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right], \shortintertext{and from \ref{thm:reshuffling95expected95smoothness},} &\;\leq 4 \left( \mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}} \right) \mathrm{D}_{f_{\rvP_i}}(\vx_k^i, \vx_*) + 2 \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] \end{align}\] Notice the variance term \(\mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right]\). This quantifies the amount of deviation from the trajectory of singly stochastic random reshuffling. As such, it quantifies how slower we will be compared to its fast rate.
Now, we will denote the \(\sigma\)-algebra formed by the randomness and the iterates up to the \(i\)th step of the \(k\)th epoch as \(\mathcal{F}_k^i\) such that \((\mathcal{F}_k^i)_{k \geq 1, i \geq 1}\) is a filtration. Then, \[\begin{align} &\mathbb{E}_{\rvveta_{k}^{i} \sim \varphi}\left[ {\lVert \vx^{i+1}_k - \vx_*^{i+1}\rVert}_2^2 \,\middle|\, \mathcal{F}_k^i \right] \\ &\leq {\lVert \vx^{i}_k - \vx_*^i\rVert}_2^2 \\ &\quad-2 \gamma \left( \mathrm{D}_{f_{\rvP_i}}(\vx_*^i, \vx_k^i) + \mathrm{D}_{f_{\rvP_i}}(\vx_k^i, \vx_*) - \mathrm{D}_{f_{\rvP_i}}(\vx_*^i, \vx_*) \right) \\ &\quad+ 4 \gamma^2 \left(\mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}}\right) \mathrm{D}_{f_{\rvP_i}}(\vx_k^i, \vx_*) \\ &\quad+ 2 \gamma^2 \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right]. \shortintertext{Now, the \(\mu\)-strong convexity of the component functions imply \(\mathrm{D}_{f_{\rvP_i}}\left(\vx_*^i, \vx_k^i\right) \leq \frac{\mu}{2} {\lVert \vx_k^i - \vx_*^i \rVert}_2^2\). Therefore,} &\leq {\lVert \vx^{i}_k - \vx_*^i\rVert}_2^2 \\ &\quad- 2 \gamma \left( \frac{\mu}{2} {\lVert \vx_k^i - \vx_*^i \rVert}_2^2 + \mathrm{D}_{f_{\rvP_i}}(\vx_k^i, \vx_*) - \mathrm{D}_{f_{\rvP_i}}(\vx_*^i, \vx_*) \right) \\ &\quad+ 4 \gamma^2 \left( \mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}} \right) \mathrm{D}_{f_{\rvP_i}}(\vx_k^i, \vx_*) \\ &\quad+ 2 \gamma^2 \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] , \shortintertext{and reorganizing the terms,} &= \left(1 - \gamma \mu \right) {\lVert \vx^{i}_k - \vx_*^i\rVert}_2^2 \\ &\quad-2 \gamma \left(1 - 2 \gamma \left(\mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}} \right)\right) \mathrm{D}_{f_{\rvP_i}}(\vx_k^i, \vx_*) \\ &\quad+ 2 \gamma \, \mathrm{D}_{f_{\rvP_i}}(\vx_*^i, \vx_*) \\ &\quad+ \gamma^2 2 \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right]. \end{align}\]
Taking full expectation, \[\begin{align} &\mathbb{E}{\lVert \vx^{i+1}_k - \vx_*^{i+1}\rVert}_2^2 \nonumber \\ &\;\leq \left(1 - \gamma \mu \right) \mathbb{E}{\lVert \vx^{i}_k - \vx_*^i\rVert}_2^2 \nonumber \\ &\quad- 2 \gamma \left(1 - 2 \gamma \left(\mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}} \right)\right) \mathbb{E}\left[ \mathrm{D}_{f_{\rvP_i}}\left(\vx_k^n, \vx_*\right) \right] \nonumber \\ &\quad+ 2 \gamma \, \mathbb{E}\left[ \mathrm{D}_{f_{\rvP_i}}\left(\vx_*^i, \vx_*\right) \right] \nonumber \\ &\quad+ 2 \gamma^2 \mathbb{E}\left[ \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] \right], \nonumber \shortintertext{and as long as \(\gamma < {1}/{\left( 2 \left(\mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}}\right) \right)}\)} &\;\leq \left(1 - \gamma \mu \right) \mathbb{E}{\lVert \vx^{i}_k - \vx_*^i\rVert}_2^2 + 2 \gamma \underbrace{ \mathbb{E}\left[ \mathrm{D}_{f_{\rvP_i}}\left(\vx_*^i, \vx_*\right) \right] }_{T_{\text{err}}} \nonumber \\ &\quad+ 2 \gamma^2 \underbrace{ \mathbb{E}\left[ \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] \right] }_{T_{\text{var}}} . \label{eq:reshuffling95convergence95main95bound} \end{align}\tag{31}\]
3.2.3.1 Bounding \(T_{\mathrm{err}}\)
From the definition of the Bregman divergence and \(L\)-smoothness, for all \(j = 1, \ldots, n\), notice that we have \[\begin{align} \mathrm{D}_{f_j}\left( \vy, \vx \right) &= f_j\left(\vy\right) - f_j\left(\vx\right) - \left\langle \nabla f_j\left(\vx\right), \vy - \vx \right\rangle \nonumber \\ &\leq \frac{L}{2} {\left\lVert \vy - \vx \right\rVert}_2^2. \label{eq:bregman95smooth95upperbound} \end{align}\tag{32}\] for all \((\vx, \vx') \in \mathcal{X}^2\). Given this, the Lyapunov error term \[\begin{align} \mathbb{E}\left[ \mathrm{D}_{f_{\rvP_i}}\left(\vx_k^i, \vx_*\right) \right] &= \mathbb{E}\left[ \frac{1}{b} \sum_{j \in \rvP_i} \mathrm{D}_{f_{j}}\left(\vx_k^i, \vx_*\right) \right] \nonumber \shortintertext{can be bounded using \(L\)-smoothness by \ref{eq:bregman95smooth95upperbound},} &\leq \mathbb{E}\left[ \frac{1}{b} \sum_{j \in \rvP_i} \frac{L_{j}}{2} {\lVert \vx_k^i - \vx_*\rVert}_2^2 \right] \nonumber \shortintertext{and \(L_{\rm{max}} \geq L_i\) for all \(i = 1, \ldots, n\),} &\leq \frac{L_{\mathrm{max}}}{2} \mathbb{E}\left[ \frac{1}{b} \sum_{j \in \rvP_i} {\lVert \vx_k^i - \vx_*\rVert}_2^2 \right] \nonumber \\ &= \frac{L_{\mathrm{max}}}{2} \mathbb{E} {\lVert \vx_k^i - \vx_*\rVert}_2^2. \label{eq:reshuffling95convergence95Terr} \end{align}\tag{33}\]
The squared error \({\lVert \vx_k^i - \vx_*\rVert}_2^2\) is bounded in [thm:reshuffle95variance] as \[\mathbb{E} {\lVert \vx_k^i - \vx_*\rVert}_2^2 \leq \epsilon_{\mathrm{sfl}}^2 \triangleq \frac{\gamma^2 n }{4 b^2} \, \tau^2 < \infty. \label{eq:reshuffling95epsilon}\tag{34}\]
3.2.3.2 Bounding \(T_{\mathrm{var}}\)
Now, let’s take a look at the variance term. First, notice that, by the Law of Total Expectation, \[\begin{align} \mathbb{E}\left[ \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] \right] = \mathbb{E}\left[ \mathbb{E}\left[ \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] \mid \rvP \right] \right]. \end{align}\] Here, \[\mathbb{E}\left[ \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] \mid \rvP \right]\] is the variance from selecting \(b\) samples without replacement. We can thus apply [thm:expected95variance95lemma] with \(b_{\mathrm{eff}} = \frac{(n-1) b}{n-b}\) such that \[\begin{align} &\mathbb{E}\left[ \mathrm{tr} \mathbb{V}_{\varphi}\left[ \rvvg_{\rvP_i}\left(\vx_*\right) \right] \mid \rvP \right] \nonumber \\ &\;\leq \frac{n - b}{\left(n - 1\right) b} \left(\frac{1}{n} \sum^n_{j=1} \sigma^2_{j}\right) + \frac{n \left(b - 1\right)}{\left(n - 1\right) b} {\left(\frac{1}{n} \sum^n_{j=1} \sigma_{j}\right)}^2, \nonumber \shortintertext{which we will denote as} &\;= \sigma^2 \label{eq:reshuffling95convergence95variance} \end{align}\tag{35}\] for clarity. Also, notice that \(\sigma^2\) no longer depends on the partitioning.
3.2.3.3 Per-step Recurrence Equation
Applying 33 35 to 31 , we now have the recurrence equation \[\begin{align} \mathbb{E}{\lVert \vx^{i+1}_k - \vx_*^{i+1}\rVert}_2^2 &\leq \left(1 - \gamma \mu \right) \mathbb{E}{\lVert \vx^{i}_k - \vx_*^i\rVert}_2^2 \\ &\qquad+ L_{\mathrm{max}} \epsilon_{\mathrm{sfl}}^2 \, \gamma + 2 \sigma^2 \, \gamma^2. \end{align}\]
Now that we have a contraction of the Lyapunov function \(\mathbb{E}{\lVert \vx_k^{i+1} - \vx_*^{i+1} \rVert}_2^2\), it remains to convert this that the Lyapunov function bounds our objective \(\mathbb{E}{\lVert \vx_k^{i+1} - \vx_* \rVert}_2^2\). This can be achieved by noticing that, at the end of each epoch, we have \(\vx_{k+1} - \vx_* = \vx^{p}_{k} - \vx^p_*\), and equivalently, we have \(\vx_k - \vx_* = \vx^{0}_k - \vx^{0}_*\) at the beginning of the epoch. The fact that the relationship with the original objective is only guaranteed at the endpoints (beginning and end of the epoch) is related to the fact that the bias of random reshuffling starts increasing at the beginning of the epoch and starts decreasing near the end.
3.2.3.4 Per-Epoch Recurrence Equation
Nevertheless, this implies that by simply unrolling the recursion as in the analysis of regular SGD, we obtain a per-epoch contraction of \[\begin{align} \mathbb{E}{\lVert \vx_{k+1}^0 - \vx_* \rVert}_2^2 &\leq {\left(1 - \gamma \mu\right)}^{p} \mathbb{E} {\lVert \vx_{k}^0 - \vx_* \rVert}_2^2 \\ &\quad+ \left( L_{\mathrm{max}} \epsilon^2_{\mathrm{sfl}} \gamma + 2 \sigma^2 \gamma^2 \right) \left(\sum^{p-1}_{i=0} {\left(1 - \mu \gamma\right)}^{i} \right). \end{align}\] And after \(K\) epochs, \[\begin{align} &\mathbb{E}{\lVert \vx_{K+1}^0 - \vx_* \rVert}_2^2 \leq {\left(1 - \gamma \mu\right)}^{p K} \mathbb{E} {\lVert \vx_{0}^0 - \vx_* \rVert}_2^2 \\ &+ \left( L_{\mathrm{max}} \epsilon^2_{\mathrm{sfl}} \gamma + 2 \sigma^2 \gamma^2 \right) \left(\sum^{p-1}_{i=0} {\left(1 - \mu \gamma\right)}^{i} \right) \left(\sum^{pK-1}_{j=0} {\left(1 - \mu \gamma\right)}^{p j} \right). \end{align}\] Note that \(T = pK\).
As done by [41], the product of sums can be bounded as \[\begin{align} &\left(\sum^{p-1}_{i=0} {\left(1 - \mu \gamma\right)}^{i} \right) \left(\sum^{T-1}_{j=0} {\left(1 - \mu \gamma\right)}^{p j} \right) \\ &\;= \sum^{p-1}_{i=0} \sum^{T-1}_{j=0} {\left(1 - \mu \gamma\right)}^{i} {\left(1 - \mu \gamma\right)}^{p j} \\ &\;= \sum^{p-1}_{i=0} \sum^{T-1}_{j=0} {\left(1 - \mu \gamma\right)}^{i + pj} \\ &\;= \sum^{Tp-1}_{i=0} {\left(1 - \mu \gamma\right)}^{i} \\ &\;\leq \sum^{\infty}_{i=0} {\left(1 - \mu \gamma\right)}^{i} \\ &\;\leq \frac{1}{\gamma \mu}. \end{align}\] Then, \[\begin{align} &\mathbb{E}{\lVert \vx_{K+1}^0 - \vx_* \rVert}_2^2 \\ &\;\leq {\left(1 - \gamma \mu\right)}^{p K} \mathbb{E} {\lVert \vx_{0}^0 - \vx_* \rVert}_2^2 + \frac{1}{\gamma \mu} \left( L_{\mathrm{max}} \epsilon^2_{\mathrm{sfl}} \gamma + 2 \sigma^2 \gamma^2 \right) \\ &= {\left(1 - \gamma \mu\right)}^{p K} \mathbb{E} {\lVert \vx_{0}^0 - \vx_* \rVert}_2^2 + \frac{\epsilon^2_{\mathrm{sfl}} }{\mu} + \frac{2\sigma^2}{\mu} \gamma. \end{align}\] Plugging in the value of \(\epsilon_{\mathrm{sfl}}^2\) from 34 , we have \[\begin{align} \mathbb{E}{\lVert \vx_{K+1}^0 - \vx_* \rVert}_2^2 &\leq {\left(1 - \gamma \mu\right)}^{p K} \mathbb{E} {\lVert \vx_{0}^0 - \vx_* \rVert}_2^2 \\ &\qquad+ \frac{L_{\mathrm{max}} n \sigma^2_{\mathrm{sub}}}{4 b^2 \mu} \gamma^2 + \frac{2\sigma^2}{\mu} \gamma. \end{align}\] This implies \[\begin{align} &\mathbb{E}{\lVert \vx_{K+1}^{0} - \vx_* \rVert}_2^2 \leq r^{K n/b} {\lVert \vx_1^{0} - \vx_* \rVert}_2^2 + C_{\rm{var}}^{\mathrm{sub}} \, \gamma^2 + C_{\rm{var}}^{\mathrm{com}} \, \gamma, \end{align}\] where \(r = 1 - \gamma \mu\), \[\begin{align} {3} C_{\rm{var}}^{\mathrm{sub}} &= \frac{1}{4} \frac{L_{\mathrm{max}}}{\mu} \frac{n}{b^2} \left( \frac{1}{n} \sum^n_{i=1} {\left\lVert \nabla f_i\left(\vx_*\right) \right\rVert}_2^2 \right), \text{and} \\ C_{\mathrm{var}}^{\rm{com}} &= { \frac{2}{\mu} \frac{n-b}{(n-1) b} {\left( \frac{1}{n} \sum^n_{i=1} \sigma^2_{i} \right)} + \frac{2}{\mu} \frac{n \left(b-1\right)}{\left(n-1\right) b} {\left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}\right)}^2 }. \end{align}\] Applying the fact that \((n-b)/n \leq (n-1)/n \leq 2\) for all \(n \geq 2\) yields the simplified constants in the statement.
theorem Let the objective \(F\) satisfy 1 2, where each component \(f_i\) is additionally \(\mu\)-strongly convex, and 6 (\(\rm{A}^{\rm{CVX}}\)), hold. Then, the last iterate \(\vx_T\) of doubly SGD-RR is \(\epsilon\)-close to the global optimum \(\vx_* = \mathop{\mathrm{arg\,max}}_{\vx \in \mathcal{X}} F\left(\vx\right)\) such that \(\mathbb{E}{\left\lVert \vx_T - \vx_* \right\rVert}_2^2 \leq \epsilon\) if \[\begin{align}T \,\geq \; & {\textstyle \max\left( 4 C_{\rm{var}}^{\mathrm{com}} \frac{1}{\epsilon} + C_{\rm{var}}^{\mathrm{sub}} \frac{1}{\sqrt{\epsilon}}, \; C_{\mathrm{bias}} \right) \log \left(2 \, {\small\lVert \vx_1^0 - \vx_* \rVert}_2^2 \frac{1}{\epsilon} \right) } \end{align}\] for some fixed stepsize, where \(T = K p = K {n}/{b}\), \[\begin{align} C_{\mathrm{bias}} &= \left(\mathcal{L}_{\mathrm{max}} + L\right)/{\mu} \\ C_{\rm{var}}^{\mathrm{com}} &= { \frac{2}{b} \left(\frac{1}{n} \sum^n_{i=1} \frac{\sigma_{i}^2}{\mu^2} \right) + { 2 \left( \frac{1}{n}\sum^n_{i=1} \frac{\sigma_{i}}{\mu} \right)}^2 }, \\ C_{\rm{var}}^{\mathrm{sub}} &= \sqrt{\frac{L_{\mathrm{max}}}{\mu}} \frac{ \sqrt{n}}{b} \frac{\tau}{\mu}. \end{align}\]
From the result of [thm:strongly95convex95reshuffling95sgd95convergence], we can invoke [thm:geometric95complexity95squared] with \[\begin{align} {3} A &= \frac{L_{\mathrm{max}} n}{4 b^2 \mu} \tau^2,\quad \\ B &= \frac{2}{\mu}\left( \frac{n-b}{(n-1) b} \left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}^2\right) + \frac{n \left(b-1\right)}{\left(n-1\right) b} {\left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}\right)}^2 \right), \\ C &= \mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}}. \end{align}\] Then, an \(\epsilon\) accurate solution in expectation can be obtained after \[\begin{align} T &\geq \max\Bigg( \underbrace{\frac{2 B}{\mu}}_{\triangleq C_1} \frac{1}{\epsilon} + \underbrace{\frac{\sqrt{2 A}}{\mu}}_{\triangleq C_2} \frac{1}{\sqrt{\epsilon}}, \; \frac{\mathcal{L}_{\mathrm{max}} + L_{\mathrm{max}}}{\mu} \Bigg) \log \left(2 r_0^2 \frac{1}{\epsilon} \right) \end{align}\] iterations with a stepsize of \[\begin{align} \gamma = \min\left( \frac{- B + \sqrt{B^2 + 2 A \epsilon }}{2 A}, \frac{1}{C} \right). \end{align}\] To make the iteration complexity more precise, the terms \(C_1, C_2\) can be organized as \[\begin{align} C_1 &= \frac{2 B}{\mu} = \frac{2}{\mu} \Bigg( \, \frac{2}{\mu} \, \Bigg\{\, \frac{n-b}{(n-1) b} \left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}^2\right) \\ &\qquad\qquad\qquad+ \frac{n \left(b-1\right)}{\left(n-1\right) b} {\left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}\right)}^2 \,\Bigg\} \,\Bigg) \\ &= \frac{4}{\mu^2} \Bigg( \frac{n-b}{(n-1) b} \left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}^2\right) + \frac{n \left(b-1\right)}{\left(n-1\right) b} \, {\left(\frac{1}{n} \sum^n_{i=1} \sigma_{i}\right)}^2 \Bigg) \\ C_2 &= \frac{\sqrt{2 A}}{\mu} \\ &= \sqrt{2 \frac{L_{\mathrm{max}} n}{4 b^2 \mu} \tau^2 } \frac{1}{\mu^2} \\ &= \frac{\sqrt{L_{\mathrm{max}}} \, \tau \, \sqrt{n}}{\sqrt{2} b \mu^{\frac{3}{2}}} \\ &\leq \frac{\sqrt{L_{\mathrm{max}}}}{\mu^{\frac{3}{2}}} \frac{ \sqrt{n}}{b} \, \tau. \end{align}\] Applying the fact that \((n-b)/n \leq (n-1)/n \leq 2\) for all \(n \geq 2\) yields the simplified constants in the statement.
When \(\sigma_{i} = 0\) for all \(i = 1, \ldots, n\), the anytime convergence bound [thm:strongly95convex95reshuffling95sgd95convergence] in the Appendix reduces exactly to Theorem 1 of [41]. Therefore, [thm:strongly95convex95reshuffling95sgd95convergence] is a strict generalization of SGD-RR to the doubly stochastic setting.
Using \(m\)-sample Monte Carlo improves the constants as follows:
Corollary 2. Let the assumptions of [thm:strongly95convex95reshuffling95sgd95complexity] hold. Then, for \(1 < b < n\) and \(m\)-sample Monte Carlo, the same guarantees hold with the constant \[\begin{align} C_{\rm{var}}^{\mathrm{com}} = \frac{2}{m b} {\left(\frac{1}{n} \sum^n_{i=1} \frac{\sigma_{i}^2}{\mu^2} \right)} + \frac{2}{m} \, {\left( \frac{1}{n}\sum^n_{i=1} \frac{\sigma_{i}}{\mu} \right)}^2. \end{align}\]
Compared to doubly SGD, doubly SGD-RR improves the dependence on the subsampling noise \(\tau^2\) from \(\mathcal{O}\left(1/\epsilon\right)\) to \(\mathcal{O}(1/\sqrt{\epsilon})\). Therefore, random reshuffling does improve the complexity of doubly SGD. Unfortunately, it also means that it does not achieve a better asymptotic complexity as in the finite sum setting. However, non-asymptotically, if the subsampling noise dominates component estimation noise, doubly SGD-RR will behave closely to an \(\mathcal{O}(1/\sqrt{\epsilon})\) (or equivalently, \(\mathcal{O}(1/T)\)) algorithm.
As was the case with independent subsampling, increasing \(b\) also reduces component estimation noise for RR-SGD. However, the impact on the complexity is more subtle. Consider that the iteration complexity is \[{\textstyle \mathcal{O}\left( \kappa_{\sigma}^2 \left(\frac{1}{mb} + \frac{1}{m}\right) \frac{1}{\epsilon} + \kappa \, \kappa_{\tau} \, \frac{\sqrt{n}}{b} \frac{1}{\sqrt{\epsilon}} \right), } \label{eq:doublysgdrr95complexity}\tag{36}\] where \(\kappa_{\sigma} = { \max_{i=1,\ldots, n} \sigma_{i}}/{\mu}\), \(\kappa_{\tau} = \tau/{\mu}\) and \(\kappa = {\max_{i=1,\ldots, n} L_{i}}/{\mu}\). The \(1/\epsilon\) term decreases the fastest with \(m\). Therefore, it might seem that increasing \(m\) is advantageous. However, the \(1/\sqrt{\epsilon}\) term has a \(\mathcal{O}\left(\sqrt{n}\right)\) dependence on the dataset size, which would be non-negligible for large datasets. As a result, in the large \(n\), large \(\epsilon\) regime, increasing \(b\) over \(m\) should be more effective.
36 also implies that, for dependent estimators, doubly SGD-RR achieves an asymptotic speedup of \({n}/{b}\) compared to full-batch SGD with only component estimation noise. Assume that the sample complexity of a single estimate is \(\Theta(m b)\) (\(\Theta(m n)\) for full-batch). Then, the sample complexity of doubly SGD-RR is \(\mathcal{O}\left(b\frac{1}{\epsilon}\right)\) and \(\mathcal{O}\left(n\frac{1}{\epsilon}\right)\) for full-batch SGD. However, the \(n/b\) seed-up comes from correlations. Therefore, for independent estimators, the asymptotic complexity of the two is equal.
4 Simulation↩︎
4.0.0.1 Setup
We evaluate the insight on the tradeoff between \(b\) and \(m\) for correlated estimators on a synthetic problem. In particular, we set \[f_i\left(\vx; \rvveta\right) = \frac{L_i}{2} {\left\lVert \vx - \vx_i^* + \rvveta \right\rVert}_2^2,\] where the smoothness constants \(L_i \sim \text{Inv-Gamma}(\frac{1}{2}, \frac{1}{2})\) and the stationary points \(\vx_i^* \sim \mathcal{N}(\symbfup{0}_d, s^2\symbfup{I}_d)\) are sampled randomly, where \(\symbfup{0}_d\) is a vector of \(d\) zeros and \(\symbfup{I}_d\) is a \(d \times d\) identity matrix. Then, we compute the gradient variance on the global optimum, corresponding to computing the BV (2) constant. Note that \(s^2\) here corresponds to the “heterogeneity” of the data. We make the estimators dependent by sharing \(\rvveta_1, \ldots, \rvveta_m\) across the batch.
4.0.0.2 Results
The results are shown in 1. At low heterogeneity, there exists a “sweet spot” between \(m\) and \(b\). However, this sweet spot moves towards large values of \(b\), where, at high heterogeneity levels, the largest values of \(b\) are more favorable. Especially in the low budget regime where \(m b \ll n\), the largest \(b\) values appear to achieve the lowest variance. This confirms our theoretical results that a large \(b\) should be preferred on challenging (large number of datapoints, high heterogeneity) problems.
5 Discussions↩︎
5.1 Applications↩︎
In 8, we establish 6 7 on the following applications:
ERM with Randomized Smoothing: In this problem, we consider ERM, where the model weights are perturbed by noise. This variant of ERM has recently gathered interest as it is believed to improve generalization performance [21], [22]. In 8.1, we establish 6 \((\rm{A}^{\rm{ITP}})\) under the interpolation assumption.
Reparameterization Gradient: In certain applications, e.g., variational inference, generative modeling, and reinforcement learning (see [42, Sec. 5]), the optimization problem is over the parameters of some distribution, which is taken expectation over. Among gradient estimators for this problem, the reparameterization gradient is widely used due to lower variance [43]. For this, in 8.2, we establish 6 (\(\rm{A}^{\rm{CVX}}\)) and (B) by assuming a convexity and smooth integrand.
5.2 Related Works↩︎
Unlike SGD in the finite sum setting, doubly SGD has received little interest. Previously, [7], [8], [23] have studied the convergence of variance-reduced gradients [44] specific to the doubly stochastic setting under the uniform Lipchitz integrand assumption (\(\vg_i(\cdot; \veta)\) is \(L\)-Lipschitz for all \(\veta\)). Although this assumption has often been used in the stochastic optimization literature [3], [24], [25], [45], it is easily shown to be restrictive: for some \(L\)-smooth \(\widehat{f_i}\left(\vx\right)\), \(\nabla f_i\left(\vx; \rvveta\right) = \nabla \widehat{f}_i\left(\vx\right) + x_1 \rvveta\) is not \(L\)-Lipschitz unless the support of \(\veta\) is compact. In contrast, we established results under weaker conditions. We also provide a discussion on the relationships of different conditions in 6.
Furthermore, we extended doubly SGD to the case where random reshuffling is used in place of sampling independent batches. In the finite-sum setting, the fact that SGD-RR converges faster than independent subsampling (SGD) has been empirically known for a long time [46]. While [47] first demonstrated that SGD-RR can be fast for quadratics, a proof under general conditions was demonstrated recently [48]: In the strongly convex setting, [41] [49], [50] establish a \(\mathcal{O}\left(1/\sqrt{\epsilon}\right)\) complexity to be \(\epsilon\)-accurate, which is tight in terms of asymptotic complexity [51]–[53].
Lastly, [19], [20], [54] provided convergence guarantees for doubly SGD for ERM of random feature kernel machines. However, these analyses are based on concentration arguments that doubly SGD does not deviate too much from the optimization path of finite-sum SGD. Unfortunately, concentration arguments require stronger assumptions on the noise, and their analysis is application-specific. In contrast, we provide a general analysis under the general ER assumption.
5.3 Conclusions↩︎
In this work, we analyzed the convergence of SGD with doubly stochastic and dependent gradient estimators. In particular, we showed that if the gradient estimator of each component satisfies the ER and BV conditions, the doubly stochastic estimator also satisfies both conditions; this implies convergence of doubly SGD.
5.3.0.1 Practical Recommendations
An unusual conclusion of our analysis is that when Monte Carlo is used with minibatch subsampling, it is generally more beneficial to increase the minibatch size \(b\) instead of the number of Monte Carlo samples \(m\). That is, for both SGD and SGD-RR, increasing \(b\) decreases the variance in a rate close to \(1/b\) when
the gradient variance of the component gradient estimators varies greatly such that \({\left( \frac{1}{n} \sum^n_{i=1} \sigma_{i} \right)}^2 \ll \frac{1}{n} \sum^n_{i=1} \sigma_{i}^2\) or when
the estimators are independent as \(\rho = 0\).
Surprisingly, such a benefit persists even in the interpolation regime \(\tau^2 = 0\). On the contrary, when the estimators are both dependent and have similar variance, it is necessary to increase both \(m\) and \(b\), where a sweet spot between the two exists. However, such a regime is unlikely to occur in practice; in statistics and machine learning applications, the variance of the gradient estimators tends to vary greatly due to the heterogeneity of data.
Acknowledgements↩︎
The authors would like to thank the anonymous reviewers for their comments, Jason Altschuler (UPenn) for numerous suggestions that strengthened the work, and Trevor Campbell (UBC) for pointing out a typo.
K. Kim was supported by a gift from AWS AI to Penn Engineering’s ASSET Center for Trustworthy AI; Y.-A. Ma was funded by the NSF Grants [NSF-SCALE MoDL-2134209] and [NSF-CCF-2112665 (TILOS)], the U.S. Department Of Energy, Office of Science, as well as the DARPA AIE program; J. R. Gardner was supported by NSF award [IIS-2145644].
Impact Statement↩︎
This paper presents a theoretical analysis of stochastic gradient descent under doubly stochastic noise to broaden our understanding of the algorithm. The work itself is theoretical, and we do not expect direct societal consequences, but SGD with doubly stochastic gradients is widely used in various aspects of machine learning and statistics. Therefore, we inherit the societal impact of the downstream applications of SGD.
6 Gradient Variance Conditions↩︎
![Figure 2: Implications between general gradient variance conditions for some unbiased estimator \rvvg\left(\vx\right) = \nabla f \left(\vx; \rvveta\right) of \nabla f\left(\vx\right) = \mathbb{E} \rvvg\left(\vx\right). The dashed arrows () hold if f is further assumed to be QFG; the dotted arrow () holds if the integrand f(\vx; \veta) is uniformly convex such that it is convex with respect to \vx for any fixed \veta. (1), (5), (9), (13) are established by [34]; (2) is proven in [thm:sgisqv]; (3) is proven in [thm:qviswg]; (4) is proven in [thm:qvisqes]; (7) is proven in [thm:usqes]; (8) is proven in [thm:usces]; (6) is proven by [24] but we restate the proof in [thm:qesises]; (11) is proven in [thm:cescer]; (10), (12) hold trivially if \vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} f\left(\vx\right) are all stationary points.](/article/2406.00920/_pdflatex/htrsfjln.png "fig:")
In this section, we will discuss some additional aspects of the ER and ES conditions introduced in 2.3. We will also look into alternative gradient variance conditions that have been proposed in the literature and their relationship with the ER condition.
6.1 Definitions↩︎
For this section, we will use the following additional definitions:
Definition 3 (Quadratic Functional Growth; QFG). We say \(f : \mathcal{X} \to \mathbb{R}\) satisfies \(\mu\)-quadratic functional growth if there exists some \(\mu > 0\) such that \[\frac{\mu}{2}{\left\lVert \vx - \vx_* \right\rVert}_2^2 \leq f\left(\vx\right) - f\left(\vx_*\right)\] holds for all \(\vx \in \mathcal{X}\), where \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} f\left(\vx\right)\).
This condition implies that \(f\) grows at least as fast as some quadratic and is weaker than the Polyak-Łojasiewicz. However, for any convex function \(f\) that satisfies this condition also means that \(f\) is \(\mu\)-strongly convex [55].
Definition 4 (Uniform Smoothness). For the unbiased estimator \(\rvvg\left(\vx\right) = \nabla f\left(\vx; \rvveta\right)\) of \(\nabla F(\vx) = \mathbb{E} \nabla f\left(\vx; \rvveta\right) = \mathbb{E} \nabla f\left(\vx; \rvveta\right)\), we say the integrand \(\nabla f_i\left(\vx; \veta\right)\) satisfies uniform \(L\)-smoothness if there exist some \(L < \infty\) such that, for any fixed \(\veta\), \[{\left\lVert \nabla f\left(\vx; \veta\right) - \nabla f\left(\vx'; \veta\right) \right\rVert}_2 \leq L {\left\lVert \vx - \vx' \right\rVert}_2\] holds for all \((\vx, \vx') \in \mathcal{X}^2\) simultaneously.
As discussed in [section:introduction] [section:related], this condition is rather strong: it does not hold for multiplicative noise unless the support is bounded.
Definition 5 (Uniform Convexity). For the unbiased estimator \(\rvvg\left(\vx\right) = \nabla f\left(\vx; \rvveta\right)\) of \(\nabla F(\vx) = \mathbb{E}\nabla f\left(\vx; \rvveta\right)\), we say the integrand \(f\left(\vx; \veta\right)\) is uniformly convex if it is convex for any \(\veta\) such that, for any fixed \(\veta\), \[f\left(\vx; \veta\right) - f\left(\vx'; \veta\right) \leq \left\langle \nabla f\left(\vx; \veta\right), \vx - \vx' \right\rangle\] holds for all \((\vx, \vx') \in \mathcal{X}^2\) simultaneously.
6.2 Additional Gradient Variance Conditions↩︎
For some estimator \(\rvvg\) of \(\nabla f\), the following conditions have been considered in the literature:
Strong growth condition (SG): \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) \right\rVert}_2^2 \leq \rho {\left\lVert \nabla f\left(\vx\right) \right\rVert}_2^2 \end{align}\]
Weak growth condition (WG): \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) \right\rVert}_2^2 \leq \rho \left( f\left(\vx\right) - f\left(\vx_*\right)\right) \end{align}\]
Quadratic variance (QV): \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) \right\rVert}_2^2 \leq \alpha \, {\left\lVert \vx - \vx_* \right\rVert}_2^2 + \beta \end{align}\]
Convex expected smoothness (CES): \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vy\right) \right\rVert}_2^2 \leq 2 \mathcal{L} \mathrm{D}_{f}\left(\vx, \vy\right) \end{align}\]
Convex expected residual (CER): \[\begin{align} \mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vy\right) \right] \leq 2 \mathcal{L} \mathrm{D}_{f}\left(\vx, \vy\right) \end{align}\]
Quadratic expected smoothness (QES): \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vy\right) \right\rVert}_2^2 \leq \mathcal{L}^2 {\left\lVert \vx - \vy \right\rVert}_2^2 \end{align}\]
ABC: \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) \right\rVert}_2^2 \leq A \left( f\left(\vx\right) - f\left(\vx_*\right)\right) + B {\left\lVert \nabla f\left(\vx\right) \right\rVert}_2^2 + C \end{align}\]
Here, \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} f\left(\vx\right)\) is any stationary point of \(f\) and the stated conditions should hold for all \((\vx, \vy) \in \mathcal{X}^2\).
SG was used by [56] to establish the linear convergence of SGD for strongly convex objectives, and \(\mathcal{O}(1/T)\) convergence for convex objectives; WG was proposed by [39] to establish similar guarantees to SG under a verifiable condition; QV was used to establish the non-asymptotic convergence on strongly convex functions by [57], while convergence on general convex functions was established by [58], including stochastic proximal gradient descent; QES was used by [25] to establish one of the earliest general non-asymptotic convergence results for SGD on strongly convex objectives; ABC was used by [35] to establish convergence of SGD for non-convex smooth functions. (See also [35] for a comprehensive overview of these conditions.) The relationship of these conditions with the ER condition are summarized in 2.
As demonstrated in 2 and discussed by [35], the ABC condition is the weakest of all. However, the convergence guarantees for problems that exclusively satisfy the ABC condition are weaker than others. (For instance, the number of iterations \(T\) has to be fixed a priori.) On the other hand, the ER condition retrieves most of the strongest known guarantees for SGD; some of which were listed in 2.4.
![Figure 3: Implications of assumptions on the components f_1, \ldots, f_n to the minibatch subsampling gradient estimator \nabla f_{\rvB} of F = \frac{1}{n}\left(f_1 + \ldots + f_n\right). (1), (4) are established by [34], while (3) trivially follows from the fact that \vx_*-convexity is strictly weaker than (global) convexity, and (2) was established by [29].](/article/2406.00920/_pdflatex/vxhrlqbi.png "fig:")
6.3 Establishing the ER Condition↩︎
For subsampling estimators, it is possible to establish some of the gradient variance conditions through general assumptions on the components. See some examples in 3. Here, we use the following definitions:
Definition 6. For the finite sum objective \(F = \frac{1}{n}\left(f_1 + \ldots + f_n\right)\), we say interpolation holds if, for all \(i = 1, \ldots, n\), \[f_i\left(\vx_*\right) \leq f_i\left(\vx\right),\] holds for all \(\vx \in \mathcal{X}\), where \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\).
Definition 7. For the finite sum objective \(F = \frac{1}{n}\left(f_1 + \ldots + f_n\right)\), we say the components are \(\vx_*\)-convex if, for all \(i = 1, \ldots, n\), \[f_i\left(\vx_*\right) - f_i\left(\vx\right) \leq \left\langle \nabla f_i\left(\vx_*\right), \vx_* - \vx \right\rangle\] holds for all \(\vx \in \mathcal{X}\), where \(\vx_* \in \mathop{\mathrm{arg\,min}}_{\vx \in \mathcal{X}} F\left(\vx\right)\).
This assumption is a weaker version of convexity; convexity needs to hold with respect to \(\vx_*\) only. It is closely related to star [59] and quasar convexity [36], [60].
6.4 Proofs of Implications in 2↩︎
We prove new implication results between some of the gradient variance conditions discussed in 6.2. In particular, the relationship between the QES and QV against other conditions has not been considered before.
Let \(\rvvg\) be an unbiased estimator of \(\nabla f\). Then,
Figure 4: 
.
The result immediately follows from the fact that \[\mathrm{tr}\mathbb{V}\left[ \rvvg\left(\vx\right) - \rvvg\left(\vx'\right) \right] \leq \mathbb{E}{\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vx'\right) \right\rVert}_2^2\] holds for all \(\vx, \vx' \in \mathcal{X}\).
Let \(\rvvg\) be an unbiased estimator of \(\nabla f\). Then,
Figure 5: 
.
Notice that, by definition, \(\nabla f\left(\vx_*\right) = 0\). Then, \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) \right\rVert}_2^2 &\leq \rho {\left\lVert \nabla f\left(\vx\right) \right\rVert}_2^2 \\ &= \rho {\left\lVert \nabla f\left(\vx\right) - \nabla f\left(\vx_*\right) \right\rVert}_2^2, \shortintertext{applying \(L\)-smoothness of \(f\),} &\leq L^2 \rho {\left\lVert \vx - \vx_* \right\rVert}_2^2. \end{align}\]
Let \(\rvvg\left(\vx\right) = \nabla f\left(\vx; \rvveta\right)\) be an unbiased estimator of \(\nabla f\left(\vx\right) = \mathbb{E} \nabla f\left(\vx; \rvveta\right)\). Then,
Figure 6: 
.
The result immediately follows from the fact that the integrand \(f\left(\vx; \veta\right)\) is \(L\)-smooth with respect to \(\vx\) uniformly over \(\veta\) as \[\begin{align} \mathbb{E}{\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vx'\right) \right\rVert}_2^2 &= \mathbb{E}{\left\lVert \nabla f\left(\vx; \rvveta\right) - \nabla f\left(\vx'; \rvveta\right) \right\rVert}_2^2 \\ &\leq L^2 {\left\lVert \vx - \vx' \right\rVert}_2^2. \end{align}\]
Let \(\rvvg\left(\vx\right) = \nabla f\left(\vx; \rvveta\right)\) be an unbiased estimator of \(\nabla f\left(\vx\right) = \mathbb{E} \nabla f\left(\vx; \rvveta\right)\). Then,
Figure 7: 
.
Since the integrand \(f\left(\vx; \veta\right)\) is both uniformly smooth and convex with respect to \(\vx\) for a any fixed \(\veta\), we have \[\begin{align} &{\left\lVert \nabla f\left(\vx; \veta\right) - \nabla f\left(\vx'; \veta\right) \right\rVert}_2 \\ &\qquad\leq 2 L \left(f\left(\vx; \veta\right) - f\left(\vx'; \veta\right) - \left\langle \nabla f\left(\vx'; \veta\right), \vx - \vx' \right\rangle \right). \end{align}\] Then, \[\begin{align} &\mathbb{E}{\left\lVert \rvvg\left(\vx\right) - \rvvg\left(\vx_*\right) \right\rVert}_2^2 \\ &\;= \mathbb{E}{\left\lVert \nabla f\left(\vx; \rvveta\right) - \nabla f\left(\vx_*; \rvveta\right) \right\rVert}_2^2 \\ &\;\leq 2 L \, \mathbb{E}\left(f\left(\vx; \veta\right) - f\left(\vx_*; \veta\right) - \left\langle \nabla f\left(\vx_*; \veta\right), \vx - \vx' \right\rangle \right) \\ &\;= 2 L \, \left(f\left(\vx\right) - f\left(\vx_*\right) - \left\langle \nabla f\left(\vx_*\right), \vx - \vx' \right\rangle \right) \\ &\;= 2 L \, \left( f\left(\vx\right) - f\left(\vx_*\right)\right) \end{align}\] holds for all \(\vx \in \mathcal{X}\).
Let \(\rvvg\) be an unbiased estimator of \(\nabla F\). Then,
Figure 8: 
.
From the classic inequality \({(a + b)}^2 \leq 2 a^2 + 2 b^2\), we have \[\begin{align} \mathbb{E}{\left\lVert \rvvg(\vx) - \rvvg\left(\vx_*\right) \right\rVert}_2^2 \leq 2 \, \mathbb{E}{\left\lVert \rvvg(\vx) \right\rVert}_2^2 + 2 \, \mathbb{E}{\left\lVert \rvvg(\vx_*) \right\rVert}_2^2. \end{align}\] Now, since QV holds with \(\beta = 0\), we have \(\mathbb{E}{\left\lVert \rvvg(\vx_*) \right\rVert}_2^2 = 0\). Therefore, \[\begin{align} \mathbb{E}{\left\lVert \rvvg(\vx) - \rvvg\left(\vx_*\right) \right\rVert}_2^2 \leq 2 \, \mathbb{E}{\left\lVert \rvvg(\vx) \right\rVert}_2^2 \leq 2 \, \alpha {\left\lVert \vx - \vx_* \right\rVert}_2^2, \end{align}\] where we have applied QV at the last inequality.
Let \(\rvvg\) be an unbiased estimator of \(\nabla f\). Then,
Figure 9: 
.
The result immediately follows from QV as \[\begin{align} \mathbb{E}{\left\lVert \rvvg(\vx) \right\rVert}_2^2 &\leq \alpha {\left\lVert \vx - \vx_* \right\rVert}_2^2, \shortintertext{applying \(\mu\)-quadratic functional growth,} &\leq \frac{2\alpha}{\mu} \left(f\left(\vx\right) - f\left(\vx_*\right) \right). \end{align}\]
Let \(\rvvg\) be an unbiased estimator of \(\nabla f\). Then,
Figure 10: 
.
From QV, we have \[\begin{align} \mathbb{E}{\left\lVert \rvvg(\vx) - \rvvg\left(\vx_*\right) \right\rVert}_2^2 &\leq \mathcal{L}^2 {\left\lVert \vx - \vx_* \right\rVert}_2^2 \shortintertext{and \(\mu\)-quadratic functional growth yields} &\leq \frac{2 \mathcal{L}^2}{\mu} \left(f\left(\vx\right) - f\left(\vx_*\right) \right). \end{align}\]
The strategy applying QFG when proving [thm:qesises] [thm:qviswg] establishes the stronger variant of the ER condition: 6 (B). However, the price for this is that one has to pay for an excess \(\kappa = \mathcal{L} / \mu\) factor, and this strategy works only works for quadratically growing objectives.
7 Proofs↩︎
7.1 Auxiliary Lemmas ([thm:variancewithoutreplacement] [thm:varianceofsum] [thm:geometric95complexity] [thm:geometric95complexity95squared] [thm:average95bregman])↩︎
7.2 Convergence of SGD ([thm:expected95residual95gradient95variance95bound] [thm:strongly95convex95sgd95convergence] [thm:strongly95convex95sgd95complexity])↩︎
7.3 General Variance Bound ([thm:doubly95stochastic95variance])↩︎
lemma expectedvariancelemma
From the formula for the variance of sums, \[\begin{align} {2} &\mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^b \rvvx_i \,\middle|\, \rvB \right] \nonumber \\ &\;= \sum_{i=1}^b \mathrm{tr}\, \mathbb{V}\left[ \rvvx_i \mid \rvB \right] + \sum_{i=1}^b \sum_{j \neq i} \mathrm{tr}\, \mathrm{Cov}\left( \rvvx_i, \rvvx_j \mid \rvB \right). \nonumber \\ &\;\leq \sum_{i=1}^b \mathrm{tr}\, \mathbb{V}\left[ \rvvx_i \mid \rvB \right] + \sum_{i=1}^b \sum_{j \neq i} \rho \, \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \mid \rvB \right]} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_j \mid \rvB \right]} \nonumber \\ &\;= \left(1 - \rho\right) \sum_{i=1}^b \mathrm{tr}\, \mathbb{V}\left[ \rvvx_i \mid \rvB \right] \nonumber \\ &\;\quad+ \rho \sum_{i=1}^b \sum_{j=1}^b \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \mid \rvB \right]} \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_j \mid \rvB \right]} \nonumber \\ &\;= \left(1 - \rho\right) \, \sum_{i=1}^b \mathrm{tr}\, \mathbb{V}\left[ \rvvx_i \mid \rvB \right] + \rho \, {\textstyle\left( \sum_{i=1}^b \sqrt{\mathrm{tr}\mathbb{V}\left[ \rvvx_i \mid \rvB \right]} \right)}^2 \nonumber \\ &\;= \left(1 - \rho\right) \, \rvV + \rho \, \rvS^2. \nonumber \end{align}\] Then, it follows that \[\begin{align} \mathbb{E}\left[ \mathrm{tr}\mathbb{V}\left[ \sum_{i=1}^b \rvvx_i \,\middle|\, \rvB \right] \right] &\leq \rho \mathbb{E}\left[ \rvS^2 \right] + \left(1 -\rho\right) \mathbb{E}\left[ \rvV \right] \\ &= \rho \mathbb{V}\left[ \rvS \right] + \rho {(\mathbb{E}\rvS)}^2 + \left(1 -\rho\right) \mathbb{E}\left[ \rvV \right], \end{align}\] from the basic property of the variance: \[\mathbb{V}\left[ \rvS \right] = \mathbb{E}\left[ \rvS^2 \right] - {(\mathbb{E}\rvS)}^2.\]
Since 3 is the only inequality we use, the equality in the statement holds whenever the equality in 3 holds.
7.4 Doubly Stochastic Gradients↩︎
7.4.1 Expected Residual Condition ([thm:general95conditions])↩︎
7.4.2 Bounded Variance Condition ([thm:bounded95variance])↩︎
7.4.3 Complexity Analysis ([thm:strongly95convex95sgd95sampling95without95replacement])↩︎
7.5 Random Reshuffling of Stochastic Gradients↩︎
7.5.1 Gradient Variance Conditions ([thm:reshuffling95expected95smoothness], [thm:reshuffle95variance])↩︎
7.5.2 Convergence Analysis ([thm:strongly95convex95reshuffling95sgd95convergence])↩︎
7.5.3 Complexity Analysis ([thm:strongly95convex95reshuffling95sgd95complexity])↩︎
8 Applications↩︎
8.1 ERM with Randomized Smoothing↩︎
8.1.1 Description↩︎
Randomized smoothing was originally considered by [61]–[63] in the nonsmooth convex optimization context, where the function is “smoothed” through random perturbation. This scheme has recently renewed interest in the non-convex ERM context as it has been found to improve generalization performance [21], [22]. Here, we will focus on the computational aspect of this scheme. In particular, we will see if we can obtain similar computational guarantees already established in the finite-sum ERM setting, such as those by [34].
Consider the canonical ERM problem, where we are given a dataset \(\mathcal{D} = {\left\{(\vx_i, y_i)\right\}}_{i = 1}^n \in {(\mathcal{X} \times \mathcal{Y})}^{n}\) and solve \[\begin{align} \mathop{\mathrm{minimize}}_{\vw \in \mathcal{W}}\;\; L\left(\vw\right) &= \frac{1}{n} \sum^n_{i=1} \ell\left(f_{\vw}\left(\vx_i\right), y_i\right) + h\left(\vw\right), \end{align}\] where \((\vx_i, y_i) \in \mathcal{X} \times \mathcal{Y}\) are the feature and label of the \(i\)th instance, \(f_{\vw} : \mathcal{X} \to \mathcal{Y}\) is the model, \(\ell : \mathcal{Y} \times \mathcal{Y} \to \mathbb{R}_{\geq 0}\) is a non-negative loss function, and \(h : \mathcal{W} \to \mathbb{R}\) is a regularizer.
For randomized smoothing, we instead minimize \[L\left(\vw\right) = \frac{1}{n} \sum^n_{i=1} R_i\left(\vw\right),\] where the instance risk is defined as \[r_i\left(\vw\right) = \mathbb{E}_{\rvvepsilon \sim \varphi} \ell\left( f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right)\] for some noise distribution \(\vepsilon \sim \varphi\). The goal is to obtain a solution \(\vw^* = \mathop{\mathrm{arg\,min}}_{\vw \in \mathcal{W}} L\left(\vw\right)\) that is robust to such perturbation.
The integrand of the gradient estimator of the instance risk is defined as \[\begin{align} \vg_{i}\left(\vw; \veta\right) &= \nabla_{\vw} \ell\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) \\ &= \frac{\partial f_{\vw + \rvvepsilon}\left(\vx_i\right)}{\partial \vw} \ell'\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right), \end{align}\] where it is an unbiased estimate of the instance risk such that \[\mathbb{E}\rvvg_{i}\left(\vw\right) = \nabla R_i\left(\vw\right).\] The key challenge in analyzing the convergence of SGD in the ERM setting is dealing with the Jacobian \(\frac{\partial f_{\vw + \rvvepsilon}\left(\vx_i\right)}{\partial \vw}.\) Even for simple toy models, analyzing the Jacobian without relying on strong assumptions is hard. In this work, we will assume that it is bounded by an instance-dependent constant.
8.1.2 Preliminaries↩︎
We use the following assumptions:
Assumption 8.
Let the mapping \(\hat{y} \mapsto \ell\left(\hat{y}, y\right)\) is convex and \(L\)-smooth for any \(y_i\) \(\forall i = 1, \ldots, n\).
The Jacobian of the model with respect to its parameters for all \(i = 1, \ldots, n\) is bounded almost surely as \[{\left\lVert \frac{ f_{\vw + \rvvepsilon}\left(\vx_i\right) }{ \partial \vw } \right\rVert}_2 \leq G_i\] for all \(\vw \in \mathcal{W}\).
Interpolation holds on the solution set such that, for all \(\vw_* \in \mathop{\mathrm{arg\,min}}_{\vw \in \mathcal{W}} L\left(\vw\right)\), the loss minimized as \[\ell\left(f_{\vw_* + \rvvepsilon}\left(\vx_i\right), y_i\right) = \ell'\left(f_{\vw_* + \rvvepsilon}\left(\vx_i\right), y_i\right) = 0\] for all \((\vx_{i}, y_i) \in \mathcal{D}\).
(a) holds for the squared loss, (c) basically assumes that the model is overparameterized and there exists a set of optimal weights that are robust with respect to perturbation. The has recently gained popularity as it qualitatively explains some of the empirical phenomenons of non-convex SGD [34], [39], [40]. (b) is a strong assumption but is commonly used to establish convergence guarantees of ERM [34].
8.1.3 Theoretical Analysis↩︎
\[\begin{align} &\mathbb{E}{\left\lVert \rvvg_i\left(\vw\right) - \rvvg_i\left(\vw_*\right) \right\rVert}_2^2 \\ &\;= \mathbb{E}{\bigg\lVert} \frac{\partial f_{\vw + \rvvepsilon}\left(\vx_i\right)}{\partial \vw} \ell'\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) \\ &\quad\qquad- \frac{\partial f_{\vw + \rvvepsilon}\left(\vx_i\right)}{\partial \vw} \ell'\left(f_{\vw_* + \rvvepsilon}\left(\vx_i\right), y_i\right) {\bigg\rVert}_2^2, \shortintertext{from the interpolation assumption (\ref{assumption:erm} (c)),} &\;= \mathbb{E}{\left\lVert \frac{\partial f_{\vw + \rvvepsilon}\left(\vx_i\right)}{\partial \vw} \ell'\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) \right\rVert}_2^2 \\ &\;\leq \mathbb{E}{\left\lVert \frac{\partial f_{\vw + \rvvepsilon}\left(\vx_i\right)}{\partial \vw} \right\rVert}_2^2 {\left| \ell'\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) \right|}^2, \shortintertext{applying \ref{assumption:erm} (b),} &\;\leq G_i^2 \mathbb{E}{\left| \ell'\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) \right|}^2. \shortintertext{and then the interpolation assumption (\ref{assumption:erm} (c)),} &\;= G_i^2 \mathbb{E}{\left| \ell'\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) - \ell'\left(f_{\vw_* + \rvvepsilon}\left(\vx_i\right), y_i\right) \right|}. \shortintertext{From \ref{assumption:erm} (a),} &\;\leq 2 L G_i^2 \mathbb{E}\Big( \ell\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) - \ell\left(f_{\vw_* + \rvvepsilon}\left(\vx_i\right), y_i\right) \\ &\qquad\qquad- \left\langle \ell'\left(f_{\vw_* + \rvvepsilon}\left(\vx_i\right), y_i\right) , f_{\vw + \rvvepsilon}\left(\vx_i\right) - f_{\vw_* + \rvvepsilon}\left(\vx_i\right) \right\rangle \Big) \shortintertext{and interpolation (\ref{assumption:erm} (c)),} &\;= 2 L G_i^2 \left( \mathbb{E} \ell\left(f_{\vw + \rvvepsilon}\left(\vx_i\right), y_i\right) - \mathbb{E} \ell\left(f_{\vw_* + \rvvepsilon}\left(\vx_i\right), y_i\right) \right) \\ &\;= 2 L G_i^2 \left( R_i\left(\vw\right) - R_i\left(\vw_*\right) \right). \end{align}\]
\[\begin{align} {2} &{\textstyle\frac{1}{n}\sum_{i=1}^n} {\left\lVert \nabla R_{i}\left(\vw\right) - \nabla R_{i}\left(\vw_*\right) \right\rVert}_2^2 \\ &\;= {\textstyle\frac{1}{n}\sum_{i=1}^n} {\left\lVert \mathbb{E} \rvvg_{i}\left(\vw\right) - \mathbb{E} \rvvg_{i}\left(\vw_*\right) \right\rVert}_2^2, \shortintertext{and from Jensen's inequality,} &\;\leq {\textstyle\frac{1}{n}\sum_{i=1}^n} \mathbb{E} {\left\lVert \rvvg_i\left(\vw\right) - \rvvg_i\left(\vw_*\right) \right\rVert}_2^2. \shortintertext{We can now reuse \ref{eq:noisy95erm95er} as} &\;\leq \frac{2}{n}\sum_{i=1}^n 2 L G_i^2 \left(R_i\left(\vw\right) - R_i\left(\vw_*\right)\right) \shortintertext{and taking \(G_{\rm{max}} > G_i\) for all \(i = 1, \ldots, n\) as} &\;\leq 2 L G_{\rm{max}}^2 \frac{1}{n}\sum_{i=1}^n \left(R_i\left(\vw\right) - R_i\left(\vw_*\right)\right) \\ &\;= 2 L G_{\rm{max}}^2 \left(L\left(\vw\right) - L\left(\vw_*\right)\right). \end{align}\]
8.2 Reparameterization Gradient↩︎
8.2.1 Description↩︎
The reparameterization gradient estimator [12], [13], [17] is a gradient estimator for problems of the form of \[f_i\left(\vw\right) = \mathbb{E}_{\rvvz \sim q_{\vw}} \ell_i\left(\rvvz\right),\] where \(\ell : \mathbb{R}^{d_{\rvvz} \to \mathbb{R}}\) is some integrand, such that the derivative is taken with respect to the parameters of the distribution\(q_{\vw}\) we are integrating over. It was independently proposed by [12], [13] in the context of variational expectation maximization of deep latent variable models (a setup commonly known as variational autoencoders) and by [17] for variational inference of Bayesian models.
Consider the case where the generative process of \(q_{\vw}\) can be represented as \[\rvvz \sim q_{\vw} \quad\Leftrightarrow\quad \rvvz \stackrel{d}{=} \mathcal{T}_{\vw}\left(\rvvu\right); \quad \rvvu \sim \varphi,\] where \(\stackrel{d}{=}\) is equivalence in distribution, \(\varphi\) is some base distribution independent of \(\vw\), and \(\mathcal{T}_{\vw}\) is a reparameterization function measurable with respect to \(\varphi\) and differentiable with respect to all \(\vw \in \mathcal{W}\). Then, the reparameterization gradient is given by the integrand \[\begin{align} \vg_i\left(\vw; \vu\right) = \nabla_{\vw} \; \ell_i\left(\mathcal{T}_{\vw}\left(\vu\right)\right), \end{align}\] which is unbiased, and often results in lower variance [18], [43] compared to alternatives such as the score gradient estimator. (See [42] for an overview of such estimators.)
The reparameterization gradient is primarily used to solve problems in the form of \[\begin{align} \mathop{\mathrm{minimize}}_{\vw \in \mathcal{W}}\;\; F\left(\vw\right) &= \sum^n_{i=1} f_i\left(\vw\right) + h\left(\vw\right) \\ &= \mathbb{E}_{\rvvz \sim q_{\vw}} \ell_i\left(\rvvz\right) + h\left(\vw\right), \end{align}\] where \(h\) is some convex regularization term.
Previously, [27] established a bound on the gradient variance of the reparameterization gradient [12], [13], [17] under the doubly stochastic setting. This bound also incorporates more advanced subsampling strategies such as importance sampling [29], [32], [64]–[66]. However, he did not extend the analysis to a complexity analysis of SGD and left out the effect of correlation between components.
8.2.2 Preliminaries↩︎
The properties of the reparameterization gradient for when \(q_{\vw}\) is in the location-scale family were studied by [27].
Assumption 9. We assume the variational family \[\mathcal{Q} \triangleq \{q_{\vw} \mid \vw \in \mathcal{W}\}\] satisfies the following:
\(\mathcal{Q}\) is part of the location-scale family such that \(\mathcal{T}_{\vw}\left(\vu\right) = \mC \vu + \vm\).
The scale matrix is positive definite such that \(\mC \succ 0\).
\(\rvvu = (\rvu_1, \ldots, \rvu_{d_{\vz}})\) constitute of i.i.d. components, where each component is standardized, symmetric, and finite kurtosis such that \(\mathbb{E}\rvu_i = 0\), \(\mathbb{E}\rvu_i^2 = 1\), \(\mathbb{E}\rvu_i^3 = 0\), and \(\mathbb{E}\rvu_i^4 = k_{\varphi}\), where \(k_{\varphi}\) is the kurtosis.
Under these conditions, [27] proves the following:
Lemma 1 ([27]; Theorem 3). Let 9 hold and \(\ell_i\) be \(L_i\)-smooth. Then, the squared norm of the reparameterization gradient is bounded: \[\mathbb{E}{\left\lVert \rvvg_i\left(\vw\right) \right\rVert}_2^2 \leq \left(d + 1\right) {\left\lVert \vm - \bar{\vz}_i \right\rVert}_2^2 + \left(d + k_{\varphi}\right) {\left\lVert \mC \right\rVert}_{\mathrm{F}}^2\] for all \(\vw = (\vm, \mC) \in \mathcal{W}\) and all stationary points of \(\ell_i\) denoted with \(\bar{\vz}_i\).
Similarly, [28] establish the QES condition as part of Lemma 3 [28]. We refine this into statement we need:
Lemma 2. Let 9 hold and \(\ell_i\) be \(L_i\)-smooth. Then, the squared norm of the reparameterization gradient is bounded: \[\mathbb{E}{\left\lVert \rvvg_i\left(\vw\right) - \rvvg_i\left(\vw'\right) \right\rVert}_2^2 \leq L_i^2 \left(d + k_{\varphi}\right) {\left\lVert \vw - \bar{\vw}_i \right\rVert}_2^2\] for all \(\vw, \vw' \in \mathcal{W}\).
\[\begin{align} &\mathbb{E}{\left\lVert \rvvg_i\left(\vw\right) - \rvvg_i\left(\vw'\right) \right\rVert} \\ &\;= \mathbb{E}{\left\lVert \nabla_{\vw} \ell_i\left(\mathcal{T}_{\vw}\left(\rvvu\right)\right) - \nabla_{\vw} \ell_i\left(\mathcal{T}_{\vw'}\left(\rvvu\right)\right) \right\rVert}_2^2 \\ &\;= \mathbb{E}{\left\lVert \frac{\partial \mathcal{T}_{\vw} \left(\rvvu\right)}{\partial \vw} \nabla \ell_i\left(\mathcal{T}_{\vw}\left(\rvvu\right)\right) - \frac{\partial \mathcal{T}_{\vw'} \left(\rvvu\right)}{\partial {\vw'}} \nabla \ell_i\left(\mathcal{T}_{\vw'}\left(\rvvz\right)\right) \right\rVert}_2^2 \\ &\;= \mathbb{E} {\left( \nabla \ell_i\left(\mathcal{T}_{\vw}\left(\rvvu\right)\right) - \nabla \ell_i\left(\mathcal{T}_{\vw'}\left(\rvvu\right)\right) \right)}^{\top} {\left( \frac{\partial \mathcal{T}_{\vw} \left(\rvvu\right)}{\partial \vw} \right)}^{\top} \frac{\partial \mathcal{T}_{\vw'} \left(\rvvu\right)}{\partial \vw'} \\ &\qquad \times \left( \nabla \ell_i\left(\mathcal{T}_{\vw}\left(\rvvu\right)\right) - \nabla \ell_i\left(\mathcal{T}_{\vw'}\left(\rvvu\right)\right) \right). \shortintertext{As shown by \cite{kim_convergence_2023}, the squared Jacobian is an identity matrix scaled with a scalar-valued function independent of \(\vw\), \(J_{\mathcal{T}}\left(\vu\right) = {\left\lVert \vu \right\rVert}_2^2 + 1\), such that} &\;= \mathbb{E} J_{\mathcal{T}}\left(\rvvu\right) {\left\lVert \nabla \ell_i\left(\mathcal{T}_{\vw}\left(\rvvu\right)\right) - \nabla \ell_i\left(\mathcal{T}_{\vw'}\left(\rvvu\right)\right) \right\rVert}_2^2, \shortintertext{applying the \(L_i\)-smoothness of \(\ell_i\),} &\;= L_i^2 \mathbb{E} J_{\mathcal{T}}\left(\rvvu\right){\left\lVert \mathcal{T}_{\vw}\left(\rvvu\right) - \mathcal{T}_{\vw'}\left(\rvvu\right) \right\rVert}_2^2, \shortintertext{and \cite{kim_convergence_2023} show that,} &\;\leq L_i^2 \left(d + k_{\varphi}\right) {\left\lVert \vw - \vw' \right\rVert}_2^2. \end{align}\]
Lastly, the properties of \(\ell_i\) are known to transfer to the expectation \(f_i\) as follows:
Lemma 3. Let 9 hold. Then we have the following:
Let \(\ell_i\) be \(L_i\) smooth. Then, \(f_i\) is also \(L_i\)-smooth
Let \(\ell_i\) be convex. Then, \(f_i\) and \(F\) are also convex.
Let \(\ell_i\) be \(\mu\)-strongly convex. Then, \(f_i\) and \(F\) are also \(\mu\)-strongly convex.
is proven by [67], while a more general result is provided by [28]; are proven by [67] and follow from the fact that \(h\) is convex.
8.2.3 Theoretical Analysis↩︎
We now conclude that the reparameterization gradient fits the framework of this work:
The result follows from combining 3 and [thm:expected95residual95without95replacement].
Let 9 hold, \(\ell_i\) be \(L_i\)-smooth, the solutions \(\vw_* \in \mathop{\mathrm{arg\,min}}_{\vw \in \mathcal{W}} F\left(\vw\right)\) and the stationary points of \(\ell_i\), \(\bar{\vz}\), be bounded such that \({\left\lVert \vw_* \right\rVert}_2 < \infty\) and \({\left\lVert \bar{\vz} \right\rVert}_2 < \infty\). Then, 7 holds.
1 implies that, as long as the \(\vw_*\) and \(\bar{\vz}\) are bounded, we satisfy the component gradient estimator part of 7, where the constant is given as \[\sigma_i^2 = L_i^2 \left(d + 1\right) {\left\lVert \vm_* - \bar{\vz}_i \right\rVert}_2^2 + L_i^2 \left(d+k_{\varphi}\right) {\left\lVert \mC_* \right\rVert}_{\mathrm{F}}^2,\] where \(\vw_* = (\vm_*, \mC_*)\).
From 2, we can conclude that the reparameterization gradient satisfies 6:
Let 9 hold and \(\ell_i\) be \(L_i\)-smooth and \(\mu\)-strongly convex. Then, 6 (\(\rm{A}^{\rm{CVX}}\)) and 6 (B) hold.
Notice the following:
4 always holds for \(\rho = 1\).
From the stated conditions, 3 establishes that both \(f_i\) and \(F\) are \(\mu\)-strongly convex.
\(\mu\)-strong convexity of \(f\) and \(F\) implies that both are \(\mu\)-QFG [55].
The reparameterization gradient satisfies the QES condition by 2.
Item 1, 2 and 3 combined imply the ES condition by [thm:qesises], which immediately implies the ER condition with the same constant. Therefore, we satisfy both 6 (\(\rm{A}^{\rm{CVX}}\)), 6 (B) where the ER constant \(\mathcal{L}_i\) is given as \[\mathcal{L}_i = \frac{L_i^2}{\mu} \left(d + k_{\varphi}\right).\]
References↩︎
See the implementation at https://github.com/zixu1986/Doubly_Stochastic_Gradients↩︎