Generalization Bounds for Stochastic Gradient Descent via Localized ε -Covers
Anonymous Author(s)
AffiliationAddress email
Abstract
In this paper, we propose a new covering technique localized for the trajectoriesof SGD.
This localization provides an algorithm-specific complexity measured bythe covering number, which can have dimension-independent cardinality in contrast to standard uniform covering arguments that result in exponential dimensiondependency.
Based on this localized construction, we show that if the objectivefunction is a finite perturbation of a piecewise strongly convex and smooth functionwith P pieces, i e., non-convex and non-smooth in general, the generalizationerror can be upper bounded by O ( (cid:112) (log n log( nP )) /n ) , where n is the numberof data samples.
In particular, this rate is independent of dimension and does notrequire early stopping and decaying step size.
Finally, we employ these results invarious contexts and derive generalization bounds for multi-index linear models,multi-class support vector machines, and K -means clustering for both hard andsoft label setups, improving the previously 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 apossibly non-convex loss incurred by a single data point z ∈ Z , and Z is a random variable on Zfollowing the data distribution.
Since the distribution of Z is unknown in general, the following proxybased 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 problemin learning theory.
Classical algorithm-independent results rely on uniform convergence over theentire domain Θ ⊂ R d ; thus, they apply to any learning algorithm.
However, these bounds oftenincrease with the dimension d [SSSSS09, SSBD14, Fel16], becoming vacuous in the modern overparameterized regime [ZBH + 21 ].
To derive dimension-independent bounds, researchers have beeninvestigating algorithm-dependent generalization properties, especially for commonly used trainingmethods such as stochastic gradient descent (SGD)
[HRS15, SHN + 18, YKM21].
Notably, algorithmic stability is a technique for deriving generalization bounds based on the propertiesof a specific learning algorithm, which leverages that if a parameter learned by an algorithm isrobust under a perturbation of samples z 1 , . .
, z n , then the generalization error at that parametermust be small [BE02].
Based on this principle, several works proposed dimension-independentgeneralization 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 forstrongly convex and smooth functions [SSSSS09] and convex and non-smooth functions [BFGT20,
AKL21].
Nevertheless, without global (strong) convexity, early stopping, and/or decaying stepsize, generalization bounds based on algorithmic stability often diverge with the number of SGDiterations [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 iterateswhose complexity can be measured by a notion called the Hausdorff dimension [Fal14].
In thiscontext, [¸SSDE20] showed that, under a continuous-time surrogate for SGD, the generalization errorcan be bounded by (cid:101)
O ( (cid:112) d H /n ) , where d H denotes the Hausdorff dimension of the optimizationtrajectory.
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], ultimatelyproviding 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 bequantified under standard verifiable conditions.
Our contributions are as follows.
• Localized ε -covers for SGD.
Our first principle contribution is a covering technique localized forthe trajectories of SGD.
This localization provides an algorithm-specific complexity measured bythe 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 dimension-
independent 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 ) − ∇ h
( θ ; 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 Pdenotes the number of strongly convex pieces needed to approximate f .
We further show thatthe gradient of any (piecewise) smooth function f can be approximated with that of a piecewisestrongly convex and smooth function, demonstrating the wide applicability of the bound (1.3).Finally in the special case P = 1 and ξ = 0 , our result reduces to a non-asymptotic bound where thecomplexity is captured by the Hausdorff dimension of the invariant measure of the SGD algorithm.• Improved bounds in specific models.
We employ the above result to derive generalization boundsin several statistical models trained by SGD, including multi-index linear models, multi-classsupport vector machines, and K -means clustering with both hard and soft label setups, improvingthe previously known state-of-the- art generalization error bounds in this context.
Notation and problem setup.
For k ∈ N and a set S , we denote [ k ] := { 1 , . . .
, k } and (cid:0) S k (cid:1) :
={T ⊆ S : |T
| = k } , where |T | denotes the cardinality of T .
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 ε ⊆ S 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 ] , thet -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 in the presence of an ℓ 2 -regularizer and Lipschitz continuity; see Section 4 for more details.
Throughout the paper, weuse θ ( 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] and ∥∇
[Equation]
θ
[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 smoothcase in which the localization argument can be simplified by the contractivity of SGD.
We presentour 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 basedon 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 highprobability by simply applying the union bound over C ε , which yields a generalization error bounddepending on the cardinality of the cover (cid:112) log |C ε | .
However, C ε is typically independent of thealgorithm being used, and standard ε -covers for Θ yield |C ε | = ( R/ε ) Ω( d ) ; thus, bounds based oncovering numbers often grow with √ d , which can be loose if d is large.
To overcome this issue, weuse 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 thepoints that can be reached by SGD trajectories, resulting in a covering number that is independent ofthe 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]
Θ ,
[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 distributioninstead of a single point [DDB20], but their contractivity can still provide the following localization.
Lemma 2.1.
Suppose that g 1 , g 2 , . .
are γ -contractive for some γ ∈ (0 , 1) .
Then, for any initialization θ (0) ∈ Θ ⊆ B dR (0) , for some ε ∈ (0 , 1) , and T := T ε = (cid:108) log( R/ε ) log(1 /γ ) (cid:109) , we have
[Equation]
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.
We have
[Equation]
since each g i is γ -contractive, which in turn implies θ ( t ) ∈ (cid:83) ϕ ∈ Ψ T (0)
B dε ( ϕ ) .
For strongly convex and smooth f ( · ; z ) , an SGD update g i with a sufficiently small step size iscontractive [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 thatapplying T synchronously coupled SGD updates that use the same sample at each iteration can makethe distance between the initial points smaller than ε ; see Figure 1 (left).
This observation impliesthat the set Ψ T
0) of all parameters that can be generated by T SGD updates when initialized at the ε the set Ψ ≥ T ( θ 0) ) of all parameters that can be generated by any t ≥ T SGD for an arbitrary initialization θ (0) see Figure 1 (right). In contrast to algorithm independent covers of
Θ that scale with |C ε | = ( R/ε ) Ω( d ) = e (cid:101) Ω( d ) , we obtain | Φ T (0) | ≤ n T =
e (cid:101) O
(1) which is independent of the dimension d and only polynomial in the number of samples n .
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 , we have
[Equation]
Notice that h = 0 reduces to the classical Lipshitz continuity, but the above condition is more general.
Assumption 2 (Bounded deviation) .
For B > 0 and for all θ ∈ Θ , we have
[Equation]
We note that Assumptions 1 & 2 (or their variants) both appear in several algorithmic stability-basedresults; see e.g. [HRS15, Thm 3.10].
We also highlight that these conditions are invariant to addingregularizers (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 αη + αβη 2 , T := (cid:108) log(2 LRn ) log(1 /γ ) (cid:109)
if γ ̸ = 0 , and T := 1 otherwise.
Then, with probabilityat 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 boundvia uniform convergence over Θ is worse than the algorithmic stability-based bounds for SGD by afactor 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 stillnot directly comparable to most stablity-based bounds which are given in expectation.
An exceptionis [FV19, Thm 4.5] which obtains near-optimal bounds for SGD in the convex case by tuning thenumber of iterations to be taken, obtaining a rate of O (log( n/δ ) 2 log( n ) / √ n ) .
In contrast, our boundreads O ( (cid:112) log( n ) 2 log(1 /δ ) /n ) which is better by a logarithmic factor and holds for any sufficientlylarge number of SGD iterations.
Lastly, we emphasize that the bound (2.1) holds for the union of alltrajectories of SGD generated by t ≥ T iterations, whereas the stability-based results often consider asingle parameter generated by SGD.
A similar setup is considered in Corollary 2.1 where we establishimproved bounds on the generalization error removing the logarithmic factor.
The proof of Theorem 2.1 follows from three steps.
i )
Lemma 2.1 implies that any θ ( t ) for t ≥ T canbe 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 onat most T
=
O (log( n )) of n samples.
iii )
The empirical process | ˆ F ( θ ( T ) ) − F ( θ ( T ) ) | uniformlyconverges 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 tightenedby considering a single realization of θ (0) in the following corollary.
Corollary 2.1.
Assume the setup in Theorem 2.1.
Then, given θ (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 ata fixed θ ∈ Θ ; however, it has the limitation that at least t ≥ T SGD updates are required.
This isremedied in the next result which also does not require 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 proofsof 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.1and the Hausdorff dimension of the support µ of the stationary distribution of (projected) SGD.
Forexample, suppose that η ∈ (0 , 1) , Θ = B dR , f ( θ ; z ) =
∥ θ − θ ∗ z ∥ 2 for some θ ∗ z ∈ B dR (i.e. α = β = 1and γ = 1 − η ), and ∥ θ ∗ z i − θ ∗ z j ∥ ≥ 2 γR for all i ̸
= j .
Here, the last assumption can be satisfiedwith high probability if d, η are large enough and θ ∗ z ∼ N (0 , σ 2 I ) , i.e. zero-mean Gaussian withcovariance σ 2 I ; see e.g. [PLYS21, App 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 derivationof (2.2) are presented in Appendix B.5.
In (2.2), d H replaces the ambient dimension d that appearin algorithm-independent bounds [SSBD14].
We note that such a connection between the fractaldimension and generalization bounds has been also studied in [¸SSDE20, CDE + 21]; however, theirbounds are asymptotic and require non-trivial assumptions that are not easy to verify in practice.
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 SGDiterations [HRS15, LLQ20].
The localized covering construction introduced in the previous sectionremedies this issue, providing more stable generalization bounds.
Specifically, we establish adimension-independent generalization bound for functions that are finite perturbations of piecewisestrongly convex and smooth functions.
We further prove an approximation result and show that anysmooth non-convex function can be approximated with a piecewise strongly convex function.
Definition 2.
f is “piecewise α -strongly convex and β -smooth with P pieces on Θ ” if there exists apartition {P 1 , . . .
, P P } of Θ and α -strongly convex and β -smooth functions h 1 , . . .
, h P on Θ suchthat f = h p on P p for all p ∈ [ P ] .
We also define ∇ f ( θ ) :
= ∇ h p ( θ ) if θ ∈ P p .
Note that since piecewise strongly convex and smooth functions may not be differentiable on theentire Θ , we will use the auxiliary gradient in Definition 2 as a surrogate in the SGD update (1.4).
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 ispiecewise strongly convex and smooth.
Furthermore, the objective of learning an entire ReLU networkis also piecewise strongly convex and smooth on small loss regions [Mil19].
Such observationseasily extend to more general settings, e.g. an objective function defined as the summation of apiecewise linear loss (e.g. hinge loss) and ℓ 2 -regularization.
However, if we further allow for finiteperturbations, 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 for any θ ∈ Θ , we have [Equation]
For this general class of non-convex functions, we have 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 /β ) and T ∈ N , let γ :=
(cid:112) 1 − 2 αη + αβη 2 .
Then, with probabilityat 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 andany choice of T ≥ 1 .
However, to obtain a meaningful generalization bound, one may choose, forexample T
=
O (log( nR ) / log( γ − 1 )) .
In the case that γ, B, L = Θ(1) , the bound simplifies to
[Equation]
Here, we observe that the first term in the bound is logarithmic in the number of pieces P and the lastterm scales linearly with the approximation error ξ .
Thus, the generalization bound depends on thetrade-off between the complexity of h through P , and how well ∇ h approximates ∇ f through ξ .
Inthis regime, in contrast to algorithmic stability-based bounds, the above result does not grow with thenumber of SGD iterations; thus, early stopping is not required for generalization in Theorem 2.2.
In light of Proposition 2.1, any piecewise smooth function can be approximated by a piecewisestrongly 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 coveringbound (cid:101) O
( B (cid:112) d/n ) by choosing ξ = O (1 / ( L √ n )) and T = Θ(log( LRn )) .
However, any valueof P that is sub-exponential in dimension yields improved generalization bounds.
In particular inthe next section, we consider certain (non-convex) statistical models and carefully design h so that Theorem 2.2 improves the existing generalization error bounds.
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 concentrationinequality 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 the our localized covering to prove generalization bounds for multi-indexlinear models, multi-class support vector machines, and K -means clustering with both hard and softlabel setups trained by SGD, improving the state-of-the-art results known for these models.
3.1 Multi-index linear models
Given a sample ( y, x ) ∈ Y × B dR x (0) , consider the ℓ 2 -regularized loss in a multi-index model with Kindices parameterized by θ = ( θ j
[Equation]
Kj =1 ∈ R
d
× K
[Equation]
Characterizing the generalization properties of multi-index models is an important problem with manyapplications including regression, classification, dimension reduction, and learning a single-layerof a neural network.
In the following theorem, we derive a generalization bound for multi-indexmodels by approximating each f ( · ; z ) with a piecewise strongly convex and smooth function, andthen 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 thatf satisfies (3.1) , Assumptions 1 & 2 hold, and ℓ ( · , . . .
, · ; y ) is piecewise β -smooth with Q pieceson Θ for all y ∈ Y .
For any η ∈ (0 , 2 /λ ) , let γ :=
| 1 − ηλ | , T := (cid:108) log(3 LRn ) log(1 /γ ) (cid:109) if γ ̸ = 0 , T :=
1otherwise, 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 notableresults include [Gue02, Zha04, PXXZ08, JKZ + 12, CMR13, LDZK19].
Under the same conditionsof Theorem 3.1, existing state-of-the-art bounds scale at least linearly in √ d or K , while our result isdimension-free and scales with √ K .
Specifically, the result in [LDZK19] translated to our settingreads ( ℓ 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 thebound in (3.2) is linear in KLRR x , our bound is linear in √ K and logarithmic in L, R, R x .
Thesignificance of this improvement can be better seen, for example, when learning the first layer ofneural networks, for which L can be very large.
Moreover, if x ∼ N (0 , I ) , then ∥ x ∥ = (cid:101) Ω( √ d ) withhigh probability; that is, (3.2) is linear in R x = Ω( √ d ) but ours only scales with √ log d , whichcan make a significant difference especially in the overparameterized regime.
We should note thatcompared 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 boundswhich 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 similarly provide generalization error bounds for the problem of K -means clustering in bothhard and soft label setups.
As before, we consider SGD as an optimizer which has been the focusof many works in this context [BB94, Scu10, TM17].
We assume throughout this section that thesamples 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 abound 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 eachpoint using one-hot encoding, the soft K -means clustering allows soft labels within the probabilitysimplex.
Specifically, given ζ > 0 and the samples z 1 , . . .
, z n , the soft K -means clustering algorithmalternates between updating the soft labels ( w j ( z i ; θ ))
Kj =1 of z
i for all i ∈ [ n ] , and estimating thecluster 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, whichconverges 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 SGDand 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 2 ,
L := 4 R √ K e ζB , T := (cid:108) log(3 LRn ) log(1 /γ ) (cid:109) , and P := poly ( η, ζ, B, K, R, T, n, e ζB ) .
Then, with probabilityat 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 theobjective 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 SGDiterations 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 := (cid:108) log(16 √ KRn ) log(1 /γ )
(cid:109) if γ ̸ = 0 , and
T := 1 otherwise.
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 remark that under η =Θ(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 weutilized is essentially based on the (piecewise) contractivity of SGD updates for (piecewise) stronglyconvex and smooth functions.
However, local strong convexity is by no means necessary and thisanalysis 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öldercontinuous gradient [Nes15].
Then, for any ε > 0 , there exists a step size η > 0 and T = T ε suchthat applying T synchronously coupled SGD updates on any two points makes the distance betweenthem smaller than ε .
Consequently, an analog of Lemma 2.1 can be established under this setup aswell, for a properly chosen η, ε, T .
We highlight that this class of functions is already covered in ourframework via Proposition 2.1; however, an analysis based on the actual curvature (as opposed to apiecewise approximation) of the objective may provide tighter generalization bounds.
Extension to SGD without projection.
Projection operation is only needed to ensure that the SGDiterates 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].
For example,if f ( · ; z ) is Lipschitz continuous, it is a straightforward exercise to show that SGD iterations arebounded in the presence of ℓ 2 -regularization and bounded initialization.
Extension to mini-batch setup and different sampling schemes.
The results we presented canbe 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 nare contractive, then an average of their subset is also contractive.
Moreover, our generalizationbounds 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 possibletrajectories of SGD , which is independent of the underlying sampling scheme.
Extension to contractive stochastic optimization methods.
In this paper, our main focus wasthe SGD algorithm and its generalization properties.
However, it is straightforward to adapt ourframework to a general stochastic optimization setup.
For example, the next result follows from theidentical steps leading to Theorem 2.2; hence, its proof is omitted.
Assumption 3.
Given samples z 1 , . . . , z n and the initial parameter θ (0) ∈ Θ , an iterative stochasticoptimization method performs the update θ ( t ) = g
i t ( θ ( t − 1) ) at iteration t , for a random samplei t ∈ [ n ] and some map g i t : Θ → Θ depending on z i t and f ( · ; z i t ) .
We assume that for any i ∈ [ n ] ,there exists a partition {P i, 1 , . . . , P i,P } of Θ such that g
i = g i,p on P i,p for all p ∈ [ P ] where eachmap g i,p is γ -contractive on the domain Θ for some γ ∈ (0 , 1) .
Theorem 4.1.
Suppose that Assumptions 1 & 2 hold and Θ ⊂ B dR (0) .
Then for any iterativestochastic optimization algorithm satisfying Assumption 3, for T := (cid:108) log(2 LRn ) log(1 /γ )
(cid:109) , with probability atleast 1 − δ , for any θ (0) ∈ Θ , t ≥ T , and i 1 , . . .
, i t ∈ [ n ] , we have
[Equation]
We note that a concurrent work by [KWS22] also investigated generalization bounds for contractiveoptimizers based on algorithmic stability. However, their results require global contractivity and theirbounds are in expectation, whereas Theorem 4.1 is a high-probability statement which only requirespiecewise contractivity as stated in Assumption 3.
Limitation of our results.
We outline a few limitations of our current analysis.
In this paper, we onlyconsidered the generalization error (also referred to as the generalization gap), and our main resultin the non-convex regime cannot be easily translated to a bound on the excess risk.
This additionalstep would require a bound on the optimization error; nevertheless, establishing such bounds in thenon-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 methodcannot be easily extended to varying step size.
This is indeed a limitation of our analysis as varyingstep size schedules are oftentimes used in practice.
We leave this extension as a future work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[No] This paper
only contains theoretical results on the generalization properties of SGD.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[N/A]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[N/A]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[N/A]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[N/A]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]