Generalization Bounds for Stochastic Gradient Descent via Localized ε -Covers
Abstract
In this paper, we propose a new covering technique localized for the trajectories of SGD.
This localization provides an algorithm-specific complexity measured by the covering number, which can have dimension-independent cardinality in contrast to standard uniform covering arguments that result in exponential dimension dependency.
Based on this localized construction, we show that if the objective function is a finite perturbation of a piecewise strongly convex and smooth function with P pieces, i.e. non-convex and non-smooth in general, the generalization error can be upper bounded by O ( (cid:112) (log n log( nP )) /n ) , where n is the number of data samples.
In particular, this rate is independent of dimension and does not require early stopping and decaying step size.
Finally, we employ these results in various contexts and derive generalization bounds for multi-index linear models, multi-class support vector machines, and K -means clustering for both hard and soft label setups, improving the known state-of-the-art rates.
Introduction
We consider the following stochastic optimization problem
[Equation]
where θ represents the optimization parameter, Θ ⊂ R d is a convex parameter domain, f ( · ; z ) is a possibly non-convex loss incurred by a single data point z ∈ Z and Z is a random variable on Z following the data distribution.
Since the distribution of Z is unknown in general, the following proxy based on independent and identically distributed (i i d.) samples z 1 , . .
, z n of Z is optimized instead
[Equation]
Given a learning algorithm A ( · ) mapping samples z 1 , . . .
, z n to an approximate solution of (1.2), bounding the generalization error 1 ˆ F ( A ( z 1 , . . z n ))
− F ( A ( z 1 , . . . , z n )) is a fundamental problem in learning theory.
Classical algorithm-independent results rely on uniform convergence over the entire domain Θ ⊂ R d ; thus, they apply to any learning algorithm.
However, these bounds often increase with the dimension d [SSSSS09, SSBD14, Fel16], becoming vacuous in the modern overparameterized regime [ZBH + 21].
To derive dimension-independent bounds, researchers have been investigating algorithm-dependent generalization properties, especially for commonly used training methods such as stochastic gradient descent (SGD)
[HRS15, SHN + 18, YKM21].
Notably, algorithmic stability is a technique for deriving generalization bounds based on the properties of a specific learning algorithm, which leverages that if a parameter learned by an algorithm is robust under a perturbation of samples z 1 , . .
, z n , then the generalization error at that parameter must be small [BE02].
Based on this principle, several works proposed dimension-independent generalization bounds for SGD and its variants under various setups [HRS15, Lon17, FV19, LLQ20,
BFGT20, LY20, FO21, LLY21, KWS22].
Bounds derived using algorithmic stability is optimal for strongly convex and smooth functions [SSSSS09] and convex and non-smooth functions [BFGT20,
AKL21].
Nevertheless, without global (strong) convexity, early stopping, and/or decaying step size, generalization bounds based on algorithmic stability often diverge with the number of SGD iterations [HRS15, LLQ20], failing to explain the empirical observations.
To obtain (ambient) dimension-independent bounds that do not diverge with the number of iterations, recent works proposed to utilize the low-dimensional fractal structures generated by the SGD iterates whose complexity can be measured by a notion called the Hausdorff dimension [Fal14].
In this context, [¸SSDE20] showed that, under a continuous-time surrogate for SGD, the generalization error can be bounded by (cid:101)
O ( (cid:112)
d H /n ) , where d H denotes the Hausdorff dimension of the optimization trajectory.
This result was later extended to discrete-time iterated function systems by [CDE + 21].
Here, the Hausdorff dimension can be smaller than the ambient dimension [CDE + 21], ultimately providing improved generalization bounds.
However, both of these results are inherently asymptotic, and rely on opaque assumptions that are hard to verify in practice.
In this paper, we propose a new framework for deriving generalization bounds for the projected
SGD with a constant step size and without requiring early stopping.
Inspired by the works [¸SSDE20,
CDE + 21], our framework is based on a complexity measure of the trajectory of SGD, which can be quantified under standard verifiable conditions.
Our contributions are as follows.
• Localized ε -covers for SGD.
Our first principle contribution is a covering technique localized for the trajectories of SGD.
This localization provides an algorithm-specific complexity measured by the covering number, which can have dimension-independent cardinality in contrast to standard covering arguments that result in exponential dimension dependency.
• Generalization bounds for SGD.
Based on this localized covering, we establish dimensionindependent generalization bounds for SGD, for a rich class of non-convex loss functions
f whose gradients can be approximated by that of a piecewise strongly convex and smooth function h ,
i e. ∥∇ f
( θ ; z [Equation] − ∇ h
[Equation] θ ; z ) ∥ ≤ ξ for some ξ .
In particular, with high probability, we prove the bound
[Equation] where
θ ( t ) denotes the parameter generated by t SGD iterations for a sufficiently large t , and P denotes the number of strongly convex pieces needed to approximate f .
We
further show that the gradient of any (piecewise) smooth function f can be approximated with that of a piecewise strongly convex and smooth function, demonstrating the wide applicability of the bound (1.3). Finally in the special case P = 1 and ξ = our result reduces
to a
non-asymptotic bound where the complexity is captured by the Hausdorff dimension of the invariant measure of SGD.
• Improved bounds in specific models.
We employ the above result to derive generalization bounds in several statistical models trained by SGD, including multi-index linear models, multi-class support vector machines and K -means clustering with both hard and soft label setups, improving the previously known state-of-the-art generalization error bounds in this context.
Notation and problem setup.
For k ∈ N , we denote [ k ] := { 1 , . . .
, k } We use ∥ · ∥ to denote the ℓ 2 -norm.
For ε > 0 and θ ∈ R d , we use B dε ( θ ) to denote the d -dimensional closed ℓ 2 -ball of radius ε , centered at θ .
Given a set S ⊆ R d and ε > 0 , we say C ε ⊆ R d is an “ ε -cover” of S if S ⊆ (cid:83)
θ ∈C ε B dε ( θ Given a step size η > 0 , an initial parameter θ (0) ∈ Θ , and a randomly sampled index i t ∈ [ n ] , the t -th iteration of the projected SGD performs the following update on the parameters
[Equation]
where Π Θ ( θ ) :
= arg min θ
′ ∈
Θ ∥ θ ′ − θ ∥ denotes the Euclidean projection.
The domain Θ is convex; thus, the projection operation is unique.
We note that the projection is not needed under the presence of an ℓ 2 -regularizer and Lipschitz continuity; see Section 4 for more details.
Throughout the paper, we use θ ( t ) :
=
g i t ◦ · · · ◦ g i 1 ( θ (0) ) for possibly random indices i 1 , . .
, i t .
Lastly, we recall a few standard notions.
f : Θ → R is called “ α -strongly convex” and “ β -smooth” respectively if for all θ, θ ′ ∈ Θ , the following conditions are satisfied
[Equation]
The function f is called “convex” if it is 0 -strongly convex.
Main results
We demonstrate our covering construction in Section 2.1 in the classical strongly convex and smooth case in which the localization argument can be simplified by the contractivity of SGD.
We present our main generalization result on non-convex losses in Section 2.2, and its implications in Section 3.
2.1 A localized covering construction: Strongly convex and smooth case
To motivate our approach, let us first briefly discuss the limitations of prior methods that are based on uniform convergence of empirical processes.
Given a set of parameters Θ ⊆ B dR (0) , let C ε be an ε -cover of Θ , i.e. Θ ⊆ (cid:83) ϕ ∈C ε B dε ( ϕ ) .
Uniform convergence over Θ can be established with high probability by simply applying the union bound over C ε , which yields a generalization error bound depending on the cardinality of the cover (cid:112) log |C ε | .
However, C ε is typically independent of the algorithm being used, and standard ε -covers for Θ yield |C ε | =
( R/ε ) Ω( d ) ; thus, bounds based on covering numbers often grow with √ d , which can be loose if d is large.
To overcome this issue, we use the contractive properties of SGD in the strongly convex and smooth case and localize the ε -cover.
Namely, instead of covering the entire feasible set Θ , we construct a cover that contains only the points that can be reached by SGD trajectories, resulting in a covering number that is independent of the ambient dimension d .
We introduce the following sets produced by SGD trajectories.
Definition 1.
We define the following two subsets of Θ .
• The set of points that can be reached by T SGD iterations initialized at θ (0) ∈ Θ ,
[Equation]
• The set of points that can be reached by any t ≥ T SGD iterations initialized at θ (0) ∈ Θ ,
[Equation]
For γ ∈ (0 , 1) , a function g : Θ → Θ is called “ γ -contractive” if for all θ, θ ′ ∈ Θ , it satisfies ∥ g ( θ ) − g
( θ ′ )
∥ ≤ γ ∥ θ − θ ′ ∥ .
In the constant step-size case, SGD iterates converge to a distribution instead of a single point [DDB20], but their contractivity can still provide the following localization: all possible SGD iterates after sufficiently many iterations can be ε -covered by n O (log(1 /ε )) points.
Lemma 2.1.
Suppose that g 1 , g 2 , . .
are γ -contractive for some γ ∈ (0 , 1) .
Then, for any initialization θ (0) ∈ Θ ⊆ B dR (0) and for any ε > 0 , for T :=
T ε = max (cid:110)(cid:108) log( R/ε ) log(1 /γ )
(cid:109) , 0 (cid:111) , we have Ψ ≥ T ( θ (0) ) ⊆ (cid:91) ϕ
∈ Ψ T (0) B dε ( ϕ ) .
Proof.
For t ≥ T , let θ ( t ) ∈ Ψ ≥ T ( θ (0) )
such that θ ( t ) :
=
g i t ◦· · ·◦ g i 1 ( θ (0) ) for some
i 1 , . .
, i t ∈ [ n ] .
Let ϕ :
=
g i t ◦ · · · ◦ g i t − T +1 (0) and notice that ϕ ∈ Ψ T (0) by construction. Then for T : = T ε and θ ( t − T ) :
= g i t − T ◦ · · · ◦ g i 1 , we have ∥ θ ( t ) − ϕ ∥ ≤ γ T ∥ θ ( t − T ) − 0 ∥
≤ γ T R ≤ ε, since each g i is γ -contractive and ∥ θ ( t − T ) ∥ ≤ R by (1.4). This implies θ ( t ) ∈ (cid:83) ϕ ∈ Ψ T (0)
B dε ( ϕ ) .
and z SGD a sufficiently small contractive [DDB20] that is, applying g i to any two points decreases the distance between them (see Appendix A
1 for a formal derivation).
Therefore, for any ε > 0 , there exists T
= T ε such that applying T synchronously coupled SGD updates that use the same sample at each iteration can make the distance between the initial points smaller than ε see
Ψ all be generated by T SGD updates when initialized see contrast is of the dimension d
We make the following additional assumptions on the loss function f .
Assumption 1 (Weak Lipschitz continuity) .
For L > 0 , there exists h : Θ → R such that for all θ, θ ′ ∈ Θ and z ∈ Z , | f ( θ ; z ) − h ( θ ) − ( f ( θ ′ ; z ) − h ( θ ′ ))
| ≤ L ∥ θ − θ ′ ∥ .
Notice that h = 0 reduces to the classical Lipshitz continuity, but the above condition is more general.
Assumption 2 (Bounded deviation) .
For B > 0 , for all θ ∈ Θ , (cid:12)(cid:12) sup z ∈Z f
( θ ; z ) − inf z ∈Z f
( θ ; z ) (cid:12)(cid:12) ≤ B. We note that Assumptions 1 & 2 (or their variants) both appear in several algorithmic stability-based results; see e.g. [HRS15, Thm 3.10].
We also highlight that these conditions are invariant to adding regularizers (e.g. consider f ( θ ; z ) ← f
( θ ; z ) + λ 2 ∥ θ ∥ 2 ) which will be useful in Section 3.
Theorem 2.1.
Suppose that Assumptions 1 & 2 hold and there exist α, β > 0 such that f ( · ; z ) is α -strongly convex and β -smooth on Θ ⊆ B dR (0) for all z ∈ Z .
For any η ∈ (0 , 2 /β ) , let γ :=
(cid:112) 1 − 2 αη +
[Equation] 2 , T := max (cid:110)(cid:108) log(2 LRn ) log(1 /γ )
(cid:109) , 0 (cid:111) .
Then, with probability at least 1 − δ , for any θ (0) ∈ Θ , t ≥ T , and i 1 , . .
, i t ∈ [ n ] , SGD iterate θ ( t ) satisfies
[Equation]
Remark.
A few remarks are in order.
For stochastic convex optimization, the best achievable bound via uniform convergence over Θ is worse than the algorithmic stability-based bounds for SGD by a factor of √ d [SSSSS09, Fel16, FV19]; however, by localizing the uniform convergence argument, we are able to obtain a dimension-independent bound.
We note that this high-probability result is still not directly comparable to most stablity-based bounds which are given in expectation.
An exception is [FV19, Thm 4.5] which obtains near-optimal bounds for SGD in the convex case by tuning the number of iterations to be taken, obtaining a rate of O (log( n/δ ) 2 log( n ) / √ n ) .
In contrast, our bound reads O ( (cid:112) log( n ) 2 log(1 /δ )
/n ) which is better by a logarithmic factor and holds for any sufficiently large number of SGD iterations.
We emphasize that the bound (2.1) holds for the union of all trajectories of SGD generated by t ≥ T iterations, whereas the stability-based results often consider a single parameter generated by SGD.
A similar setup is considered in Corollary 2.1 where we establish improved bounds on the generalization error removing the logarithmic factor. Lastly, we note that our localized covering can also be used for deriving bounds in expectation; see Appendix D.
The proof of Theorem 2.1 follows from three steps.
i )
Lemma 2.1 implies that any θ ( t ) for t ≥ T can be approximated by some parameter θ ( T ) generated by T SGD updates initialized at 0.
ii ) ˆ F ( θ ( T ) ) concentrates around F ( θ ( T ) )
since θ ( T ) is almost independent of the samples, i.e. it depends on at most T
=
O (log( n )) of n samples.
iii )
The empirical process | ˆ F ( θ ( T ) )
− F ( θ ( T ) ) | uniformly converges over the set Ψ T (0) which has a dimension-free cardinality.
See Appendix B.2 for details.
While the bound in Theorem 2.1 holds uniformly over all possible initializations, it can be tightened by considering a single realization of θ (0) in the following corollary.
Corollary 2.1.
Assume the setup in Theorem 2.1.
Then, for any θ (0) ∈ Θ , t ≥ T , and i 1 , . .
, i t ∈ [ n ] , with probability at least 1 − δ , we have
[Equation]
For sufficiently large n , the bound in Corollary 2.1 coincides with the tight concentration bound at a fixed θ ∈ Θ ; however, it has the limitation that at least t ≥ T SGD updates are required.
This is remedied in the next result which does not require strong convexity and smoothness.
Corollary 2.2.
Suppose that Assumption 2 holds.
Then, for any θ (0) ∈ Θ , t ≥ 0 , and i 1 , . . .
, i t ∈ [ n ] , with probability at least 1 − δ , we have
[Equation]
The bound in Corollary 2.2 requires early stopping since it diverges as t → ∞ .
However, since
Corollary 2.2 only requires Assumption 2, it can be easily combined with other results, e.g. Corollary 2.1, to provide a generalization bound that holds for any number of SGD iterations.
The proofs of Corollaries 2.1 and 2.2 are presented in Appendices B.3 and B.4, respectively.
Relation to fractal dimension.
There is an interesting connection to be made between Theorem 2. and the Hausdorff dimension of the support µ of the stationary distribution of (projected) SGD.
For example, suppose that η ∈ (0 , 1) , Θ = B dR (0) , f ( θ ; z ) = 12 ∥ θ − θ ∗ z ∥ 2 for some θ ∗ z ∈ B dR (0) , and ∥ θ ∗ z i − θ ∗ z j ∥ ≥ 2 γR for all i ̸
= j .
Here, the last assumption can be satisfied with high probability if d, η are large enough and θ ∗ z ∼ N (0 , σ 2 I ) , i.e. zero-mean Gaussian with covariance σ 2 I ; see e.g. [PLYS21, Appendix C].
Under these assumptions, the bound in Theorem 2.1 can be reformulated as
[Equation]
where d H denotes the Hausdorff dimension of µ .
The precise definition of d H and the derivation of (2.2) are presented in Appendix B.5.
In (2.2), d H replaces the ambient dimension d that appear in algorithm-independent bounds [SSBD14].
We note that such a connection between the fractal dimension and generalization bounds has been also studied in [¸SSDE20, CDE + 21]; however, their bounds are asymptotic and require non-trivial assumptions that are not easy to verify in practice. Generalization due to contractivity. Contractivity of SGD has also been used to derive generalization bounds in the concurrent work [KWS22]. Although the localized covering construction in
Lemma 2.1 relies on the same principle, our results differ from those in [KWS22] on two key aspects. First and foremost, as we shall see in Section 2.2, the localized covering argument used in Lemma 2. can also be applied to the non-convex case in which the uniform stability-based argument used in [KWS22] provably breaks down (Appendix F); thus, extending the results of [KWS22] to cover nonconvex objectives is highly non-trivial. Further, the generalization bounds in [KWS22] are provided in expectation, i.e. | E [ ˆ F ( θ ( t ) )
− F ( θ ( t ) )] | ≤ O (1 /n ) , whereas we provide high-probability bounds, i.e. | ˆ F ( θ ( t ) )
− F ( θ ( t ) ) | ≤ (cid:101) O
(1 / √ n ) .
When translated to bounds in expectation, our results read | E
[ ˆ F ( θ ( t ) ) − F ( θ ( t ) )] | ≤ (cid:101) O (1 /n ) and E [ | ˆ F ( θ ( t ) )
− F ( θ ( t ) ) | ] ≤ (cid:101) O (1 / √ n ) ; see Appendices D & E. 2.2 Non-convex case: Perturbations of piecewise strongly convex and smooth functions Algorithmic stability technique yields (near) optimal rates for strongly convex objectives; however, when applied to non-convex functions, the resulting bounds often diverge with the number of SGD iterations [HRS15, LLQ20].
The localized covering construction introduced in the previous section remedies this issue, providing more stable generalization bounds.
Specifically, we establish a dimension-independent generalization bound for functions that are finite perturbations of piecewise strongly convex and smooth functions.
We further prove an approximation result and show that any smooth non-convex function can be approximated with a piecewise strongly convex function.
Since piecewise strongly convex and smooth functions may not be differentiable on the entire Θ , we will use the auxiliary gradient in Definition 2 as a surrogate for the SGD update (1.4).
Definition 2.
f is “piecewise α -strongly convex and β -smooth with P pieces on Θ ” if there exists a partition {P 1 , . . . , P P } of Θ and α -strongly convex and β -smooth functions h 1 , . . . , h P on Θ such that f = h p on P p for all p ∈ [ P ] .
We also define ∇ f ( θ ) : = ∇ h p ( θ ) if θ ∈ P p . Piecewise strongly convex and smooth objectives are widely used in machine learning applications. For example, the objective of learning a single layer of a ReLU network with ℓ 2 -regularization is piecewise strongly convex and smooth.
Furthermore, the objective of learning an entire ReLU network is also piecewise strongly convex and smooth on small loss regions [Mil19].
Such observations easily extend to more general settings, e.g. an objective function defined as the summation of a piecewise linear loss (e.g. hinge loss) and ℓ 2 -regularization.
However, if we further allow for finite perturbations, any piecewise smooth non-convex function can be covered within this framework.
Proposition 2.1.
For any piecewise β ′ -smooth f with Q pieces on Θ ⊂ B dR (0) , and for any ξ > 0 and 0 < α ≤ β , there exists a piecewise α -strongly convex and β -smooth h with at most
Q (3( β + β ′ )
R/ξ ) d pieces such that ∥∇ f ( θ ) − ∇ h ( θ ) ∥ ≤ ξ for all θ ∈ Θ .
For this general class of non-convex functions, we derive the following generalization bound.
Theorem 2.2.
Suppose that Assumptions 1 & 2 hold and Θ ⊆ B dR (0) .
Suppose further that there exists h : Θ × Z
→ R such that for any z ∈ Z , h ( · ; z ) is piecewise α -strongly convex and β -smooth with P pieces on Θ satisfying
[Equation]
for all θ ∈ Θ .
For any η ∈ (0 , 2 /β ) , let γ :=
(cid:112) 1 − 2 αη
+ αβη 2 .
Then given T ∈ N , with probability at least 1 − δ , for any θ (0) ∈ Θ , t ≥ T , and i 1 , . .
, i t ∈ [ n ] , we have
[Equation]
Remark.
The bound above is stated in full generality and holds for any SGD iterate t ≥ T and any choice of T ≥ 1 .
However, to obtain a meaningful generalization bound, one may choose, for example T
=
O (log( nR ) / log( γ − 1 )) .
In the case that γ, B, L = Θ(1) , the bound simplifies to
[Equation]
Here, the first term in the bound is logarithmic in the number of pieces P and the last term scales linearly with the approximation error ξ .
Thus, the generalization bound depends on the trade-off between the complexity of h through P , and how well ∇ h approximates ∇ f through ξ .
In this regime, in contrast to algorithmic stability-based bounds, the above result does not grow with the number of SGD iterations, i.e. early stopping is not required for generalization in Theorem 2.2.
We finally note that Theorem 2.2 holds for any f with an auxiliary gradient ∇ f satisfying (2.3).
In light of Proposition 2.1, any piecewise smooth function can be approximated by a piecewise strongly convex and smooth function; in the worse case, the number of pieces P is at most e (cid:101) Ω( d ) .
Therefore in this pessimistic case, Theorem 2.2 recovers the classical algorithm-independent covering bound (cid:101) O
( B (cid:112) d/n ) by choosing ξ
= O (1 / ( L √ n )) and T = Θ(log( LRn )) .
However, any value of P that is sub-exponential in dimension yields improved generalization bounds.
In particular in the next section, we consider certain (non-convex) statistical models and carefully design h so that
Theorem 2.2 improves the existing generalization error bounds.
In contrast, uniform stability-based bounds for piecewise strongly convex and smooth functions are in general Ω(1) after sufficiently many SGD iterations; see Appendix F.
That is, the contractivity-based bounds in [KWS22] cannot be directly extended to piecewise contractivity (2.3).
The proof of Theorem 2.2 relies on a modified version of the covering construction presented in
Section 2.1.
First, we define the auxiliary parameter update g i,p ( θ ) :
= θ −∇ h p ( θ ; z i ) where h p ( · ; z i ) denotes the strongly convex and smooth function satisfying ∇ f ( · ; z i )
≈ ∇ h p ( · ; z i ) on the p -th piece.
We show that Ψ ′ T (0) :
= { g i T ,p T ◦ · · · ◦ g i 1 ,p 1 (0) : i 1 , . .
, i T ∈ [ n ] , p 1 . . . , p T ∈ [ P ] } εcovers Ψ ≥ T ( θ (0) ) for any θ (0) ∈ Θ
as in Lemma 2.1.
Since | Ψ ′ T | ≤ ( nP ) T , applying a concentration inequality and the union bound over the localized cover yields Theorem 2.2.
Formal proofs of
Proposition 2.1 and Theorem 2.2 are provided in Appendices B.6 and B.7, respectively.
Applications
In this section, we use our localized covering to prove generalization bounds for multi-index linear models, multi-class support vector machines, and K -means clustering for both hard and soft label setups, trained by SGD, improving the state-of-the-art results known for these models.
3.1 Multi-index linear models
Given a sample z
=
( y, x ) ∈ Y × B dR x (0) , consider the ℓ 2 -regularized loss in a multi-index model with K indices parameterized by θ =
( θ j ) Kj =1 ∈ R
d
× K
[Equation]
Characterizing the generalization properties of multi-index models is an important problem with many applications including regression, classification, dimension reduction, and learning a single-layer of a neural network.
In the following theorem, we derive a generalization bound for multi-index models by approximating each f ( · ; z ) with a piecewise strongly convex and smooth function, and then applying Theorem 2.2.
The proof is given in Appendix C.1.
Theorem 3.1.
Let θ =
( θ j ) Kj =1 and Θ = (cid:81)
Kj =1 Θ j for some convex Θ j ⊆ B dR (0) .
Suppose that f satisfies (3.1) , Assumptions 1 & 2 hold, and ℓ ( · , . . .
, · ; y ) is piecewise β -smooth with Q pieces on Θ for all y ∈ Y .
For any η ∈ (0 , 2 /λ ) , let γ :=
| 1 − ηλ | , T := max (cid:110)(cid:108) log(3 LRn ) log(1 /γ )
(cid:109) , 0 (cid:111) , and
P := poly ( β, η, K, L, R, R x , T, n ) .
Then, with probability at least 1 − δ , for any θ (0) ∈ Θ , t ≥ T , and i 1 , . .
, i t ∈ [ n ] , we have
[Equation]
Generalization behavior of multi-index models has received considerable attention; a subset of notable results include [Gue02, Zha04, JKZ + 12, CMR13, LDZK19].
Under the same conditions of
Theorem 3.1, existing state-of-the-art bounds scale at least linearly in √ d or K , while our result is dimension-free and scales with √ K .
Specifically, the result in [LDZK19] translated to our setting reads ( ℓ is L -Lipschitz, and f ( θ ; z ) ∈ [0 , B ] which are stronger conditions than Assumptions 1 & 2)
[Equation]
for all θ ∈ Θ , which is obtained via [LDZK19, Cor 3 & 9] with Λ ← R and p ← ∞ .
While the bound in (3.2) is linear in KLRR x , our bound is linear in √ K and logarithmic in L, R, R x .
The significance of this improvement can be better seen, for example, when learning the first layer of neural networks, for which L can be very large.
Moreover, if x ∼ N (0 , I ) , then ∥ x ∥ = (cid:101) Ω( √ d ) with high probability; that is, (3.2) is linear in R x = Ω( √ d ) but ours only scales with √ log d , which can make a significant difference especially in the overparameterized regime.
We should note that compared to our bound, (3.2) does not require the smoothness of ℓ .
For a specific application of Theorem 3.1, consider the multi-class support vector machines, i.e. for
Y =
[ K ] and ρ : [ − RR x , RR x ] → [0 , B ] is L -Lipschitz and β -smooth, the objective is given as
[Equation]
Theorem 3.1 provides a generalization bound of (cid:101)
O
( (cid:112) K/n ) that improves the existing bounds which scale at least linearly with K and/or √ d [Zha04, DSBDSS15, LDBK15, LDZK19]; see e.g.
[LDZK19, Sec 2.1] for similar related results.
3.2 K -means clustering
We can also provide generalization error bounds for the problem of K -means clustering in both hard and soft label setups.
As before, we consider SGD as an optimizer which has been the focus of many works in this context [BB94, Scu10, TM17].
We assume throughout this section that the samples are supported on a bounded domain, and the SGD iterations are applied without projection.
Generalization properties of K -means clustering has been studied for decades [Ant05, Lev13,
TTJC15, TM16].
While most existing bounds are at least linear in K , [LL21] recently provided an improved bound in the hard label setup, which is of order ( K log 3 n/n ) 1 / 2 for bounded inputs.
However, state-of-the-art bounds in the soft label setup were still linear in K .
Below, we establish a bound of order (cid:112) K log n/n in the soft label setup, which is the first bound that is sublinear in K .
We further improve the bounds of [LL21] in the hard label setup, but by a logarithmic factor.
Soft K -means clustering.
While the hard K -means clustering assigns exactly one cluster to each point using one-hot encoding, the soft K -means clustering allows soft labels within the probability simplex.
Specifically, given ζ > 0 and the samples z 1 , . . .
, z n , the soft K -means clustering algorithm alternates between updating the soft labels ( w j ( z i ; θ ))
Kj =1 of z
i for all i ∈ [ n ] , and estimating the cluster centers θ = ( θ j ) Kj =1 ∈ ( B dR (0))
K using the classical update rule [Mac03]
[Equation]
Here, this procedure is equivalent to a special case of expectation-maximization algorithm, which converges to a local minimum of the following objective (see Appendix C.2 for details)
[Equation]
Instead of running the standard alternating procedure (3.3), we directly minimize (3.4) using SGD and derive generalization bounds for the soft K -means clustering by approximating the objective (3.4) with piecewise strongly convex and smooth functions.
The proof is presented in Appendix C.3.
Theorem 3.2.
For B := 4( R + 1) 2 and for any η ∈ (0 , Ke − ζB ) , let γ := (cid:113) 1 − 4 ηe − ζB K + 4 η 2 K
[Equation]
L := 4 R √ K e ζB , T := max (cid:110)(cid:108) log(3 LRn ) log(1 /γ )
(cid:109) , 0 (cid:111) , and P :=
poly ( η, ζ, B, K, R, T, n, e ζB ) .
Then, with probability at least 1 − δ , for any θ (0) ∈ Θ , t ≥ T , and i 1 , . .
, i t ∈ [ n ] , we have
[Equation]
The above generalization bound scales with √ K , which improves the previously known bound (cid:101) O ( K/ √ n ) for clustering with soft labels [LL21].
Theorem 3.2 also implies (cid:101) O
( (cid:112) K/n ) bound under
R, ζ = Θ(1) and η = Θ( K ) , which coincides with the best known rate in the hard label setup.
Here, the choice of η = Θ( K ) may look odd; however, it is indeed practical since the derivative of the objective function f scales with 1 /K , i.e. ∥ ∂f ( θ ; z ) /∂θ j ∥ ≤ 4 Re ζB /K .
Hard K -means clustering.
Under the same setup as before, we minimize the objective function
[Equation]
Note that (3.5) coincides with the soft K -means objective (3.4) as ζ → ∞ if ∥ θ j − z i ∥ ̸ = ∥
θ k − z i ∥ for all i, j, k .
Since f ( θ ; z ) =
min j ∈ [ K ] ∥ θ j − z ∥ 2 may not be differentiable, we use an auxiliary gradient at a non-differentiable θ , i.e. for a randomly or deterministically chosen S ⊆ arg min j ∈ [ K ] ∥ θ j − z ∥ 2 ,
[Equation]
For example, one may choose a single cluster index ( |S| =1 ) and compute the corresponding gradient.
The following result characterizes the generalization of hard K -means clustering trained by SGD using this auxiliary gradient.
Its proof is deferred to Appendix C.4.
Theorem 3.3.
For any η ∈ (0 , 1) , let B := 4 R 2 , γ :=
| 1 − 2 η | , T := max (cid:110)(cid:108) log
[Equation]
16 √ KR 2 n ) log(1 /γ )
[Equation]
0 (cid:111) .
Then, given θ (0) ∈ Θ , the following bound holds with probability at least 1 − δ , for any t ≥ 0 and i 1 , . . .
, i t ∈ [ n ] , we have
[Equation]
We note that under η, R =Θ(1) and K = O ( n ) , Theorem 3.3 provides O (( K log 2 n/n ) 1 / 2 ) generalization bound which improves the bound given by [LL21], but this time by a logarithmic factor.
Discussions
Beyond (piecewise) strong convexity and smoothness.
The localized covering construction we utilized is essentially based on the (piecewise) contractivity of SGD updates for (piecewise) strongly convex and smooth functions.
However, local strong convexity is by no means necessary and this analysis can be extended to a broader class of objective functions.
Consider, for example, an objective function f that is uniformly convex [DN21] and has Hölder continuous gradient [Nes15].
Then, for any ε > 0 , there exists a step size η > 0 and T =
T ε such that applying T synchronously coupled SGD updates on any two points makes the distance between them smaller than ε .
Consequently, an analog of Lemma 2.1 can be established under this setup as well, for a properly chosen η, ε, T .
We highlight that this class of functions is already covered in our framework via Proposition 2.1; however, an analysis based on the actual curvature (as opposed to a piecewise approximation) of the objective may provide tighter generalization bounds.
Extension to SGD without projection.
Projection operation is only needed to ensure that the SGD iterates stay bounded.
Indeed, the projection is not needed in the presence of explicit regularization, or under a dissipativity-type condition on the objective [RRT17, EMS18, YBVE21, EH21].
For example, if f ( · ; z ) is Lipschitz continuous, it is a straightforward exercise to show that SGD iterations are bounded in the presence of ℓ 2 -regularization and bounded initialization.
Extension to mini-batch setup and different sampling schemes.
The results we presented can be easily extended to the mini-batch setting.
To see this, note that the localization is based on the (piecewise) contractivity of a single iteration of the algorithm, and we have that if all g 1 , . .
, g n are contractive, then an average of their subset is also contractive.
Moreover, our generalization bounds hold for any sampling scheme, e.g. sampling without replacement, random shuffling, datadependent sampling, since the covering construction is based on an ε -cover of the union of all possible trajectories of SGD , which is independent of the underlying sampling scheme.
Extension to contractive stochastic optimization methods.
In this paper, our main focus was the SGD algorithm and its generalization properties.
However, it is straightforward to adapt our framework to a general stochastic optimization setup.
For example, the next result follows from the identical steps leading to Theorem 2.2; hence, its proof is omitted.
Definition 3. An “iterative stochastic algorithm” using g : Θ × Z → Θ performs the update θ ( t ) = g i t ( θ ( t − 1) ) : = g ( θ ( t − 1) ; z i t ) at iteration t , for a random sample i t ∈ [ n ] .
Definition 4. g : Θ → Θ is “piecewise γ -contractive with P pieces on Θ ” if there exists a partition
P 1 , . .
P P of Θ and γ -contractive h 1 , . .
, h P on Θ such that g = h p on P p for all p ∈
[ P ] .
Theorem 4.1.
Suppose that Assumptions 1 & 2 hold, Θ ⊂ B dR (0) , and an iterative stochastic algorithm g is given.
Suppose further that there exists h : Θ × Z → Θ such that for any z ∈ Z , h ( · ; z ) is piecewise γ -contractive with P pieces on Θ satisfying
[Equation]
for all θ ∈ Θ . Choose T ∈ N .
Then with probability at least 1 − δ , for any θ (0) ∈ Θ , t ≥ T , and i 1 , . . . , i t ∈ [ n ] , we have [Equation]
We note that a concurrent work by [KWS22] requires global contractivity and their bounds are in expectation, whereas Theorem 4.1 is a high-probability statement which only requires piecewise contractivity as stated in Definition 4.
Limitation of our results.
We outline a few limitations of our current analysis.
In this paper, we only considered the generalization error (also referred to as the generalization gap), and our main result in the non-convex regime cannot be easily translated to a bound on the excess risk.
This additional step would require a bound on the optimization error; nevertheless, establishing such bounds in the non-convex regime is highly non-trivial.
It is worth highlighting that without assuming convexity, algorithmic stability-based bounds also suffer from this limitation [HRS15, FV19, BKZ20].
Because of the shared contractivity parameter γ across different iterations and pieces, our method cannot be easily extended to varying step size.
This is indeed a limitation of our analysis as varying step size schedules are oftentimes used in practice.
We leave this extension as a future work.
SP was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2019-0-00079, Artificial Intelligence Graduate School Program, Korea University).
US was funded in part by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute), and the European Research Council Starting
Grant DYNASTY – 101039676.
MAE was supported by NSERC Grant [2019-06167], CIFAR AI
Chairs program, and CIFAR AI Catalyst grant.
A Technical results
All results presented in this section are standard.
However, we provide their proofs for convenience.
Lemma A.1.
Let S be a Hibert space associated with the norm ∥ · ∥ induced by the inner product ⟨· , ·⟩ .
Let C ⊂ S be a convex set and Π C ( x ) :=
arg min z ∈C ∥ x − z ∥ be a projection of x onto C .
Then, [Equation] Proof.
Since C is convex, Π C ( x ) is well-defined.
From the definition of the projection and the convexity of C , we have
[Equation]
for all t ∈ [0 , 1] .
Since we have ∂
[Equation]
and ∂∂t ∥ Π C ( x ) − x + t (Π C ( y ) − Π C ( x )) ∥ 2 ≥ 0 at t = 0 by (A.1), we have [Equation]
A
[Equation] Likewise, we have [Equation] (A
[Equation]
Now, consider the function
[Equation]
i.e. f (0) =
∥ Π C ( x ) − Π C ( y ) ∥ 2 and f (1) =
∥ x − y ∥ 2 .
Then f (0) ≤ f (1) since the following inequality holds:
[Equation]
where the inequality follows from (A.3) and (A.4).
This completes the proof of Lemma A.1.
Lemma A.2 (Hoeffding’s inequaltiy [Hoe63]) .
Let X be a random variable on [ a, b ] and X 1 , . . .
, X n be independent copies of X .
Then, it holds that
[Equation]
Lemma A.3.
Suppose that f : R d → R is convex and differentiable.
Then, the following conditions are equivalent.
{+f is β -smooth, i.e. ∥∇ f ( θ ) − ∇ f ( θ ′ ) ∥ ≤ β ∥ θ − θ ′ ∥ for all θ, θ ′ ∈ R d . 2. β 
 ∥ θ ∥ 2 − f ( θ ) is convex for all θ ∈ R d . 3.+}
f ( θ ) − f ( θ ′ )
− ∇ f ( θ ′ ) ⊤ ( θ − θ ′ ) ≤ β 2 ∥ θ − θ ′ ∥ 2 for all θ, θ ′ ∈ R d . 4. ( ∇ f ( θ ) − ∇ f ( θ ′ )) ⊤ ( θ − θ ′ ) ≥
β
∥∇ f ( θ ) − ∇ f ( θ ′ )
∥ 2 for all θ, θ ′ ∈ R d . Proof. 1 ⇒ 2. Let g ( θ ) :
=
{+β 
+}
∥ θ ∥ 2 − f ( θ ) . Then for any θ ̸ = θ ′ , [Equation] In addition, by the mean value theorem, there exists s ∈ (0 , 1) such that for θ s = sθ + (1 − s ) θ ′ ,
[Equation]
Using this, we can derive the following inequality [Equation] which implies that g is convex. 2 ⇔ 3.
The following relation shows the equivalence of the second and the third statements:
[Equation]
2,3 ⇒ 4.
We first introduce the following claim. Claim 1. Suppose that f : R d → R is convex and differentiable and have a global minimum θ ∗ . Then, 12 β ∥∇ f ( θ ) ∥ 2 ≤ f ( θ ) − f ( θ ∗ ) .
Proof.
The statement of Claim 1 is a consequence of the following relation
[Equation]
Let f θ ( ϕ ) :
= f ( ϕ ) − ∇ f ( θ ) ⊤ ϕ . Since β 2 ∥ ϕ ∥ 2 − f θ ( ϕ ) is convex and ϕ = θ minimizes f θ , from
Claim 1, we have
[Equation]
Adding two above inequalities derive the fourth statement.
4 ⇒ 1.
The following inequality is sufficient for deriving the first statement [Equation] A.1 Contractivity of projected SGD for strongly convex and smooth objectives Lemma A.4.
Let f : R d → R be α -strongly convex and β -smooth. Then for any η ∈ (0 , 2 /β ) and for any convex Θ ⊂ R d , θ (cid:55)→ Π Θ ( θ − η ∇ f ( θ )) is (cid:112) 1 − 2 αη + αβη 2 -contractive.
Proof.
Let g ( θ ) :
= θ − η ∇ f ( θ ) . Then for any θ, θ ′ ∈ R d , [Equation] Here, the inequality is from the α -strong convexity and Lemma A.3, i.e.
[Equation]
Using Lemma A.1 completes the proof of Lemma A.4.
In Lemma 2.1, we show that the union of all trajectories generated by {+t ≥ T = T ε SGD iterations, i.e. (cid:83) θ (
 ) ∈+}
Θ Ψ ≥ T ( θ (0) ) , can be ε -covered by a set of points generated by exactly T SGD iterations, i.e.
Ψ T (0) . Here, each g i T ◦ · · · ◦ g i 1 (0) ∈ Ψ T (0) can be viewed as a deterministic algorithm mapping samples z 1 , . . . , z n to a parameter in Θ , where each g i T ◦ · · · ◦ g i 1 (0) only depends on at most T samples z
i 1 , . . . , z i T . Under this observation, in this section, we generalize our localized cover so that its elements are general deterministic algorithms depending on a small number of samples, i.e. not restricted to possible instances of SGD. Then, we derive generalization bounds using our (generalized) localized cover.
To this end, we define algorithms depending on samples. Definition 5. ϕ : Z n → Θ is a “(deterministic) algorithm depending on at most T samples” if there exists I ∈ {S ⊂ [ n ] : |S| = T } satisfying the following: for any z 1 , . . . , z n , z ′ 1 , . . .
, z ′ n ∈ Z such that z i = z ′ i for all i ∈ I , ϕ ( z 1 , . . . , z n ) =
ϕ ( z ′ 1 , . .
, z ′ n ) .
Here, we refer to I as the “set of sample indices determining ϕ .”
The collection of all algorithms depending on at most T samples is denoted by A T .
In Definition 5, we define ϕ to be an algorithm depending on at most T samples if the value of ϕ can be fully determined by a subset of the samples { z i } i ∈I for some I ∈ {S ⊂ [ n ] : |S| = T } .
We note that Ψ T ( θ (0) )
⊂
A T for any T ≥ 0 and θ (0) ∈ Θ .
Now, we introduce the following theorem for deriving generalization bounds using a localized cover where each element in the cover is an algorithm depending on at most T samples.
Theorem B.1.
Suppose that Assumptions 1 & 2 hold and Θ ⊂ R d .
Let ε > 0 , Φ n ⊆
A n , and Φ T,ε ⊆
A T such that
[Equation]
Let µ be a distribution over Z and z 1: n = ( z 1 , . . .
, z n ) is such
that z
i ’s are i.i.d.
samples from µ .
Then, with probability at least 1 − δ over the sampling distribution of z 1: n , for any ψ ∈ Ψ ,
[Equation]
Theorem B.1 implies that if (i)
Φ n can be covered by a set Φ T,ε of algorithms depending on at most T samples and (ii) B , L , T , and | Φ T,ε | are independent of d , then a dimension-independent generalization bound can be derived.
We note that the assumption (B.1) is a generalization of the observation in Lemma 2.1 since Ψ ≥ T ⊂
A n and Ψ T ⊂
A T .
Proof of Theorem B.1.
From Assumption 1 and the assumption (B.1), for any ψ ∈ Φ n , we have
[Equation]
Namely, if we bound | ˆ F ( ϕ ( z 1: n )) − F ( ϕ ( z 1: n )) | for all ϕ ∈ Φ T,ε , the bound in Theorem B.1 follows.
Since the statement of Theorem B.1 is trivial if T = n or | Φ T,ε | = ∞ , we assume T < n and | Φ T,ε | < ∞ .
Now, we derive the target bound: for ϕ ∈
Φ T,ε , θ := ϕ
( z 1: n ) ∈ Θ , and the set of indices I ϕ determining ϕ with |I ϕ
| ≤ T ,
[Equation]
≤ B |I
ϕ | n
[Equation]
n
[Equation]
≤ BT
[Equation]
[Equation]
[Equation]
[Equation]
For the first inequality in (B.3), we use the triangle inequality to upper bound
| ˆ F ( θ ) − F ( θ ) | using two terms: we only utilize the samples determining θ in the first term while remaining samples independent of θ are considered in the second term.
The second inequality directly follows from
Assumption 2.
To bound the second term in RHS of the second inequality, one can apply the
Hoeffding’s inequality (see Lemma A.2), which leads us to the third inequality in (B.3).
Here, the last inequality naturally follows.
By using (B.2), (B.3), and by applying the union bound for all ϕ ∈ Φ T,ε , we obtain the bound in Theorem B.1.
B.2 Proof of Theorem 2.
First, observe that g 1 , . .
. , g n are γ -contractive by Lemma A.4.
Let ε = 1 / (2 Ln ) ; then by Lemma 2.1, we have
[Equation]
Now, we apply Theorem B.1 with
[Equation]
where the assumption (B.1) in Theorem B.1 is satisfied by (B.4). This completes the proof of Theorem 2.1. B.3 Proof of Corollary 2.
Let θ ( t ) :
=
g i t ◦ · · · ◦ g i 1 ( θ (0) ) and ϕ :=
g i t ◦ · · · ◦ g t − T +1 (0) , i.e. ϕ is an algorithm depending on at most T samples.
Since each g i is γ -contractive by Lemma A.4, one can observe that
[Equation]
Now, we apply Theorem B.1 with
[Equation]
← { θ ( t ) } , Φ T,ε ← { ϕ } , ε ← ε, T ← T, δ ← δ.
where the assumption (B.1) is satisfied by (B.5) and | Φ ε | = 1 .
This provides the bound in Corollary 2.1.
B.4 Proof of Corollary 2.
The proof of Corollary 2.2 is simple.
Since θ (0) , t , and i 1 , . .
, i t are fixed, θ ( t ) =
g i t ◦ · · · ◦ g i 1 ( θ (0) ) is an algorithm depending on at most t samples.
Then by using Theorem B.1 with [Equation] , we obtain the bound in Corollary 2.2. B.5 Derivation of (2.2)
We first formally define the Hausdorff dimension.
Definition 6.
Given µ ⊂ R d , d H defined below is the “Hausdorff dimension” of µ :
[Equation]
such that there exists ( θ i ) ki
=1 satisfying µ ⊂ (cid:91)
i
B r
[Equation]
( θ i ) (cid:111) .
From the assumption on f ( θ ; z ) , one can observe that for all i ∈ [ n ] and θ, θ ′ ∈ B dR (0) , ∥ g i ( θ ) −
g i ( θ ′ )
∥ ≤ γ ∥ θ − θ ′ ∥ where γ = 1 − η since α = β = 1 .
Then, the assumption ∥ θ ∗ z i − θ ∗ z j ∥ ≥ 2 γR guarantees
that g i ( int ( B R )) ∩ g j ( int ( B R ))
= ∅ where int ( S ) denotes the interior of a set S .
Finally, the following theorem shows that d H = log n log(1 /γ ) .
Substituting d H to the bound in Theorem 2.1 results in (2.2).
Theorem B.2 (Theorem 9.3 in [Fal14]) .
Suppose that g i ( int ( B R )) ∩ g j ( int ( B R )) = ∅ for all i ̸ = j and ∥ g i ( θ ) − g i ( θ ′ )
∥
= γ ∥ θ − θ ′ ∥ for all i . Then d H = log n log(1 /γ ) . B.6 Proof of Proposition 2.  In this proof, we explicitly construct a piecewise β -strongly convex and β -smooth function, i.e. quadratic, which is always piecewise α -strongly convex and β -smooth for any α ∈ [0 , β ] .
Let {Q 1 , . . .
, Q Q } be the partition of B dR (0) such that f = ℓ q on Q q for some smooth ℓ q on B dR (0) .
Given q ∈ [ Q ] and ε > 0 , let C q,ε = { ϕ q, 1 , . . .
, ϕ q,P q } ⊂
Θ be an ε -cover of Θ with the minimum cardinality, i.e. P q ≤ (3 R/ε ) d for ε ≤ R .
Let {P q, 1 , . . .
, P q,P q } be a partition of Q q where each
P q,p is defined as follows:
[Equation]
For each q ∈ [ Q ] and p ∈ [ P q ] , we also define h ( θ ) :
= h q,p ( θ ) on P q,P q where
[Equation]
and ∇ j denotes the partial derivative with respect to the j -th entry.
Then, one can observe that h is piecewise α -strongly convex and β -smooth with Q (3 R/ε ) d pieces for any α ≤ β .
Furthermore, for any θ ∈ P q,p , we have
[Equation]
Choosing ε := ξ/ ( β + β ′ ) completes the proof of Proposition 2.1.
B.7 Proof of Theorem 2.
Given z ∈ Z , let P z, 1 , . .
, P z,P be a partition of Θ and h 1 ( · ; z ) , . . .
, h P ( · ; z ) be α -strongly convex and β -smooth functions
such that h ( · ; z ) = h p ( · ; z ) on P z,p for all p ∈ [ P ] . Let g i,p ( θ ) :
= θ − η ∇ h p ( θ ; z i ) , i.e. each g i,p is γ -contractive by Lemma A.4.
Given ψ (0) ∈ Θ , t ≥ T , and i 1 , . . . , i t ∈ [ n ] , let ψ ( s ) : = g i s ◦ · · · ◦ g i 1 ( ψ (0) ) for all s ∈ [ t ] and let p s ∈
[ P ] be an index satisfying ψ ( s − 1)
∈
P z is ,p s . Let ϕ : = g i t ,p t ◦ · · · ◦ g i t − T +1 ,p t − T +1 (0) , i.e. ϕ is an algorithm depending on at most T samples.
Then for any θ ∈ Θ , we have
[Equation]
This implies that
[Equation] Now, we apply Theorem B.1 with [Equation] where the assumption (B.1) in Theorem B.1 is satisfied by (B.6) and | Φ T,ε | ≤ ( nP ) T .
This leads us to the bound in Theorem
[Equation]
[Equation]
C Proofs of results in SectionC.1 Proof of Theorem 3.
Given z = ( y, x ) ∈ Y × B dR x (0) = Z , let {Q z, 1 , . . .
, Q z,Q } be a partition of Θ such that ℓ ( · , . . .
, · ; y ) is β -smooth on Q z,q for all q ∈ [ Q ] .
For proving Theorem 3.1, we utilize Theorem 2.2 by approximating each f ( · ; z ) using a piecewise strongly convex and smooth function h ( · ; z ) with P K Q pieces.
To this end, we first define P by [Equation] where κ := [Equation]
Using P defined as above, for each z = ( y, x
) ∈
Y ×B dR x (0) , we construct a partition {P z,q,p 1 ,...,p K : p 1 , . . . , p K ∈ [ P ] , q ∈ [ Q ] } of Θ as follows [Equation]
where [Equation]
For the notational simplicity, let u := ( z, q, p 1 , . . . , p K ) ∈
Z × [ Q ] × [ P ] K . Now, we define our approximation as h ( · ; z ) = h u ( · ) on P u if P u ̸ = ∅ . Here, h u ( · ) is defined as follows: for some fixed ν u := ( ν u, 1 , . . .
, ν u,K ) ∈ P u , h u ( θ ) := ℓ ( ν ⊤ u
[Equation] [Equation]
. . . , ν ⊤ u
[Equation] K x [Equation] y
[Equation]
where ∇ j denotes the partial derivative with respect to the j -th argument.
Namely, h u is a first order approximation of ℓ at θ with ℓ 2 -regularization; hence, h u is λ -strongly convex and λ -smooth.
Then given z ∈ Z and θ ∈ Θ , for q and p 1 , . . .
, p K satisfying θ ∈ P u for u =
( z, q, p 1 , . . .
, p K ) , one can observe that
[Equation]
Now, we apply Theorem 2.2 with
[Equation]
Then in the bound of Theorem 2.2, we have
[Equation]
since γ T R ≤ 1 / (3 Ln ) and 1 − γ T 1 − γ = (cid:80) T − 1 j =0 γ j ≤ T .
This completes the proof of Theorem 3.1.
C.2 Equivalence between soft K -means algorithm and
expectation-
maximization
for
We first observe that applying the affine transformation x (cid:55)→ − nKx + log( ( ξ/π ) d/ 2 K ) to (3.4) results in the following objective:
[Equation]
i.e. the expectation-maximization algorithm for (C.1) aims to find a local minimum of (3.4). Since (C.1) is the log-likelihood for the mixture of Gaussians N ( θ 1 , 12 ζ I ) , . . . , N ( θ K , 12 ζ I ) under the same cluster density 1 /K with observations z 1 , . . . , z n , the expectation-maximization algorithm for (C.1) is identical to the alternative procedure (3.3)
[Pri12].
C.3 Proof of Theorem 3.
First, observe that for θ ∈
( B dR (0)) K , z ∈ B dR (0) , we have
[Equation]
where 1 k = k ′ is one if k = k ′ and zero otherwise.
Using this, we bound f and its first and second derivatives as follows:
[Equation]
To utilize Theorem 2. [Equation] we approximate each f ( · ; z ) using a piecewise strongly convex and smooth function h ( · ; z ) with P K pieces. To this end
[Equation] we first construct a partition {P z,p 1 ,...,p K : p 1 , . . p K ∈ [ P ] } of Θ for each z ∈ Z as follows:
[Equation] where [Equation] 1] [ µ P , µ P +1 ] if p = P .
We define h ( · ; z ) : = h z,p 1 ,...,p K ( · ) on P z,p 1 ,...,p K if P z,p 1 ,...,p K ̸ = ∅ . Here, note that P z,p 1 ,...,p K \ { z } ̸ = ∅ if P z,p 1 ,...,p K ̸ = ∅ from the definition of P z,p 1 ,...,p K . Now, we define h z,p 1 ,...,p K ( · ) as follows: for some fixed ν z,p 1 ,...,p K =
(( ν z,p 1 ,...,p K ) j ) Kj =1 ∈ P z,p 1 ,...,p K \ { z } , [Equation] Our construction of h z,p 1 ,...,p K has some nice properties. For example, for
[Equation] [Table] h z,p 1 ,...,p K is a z,p 1 ,...,p K -strongly convex and b z,p 1 ,...,p K -smooth and α ≤ a z,p 1 ,...,p K ≤ b z,p 1 ,...,p K ≤ β . Furthermore, for any ν ′ =
( ν ′ j ) Kj =1 satisfying [Equation] we have [Equation]
from the symmetry of f and h z,p 1 ,...,p K . Now, we bound ∥∇ f
( θ ; z ) − ∇ f ( θ ; z ) ∥ to utilize Theorem 2.2.
For θ = ( θ j ) Kj =1 ∈ ( B dR (0))
K , let p 1 , . .
. , p K be indices in [ P ] satisfying θ ∈ P z,p 1 ,...,p K .
Let ν ′ be a point on the line connecting z and θ such that ν ′ satisfies (C.2), i.e. ∥ ν ′ j − θ j ∥ ≤ κ for all j ∈ [ K ] . Then, we have
[Equation]
The first inequality holds since ∇ f ( ν ′ ; z )
= ∇ h z,p 1 ,...,p K ( ν ′ ) by (C.3) and other inequalities holds from the definitions of β and β ′ . Now, we apply Theorem 2.2 with [Equation] Then in the bound of Theorem 2.2, we have
[Equation]
since γ T R ≤ 1 / (3 Ln ) and 1 − γ T 1 − γ = (cid:80) T − 1 j =0 γ j ≤ T . This completes the proof of Theorem 3.2.
For i ∈ [ n ] and j ∈ [ K ] , we first define g i,j as g i,j ( θ ) :=( θ ′ k ) Kk = where θ ′ k = (cid:26) θ k − 2 η ( θ k − z i ) if k =
j θ k if k ̸ = j .
For ε := 1 8 Rn , We define Φ ε := { ( θ ′ j ) Kj =1 : θ ′ j = (cid:0) g i j,tj ,j ◦ · · · ◦ g i j, 1 ,j ( θ (0) ) (cid:1) j , t j ∈
[ T ] ∪ { 0 } , i j, 1 , . . .
, i j,t j ∈ [ n ] , ∀ j ∈ [ K ] } . Given θ (0) , since g i,j is γ -contractive by Lemma A.4, we have for any t ≥ 0 , i 1 , . . , i t ∈ [ n ] , and j ∈ [ K ] ,
[Equation]
i.e.
[Equation] Namely, Φ ε contains all possible SGD parameters starting from θ (0) where each element in Φ ε is an algorithm depending on at most KT samples. Since each f ( · ; z ) is (4 R ) -Lipschitz on Θ , we apply
Theorem B.1 with
[Equation]
where
[Equation]
and the assumption (B.1) in Theorem B.1 is satisfied by (C.4). This leads us to the bound in
Theorem
[Equation]
[Equation]
In this section, we bound the expected generalization gap using our localized cover for a contractive iterative stochastic optimizer g (see Definition 3), i.e. each g ( · ; z ) is contractive for all z ∈ Z . Here, we use S for denoting the set of samples { z 1 , . . .
, z n } and g i for denoting g ( · ; z i ) .
Theorem D.1.
Suppose that Assumptions 1 & 2 hold and there exists γ ∈ (0 , 1) such that g ( · ; z ) is γ -contractive on Θ ⊆ B dR (0) for all z ∈ Z .
Let T := max (cid:110)(cid:108) log(2 LRn ) log(1 /γ )
(cid:109) , 0 (cid:111) .
Then for any θ (0) ∈ Θ , t ≥ 0 , and i 1 , . .
, i t ∈ [ n ] , the t -th iterate θ ( t ) : =
g i t ◦ · · · ◦ g i 1 ( θ (0) ) satisfies
[Equation]
Compared to the uniform stability-based bound {+for γ -contractive 1 n (cid:80) ni =
 g i under L ′ -Lipschitz f ( · ; z ) and ∥ g ( θ ; z ) ∥ ≤+} K [KWS22]
[Equation] our bound in Theorem D.1 is linear in B , has an additional log n factor, and requires contractive g ( · ; z ) but does not depend on L † and K .
Here, note that our bound uniformly holds for all parameters generated by SGD while the bound (D.1) requires the expectation over θ (0) and the indices i 1 , . . , i t used for the t updates using the iterative stochastic optimizer. Although the bounds in Theorem D.1 and (D.1) provide some information about the expected generalization gap, these bounds cannot be used for understanding the absolute deviation | ˆ F ( θ ( t ) ) −
F ( θ ( t ) ) | , which can be of practical interest. In the following theorem, we provide a bound on the expectation of the absolute generalization gap where Ψ ≥ T ( θ ) and Ψ T ( θ ) are defined for general g i , analogous to Definition 1: [Equation] Theorem D.2.
Assume the setup in Theorem D.1. Then there exists some absolute constant C >such that [Equation] To our knowledge, the bound in Theorem D.2 is the first bound on E S [ | ˆ F ( θ ( t ) )
− F ( θ ( t ) ) | ] for the contractive case, including strongly convex and smooth functions. Compared to Theorem D.1, the bound in Theorem D.2 looses O ( (cid:112) T/ ( n log n )) factor for taking the absolute value inside the expectation. Nevertheless, this bound can be improved by removing the supremum inside the expectation.
Corollary D.1.
Assume the setup in Theorem 2.1. Then there exists some absolute constant C >such that for any θ (0) ∈ Θ , t ≥ 0 , and i 1 , . . , i t ∈ [ n ] , we have
[Equation]
(cid:114)n. Compared to Theorem D.2, the bound in Corollary D.1 does not have √ T log n factor but the LHS in the bound is weaker. D.1 Proof of Theorem D.  For any θ ( t ) ∈ (cid:83) θ (0) ∈
{+Θ Ψ ≥ 
+}
( θ (0) ) , let θ ( t ) =
g i t ◦ · · · ◦ g i 1 ( θ (0) ) .
Let ϕ
= g i t ◦ · · · ◦ g i t − T +1 (0) and I = { i t − T +1 , . . . , i t } ⊂
[ n ] if t > T and ϕ = g i t ◦ · · · ◦ g i 1 (0) and I = { i 1 , . . , i t } ⊂
[ n ] otherwise. Then as in the proof of Lemma 2.1, we have [Equation] for ε = 1 / (2 Ln ) , regardless of the choice of S .
Using this, we have
[Equation]
The first inequality is from the triangle inequality and the second inequality is from (D.3). The third inequality is again from the triangle inequality and the first equality is from the fact that ϕ is independent of { z i : i ∈ I} . The last inequality is from Assumption 2. Since the above bound holds for any θ (0)
∈ Θ and I t , this completes the proof of Theorem D.1. D.2 Proof of Theorem D.  In this proof, we assume T < n since the statement trivially follows otherwise.
As in the statement of Lemma 2.1, one can observe that
[Equation]
for ε = 1 / (2 Ln ) . For each ϕ ∈ Ψ T (0) , let I ϕ : = { i 1 , . .
, i T } such that ϕ
= g i T ◦ · · · ◦ g i 1 (0) .
Using this we can derive the following bound: for Π Ψ T (0) ( θ ( t ) ) : = arg min ϕ ∈ Ψ T (0) ∥ θ ( t ) − ϕ ∥ ,
E S (cid:34) sup
[Equation]
+ E S (cid:34) sup [Equation] (cid:35) ≤ 2 Lε + E S [Equation]
sup [Equation] ≤n + [Equation]
S [Equation]
n [Equation] i ∈I ϕ ( [Equation]
)) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) + E S   sup ϕ ∈ Ψ T (0) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) n (cid:88) i ∈ [ n ] \I ϕ ( f ( ϕ ; z i ) − F ( ϕ )) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ≤ BT +n [Equation]
sup
[Equation] T (0) (cid:12) (cid:12) (cid:12)(cid:12)(cid:12)(cid:12) | [ n ] \ I ϕ | n · | [ n ] \ I ϕ | (cid:88) i ∈ [ n ] \I ϕ ( f ( ϕ ; z i ) − F ( ϕ ))
(cid:12) (cid:12) (cid:12)(cid:12)(cid:12)(cid:12) ≤ BT +n + CB (cid:114) T log n n .
The first inequality is from the triangle inequality and the second inequality is from (D.4).
The third inequality is again from the triangle inequality and the definition of ε , and the fourth inequality is from Assumption 2.
Since each 1 | [ n ] \I ϕ |
[Equation]
i ∈ [ n ] \I ϕ ( f ( ϕ ; z i ) − F ( ϕ )) is sub-Gaussian with the sub-Gaussian norm bounded by cB/ (cid:112) | [ n ] \ I ϕ | for some absolute constant c (see Proposition 2.5.
in [Ver18]), the last inequality follows from a standard upper bound for empirical processes (see Exercise 2.5.10 in [Ver18]). This completes the proof of Theorem D.2.
D.3 Proof of Corollary D.
As in the proof of Theorem D.2, we
For any θ ( t ) ∈ (cid:83) θ (0) ∈
{+Θ Ψ ≥ 
+}
( θ (0) ) , let θ ( t ) =
g i t ◦ · · · ◦ g i 1 ( θ (0) ) . Let ϕ
=
g i t ◦ · · · ◦ g i t − T +1 (0) and I =
{ i t − T +1 , . . .
, i t } ⊂ [ n ] if t > T and ϕ = g i t ◦ · · · ◦ g i 1 (0) and I =
{ i 1 , . .
, i t } ⊂ [ n ] otherwise. Then for ε = 1 / (2 Ln ) , we have
[Equation]
Using this, we derive the following inequality:
[Equation]
The first and second inequality follows from the triangle inequality and (D.5), respectively.
The third inequality is again from the triangle inequality while Assumption 2 gives us the fourth inequality. The sixth inequality is from Proposition 2.5.2 and Exercise 2.5.10 in [Ver18] where the last inequality naturally follows.
Since the above bound holds for any θ (0)
∈ Θ and I t , this completes the proof of Corollary D.
E Comparison with existing bounds
[Table]
[Table]
In this section, we show that the uniform stability-based bound [HRS15] is Ω(1) for piecewise strongly convex and smooth functions after sufficiently many SGD iterations in general.
To this end, we first introduce the formal definition of the uniform stability and a standard tool for showing that the stability implies generalization, which are from Definition 2.1 and Theorem 2.2 in [HRS15]. Definition 7 (Uniform stability) . A randomized algorithm ρ is “ ε -uniformly stable” if for all data sets S , S
′ ∈
Z n such that S and S ′ differ in at most one example, we have
[Equation]
Theorem F.1. If ρ is ε -uniformly stable, then
[Equation]
For the remaining section, we provide an example that
[Equation]
regardless of n where ρ denotes sufficiently many SGD updates. Namely, the uniform stability-based bound based on Theorem F.1 is Ω(1) . Let Θ = [0 , 4] , Z = { 0 , 1 } , θ (0) ∼ Unif ([0 , 4]) , P ( z = 0) =
P ( z = 1) = 12 , f ( · ; 0) = min { ( x − 1) 2 , 12 + 12 ( x − 3) 2 } , f ( · ; 1) =
( x − 1) 2 , and the auxiliary gradient ∇ f (2; 0) = 2 at the nondifferentiable point θ = 2 , i.e., f ( · ; z ) is piecewise 1 -strongly convex and 2 -smooth with the partition P = { [0 , 2] , (2 , 4] } . For simplicity, choose η = 1 / 3 which can be generalized to arbitrary η ∈ (0 , 1) . Under this setup, using Theorem 2.2, one can easily derive a generalization bound that does not increase with the number of SGD iterations and converges to zero as n grows.
However, under the same setup and sufficiently many SGD iterations, the uniform stability-based bound is lower bounded by a constant regardless of n .
To see this, let S
= . , 0) and S
′ =
Then, one can observe that f ([0 , 2]; 0) ⊂ [0 , 2] and f ((2 , 4]; 0) ⊂ (2 , 4] , i.e. SGD iterates for S converge to either 1 or 3 depending on whether θ (0) ∈ [0 , 2] or θ (0) ∈ (2 , 4] .
Since we assumed θ (0) ∼ Unif ([0 , 4]) , we have
[Equation]
where θ ( t ) denotes a random parameter generated by t SGD updates for S , from θ (0) ∼ Unif ([0 , 4]) .
Furthermore, we have f (Θ; 1) ⊂ [0 , 2] and f ([0 , 2]; 0) ⊂ [0 , 2] .
This implies that if a single SGD update for S ′ use the first sample in S ′
(i.e. z = 1 ), which occurs with high probability under sufficiently many SGD iterations, then the SGD iterates will converge to 1 almost surely.
In other words, we have
[Equation]
where ϕ ( t ) denotes a random parameter generated by t SGD updates for S ′ , from ϕ (0) ∼ Unif ([0 , 4]) .
Combining (F.1) and (F.2) implies a constant lower bound on the uniform stability-based bound, regardless of n , under sufficiently many SGD iterations.
We note that the same conclusion can also be derived for any η ∈ (0 , 1) as