SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lam M. Nguyen*
Thomas J. Watson Research Center
Trang H. Tran*
School of ORIE
Ithaca, NY, USA
http27@cornell.edu
Stochastic gradient descent (SGD) algorithm is the method of choice in many machine learning tasks thanks to its scalability and efficiency in dealing with large-scale problems. In this paper, we focus on the shuffling version of SGD which matches the mainstream practical heuristics. We show the convergence to a global solution of shuffling SGD for a class of non-convex functions under over-parameterized settings. Our analysis employs more relaxed non-convex assumptions than previous literature. Nevertheless, we maintain the desired computational complexity as shuffling SGD has achieved in the general convex setting.
In the last decade, neural network-based models have shown great success in many machine learning applications such as natural language processing [Collobert and Weston, 2008, Goldberg et al., 2018], computer vision and pattern recognition [Goodfellow et al., 2014, He and Sun, 2015]. The training task of many learning models boils down to the following finite-sum minimization problem:
$$
\min  _ {w \in \mathbb {R} ^ {d}} \left\{F (w) := \frac {1}{n} \sum_ {i = 1} ^ {n} f (w; i) \right\}, \tag {1}
$$
where  $f(\cdot ;i):\mathbb{R}^d\to \mathbb{R}$  is smooth and possibly non-convex for  $i\in [n]\coloneqq \{1,\dots ,n\}$ . Solving the empirical risk minimization (1) had been a difficult task for a long time due to the non-convexity and the complicated learning models. Later progress with stochastic gradient descent (SGD) and its variants [Robbins and Monro, 1951, Duchi et al., 2011, Kingma and Ba, 2014] have shown great performance in training deep neural networks. These stochastic first-order methods are favorable thanks to its scalability and efficiency in dealing with large-scale problems. At each iteration SGD samples an index  $i$  uniformly from the set  $\{1,\ldots ,n\}$ , and uses the individual gradient  $\nabla f(\cdot ;i)$  to update the weight.
While there has been much attention on the theoretical aspect of the traditional i.i.d. (independently identically distributed) version of SGD [Nemirovski et al., 2009, Ghadimi and Lan, 2013, Bottou et al., 2018], practical heuristics often use without-replacement data sampling schemes. Also known as shuffling sampling schemes, these methods generate some random or deterministic permutations of the index set  $\{1,2,\dots,n\}$  and apply gradient updates using these permutation orders. Intuitively, a collection of such  $n$  individual updates is a pass over all the data, or an epoch. One may choose to create a new random permutation at the beginning of each epoch (in Random Reshuffling scheme) or use a random permutation for every epoch (in Single Shuffling scheme). Alternatively, one may use a
*Equally contributed. Corresponding author: Lam M. Nguyen
Incremental Gradient scheme with a fixed deterministic order of indices. In this paper, we use the term unified shuffling SGD for SGD method using any data permutations, which includes the three special schemes described above.
Although shuffling sampling schemes usually show a better empirical performance than SGD [Bottou, 2009], the theoretical guarantees for these schemes are often more limited than vanilla SGD version, due to the lack of statistical independence. Recent works have shown improvement in computational complexity for shuffling schemes over SGD in various settings [Gurbüzbalaban et al., 2019, Haochen and Sra, 2019, Safran and Shamir, 2020, Nagaraj et al., 2019, Nguyen et al., 2021, Mishchenko et al., 2020, Ahn et al., 2020]. In particular, in a general non-convex setting, shuffling sampling schemes improve the computational complexity in terms of  $\hat{\varepsilon}$  for SGD from  $\mathcal{O}\left(\sigma^2/\hat{\varepsilon}^2\right)$  to  $\mathcal{O}\left(n\sigma/\hat{\varepsilon}^{3/2}\right)$ , where  $\sigma$  is the bounded variance constant [Ghadimi and Lan, 2013, Nguyen et al., 2021, Mishchenko et al., 2020]. We summarize the detailed literature for multiple settings later in Table 1.
While global convergence is a desirable property for neural network training, the non-convexity landscape of complex learning models leads to difficulties in finding the global minimizer. In addition, there is little to no work studying the convergence to a global solution of shuffling-type SGD algorithms for a general non-convex setting. The closest line of research investigates the Polyak-Lojasiewicz (PL) condition (a generalization of strong-convexity), which demonstrates similar convergence rates as the strongly convex rates for shuffling SGD methods [Haochen and Sra, 2019, Ahn et al., 2020, Nguyen et al., 2021]. In another direction, Gower et al. [2021] and Khaled and Richtárik [2020] investigates the global convergence for some class of non-convex functions, however for vanilla SGD method. Beznosikov and Takáč [2021] investigate a random shuffle version of variance reduction methods (e.g. SARAH algorithm Nguyen et al. [2017]), but this approach only can show convergence to stationary points. With a target on shuffling SGD methods and specific learning architectures, we come up with the central question of this paper:
How can we establish the convergence to global solution for a class of non-convex functions using shuffling-type SGD algorithms? Can we exploit the structure of neural networks to achieve this goal?
We answer this question affirmatively, and our contributions are summarized below.
- We investigate a new framework for the convergence of a shuffling-type gradient algorithm to a global solution. We consider a relaxed set of assumptions and discuss their relations with previous settings. We show that our average-PL inequality (Assumption 3) holds for a wide range of neural networks equipped with squared loss function.
- Our analysis generalizes the class function called star- $M$ -smooth-convex. This class contains non-convex functions and is more general than the class of star-convex smooth functions with respect to the minimizer (in the over-parameterized settings). In addition, our analysis does not use any bounded gradient or bounded weight assumptions.
- We show the total complexity of  $\mathcal{O}\left(\frac{n}{\hat{\varepsilon}^{3/2}}\right)$  for a class of non-convex functions to reach an  $\hat{\varepsilon}$ -accurate global solution. This result matches the same gradient complexity to a stationary point for unified shuffling methods in non-convex settings, however, we are able to show the convergence to a global minimizer.
In recent years, there have been different approaches to investigate the global convergence for machine learning optimization. This includes a popular line of research that studies some specific neural networks and utilizes their architectures. The most early works show the global convergence of Gradient Descent (GD) for simple linear networks and two-layer networks [Brutzkus et al., 2018, Soudry et al., 2018, Arora et al., 2019, Du et al., 2019b]. These results are further extended to deep learning architectures [Allen-Zhu et al., 2019, Du et al., 2019a, Zou and Gu, 2019]. This line of research continues with Stochastic Gradient Descent (SGD) algorithm, which proves the global convergence of SGD for deep neural networks for some probability depending on the initialization process and the number of input data [Brutzkus et al., 2018, Allen-Zhu et al., 2019, Zou et al.,
The computational complexity is the number of (individual) gradient computations needed to reach an  $\hat{\varepsilon}$ -accurate stationary point (i.e. a point  $\hat{w} \in \mathbb{R}^d$  that satisfies  $\| \nabla F(\hat{w})\|^2 \leq \hat{\varepsilon}$ ).
2018, Zou and Gu, 2019]. The common theme that appeared in most of these references is the overparameterized setting, which means that the number of parameters in the network are excessively large [Brutzkus et al., 2018, Soudry et al., 2018, Allen-Zhu et al., 2019, Du et al., 2019a, Zou and Gu, 2019]. This fact is closely related to our setting, and we will discuss it throughout our paper.
Polyak-Lojasiewicz (PL) condition and related assumptions. An alternative approach is to investigate some conditions on the optimization problem that may guarantee global convergence. A popular assumption is the Polyak-Lojasiewicz (PL) inequality, a generalization of strong-convexity [Polyak, 1964, Karimi et al., 2016, Nesterov and Polyak, 2006]. Using this PL assumption, it can be shown that (stochastic) gradient descent achieves the same theoretical rate as in the strongly convex setting (i.e. linear convergence for GD and sublinear convergence for SGD) [Karimi et al., 2016, De et al., 2017, Gower et al., 2021]. Recent works demonstrate similar results for shuffling type SGD [Haochen and Sra, 2019, Ahn et al., 2020, Nguyen et al., 2021], both for unified and randomized shuffling schemes. On the other hand, [Schmidt and Roux, 2013, Vaswani et al., 2019] propose to use a new assumption called the Strong Growth Condition (SGC) that controls the rates at which the stochastic gradients decay comparing to the full gradient. This condition implies that the stochastic gradients and their variances converge to zero at the optimum solution [Schmidt and Roux, 2013, Vaswani et al., 2019]. While the PL condition for  $F$  implies that every stationary point of  $F$  is also a global solution, the SGC implies that such a point is also a stationary point of every individual function. However, complicated models as deep feed-forward neural networks generally have non-optimal stationary points [Karimi et al., 2016]. Thus, these assumptions are somewhat strong for non-convex settings.
Although there are plenty of works investigating the PL condition for the objective function  $F$  [De et al., 2017, Vaswani et al., 2019, Gower et al., 2021], not many materials devoted to study the PL inequality for the individual functions  $f(\cdot; i)$ . A recent work [Sankararaman et al., 2020] analyzes SGD with the specific notion of gradient confusion for over-parameterized settings where the individual functions satisfy PL condition. They show that the neighborhood where SGD converges linearly depends on the level of gradient confusion (i.e. how much the individual gradients are negatively correlated). Taking a different approach, we investigate the PL property for individual functions and further show that our condition holds for a general class of neural networks with quadratic loss.
Over-paramaterized settings for neural networks. Most of the modern learning architectures contain deep and large networks, where the number of parameters are often far more than the number of input data. This leads to the fact that the objective loss function is trained closer and closer to zero. Understandably, in such settings all the individual functions  $f(\cdot; i)$  are minimized simultaneously at 0 and they share a common minimizer. This condition is called the interpolation property (see e.g. [Schmidt and Roux, 2013, Ma et al., 2018, Meng et al., 2020, Loizou et al., 2021]) and is studied well in the literature (see e.g. [Zhou et al., 2019, Gower et al., 2021]). For a comparison, functions satisfying the strong growth condition necessarily satisfy the interpolation property. This property implies zero variance of individual gradients at the global minimizer, which allows good behavior for SGD near the solution. In this work, we slightly change this assumption which requires a small variance up to some level of the threshold  $\varepsilon$ . Note that when letting  $\varepsilon \to 0$ , our assumption exactly recovers the interpolation property.
Star-convexity and related conditions. There have been many attentions to a class of structured non-convex functions called star-convex [Nesterov and Polyak, 2006, Lee and Valiant, 2016, Bjorck et al., 2021]. Star-convexity can be understood as convexity between an arbitrary point  $w$  and the global minimizer  $w_{*}$ . The name star-convex comes from the fact that each sublevel set is star-shaped [Nesterov and Polyak, 2006, Lee and Valiant, 2016]. Zhou et al. [2019] shows that if SGD follows a star-convex path and there exists a common global minimizer for all component functions, then SGD converges to a global minimum.
In recent progress, Hinder et al. [2020] considers the class of quasar-convex functions, which further generalizes star-convexity. This property was introduced originally in [Hardt et al., 2018] under the name 'weakly quasi-convex', and investigated recently in literature [Hinder et al., 2020, Jin, 2020, Gower et al., 2021]. This class uses a parameter  $\zeta \in (0,1]$  to control the non-convexity of the function, where  $\zeta = 1$  yields the star-convexity and  $\zeta$  approaches 0 indicates more non-convexity [Hinder et al., 2020]. Intuitively, quasar-convex functions are unimodal on all lines that pass through a global minimizer. Gower et al. [2021] investigates the performance of SGD for smooth and quasar-
convex functions using an expected residual assumption (which is comparable to the interpolation property). They show a convergence rate of  $\mathcal{O}(1 / \sqrt{K})$  for i.i.d. sampling SGD with the number of total iterations  $K$ , which translates to the computational complexity of  $\mathcal{O}\left(1 / \hat{\varepsilon}^2\right)$ . To the best of our knowledge, this paper is the first work studying the relaxation of star-convexity and global convergence for SGD with shuffling sample schemes, not for the i.i.d. version.
2 Theoretical Setting
We first present the shuffling-type gradient algorithm below. Our convergence results hold for any permutation of the training data  $\{1,2,\dots,n\}$ , including deterministic and random ones. Thus, our theoretical framework is general and applicable for any shuffling strategy, including Incremental Gradient, Single Shuffling, and Random Reshuffling.
Algorithm 1 (Shuffling-Type Gradient Algorithm for Solving (1))
1: Initialization: Choose an initial point  $\tilde{w}_0\in \mathrm{dom}(F)$
2: for  $t = 1,2,\ldots ,T$  do
3: Set  $w_0^{(t)} \coloneqq \tilde{w}_{t-1}$ ;
4: Generate any permutation  $\pi^{(t)}$  of  $[n]$  (either deterministic or random);
5: for  $i = 1,\dots ,n$  do
6: Update  $w_{i}^{(t)}\coloneqq w_{i - 1}^{(t)} - \eta_{i}^{(t)}\nabla f(w_{i - 1}^{(t)};\pi^{(t)}(i))$
7: end for
8: Set  $\tilde{w}_t\coloneqq w_n^{(t)}$
9: end for
We further specify the choice of learning rate  $\eta_i^{(t)}$  in the detailed analysis. Now we proceed to describe the set of assumptions used in our paper.
Assumption 1. Suppose that  $f_{i}^{*} \coloneqq \min_{w \in \mathbb{R}^{d}} f(w; i) > -\infty, i \in \{1, \dots, n\}$ .
Assumption 2. Suppose that  $f(\cdot; i)$  is  $L$ -smooth for all  $i \in \{1, \dots, n\}$ , i.e. there exists a constant  $L \in (0, +\infty)$  such that:
$$
\| \nabla f (w; i) - \nabla f \left(w ^ {\prime}; i\right) \| \leq L \| w - w ^ {\prime} \|, \quad \forall w, w ^ {\prime} \in \mathbb {R} ^ {d}. \tag {2}
$$
Assumption 1 is required in any algorithm to guarantee the well-definedness of (1). In most applications, the component losses are bounded from below. By Assumption 2, the objective function  $F$  is also  $L$ -smooth. This Lipschitz smoothness Assumption is widely used for gradient-type methods. In addition, we denote the minimum value of the objective function  $F_{*} = \min_{w\in \mathbb{R}^{d}}F(w)$ . It is worthy to note the following relationship between  $F_{*}$  and the component minimum values:
$$
F _ {*} = \min  _ {w \in \mathbb {R} ^ {d}} F (w) = \frac {1}{n} \min  _ {w \in \mathbb {R} ^ {d}} \left(\sum_ {i = 1} ^ {n} f (w; i)\right) \geq \frac {1}{n} \sum_ {i = 1} ^ {n} \min  _ {w \in \mathbb {R} ^ {d}} \left(f (w; i)\right) = \frac {1}{n} \sum_ {i = 1} ^ {n} f _ {i} ^ {*}. \tag {3}
$$
We are interested in the case where the set of minimizers of  $F$  is not empty. The equality  $F_{*} = \frac{1}{n}\sum_{i=1}^{n}f_{i}^{*}$  attains if and only if a minimizer of  $F$  is also the common minimizer for all component functions. This condition implies that the variance of individual functions is 0 at the common minimizer.
2.1 PL Condition for Component Functions
Now we are ready to discuss the Polyak-Lojasiewicz condition as follows.
Definition 1 (Polyak-Lojasiewicz condition). We say that  $f$  satisfies Polyak-Lojasiewicz (PL) inequality for some constant  $\mu > 0$  if
$$
\left\| \nabla f (w) \right\| ^ {2} \geq 2 \mu [ f (w) - f ^ {*} ], \quad \forall w \in \mathbb {R} ^ {d}, \tag {4}
$$
where  $f^{*} \coloneqq \min_{w \in \mathbb{R}^{d}} f(w)$ .
The PL condition for the objective function  $F$  is sufficient to show a global convergence for (stochastic) gradient descent [Karimi et al., 2016, Nesterov and Polyak, 2006, Polyak, 1964]. It is well known that a function satisfying the PL condition is not necessarily convex [Karimi et al., 2016]. However, this assumption on  $F$  is somewhat strong because it implies that every stationary point of  $F$  is also a global minimizer. Our goal is to consider a class of non-convex function which is more relaxed than the PL condition on  $F$ , while still having the good global convergence properties. In this paper, we formulated an assumption called "average PL inequality", specifically for the finite sum setting:
Assumption 3. Suppose that  $f(\cdot; i)$  satisfies average PL inequality for some constant  $\mu > 0$  such that
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \| \nabla f (w; i) \| ^ {2} \geq 2 \mu \frac {1}{n} \sum_ {i = 1} ^ {n} [ f (w; i) - f _ {i} ^ {*} ], \quad \forall w \in \mathbb {R} ^ {d}. \tag {5}
$$
where  $f_{i}^{*}\coloneqq \min_{w\in \mathbb{R}^{d}}f(w;i)$
Comparisons. Assumption 3 is weaker than assuming the PL inequality for every component function  $f(\cdot ;i)$ . In general setting, Assumption 3 is not comparable to assuming the PL inequality for  $F$ . Formally, if  $F$  satisfies PL the condition for some parameter  $\tau >0$ , then we have:
$$
2 \tau [ F (w) - F _ {*} ] \leq \| \nabla F (w) \| ^ {2} \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \| \nabla f (w; i) \| ^ {2}. \tag {6}
$$
However, by equation (3) we have that  $[F(w) - F_{*}] \leq \frac{1}{n} \sum_{i=1}^{n} [f(w; i) - f_{i}^{*}]$ . Therefore, the PL inequality for each function  $f(\cdot; i)$ , cannot directly imply the PL condition on  $F$  and vice versa.
In the interpolated setting where there is a common minimizer for all component function  $f(\cdot; i)$ , it can be shown that the PL condition on  $F$  is stronger than our average PL assumption:
$$
2 \tau \frac {1}{n} \sum_ {i = 1} ^ {n} [ f (w; i) - f _ {i} ^ {*} ] = 2 \tau [ F (w) - F _ {*} ] \leq \| \nabla F (w) \| ^ {2} \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \| \nabla f (w; i) \| ^ {2}.
$$
On the other hand, our assumption cannot imply the PL inequality on  $F$  unless we impose a strong relationship that upper bound the sum of individual squared gradients  $\frac{1}{n}\sum_{i=1}^{n}\|\nabla f(w;i)\|^2$  in terms of the full squared gradient  $\|\nabla F(w)\|^2$ , for every  $w \in \mathbb{R}^d$ . For these reasons, the average PL Assumption 3 is arguably more reasonable than assuming the PL inequality for the objective function  $F$ . Moreover, we show that Assumption 3 holds for a general class of neural networks with a final bias layer and squared loss function. We have the following theorem.
Theorem 1. Let  $\{(x^{(i)},y^{(i)})\}_{i = 1}^{n}$  is a training data set where  $x^{(i)}\in \mathbb{R}^m$  is the input data and  $y^{(i)}\in \mathbb{R}^c$  is the output data for  $i = 1,\ldots ,n$ . We consider an architecture  $h(w;i)$  with  $w$  be the vectorized weight and  $h$  consists of a final bias layer  $b$ :
$$
h (w; i) = W ^ {T} z (\theta ; i) + b,
$$
where  $w = \mathbf{vec}(\{\theta, W, b\})$  and  $z(\theta; i)$  are some inner architectures, which can be chosen arbitrarily. Next, we consider the squared loss  $f(w; i) = \frac{1}{2} \| h(w; i) - y^{(i)} \|^2$ . Then
$$
\left\| \nabla f (w; i) \right\| ^ {2} \geq 2 \left[ f (w; i) - f _ {i} ^ {*} \right], \quad \forall w \in \mathbb {R} ^ {d}, \tag {7}
$$
Therefore, for this application, Assumption 3 holds with  $\mu = 1$ .
2.2 Small Variance at Global Solutions
In this section, we change the interpolation property in previous literature [Ma et al., 2018, Meng et al., 2020, Loizou et al., 2021] by a small threshold. For any global solution  $w_{*}$  of  $F$ , let us define
$$
\sigma_ {*} ^ {2} := \inf  _ {w _ {*} \in \mathcal {W} _ {*}} \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \| \nabla f (w _ {*}; i) \| ^ {2}\right). \tag {8}
$$
We can show that when there is a common minimizer for all component functions (i.e. when the equality  $F_{*} = \frac{1}{n}\sum_{i = 1}^{n}f_{i}^{*}$  holds), the best variance  $\sigma_{*}^{2}$  is 0. It is sufficient for our Theorem to impose a  $\mathcal{O}(\varepsilon)$ -level upper bound on the variance  $\sigma_{*}^{2}$ :
Assumption 4. Suppose that the best variance at  $w_{*}$  is small, that is, for  $\varepsilon > 0$
$$
\sigma_ {*} ^ {2} \leq P \varepsilon , \tag {9}
$$
for some  $P > 0$
It is important to note that in current non-convex literature, Assumption 4 alone (or, assuming  $\sigma_{*}^{2} = 0$  alone) is not sufficient enough to guarantee a global convergence property for SGD. Typically, some other conditions on the good landscape of the loss function are needed to complement the over-parameterized setting. Thus, we have motivation to introduce our next assumption.
2.3 Generalized Star-Smooth-Convex Condition for Shuffling Type Algorithm
We introduce the definition of star-smooth-convex function as follows.
Definition 2. The function  $g$  is star-  $M$ -smooth-convex with respect to a reference point  $\hat{w} \in \mathbb{R}^d$  if
$$
\left\| \nabla g (w) - \nabla g (\hat {w}) \right\| ^ {2} \leq M \langle \nabla g (w) - \nabla g (\hat {w}), w - \hat {w} \rangle , \quad \forall w \in \mathbb {R} ^ {d}. \tag {10}
$$
It is well known that when  $g$  is  $L$ -smooth and convex [Nesterov, 2004], we have the following general inequality for every  $w, w' \in \mathbb{R}^d$ :
$$
\left\| \nabla g (w) - \nabla g \left(w ^ {\prime}\right) \right\| ^ {2} \leq L \langle \nabla g (w) - \nabla g \left(w ^ {\prime}\right), w - w ^ {\prime} \rangle \tag {11}
$$
Our class of star-smooth-convex function requires a similar inequality to hold only for the special point  $w' = \hat{w}$ . Interestingly, this is related to a class of star-convex functions, which satisfies the convex inequality for the minimizer  $\hat{w}$ :
$$
\text {(s t a r - c o n v e x i t y w . r . t} \hat {w}) \quad g (w) - g (\hat {w}) \leq \langle \nabla g (w), w - \hat {w} \rangle , \forall w \in \mathbb {R} ^ {d}, \tag {12}
$$
This class of functions contains non-convex objective losses and is well studied in the literature (see e.g. [Zhou et al., 2019]). Our Lemma 1 shows that the class of star-smooth-convex function is broader than the class of  $L$ -smooth and star-convex functions. Therefore, our problem of interest is non-convex in general.
Lemma 1. The function  $g$  is star- $M$ -smooth-convex with respect to  $\hat{w}$  for some constant  $M > 0$  if one of the two following conditions holds:
1.  $g$  is  $L$ -smooth and convex.
2.  $g$  is  $L$ -smooth and  $g$  is star-convex with respect to  $\hat{w}$ .
Proof. The first statement of Lemma 1 follows directly from equation (11). We have that  $g$  is star- $M$ -smooth-convex with respect to any reference point and  $M = L$ .
Now we proceed to the second statement. From the star-convex property of  $g$  with respect to  $\hat{w}$ , we have
$$
g (w) - g (\hat {w}) \leq \langle \nabla g (w), w - \hat {w} \rangle , \forall w \in \mathbb {R} ^ {d},
$$
and  $\nabla g(\hat{w}) = 0$  since  $\hat{w}$  is the global minimizer of  $g$ . On the other hand,  $g$  is  $L$ -smooth and we have
$$
g (\hat {w}) \leq g \left(w - \frac {1}{L} \nabla g (w)\right) \leq g (w) - \frac {1}{2 L} \| \nabla g (w) \| ^ {2},
$$
which is equivalent to  $\| \nabla g(w) \|^2 \leq 2L[g(w) - g(\hat{w})]$ ,  $i \in [n]$ . Since  $\nabla g(\hat{w}) = 0$ ,  $i \in [n]$ , we have for  $\forall w \in \mathbb{R}^d$
$$
\| \nabla g (w) - \nabla g (\hat {w}) \| ^ {2} \leq 2 L [ g (w) - g (\hat {w}) ] \stackrel {(1 2)} {\leq} 2 L \langle \nabla g (w) - \nabla g (\hat {w}), w - w _ {*} \rangle .
$$