Stability and Generalization for Randomized Coordinate Descent

1 Introduction

Randomized coordinate descent (RCD) is a popular method for solving large-scale optimization problems which are ubiquitous in the big-data era [Nesterov, 2012; Richt´arik and Tak´aˇc, 2014; Richt´arik and Tak´aˇc, 2016; Zeng et al., 2019]. As an iterative algorithm, it iteratively updates a single randomly chosen coordinate along the negative direction of the derivative while keeping the other coordinates fixed. Due to its ease of implementation and high efficiency, RCD has found wide applications in various areas such as compressed sensing, network problems and optimization [Richt´arik and Tak´aˇc, 2014]. In particular, the conceptual and algorithmic simplicity makes it especially useful for large-scale problems for which even the simplest full-dimensional vector operations are very expensive [Nesterov, 2012].

The popularity of RCD motivates a lot of theoretical analysis to understand its empirical behavior. Specifically, iteration complexities of RCD are well studied in the literature under different settings (e.g., convex/strongly convex [Nesterov, 2012], smooth/nonsmooth cases [Richt´arik and Tak´aˇc, 2014; Richt´arik and Tak´aˇc, 2016; Lu and Xiao, 2015]) for different variants (e.g., distributed RCD [Richt´arik and Tak´aˇc, 2016], accelerated RCD [Nesterov, 2012; Ren and Zhu, 2017;

∗Corresponding Author

Gu et al., 2018; Li and Lin, 2020; Chen and Gu, 2016] and RCD for primal-dual problems [Qu et al., 2016]). These discussions concern how the empirical risks of the models trained by RCD would decay along the optimization process. As a comparison, there is little analysis on how these models would behave on testing examples, which is what really matters in machine learning. Actually, if the models are very complicated, it is very likely that the models would admit a small empirical risk or even interpolate the training examples but meanwhile suffer from a large test error. This discrepancy between training and testing, as referred to as overfitting, is a fundamental problem in machine learning [Bousquet and Elisseeff, 2002]. The existing convergence analysis of RCD is not enough to fully understand why models trained by RCD have a good prediction performance in real applications. In particular, it is not clear how the optimization and statistical behavior of RCD would change along the optimization process, which is useful for designing efficient models in practice. For example, generalization analysis provides a principled guideline on how to stop the algorithm appropriately for a best generalization.

In this paper, we aim to bridge the generalization and optimization of RCD by leveraging the celebrated concept of algorithmic stability. We establish stability bounds of RCD as measured by several concepts, including ℓ1-argument stability, ℓ2-argument stability and uniform stability. Under standard assumptions on smoothness, Lipschitz continuity and convexity of objective functions, we show clearly how the stability and the optimization error would behave along the learning process. This suggests a principled way to earlystop the algorithm to get a best generalization behavior. We consider convex, strongly convex and nonconvex objective functions. In the convex and strongly convex cases, we develop minimax optimal generalization bounds of the order O(1/√n) and O(1/n) respectively, where n is the sample size. Our analysis not only suggests that RCD has a better stability than stochastic gradient descent (SGD), but also is able to exploit a low noise condition to get an optimistic bound O(1/n) in the convex case. Finally, we develop generalization bounds with high probability which are useful to understand the robustness and variation of the training algorithm [Feldman and Vondrak, 2019].

We survey related work in Section 2 and formulate the problem in Section 3. We give stability and generalization bounds in Section 4 and Section 5. We report experimental results in Section 6 and present some proofs in Section 7. The full proofs can be found in the arXiv version of this paper.

2 RelatedWork

2.1 Randomized Coordinate Descent

RCD was widely used to solve large-scale optimization problems in machine learning, including linear SVMs [Chang et al., 2008], ℓ1-regularized models for sparse learning [ShalevShwartz and Tewari, 2009] and low-rank matrix learning [Hu and Kwok, 2019]. The convergence rate of RCD and its accelerated variant were studied in the seminal work [Nesterov, 2012], where the advantage of RCD over deterministic algorithms is clearly illustrated. These results were extended to structure optimization problems where the objective function consists of a smooth data-fitting term and a nonsmooth regularizer [Richt´arik and Tak´aˇc, 2014; Lu and Xiao, 2015]. RCD was also adapted to distributed data analysis [Richt´arik and Tak´aˇc, 2016; Sun et al., 2017; Xiao et al., 2019], primal-dual optimization [Qu et al., 2016] and privacy-preserving problems [Damaskinos et al., 2020]. All these discussions consider the convergence rate of optimization errors for RCD. As a comparison, we are interested in the generalization behavior of models trained by RCD, which is the ultimate goal in machine learning.

