Stability Bounds for Stochastic Gradient Descent

Why Sampling Matters for Stability

A learning algorithm is stable when changing one training observation has only a small effect on its predictions or losses. For stochastic gradient descent, that effect depends not only on the step sizes and the number of iterations, but also on whether the algorithm happens to sample the changed observation. The key calculation tracks how the distance between two coupled runs evolves: updates based on the same observation do not increase that distance, while an update using the changed observation can increase it.

For convex losses that are Lipschitz but need not be differentiable, we will use a proximal stochastic gradient update. This is an implicit SGD update: it minimizes the sampled loss plus a quadratic penalty around the current iterate. The proximal formulation gives a reliable stability argument without requiring a particular choice of subgradient.

Setup and Definition

Let $Z_1,\ldots,Z_n$ be a training sample from a distribution $P$, and let $C\subseteq\mathbb{R}^d$ be nonempty, closed, and convex. Assume each loss $w\mapsto \ell(w,z)$ is finite, convex, and $L$-Lipschitz on $\mathbb{R}^d$, uniformly over $z$. Thus $\lvert \ell(w,z)-\ell(v,z)\rvert\leq L\|w-v\|$. The closedness of $C$ ensures the constrained proximal minimization below has a solution; convexity and the quadratic term make it unique.

Definition: Two samples $S=(z_1,\ldots,z_n)$ and $S'=(z'_1,\ldots,z'_n)$ are replace-one neighbors if they differ in exactly one coordinate. For a randomized learning algorithm $A$, its expected replace-one stability is $$ \beta_{\mathrm{rep}}(A)= \sup_{\substack{S,S'\text{ replace-one neighbors}\\z}} \mathbb{E}\left[\left|\ell(A(S),z)-\ell(A(S'),z)\right|\right], $$ where the expectation is over the algorithm's randomness, using a specified coupling of its random choices on $S$ and $S'$. We use the same random choices in both runs.

The supremum over test observations $z$ makes this a uniform replace-one bound. The expectation is essential: the two runs need not have small loss difference for every possible random sequence, provided their expected difference is controlled.

Proximal Updates: Two Useful Bounds

For a convex loss $f$ and step size $\eta>0$, define its constrained proximal update at $x\in C$ by $\operatorname{prox}_{\eta f,C}(x)=\arg\min_{u\in C}\{f(u)+\|u-x\|^2/(2\eta)\}$. The objective is continuous and coercive on the nonempty closed set $C$, so it attains a minimum; strict convexity of the quadratic term makes that minimum unique.

Lemma: Suppose $f$ is convex and $L$-Lipschitz on $\mathbb{R}^d$, and $C$ is nonempty, closed, and convex. Write $P_f(x)=\operatorname{prox}_{\eta f,C}(x)$. Then, for $x,y\in C$, $$ \|P_f(x)-P_f(y)\|\leq\|x-y\|, \qquad \|P_f(x)-x\|\leq\eta L. $$

Proof. Let $u=P_f(x)$. The constrained optimality condition is $x-u=\eta(g+n)$, where $g\in\partial f(u)$ and $n\in N_C(u)$; here $N_C(u)=\{q:\langle q,v-u\rangle\leq0\text{ for all }v\in C\}$. Because $f$ is $L$-Lipschitz, every subgradient $g$ has norm at most $L$. Taking the inner product of the optimality condition with $x-u$, and using $x\in C$, gives $$ \|x-u\|^2 =\eta\langle g,x-u\rangle+\eta\langle n,x-u\rangle \leq\eta L\|x-u\|. $$ If $x=u$, the displacement bound holds immediately; otherwise divide by $\|x-u\|$ to obtain $\|x-u\|\leq\eta L$.

For nonexpansiveness, let $u=P_f(x)$, $v=P_f(y)$, and write $x-u=\eta(g+n)$, $y-v=\eta(h+m)$, with $g\in\partial f(u)$, $h\in\partial f(v)$, $n\in N_C(u)$, and $m\in N_C(v)$. Monotonicity of the subdifferential gives $\langle g-h,u-v\rangle\geq0$; the definition of the normal cone gives $\langle n-m,u-v\rangle\geq0$. Hence $$ \langle x-y,u-v\rangle =\|u-v\|^2+\eta\langle g-h+n-m,u-v\rangle \geq\|u-v\|^2. $$ Cauchy–Schwarz now implies $\|u-v\|\leq\|x-y\|$ when $u\ne v$, and the result also holds if $u=v$. $\square$

Expected Stability Under Sampling With Replacement

Initialize at a fixed $w_1\in C$. At iteration $t$, draw $I_t$ uniformly from $\{1,\ldots,n\}$, independently across iterations, and update $$ w_{t+1}=\arg\min_{u\in C}\left\{\ell(u,Z_{I_t})+\frac{1}{2\eta_t}\|u-w_t\|^2\right\}, $$ where the step sizes $\eta_t>0$ are deterministic. For neighboring samples $S,S'$ differing at coordinate $k$, couple the runs by using the same indices $I_1,\ldots,I_T$. Let $w_t,w'_t$ be their iterates and $D_t=\|w_t-w'_t\|$.

Theorem: Under the conditions above, after $T$ iterations the expected replace-one stability satisfies $$ \beta_{\mathrm{rep}}(A_T)\leq \frac{2L^2}{n}\sum_{t=1}^{T}\eta_t. $$ For constant step size $\eta_t=\eta$, this becomes $\beta_{\mathrm{rep}}(A_T)\leq 2L^2\eta T/n$.

Proof. If $I_t\ne k$, both runs use the same loss at iteration $t$. Nonexpansiveness from the lemma yields $D_{t+1}\leq D_t$. If $I_t=k$, the two sampled losses can differ, and the consequence of the lemma gives $D_{t+1}\leq D_t+2\eta_tL$. Combining the cases, $$ D_{t+1}\leq D_t+2\eta_tL\,\mathbf{1}\{I_t=k\}. $$ Since $D_1=0$ and $P(I_t=k)=1/n$, taking expectations and summing over $t$ gives $$ \mathbb{E}[D_{T+1}] \leq \frac{2L}{n}\sum_{t=1}^{T}\eta_t. $$ For every test observation $z$, its loss is $L$-Lipschitz in the parameter, so $\lvert\ell(w_{T+1},z)-\ell(w'_{T+1},z)\rvert\leq L D_{T+1}$. Take expectations and then the supremum over neighboring samples and $z$ to obtain the asserted bound. $\square$

Worked Examples and Sampling Variants

Worked Example: Constant Step Size

Suppose $n=100$, $L=2$, $T=50$, and $\eta_t=0.01$ at every iteration. Then $\sum_{t=1}^{T}\eta_t=50(0.01)=0.5$. The theorem gives $$ \beta_{\mathrm{rep}}(A_T) \leq \frac{2(2)^2}{100}(0.5) =\frac{8}{100}(0.5) =0.04. $$ The bound is proportional to the total step-size budget and inversely proportional to the training-set size. It is a bound, not a claim that the actual stability equals $0.04$.

Worked Example: A One-Step Bound Can Be Exact

Take $C=\mathbb{R}$, $n=5$, one iteration with $\eta=0.2$, and training losses $\ell(w,z)=aw$ with $a\in\{-1,1\}$. These losses are convex and $1$-Lipschitz. Let neighboring samples differ at one coordinate, with slope $1$ in one sample and slope $-1$ in the other. If that coordinate is drawn, the proximal updates from $w_1=0$ are $w_2=-0.2$ and $w'_2=0.2$, because minimizing $aw+(u-w_1)^2/(2\eta)$ gives $u=w_1-\eta a$. Their distance is $0.4$. If any other coordinate is drawn, the updates agree. The changed coordinate is selected with probability $1/5$, so for a test loss with slope $1$, the expected absolute loss difference is $$ \frac{1}{5}(1)(0.4)+\frac{4}{5}(0)=0.08. $$ The theorem gives $2(1)^2(0.2)/5=0.08$, matching this example exactly.

Worked Example: Random Reshuffling

Now sample the $n$ observations in a uniformly random permutation, without replacement, and use one observation per iteration for $T\leq n$ iterations. Couple the two runs with the same permutation. At each position $t$, the changed coordinate $k$ appears with probability $1/n$. The same distance inequality applies: $D_{t+1}\leq D_t+2\eta_tL\mathbf{1}\{I_t=k\}$. Taking expectations does not require the event at position $t$ to be independent of earlier positions; linearity of expectation suffices. Therefore $$ \mathbb{E}[D_{T+1}] \leq 2L\sum_{t=1}^{T}\eta_t P(I_t=k) =\frac{2L}{n}\sum_{t=1}^{T}\eta_t, $$ and the expected replace-one stability is at most $2L^2\sum_{t=1}^{T}\eta_t/n$, just as in the with-replacement bound.

More generally, suppose sampling at iteration $t$ selects coordinate $k$ with probability $p_{t,k}$, with the sampling rule the same for both datasets and independent of the data values. The same argument yields the coordinate-specific bound $2L^2\sum_{t=1}^{T}\eta_t p_{t,k}$. Taking the largest such bound over $k$ gives a uniform replace-one guarantee. Uniform sampling has $p_{t,k}=1/n$; a sampling design that repeatedly favors one coordinate can produce a larger stability bound for replacing that coordinate.

Step Sizes, Generalization, and a Common Pitfall

The bound depends on $\sum_t\eta_t$, not on the iteration count alone. With constant step size, stability grows at most linearly in $T$. If $\eta_t=\eta_0/t$, then $\sum_{t=1}^{T}\eta_t\leq\eta_0(1+\log T)$, so the bound grows at most logarithmically. If $\eta_t=\eta_0/\sqrt{T}$ for all $T$ iterations, the sum is $\eta_0\sqrt{T}$, giving an upper bound $2L^2\eta_0\sqrt{T}/n$. These statements describe stability only; they do not by themselves show that a particular schedule optimizes prediction or estimation error.

Stability is useful because it controls sensitivity of the learned predictor to individual training observations. For an iid sample and the usual empirical risk $R_S(w)=n^{-1}\sum_{i=1}^n\ell(w,Z_i)$, a uniform expected replace-one bound also controls the absolute expected generalization gap by the same bound. Thus, in this setting, the displayed quantity gives an expected generalization guarantee as well as a sensitivity guarantee.

Pitfall: assuming any subgradient update is nonexpansive. Convexity alone does not make the explicit map $w\mapsto w-\eta g(w)$ nonexpansive when a loss is nonsmooth. Different subgradient choices at nearby points can make the distance grow, so the shared-sample step in the proof would not follow.
Correction. The theorem proved here is for the proximal stochastic update, whose nonexpansiveness follows from convex optimality conditions. For an explicit subgradient algorithm, a stability proof needs additional assumptions or a different argument; the proximal result cannot simply be transferred to it.
Pitfall: forgetting the sampling probability. A differing observation affects the iterates only when it is selected, but its effect can persist through later updates. The bound accounts for both facts: the per-selection increment is at most $2\eta_tL$, and the selection probability under uniform sampling is $1/n$.

Key takeaway: For convex $L$-Lipschitz losses, coupled proximal-SGD runs contract on shared observations and can separate by at most $2\eta_tL$ when the replaced observation is sampled. Under uniform sampling, this yields expected replace-one stability at most $2L^2\sum_t\eta_t/n$.

Check Your Understanding

Use the coupling argument and the step-size sum to answer these questions.

For $n=40$, $L=3$, and step sizes $\eta_t=0.02$ for 10 iterations, what stability bound does the theorem give?
Why does the update using a shared training observation not increase the distance between the two iterates?
For nonuniform sampling, what is the coordinate-specific bound when coordinate $k$ is selected with probability $p_{t,k}$ at iteration $t$?
Why is closedness of $C$ included in the assumptions for the proximal update?
What extra issue prevents the proximal proof from automatically establishing stability for explicit subgradient SGD with nonsmooth losses?