2.2 Stability and Generalization

We now review the related work on algorithmic stability and its application on generalization analysis. The framework of algorithmic stability was established in a seminal paper [Bousquet and Elisseeff, 2002], where the important uniform stability was introduced. This algorithmic stability was extended to study randomized algorithms in Elisseeff et al. [2005]. Other than uniform stability, several other stability measures including hypothesis stability [Bousquet and Elisseeff, 2002], on-average stability [Shalev-Shwartz et al., 2010] and argument stability [Liu et al., 2017] have been introduced in statistical learning theory, whose connection to learnability has been established [Mukherjee and Zhou, 2006; Shalev-Shwartz et al., 2010]. The uniform stability of stochastic gradient descent (SGD) was established for learning with (strongly) convex, smooth and Lipschitz loss functions [Hardt et al., 2016]. This motivates the recent work of studying generalization of stochastic optimization algorithms via several stability [Meng et al., 2017; Charles and Papailiopoulos, 2018; Kuzborskij and Lampert, 2018; Yin et al., 2018; Yuan et al., 2019; Lei and Ying, 2020; Bassily et al., 2020; Lei et al., 2020; Lei and Ying, 2021; Wang et al., 2021; Yang et al., 2021]. For example, an onaverage model stability [Lei and Ying, 2020] has been proposed to remove the smoothness assumption or Lipschitz continuity assumption in Hardt et al. [2016]. Recently, elegant concentration inequalities have been developed to get highprobability bounds via uniform stability [Feldman and Vondrak, 2019; Bousquet et al., 2020]. To our knowledge, the algorithmic stability of RCD has not been studied yet, which is the topic of this paper.

3 ProblemFormulation

Let ρ be a probability measure defined over a sample space Z = X ×Y, where X is an input space and Y ⊂ R is an output space. We aim to build a parametric model hw : X ↦→ R, where w is the model parameter which belongs to the parameter space W ⊆ Rd . The performance of the model hw on a single example z can be measured by a nonnegative loss function f(w; z). The quality of a model can be quantified by a population risk F (w) = Ez[f(w; z)], where Ez denotes the expectation w.r.t. z. We wish to approximate the best model w∗ ∈ argminw∈W F(w). However, the probability measure is often unknown and we only have access to a training sample S = {z1, z2, . . . , zn} drawn independently from ρ. The empirical behavior of hw on S can be measured by a empirical risk FS(w) = 1n ∑ni=1 f (w; zi).

We give necessary notations. For any x ∈ Rd , we denote the norm ‖x‖p = (∑di=1 |xi|p )1/p for p ≥ 1. For any m ∈ N, we denote [m] := {1, . . . ,m}. We use the notation B = O( ˜B) if there exists a constant c0 > 0 such that B ≤ c0 ˜B, and use B ≍ ˜ B if there exist constants c1, c2 > 0 such that c1 ˜ B < B ≤ c2 ˜ B. We say g : W ↦→ R is L-smooth if ‖∇g(w) -∇g(w′ )‖2 ≤ L‖w-w′ ‖2 for all w,w′ ∈ W.

We apply a randomized algorithm A to the sample S and get an output model A(S) ∈ W. We are interested in studying the excess generalization error F(A(S))-F(w∗ ). Since E[FS(w∗ )] = F(w∗ ), we have the decomposition

We refer to the first term ES,A[F(A(S)) - FS(A(S))] as the estimation error, and the second term ES,A[FS(A(S)) - FS(w∗ )] as the optimization error. A standard approach to control estimation error is to study the algorithmic stability of the algorithm A, i.e., how the model would change if we change the training sample by a single example. There are several variants of stability measures including the uniform stability, hypothesis stability, on-average stability and argument stability [Bousquet and Elisseeff, 2002; Hardt et al., 2016; Elisseeff et al., 2005], among which the uniform stability is the most popular one.

Definition 1 (Uniform Stability). A randomized algorithm A is ϵ-uniformly stable if for all datasets S, ˜ S ∈ Z n differing by one example, we have supz [f(A(S); z) - f(A( ˜ S); z)] ≤ ϵ.

We consider the on-average argument stability, an advantage of which is that it can imply better generalization bounds without a Lipschitz continuity assumption on loss functions.

Definition 2 (On-average Argument Stability). Let S = {z1, . . . , zn} and S′ = {z′1 , . . . , z ′n } be drawn independently from ρ. For any i = 1, . . . , n, define S(i) = {z1, . . . , zi-1, z′i , zi+1, . . . , zn} as the set formed from S by replacing zi with z′i . We say a randomized algorithm A is ℓ1 on-average argument ϵ-stable if ES,S′ ,A[ 1n ∑ni=1 ‖A(S)- A(S(i) )‖2 ] ≤ ϵ, and ℓ2 on-average argument ϵ-stable if ES,S′ ,A[ 1n ∑ni=1 ‖A(S) - A(S(i) )‖22 ] ≤ ϵ2 .

Lemma 1 gives a connection between on-average stability and generalization [Lei and Ying, 2020]. Assumption 1 holds for popular loss including logistic loss and Huber loss.

Assumption 1. Let G1, G2 > 0. Assume for all w ∈ W and z ∈ Z, ‖∇f(w; z)‖1 ≤ G1 and ‖∇f(w; z)‖2 ≤ G2.

Lemma 1. Let S, S′ and S(i) be constructed as Definition 2. (a) If Assumption 1 holds, then

• (b) If for any z, the function w ↦→ f(w; z) is nonnegative and L-smooth, then for any γ > 0 we have

In this paper, we consider the specific RCD method widely used in large-scale learning problems. Let w1 ∈ W be the initial point. At the t-th iteration it first randomly selects a single coordinate it ∈ [d], and then performs the update along the it-th coordinate as [Nesterov, 2012]

where ∇ig denotes the derivative of g w.r.t. the i-th coordinate and ei is a vector in Rd with the i-th coordinate being 1 and other coordinates being 0. Here {ηt} is a nonnegative stepsize sequence. It is clear that RCD sequentially updates a randomly selected coordinate while keeping others fixed. In this paper, we consider the update of only a single coordinate per iteration. Our discussions can be readily extended to randomized block coordinate descent where the coordinates are partitioned into blocks, and each block of coordinates is updated per iteration [Nesterov, 2012].

4 Stability of RCD

In this section, we present our stability bounds of RCD. To this aim, we first introduce several standard assumptions [Nesterov, 2012]. The first assumption is the convexity of the empirical risk. Note we do not require the convexity of each loss function, which is used in the stability analysis of SGD [Hardt et al., 2016; Kuzborskij and Lampert, 2018].

Assumption 2. For any training dataset set S, FS is convex.

Our second assumption is the coordinate-wise smoothness. Definition 3. We say a differentiable function g : W ↦→ R has coordinate-wise Lipschitz continuous gradients with parameter ˜L > 0 if the following inequality holds for all α ∈ R, w ∈ W, i ∈ [d] g(w+αei) ≤ g(w)+α∇ig(w)+˜Lα2 /2. Assumption 3. For any training dataset S, FS is L-smooth and has coordinate-wise Lipschitz continuous gradients with

parameter ˜L > 0.

We first consider convex problems. Part (a) of Theorem 2 considers the ℓ1 argument stability, while Part (b) considers ℓ2 argument stability. The proof is given in Section 7.

Theorem 2. Let Assumptions 2, 3 hold. Let {wt}, {w(i) t } be given by (2) with ηt ≤ 2/˜L based on S and S(i) , respectively. (a) If Assumption 1 holds, then

(b) For any p > 0 the ℓ2 on-average argument stability can be bounded by

Remark 1. We compare Theorem 2 with related work. Under Assumptions 1, 2 and 3, it was shown SGD with t iterations enjoys the ℓ1 argument stability bound O( G2 n ∑tk=1 ηk). Eq. (3) shows RCD admits a better stability since there is a d in the denominator. Note in the worst case we can choose G1 ≤ √dG2 for which our stability bounds of RCD are of the order O( G2 n√d ∑tk=1 ηk). A notable property of Part (b) is that the stability bound (4) does not require the Lipschitz condition as ‖∇f(w; z)‖2 ≤ G2, which is widely used in the existing stability analysis [Hardt et al., 2016; Charles and Papailiopoulos, 2018; Kuzborskij and Lampert, 2018]. Indeed, a key point here is that we replace the Lipschitz constant G2 by empirical/population risks F and FS. Since we are minimizing the empirical risk by RCD, it is reasonable that F and FS would be small and in this case the algorithm would be more stable. This gives an intuitive connection between stability and optimization: a small optimization error is also beneficial to improve stability.

We now consider strongly convex case. Theorem 2 shows the stability becomes worse as we run more iterations. In the following theorem, we show the stability can be further improved if we impose a strong convexity assumption.

Assumption 4. Assume for all S, i ∈ [d],w ∈ W, the function v ↦→ FS(w+ vei) is σ-strongly convex, i.e.,

Theorem 3. Let Assumptions 1, 3, 4 hold. Let {wt}, {w(i) t } be produced by (2) with ηt ≤ 1/˜L based on S and S(i) , respectively. Then we have EA[‖wt+1 - w(i) t+1 ‖2] ≤ 4G1 nσ .

Remark 2. Stability bounds of the order O(1/(nσ)) were established for SGD under a strong convexity assumption [Hardt et al., 2016], which are extended to RCD here. Another difference is that the stability bounds in Hardt et al. [2016] are established for either the constant stepsize sequence ηt ≡ η or the specific stepsize sequence ηt = 1/(tσ). As a comparison, our results apply to general stepsizes.

We further consider nonconvex case. We now present stability bounds for nonconvex problems, which are ubiquitous in the modern machine learning. We denote ∏tk=t+1 (1 + ˜Lηkd- 12 ) = 1 [Ying and Zhou, 2017].

Theorem 4. Let Assumptions 1 and 3 hold. Let {wt}, {w(i) t } be produced by (2) based on S and S(i) , respectively. Then

We finally consider almost sure bounds. The following theorem gives almost sure stability bounds, which is useful to develop high-probability generalization bounds. We need a coordinate-wise Lipschitz continuity assumption.

Assumption 5. For all S and [i] ∈ [d], assume |∇iFS(w)| ≤ ˜G for all w ∈ W.

Theorem 5. Let Assumptions 1, 2, 3, 5 hold. Then RCD with T iterations is 2G2 ˜G n ∑Tt=1 ηt-uniformly stable.

5 Generalization of RCD

In this section, we use our stability bounds to develop generalization bounds for RCD. According to (1), our analysis requires to handle optimization errors by tools in convex optimization (details are given in the arXiv version).

We first consider convex case. We use the technique of ℓ1 on-average argument stability to develop generalization bounds under a Lipschitz continuity assumption.

Theorem 6. Let Assumptions 1, 2, 3 hold. Let {wt} be produced by (2) with ηt ≡ η ≤ 2/˜L. Then

The first term on the right-hand side of Eq. (5) is related to estimation error, while the remaining two terms are related to optimization error. According to (5), we know that estimation error bounds increase as we run more and more iterations, while optimization errors decrease. This suggests that we should balance these two errors by stoping the algorithm at an appropriate iteration to enjoy a favorable generalization.

Remark 3. Under the same condition, it was shown that SGD with T ≍ n can achieve the excess generalization bounds O(1/√n) [Hardt et al., 2016]. Here we show that the same generalization bounds can be achieved by RCD.

In Theorem 6, we require the boundedness assumption of stochastic gradients (note the bounded gradient assumption does not hold for the least square loss). We now show that this boundedness assumption can be removed by using the ℓ2-on-average argument stability. A nice property is that it incorporates the information of F(w∗ ) in the generalization bounds. This suggests that better generalization bounds can be achieved if F(w∗ ) is small, which are called optimistic bounds in the literature [Srebro et al., 2010; Zhang and Zhou, 2019]. Here we introduce a parameter γ to balance different components of the generalization bounds.

Theorem 7. Let Assumptions 2, 3 hold. Let {wt} be produced by (2) with nonincreasing ηt ≤ 2/˜L. For any γ > 0 such that (1+T)(1+γ)L2 e ∑Tt=1 η2 t ≤ n2d/4, we have

The following corollary gives a quantitative suggestion on how to stop the algorithm for a good generalization.

Corollary 8. Let Assumptions 2, 3 hold and d = O(n2 ). Let {wt} be produced by (2) with ηt ≡ η ≤ 2/˜L.

Remark 4. If T ≍ √nd 34 , then (1+T)(L+n√d/T)LeTη2 ≍ nd 32 + n 32 d 54 . Then the assumption in Part (a) holds if d ≤ cn2 for some appropriate c > 0. If T ≍ n√d, then (1 + T)L2 eT ≍ n2d. In this case, the assumption (1 + T)L2 eTη2 ≤ n2d/8 in Part (b) is also easy to satisfy.

Remark 5. As compared to Theorem 6, Part (a) admits a worse dependency on the dimensionality, which is the cost we pay for removing the Lipschitz continuity assumption. Furthermore, Part (b) shows that RCD is able to achieve a generalization bound as fast as O(√d/n) if the best model has a small population risk, while Theorem 6 fails to exploit this low-noise assumption and can only imply at most the generalization bound O(1/√n).

Now, we present generalization bounds of RCD for strongly convex objective functions.

Theorem 9. Let Assumptions 1, 2, 3, 4 hold. Let {wt} be produced by (2) with ηt ≡ η ≤ 1/˜L. Then

In particular, we can set T ≍ dσ-1 log 1/(nσ) to get ES,A[F (wT +1) - F (w∗ )] = O(1/(nσ)).

Remark 6. Stability bounds of the order O(1/(nσ)) were established for SGD under a strongly convex setting [Hardt et al., 2016], which together with optimization error bounds of the order O(1/(Tσ)) [Rakhlin et al., 2012], shows that SGD with n iterations can achieve excess risk bounds O(1/(nσ)). Here we show that this optimal generalization bound can also be achieved for RCD with dσ-1 log 1/(nσ) iterations.

Finally, we present high-probability bounds, which are more challenging than bounds in expectation and are important to understand the variation of algorithms in repeated runs.

Theorem 10. Let Assumptions 1, 2, 3, 5 hold. Let {wt} be produced by (2) with ηt ≡ η ≤ 2/˜L and δ ∈ (0, 1). Assume ‖wt‖∞ ≤ R and |f(wt; z)| ≤ R for all t. If we choose T ≍ n 2 3 d 13 log- 23 n log- 13 (1/δ), then with probability 1 - δ

where ¯wT = 1 T ∑Tt=1 wt is an average of iterates.

6 Experiments

In this section, we present some experimental results to illustrate our stability bounds. We follow the set up in Hardt et al. [2016], i.e., we consider two neighboring datasets and run RCD/SGD with ηt ≡ 0.01 on these neighboring datasets to produce two iterate sequences {wt}, {w′t }. We then plot the Euclidean distance between two iterate sequences as a function of the iteration number. We consider the least square regression for two datasets: ionosphere, svmguide3 and MNIST. We repeat the experiments 100 times and report the average of results. In Figure 1 we plot the Euclidean distance as a function of the number of iterations. Experimental results show that the Euclidean distance for RCD is much smaller than that with SGD, which is consistent with our theoretical results that RCD is more stable than SGD.

7 Proof of Theorem 2

Lemma 11. Let g : Rd ↦→ R be convex and have coordinatewise Lipschitz continuous gradients with parameter ˜L > 0. Then for any η ≤ 2/˜L and any i ∈ [d] we have the following inequality for any w and ˜w

‖w- η∇ig(w)ei - ˜w+ η∇ig( ˜w)ei‖2 ≤ ‖w- ˜w‖2. (7) Furthermore, if g is σ-coordinate-wise strongly convex and η ≤ 1/˜L, then (wi denotes the i-th coordinate of w ∈ Rd )

Proof of Theorem 2. By the update rule (2), we know

where we have used Lemma 11 in the last step. Since S and S(i) differ by the i-th example, we know

Note that it is uniformly drawn from [d], we further know

where we have used Assumption 1 in the last step. Plugging the above inequality back into (10), we get

Applying the above inequality recursively gives the stated inequality. This completes the proof of Part (a).

We now prove Part (b). According to (2), we know

It then follows from Lemma 11 that

Note that S and S(i) differ by the i-th example, we can analyze analogously to (11) and get

Since it is uniformly drawn from [d], we further know

where we have used the self-bounding property according to the L-smoothness of f in the last step. Putting the above inequality back into (13) implies

It then follows that

8 Conclusions

In this paper, we initialize the generalization analysis of RCD based on the algorithmic stability. We establish upper bounds of argument stability and uniform stability for RCD, which further imply the optimal generalization bounds of the order O(1/√n) and O(1/n) in the convex and strongly convex case, respectively. We also consider nonconvex case and develop high-probability bounds. Remarkably, our analysis can leverage the low-noise assumption to yield optimistic generalization bounds O(1/n) in the convex case without a bounded gradient assumption.

There are several interesting future directions. First, it would be interesting to extend our analysis to other variants, such as distributed RCD and RCD for structure optimization. Second, here we assume the objectives are convex/strongly convex and each coordinate is sampled with the same probability during RCD updates. It is interesting to extend our discussion to nonconvex setting and importance sampling [Nesterov, 2012], which are popular in modern machine learning.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China (Grant Nos. 61903309, 61806091,

11771012, U1811461) and the Fundamental Research Funds for the Central Universities (JBK1806002).