Self-avoiding walks in a rectangle

Abstract: A celebrated problem in numerical analysis is to consider Brownian motion originating at the centre of a $10 \times 1$ rectangle, and to evaluate the ratio of probabilities of a Brownian path hitting the short ends of the rectangle before hitting one of the long sides. For Brownian motion this probability can be calculated exactly \cite{BLWW04}. Here we consider instead the more difficult problem of a self-avoiding walk in the scaling limit, and pose the same question. Assuming that the scaling limit of SAW is conformally invariant, we evaluate, asymptotically, the same ratio of probabilities. For the SAW case we find the probability ratio is approximately 200 times greater than for Brownian motion.

\section{Introduction}
In the January 2002 issue of {\it SIAM News} L N Trefethen \cite{T02} presented ten problems used in teaching numerical analysis at Oxford University. The answer to each problem was a real number, and readers were set the challenge of computing 10 significant digits of the answer to each problem, with the additional incentive of \$100 for the best entrant. In the event, some 94 teams from 25 nations submitted solutions, and twenty of these teams achieved a perfect score. A donor, W J Browning, kindly provided additional funds so that all 20  teams achieving a perfect score could be awarded the \$100 prize!

As is so often the case, the mathematics revealed by the different approaches to the solutions was at least as interesting as the original problems, and indeed spawned a book {\it The Siam 100--Digit Challenge: A study in high-accuracy Numerical Computing,} \cite{BLWW04}, which is usefully and engagingly described in J Borwein's review \cite{B05}. Indeed, the authors of the book responded by providing not 10, but 10,000-digit solutions to all the problems but one, for which ``only'' several hundred digits are available. Borwein's personal favourite, and the catalyst for the calculation reported here, is problem number 10: {\em A particle at the centre of a $10 \times 1$ rectangle undergoes Brownian motion (i.e., 2-D random walk with infinitesimal step lengths) till it hits the boundary. What is the probability that it hits at one of the ends rather than at one of the sides?}

Remarkably, it turns out \cite{BLWW04} that a closed form solution to this question can be found, {\em via} the theory of elliptic functions and invoking some results of Ramanujan on singular moduli. The result is $$p_e=\frac{2}{\pi}\arcsin\left ( (3 - 2\sqrt{2})^2(2+\sqrt{5})^2(\sqrt{10}-3)^2(5^{1/4}-\sqrt{2})^4 \right ) = 0.000000383758797925 \ldots .$$ In discussion with Nathan Clisby, the question as to the corresponding problem where Brownian motion is replaced by the scaling limit of a self-avoiding walk (SAW) arose. The rest of this paper addresses our answer to this question. In fact, we prefer to address the equivalent question, which is the ratio of the probability that the walk hits the end to the probability that it hits a side. This is just $$\frac{p_e}{1-p_e}$$ which has the numerical value $0.00000038375894519599411176841999126970034234598936\ldots$ for the random walk case.

\section{Self-avoiding walks}
A self-avoiding walk (SAW)  of length $n$
on a periodic graph or lattice                      ${\mathcal L}$
is a sequence of distinct
vertices $w_0,w_1,
\ldots ,w_n$ in ${\mathcal L}$
such that each vertex is a nearest
neighbour of its predecessor. In  Figure~\ref{fig:saw} a short SAW on the square lattice is shown

\begin{figure}[ht!]
\begin{center}
\includegraphics[width=4cm]{saw.pdf}
\caption{A self-avoiding walk on the square lattice.}\label{fig:saw}
\end{center}
\end{figure}
\noindent
while in Figure~\ref{fig:23D} a rather long walk of $2^{25}$ steps is shown (generated by a Monte Carlo algorithm \cite{C10}).

\begin{figure}[ht!]
\begin{center}

\includegraphics[width=5cm]{saw2d.pdf}

\caption{A typical two-dimensional SAW of $2^{25}$ steps on the square lattice --
 courtesy of Nathan Clisby.}\label{fig:23D}

\end{center}
\end{figure}

One of the most important questions one asks about SAWs is: how many SAWs are there of length $n,$ (typically defined up to translations) denoted $c_n$?
Frequently one rather considers the associated generating function $$C(x) = \sum_{n \ge 0} c_n x^n.$$ For SAWs on a lattice it is straightforward to show that $c_n$ grows exponentially with $n$ \cite{madras-slade}, so that $c_n \propto \mu^n,$ where $\mu = 1/x_c$ depends on the lattice, and determines the radius of convergence of the generating function $C(x)$. Another important question, and central to our investigation, is the existence and nature of the {\em scaling limit} of SAWs, described in the next section.

\section{The scaling limit}
An intuitive grasp of this concept can  be gained by looking at the  two figures above. In the first, the effect of the lattice is clear. In the second, there is no obvious lattice, and indeed no way to tell that this is not a curve in the continuum. We formalise this notion as follows: Consider a smooth (enough) closed domain $\Omega,$ on an underlying square grid, with grid spacing $\delta$ as shown in Figure \ref{fig:scal-lim}. Denote by $\Omega_\delta$ that portion of the grid contained in $\Omega.$ Take two distinct points on the boundary of $\Omega$ labelled $a$ and $b.$
Now take the nearest lattice vertex to $a,$ and label it $a_\delta,$ and similarly let $b_\delta$ be the label of the nearest lattice vertex to point $b.$ Consider the set of SAWs  $\omega(\Omega_\delta)$ on the finite domain $\Omega_\delta$ from $a_\delta$ to $b_\delta.$ Recall that $\delta > 0$ sets the scale of the grid. Now let $|\omega|$ be the length of a walk  $\omega_\delta \in \omega(\Omega_\delta),$ and weight the walk by $x^{|\omega|}.$ The reason for this is that the walks are of different lengths, making the uniform measure not particularly natural. (There is also a normalising factor, which for simplicity we ignore).

\begin{figure}[h]
\label{fig:scal-lim}
\centering
\begin{tikzpicture}[every node/.style={font=\Large},scale=0.8
\draw[help lines,step=1cm] (0.6,0.6) grid (8.4,7.4);
\coordinate (A) at (1.1,5.6);
\coordinate (Adelta) at (1,5);
\coordinate (B) at (7.9,3.3);
\coordinate (Bdelta) at (7,3);
\draw (A) node [left=2pt] {$a$} [fill] circle (2pt);
\draw (Adelta) node [below right] {$a_{\delta}$} [fill] circle (2pt);
\draw (B) node [right=2pt] {$b$} [fill] circle (2pt);
\draw (Bdelta) node [below left] {$b_{\delta}$} [fill] circle (2pt);
\filldraw[thick,smooth cycle,fill opacity=0.05] plot coordinates {(0.8, 4.5) (A) (2.3,6.5) (3.8,6.9) (5,7.2) (6,7) (7,6.3) (7.7,5) (B) (7.4,2) (5.9,0.9) (3.8,0.8) (2,1.8) (1.5,3) };
\draw[very thick] (Adelta) -- ++(1,0) -- ++(0,-1) -- ++(1,0) -- ++(0,-1) -- ++(-1,0) -- ++(0,-1) -- ++(2,0) -- ++(0,4) -- ++(1,0) -- ++(0,-2) -- ++(1,0) -- ++(0,-1) -- (Bdelta);
\draw (4.5,0) node {$\Omega$};
\end{tikzpicture}
\caption{Discretisation of domain $\Omega.$}
\end{figure}

As we let $\delta \to 0$ we expect the behaviour of the walk to depend on the value of $x.$ For $x < x_c$ it is possible to prove that $\omega_{\delta}$ goes to a straight line as $\delta \to 0.$ (Strictly speaking it converges in distribution to a straight line, with fluctuations O$(\sqrt{\delta}).$)
For $x > x_c$ it is expected that $\omega_{\delta}$ becomes (again, in distribution) space-filling as $\delta \to 0.$ But at $x = x_c$ it is {\em conjectured} that $\omega_{\delta}$ becomes (in distribution) a random continuous curve, and is conformally invariant. This describes the scaling limit. If this conjecture is correct, a second, pivotal, conjecture by Lawler, Schramm and Werner \cite{LSW04} is that this random curve converges to $\rm{SLE}_{8/3}$ (described below) from $a$ to $b$ in the domain $\Omega.$

\subsection{Schramm L\"owner Evolution}
For an approachable discussion of SLE$_\kappa,$ see Chapter 15 of \cite{G08}. Here we give a very minimal outline. Let ${\mathbb H}$ denote the upper half-plane. Consider a path $\gamma$ starting at the boundary and finishing at an internal vertex. Then ${\mathbb H}\backslash \gamma$ is the complement of this path, and is a slit upper half-plane. It follows from the Riemann Mapping Theorem that it can be conformally mapped to the upper half-plane. L\"owner \cite{L23} discovered that by specifying the map so that it approaches the identity at infinity, the conformal map so described (actually a family of maps, appropriate to each point on the curve) satisfies a simple differential equation, called the L\"owner equation. The mapping can alternatively be defined by a real function. This observation led Schramm to apply the L\"owner equation to a conformally invariant measure for planar curves. That is to say, the L\"owner equation generates a set of conformal maps, driven by a continuous real-valued function. Schramm's profound insight was to use Brownian motion $B_t$ as the driving function\footnote{It is the only process compatible with both conformal invariance and the so-called domain Markov property.}. So let $B_t,$ $t \ge 0$ be standard Brownian motion on ${\mathbb R}$ and let $\kappa$ be a real parameter. Then SLE$_\kappa$ is the family of conformal maps ${g_t:t \ge 0}$ defined by the L\"owner equation
$$ \frac{\partial}{\partial t} g_t(z) = \frac {2}{g_t(z)-\sqrt{\kappa}B_t}, \,\,\, g_0(z)=z.$$ For $\kappa \le 4$ the domain of $g_t$, which is the set of initial conditions $z$ for which the above differential equation still has a solution at time $t$, equals the half-plane minus a simple curve. The curve starts at the origin and can be shown to go to $\infty$. This is called {\em chordal} SLE$_\kappa$ as it describes paths growing from the boundary point $0$ and ending at the boundary point $\infty$. Larger values of $\kappa$ lead to more complicated behaviour.

For a general simply connected domain $D$, $\rm{SLE}_\kappa$ between two boundary points is just defined by taking the image of $\rm{SLE}_\kappa$ in the half-plane using a conformal map that takes the half-plane to $D,$ and $0$ and $\infty$ to the two prescribed boundary points. It is conjectured, and widely believed, that the scaling limit of the SAW between two fixed points is conformally invariant. Lawler, Schramm and Werner \cite{LSW04} showed that this conformal invariance and the so called restriction property for the SAW imply that the scaling limit of the SAW is described by $\rm{SLE}_{8/3}$. Note that this only describes the SAW in $D$ between two boundary points. If we want to describe the SAW which starts at a fixed point on the boundary but can end anywhere on the boundary, then we need to know the distribution of this random endpoint. We discuss this distribution in the following section.

 Hopefully this rather vague description will convey the flavour of this exciting and powerful development in studying not just two-dimensional SAWs, but a variety of other processes, such as percolation, the random cluster model, and the Ising model. We refer the reader to \cite{B10} for greater detail of both SLE and these applications.

 \subsection{Conformal mapping}

Let $D$ be a bounded, simply connected domain in the complex plane containing $0.$ We are interested in paths in $D$ starting at $0$ and ending on the boundary of the domain. Initially we will consider random walks, later self-avoiding walks. As described above in the description of the scaling limit, we can discretize the space with a lattice of lattice spacing $\delta.$ In both the random walk case and the SAW case we are interested in the scaling limit as $\delta \to 0.$ For random walks the scaling limit is Brownian motion, stopping when it hits the boundary of $D.$ The distribution of the end-point is harmonic measure. If the domain boundary is piecewise smooth, then harmonic measure is absolutely continuous with respect to arc length along the boundary \cite{R16}. Let $h_D(z)$ denote the density with respect to arc length (often called the Poisson kernel). If $f$ is a conformal map on $D$ that fixes the origin, and such that the boundary of $f(D)$ is also piecewise smooth, then the conformal invariance of Brownian motion implies that the density for harmonic measure on the boundary of $f(D)$ is related to the boundary of $D$ by
\begin{equation}\label{cirw}
h_D(z)=|f'(z)|h_{f(D)}(f(z)).
\end{equation}
In the case of the SAW, Lawler Schramm and Werner \cite{LSW04} predicted that the corresponding density of the probability measure $\rho(z)$  transforms under conformal maps as
\begin{equation}\label{cisaw}
\rho_D(z)=c|f'(z)|^{b}\rho_{f(D)}(f(z)),
\end{equation}
where $b=5/8$ and the constant $c$ is required to ensure that $\rho_D(z)$ is a probability density. If one starts the random walk or the SAW at the center of a disc, then the hitting density on the circle will be uniform. So the above equations determine the hitting density for any simply connected domain. Simulations in \cite{DGKLP11}, \cite{KL11} and \cite{K12} provide strong support for the conjectured behaviour. (Eq. (\ref{cisaw}) is correct for domains whose boundary consists only of vertical and horizontal line segments. For general domains there is a lattice effect that persists in the scaling limit that produces a factor that depends on the angle of the tangent to the boundary that must be included \cite{KL11}.)

 The solution of the original problem for random walks by conformal maps is described in \cite{BLWW04}, where a conformal map from the unit disc to an $a \times c$ rectangle is given by the Schwarz-Christoffel formula. Our approach also uses conformal maps, but we instead use a map between the upper half-plane and a rectangle, where the mapping is again given by a Schwarz-Christoffel transformation. For $\alpha > 1,$ let
$$f(z) = \int_0^z \frac{d\xi}{\sqrt{1-\xi^2}\sqrt{\alpha^2 - \xi^2}}.$$ $f(z)$
 is a Schwarz-Christoffel transformation that maps the upper half plane to a rectangle.
The rectangle has one edge along the real axis and $0$ is a midpoint of this side. So the corners
can be written as $\pm a/2$ and $ic \pm a/2$ where $a, c > 0$ are the length of the horizontal and vertical
edges, respectively. We have
$$f(1) = a/2, \,\, f(-1)=-a/2, \,\, f(\alpha) = a/2 + ic, \,\, f(-\alpha) = -a/2+ic, \,\, f(0) =0.$$

So
$$ a=\int_{-1}^1 \frac{dx}{\sqrt{1-x^2}\sqrt{\alpha^2 - x^2}}, \,\,\, c =\int_{1}^\alpha \frac{dx}{\sqrt{x^2-1}\sqrt{\alpha^2 - x^2}}.$$
We note that $$ a =\frac{2}{\alpha}{\bf K} \left ( \frac{1}{\alpha} \right ), \,\, c=\frac{1}{\alpha}{\bf K} \left ( \frac{\sqrt{\alpha^2-1}}{\alpha} \right ), \,\, \alpha > 1.$$
Here ${\bf K}(x)$ is the complete elliptic integral of the first kind. By dilation invariance we only need concern ourselves with the aspect ratio $a/c.$ So given an aspect ratio
$r$, we have to find $\alpha$ such that

$$ r = \frac{\int_{-1}^1 (1-x^2)^{-1/2}(\alpha^2 - x^2)^{-1/2} dx}{\int_{1}^{\alpha} (x^2-1)^{-1/2}(\alpha^2 - x^2)^{-1/2} dx} = \frac{2 {\bf K} \left ( \frac{1}{\alpha} \right )}{{\bf K} \left ( \frac{\sqrt{\alpha^2-1}}{\alpha} \right )}.$$

Setting $k=1/\alpha,$ then $c$ can be written ${\bf K}(k')={\bf K}'(k).$ Therefore $$\frac{c}{a}=\frac{1}{r}=\frac{{\bf K}(1/\alpha')}{2{\bf K}(1/\alpha)}=\frac{{\bf K}'(1/\alpha)}{2{\bf K}(1/\alpha)}.$$ Now the nome $q$  is defined\footnote{Amramowitz and Stegun, p591 formula 17.3.17} by $q=\exp(-\pi{\bf K}'/{\bf K}),$ so $$r=\frac{{2\bf K}(1/\alpha)}{{\bf K'}(1/\alpha)},$$ and $$e^{-2\pi/ r} = q.$$ Then we have
$$\frac{1}{\alpha}= \left ( \frac{\theta_2(q)}{\theta_3(q)} \right )^2,$$ so
$$\sqrt{\alpha} = \frac{\theta_3(e^{-2\pi/r})}{\theta_2(e^{-2\pi/r})},$$
where $\theta_j(q)=\theta_j(0,q)$ is the Jacobi theta function.
Evaluating this in one's favourite algebraic package gives the required value of $\alpha$ for any $r \ge 1$ instantly.

Alternatively, we note that for an aspect ratio of 10, as posed in the original problem, $\alpha$ is very close to 1. We can achieve very high  accuracy by expanding the
ratio of the above integrals around $\alpha=1$ from the Taylor expansion of elliptic integrals, and find
$$r=\frac{1}{\pi}\left ( 4\log(2\sqrt{2})-2\log(\alpha-1)+(\alpha-1)-\frac{3}{8}(\alpha-1)^2+\frac{5}{24}(\alpha-1)^3+O(\alpha-1)^4 \right ).$$
Solving this numerically for $r=10$ gives $\alpha=1.00000120561454706472212\ldots.$
To leading order one obtains $$\alpha\approx 1+8 e^{-\pi r/2},$$ which, for $r=10$ gives $1.0000012056138\ldots$ which is in error only in the 14th digit. Going further, we can write
\begin{equation} \label{alfa}
\alpha = 1+8 e^{-\pi r/2}+32 e^{-\pi r}+O( e^{-3\pi r/2}),
\end{equation}
 which, for $r=10$ gives 19 significant digits.

The preimage of the center of the rectangle will be on the imaginary axis, so write it as $di.$
To find $d$ in terms of $\alpha,$ consider rotating the rectangle by 180 degrees. This will correspond
(under$ f$) to a conformal automorphism of the upper half plane, i.e., a M\"obius transformation
$\phi(z).$ Since the rotation interchanges opposite corners, $\phi$ will interchange $1$ and $-\alpha$ and it will
interchange $−1$ and $\alpha$. The rotation interchanges the midpoints of the horizontal edges of the
rectangle, so $\phi$ interchanges $0$ and $\infty.$ So $\phi(z) = −\alpha /z.$ The point $di$ should be fixed by $\phi,$
giving
$d^2 = \alpha.$

For the random walk and the SAW in the half plane starting at $id,$ it follows from (\ref{cisaw}) and (\ref{cirw}) that the (unnormalized)
hitting density along the real axis is $(x^2 + d^2)^{-b} = (x^2 + \alpha)^{-b}$. So letting $\rho_R$ denote the hitting density for a rectangle with the walk starting at the center, we have by (\ref{cisaw}) and (\ref{cirw}),
\begin{equation}\label{rh}
(x^2 + \alpha)^{-b} \propto |f'(x)|^b \rho_R(f(z)).
\end{equation}

We require the ratio of the integral of $\rho(z)$ along a vertical edge to the integral along a horizontal edge,
$$\frac{\int_{0}^{c} \rho_R(a/2+iy) dy}{\int_{-a/2}^{a/2} \rho_R(x) dx}.$$
By a change of variable, setting $u = f^{-1}(x)$ in the denominator and $u=f^{-1}(a/2+iy)$ in the numerator, this becomes
 $$ \frac{\int_{1}^{\alpha} \rho_R(f(u))|f'(u)| du}{\int_{-1}^{1} \rho_R(f(u))|f'(u)| du}.$$
Recall that $f'(u) = (1-u^2)^{-1/2}(\alpha^2 - u^2)^{-1/2},$ so that the ratio of probabilities of a first hit on the vertical side to a first hit on the horizontal side, $R(\alpha,b)$ is
\begin{equation}\label{rat}
R(\alpha,b) = \frac {\int_{1}^{\alpha} (u^2+\alpha)^{-b}(u^2-1)^{(b-1)/2}(\alpha^2 - u^2)^{(b-1)/2} du}{\int_{-1}^{1} (u^2+\alpha)^{-b}(1-u^2)^{(b-1)/2}(\alpha^2 - u^2)^{(b-1)/2} du}.
\end{equation}
\section{Random walks.}
We first consider the random walk case, corresponding to $b=1.$ The integrals in (\ref{rat}) greatly simplify, giving
$$R(\alpha,1) = \frac{\arctan(\sqrt{\alpha})-\arctan(1/\sqrt{\alpha})}{2\arctan(1/\sqrt{\alpha})}.$$

It is straightforward to calculate the asymptotic expansion of both the numerator and denominator, and hence their ratio. In this way we find
$$R(\alpha,1) = \frac{8}{\pi}e^{-\pi r/2}  + O\left (e^{-\pi r}\right ).$$
For $r=10$ this evaluates to $3.83758797925134\ldots \times 10^{-7},$ which is correct to 6 significant digits.

Proceeding to the second-order term, we obtain $$R(\alpha,1) = \frac{8}{\pi}e^{-\pi r/2}  + \frac{64}{\pi^2}e^{-\pi r}+ O\left (e^{-3\pi r/2}\right ).$$
For $r=10$ this evaluates to $3.8375894519594\ldots \times 10^{-7},$ which is correct to 13 significant digits. It is straightforward to obtain further terms should one so wish.

\section{The general case, including SAW.}
Unfortunately for $b \ne 1$ we cannot evaluate the integrals exactly, but as for the random walk case, asymptotics provides us with sufficient accuracy. Referring to equation (\ref{rat}), and noting that $\alpha$ is very close to 1 when the aspect ratio is 10, the denominator integral can be accurately approximated by $$
\int_{-1}^{1} (u^2+1)^{-b}(1-u^2)^{b-1}du = \frac{\sqrt{\pi} \Gamma \left (\frac{b}{2}\right )}{2\Gamma\left (\frac{b}{2}+\frac{1}{2}\right )}.$$

The numerator integral has a very narrow range of integration, so we can approximate the factor $u^2 + \alpha$ by 2, the factor $(u^2-1)$ by $2(u-1)$ and the factor $\alpha^2-u^2$ by $2(\alpha-u).$ Thus the numerator can be approximated by $$\frac{1}{2}
\int_{1}^{\alpha} (u-1)^{(b-1)/2}(\alpha-u)^{(b-1)/2}du.$$ If we set $u=1+t(\alpha-1)$, this becomes
$$\frac{1}{2}(\alpha-1)^b\int_0^1[t(1-t)]^{(b-1)/2}dt= \frac{2^{-b-1}\sqrt{\pi}\Gamma\left (\frac{1+b}{2}\right )}{\Gamma \left (1+\frac{b}{2}\right )}(\alpha-1)^b.$$

Combining the above results, and using the result derived above, $\alpha-1 \approx 8 e^{-\pi r/2},$ we find, asymptotically
\begin{equation}
\label{finas}
{\tilde R}(r,b) \approx \frac{2^{2b} \Gamma \left ( \frac{1}{2}+\frac{b}{2} \right )^2}{\Gamma\left ( 1+ \frac{b}{2} \right ) \Gamma \left ( \frac{b}{2} \right )} e^{-\pi b r/2}.
\end{equation}
For SAW, $b = 5/8,$ and (\ref{finas}) reduces to $${\tilde R}(r,5/8) \approx 1.2263431442 e^{-5\pi r/16}.$$
For $r=10$ this gives ${\tilde R}(10,5/8) \approx 6.6824528 \times 10^{-5}.$

In order to get an estimate of the accuracy of this result, we calculated the relevant integrals numerically, which is quite straightforward in either Maple or Mathematica. In this way we found the more accurate estimate ${\tilde R}(10,5/8) \approx 6.682989935 \times 10^{-5}.$

A rather elaborate calculation gives the next term in the asymptotic expansion of the aspect ratio as

$${\tilde R}(r,b)=\frac{2^{2b+1}\Lambda}{b} e^{-b\pi r/2}\left [ 1 + \frac{\Lambda 2^{b+1}} {b \sin\left ( \frac{\pi b}{2} \right )}e^{-b\pi r/2} + 4(b-1+2\Lambda)e^{-\pi r/2}+ O(e^{-b\pi r})\right ],$$ for $0 < b < 1,$ where $\Lambda = \left (\frac{\Gamma \left ( \frac{1+b}{2} \right ) }{\Gamma \left ( \frac{b}{2} \right ) }\right )^2.$ Evaluating this for $r=10, \,\, b= \frac{5}{8}$, we find $${\tilde R}(10,5/8)= 0.00006682989679\cdots$$ which is in error by 2 in the $8^{th}$ significant digit, whereas the leading order term gives an estimate which is in error by 6 in the $5^{th}$ significant digit

We see that the corresponding ratio for the scaling limit of SAWs in a $10 \times 1$ rectangle is  some two orders of magnitude larger than the corresponding result for Brownian motion.
This higher probability is intuitively what one might expect, as SAWs tend to move in a straight trajectory rather more than do random walks.

\section{Acknowledgements} AJG would like to than Jon Borwein for introducing him to the original problem, and Nathan Clisby for a discussion that spawned the current work. This work was supported by the Australian Research Council through grant DP120100939  (AJG).

Self-avoiding walks in a rectangle


Data-Dependent Stability of Stochastic Gradient Descent

Abstract: We establish a data-dependent notion of algorithmic stability for Stochastic Gradient Descent (SGD), and employ it to develop novel generalization bounds. This is in contrast to previous distribution-free algorithmic stability results for SGD which depend on the worst-case constants. By virtue of the data-dependent argument, our bounds provide new insights into learning with SGD on convex and non-convex problems. In the convex case, we show that the bound on the generalization error depends on the risk at the initialization point. In the non-convex case, we prove that the expected curvature of the objective function around the initialization point has crucial influence on the generalization error. In both cases, our results suggest a simple data-driven strategy to stabilize SGD by pre-screening its initialization. As a corollary, our results allow us to show optimistic generalization bounds that exhibit fast convergence rates for SGD subject to a vanishing empirical risk and low noise of stochastic gradient.

\section*{Acknowledgments}
This work was in parts funded by the European Research Council (ERC)
under the European Union's Horizon 2020 research and innovation programme
(grant agreement no 637076).\\~\\
This work was in parts funded by the European Research Council under the
European Union's Seventh Framework Programme (FP7/2007-2013)/ERC grant
agreement no 308036.

\section{Conclusions and Future Work}
\label{sec:conclusions}
In this work we proved data-dependent stability bounds for \ac{SGD} and revisited its generalization ability.
We presented novel bounds for convex and non-convex smooth loss functions, partially controlled by data-dependent quantities, while previous stability bounds for \ac{SGD} were derived through the worst-case analysis.
In particular, for non-convex learning, we demonstrated theoretically that generalization of \ac{SGD} is heavily affected by the expected curvature around the initialization point.
We demonstrated empirically that our bound is indeed tighter compared to the uniform one.
In addition, our data-dependent analysis also allowed us to show optimistic bounds on the generalization error of \ac{SGD}, which exhibit fast rates subject to the vanishing empirical risk of the algorithm's output.

In future work we further intend to explore our theoretical findings experimentally and evaluate the feasibility of the transfer learning based on the second-order information.
Another direction lies in making our bounds adaptive. So far we have presented bounds that have data-dependent components, however the step size cannot be adjusted depending on the data, e.g.\ as in~\cite{zhao2015stochastic}.
This was partially addressed by~\cite{london2016generalization}, albeit in the context of uniform stability,
and we plan to extend this idea to the context of data-dependent stability.

\subsection{Convex Losses}
\label{sec:proof_convex}
In this section we prove on-average  stability for loss functions that are non-negative, $\beta$-smooth, and convex.
\begin{theorem}
\label{thm:convex_stab_decomposed}
Assume that $f$ is convex, and that \ac{SGD}'s is ran with step sizes $\cbr{\alpha_t}_{t=1}^T$.
Then, for every $t_0 \in \{0, 1, 2, \ldots m\}$, \ac{SGD} is $\epsilon(\mathcal{D}, \boldsymbol{w}_1)$-on-average  stable with
\begin{align*}
\epsilon(\mathcal{D}, \boldsymbol{w}_1)
\leq \frac{2}{m} \sum_{t=t_0+1}^{T} \alpha_t \E_{S,z}\br{ \|\nabla f(\boldsymbol{w}_t, z_{j_t})\| } + \E_{S,A}\br{R(A_S)} \frac{t_0}{m}~.
\end{align*}
\end{theorem}
\begin{proof}
For brevity denote  $\Delta_t(S,z) := \E_A\br{\delta_t(S,z) \ | \ \delta_{t_0}(S,z) = 0}$.
We start by applying Lemma~\ref{lem:stab_decompose}:
\begin{align}
\E_{S,z}\E_A\br{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)}
\leq L \E_{S,z}\br{ \Delta_T(S,z) } + \E_{S,A}\br{R(A_S)} \frac{t_0}{m}~. \label{eq:convex_t0_bound}
\end{align}
Our goal is to bound the first term on the r.h.s.\ as a decreasing function of $t_0$, so that eventually we can minimize the bound w.r.t.\ $t_0$.
At this point we focus on the first term, and the proof partially follows the outline of the proof of Theorem~3.7 in \cite{hardt2016train}.
The strategy will be to establish the bound on $\Delta_T(S,z)$ by using a recursive argument.
In fact we will state the bound on $\Delta_{t+1}(S,z)$ in terms of $\Delta_t(S,z)$ and then unravel the recursion.
Finally, we will take expectation w.r.t.\ the data after we obtain the bound by recursion.

To do so, we distinguish two cases: 1) \ac{SGD} encounters a perturbed point at step $t$, that is $t = i$, and 2) the current point is the same in $S$ and ${S^{(i)}}$, so $t \neq i$.
For the first case, we will use data-dependent boundedness of the gradient update rule, Corollary~\ref{cor:G_boundedness}, that is
\[
\|G_t(\boldsymbol{w}_{S,t}) - G_t(\boldsymbol{w}_{{S^{(i)}},t})\| \leq \delta_t(S,z) + 2 \alpha_t \|\nabla f(\boldsymbol{w}_{S,t}, z_{j_t})\|~.
\]
To handle the second case, we will use the expansiveness of the gradient update rule, Lemma~\ref{lem:worst-case-expansive}, which states that for convex loss functions, the gradient update rule is $1$-expansive, so $\delta_{t+1}(S,z) \leq \delta_t(S,z)$.
Considering both cases of example selection,
and noting that \ac{SGD} encounters the perturbation w.p.\ $\frac{1}{m}$,
we write $\E_A$ for a step $t$ as
\begin{align*}
\Delta_{t+1}(S,z) &\leq \pr{1 - \frac{1}{m}} \Delta_t(S,z)
+ \frac{1}{m} \pr{ \Delta_t(S,z) + 2 \alpha_t \|\nabla f(\boldsymbol{w}_{S,t}, z_{j_t})\| }\\
&= \Delta_t(S,z) + \frac{2 \alpha_t \|\nabla f(\boldsymbol{w}_{S,t}, z_{j_t})\|}{m}~.
\end{align*}
Unraveling the recursion from $T$ to $t_0$ and plugging the above into~\eqref{eq:convex_t0_bound} yields
\begin{align*}
\E_A\E_{S,z}[\delta_T(S,z)]
\leq \frac{2}{m} \sum_{t=t_0+1}^{T} \alpha_t \E_{S,z}\br{ \|\nabla f(\boldsymbol{w}_t, z_{j_t})\| } + \E_{S,A}\br{R(A_S)} \frac{t_0}{m}~.
\end{align*}
This completes the proof.
\end{proof}
Next statement is a simple consequence of Theorem~\ref{thm:convex_stab_decomposed} and Lemma~\ref{lem:E_path_bounded_by_R1_and_noise}.
\begin{proof}[Proof of Theorem~\ref{thm:sgd_stab_convex}.]
Consider Theorem~\ref{thm:convex_stab_decomposed} and set $t_0 = 0$.
\begin{align}
  \label{eq:convex_stab_bound_proof_eq1}
  \epsilon(\mathcal{D}, \boldsymbol{w}_1) &\leq \frac{2}{m} \sum_{t=1}^{T} \alpha_t \E_{S,z}\br{\|\nabla f(\boldsymbol{w}_{S,t}, z_{j_t})\|}~.
\end{align}
Bounding the sum using Lemma~\ref{lem:E_path_bounded_by_R1_and_noise} recalling that $\alpha_t = c/ \sqrt{t}$, we get
\begin{align*}
  \E_S\br{\sum_{t=1}^{T} \alpha_t \|\nabla f(\boldsymbol{w}_t, z_{j_t})\|}
  &\leq 2 \sqrt{ \pr{\sum_{t=1}^T \alpha_t} \pr{ R(\boldsymbol{w}_1) - R^{\star} + \frac{\beta \sigma^2}{2} \sum_{t=1}^T \alpha_t^2} }
  + \sigma \sum_{t=1}^T \alpha_t\\
  &\leq 2 \sqrt{2 c} \cdot \sqrt[4]{T} \cdot \sqrt{R(\boldsymbol{w}_1) - R^{\star} }
  + 2 c \sigma \pr{ \sqrt[4]{T} \sqrt{\frac{\beta}{2}} + \sqrt{T} }~.
\end{align*}
Combining above with~\eqref{eq:convex_stab_bound_proof_eq1} completes the proof.
\end{proof}

\section{Introduction}
\emph{Stochastic gradient descent (\ac{SGD})} has become one
of the workhorses of modern machine learning.
In particular, it is the optimization method of choice for training
highly complex and non-convex models, such as neural networks.
When it was observed that these models generalize better
(suffer less from overfitting) than classical machine learning
theory suggests, a large theoretical interest emerged to explain
this phenomenon.
Given that \ac{SGD} at best finds a local minimum of the non-convex
objective function, it has been argued that all such minima might
be equally good.
However, at the same time, a large body of empirical work and tricks
of trade, such as \emph{early stopping}, suggests that in practice one
might not even reach a minimum, yet nevertheless observes excellent
performance.

In this work we follow an alternative route that aims to \emph{directly}
analyze the generalization ability of \ac{SGD} by studying how sensitive
it is to small perturbations in the training set.
This is known as \emph{algorithmic stability} approach~\cite{BousquetE02}
and was used recently~\cite{hardt2016train} to establish generalization
bounds for both convex and non-convex learning settings.
To do so they employed a rather restrictive notion of stability that does
not depend on the data, but captures only intrinsic characteristics of the
learning algorithm and global properties of the objective function.
Consequently, their analysis results in worst-case guarantees that in some
cases tend to be too pessimistic.
As recently pointed out in~\cite{zhang2017understanding}, \emph{deep learning}
might indeed be such a case, as this notion of stability is insufficient
to give deeper theoretical insights, and a less restrictive one is desirable.

\textbf{As our main contribution} in this work we establish that a data-dependent notion of algorithmic stability, very similar to the \emph{On-Average  Stability}~\cite{shalev2010learnability}, holds for \ac{SGD}
when applied to convex as well as non-convex learning problems.
As a consequence we obtain new generalization bounds that depend  on the data-generating distribution and the initialization point of an algorithm.
For convex loss functions, the bound on the generalization error is essentially multiplicative in the risk at the initialization point when noise of stochastic gradient is not too high.
For the non-convex loss functions, besides the risk, it is also critically controlled by the expected second-order information about the objective function at the initialization point.
We further corroborate our findings empirically and show that, indeed, the data-dependent generalization bound is tighter than the worst-case counterpart on non-convex objective functions.
Finally, the nature of the data-dependent bounds allows us to state \emph{optimistic} bounds that switch to the faster rate of convergence subject to the vanishing empirical risk.

In particular, our findings justify the intuition that \ac{SGD} is more stable
in less curved areas of the objective function and link it to the generalization ability.
This also backs up numerous empirical findings in the deep learning literature that
solutions with low generalization error occur in less curved regions.
At the same time, in pessimistic scenarios, our bounds are no worse than those
of~\cite{hardt2016train}.

Finally, we exemplify an application of our bounds, and propose a simple yet
principled \emph{transfer learning} scheme for the convex and
non-convex case, which is guaranteed to transfer from the best source
of information.
In addition, this approach can also be used to select a good initialization
given a number of random starting positions. This is a theoretically sound
alternative to the purely random commonly used in non-convex learning.

The rest of the paper is organized as follows.
We revisit the connection between stability and generalization of \ac{SGD}
in Section~\ref{sec:stability} and introduce a data-dependent notion of stability in Section~\ref{sec:data-stability}.
We state the main results in Section~\ref{sec:results},
in particular, Theorem~\ref{thm:sgd_stab_convex} for the convex case,
and Theorem~\ref{thm:sgd_stab_nonconvex} for the non-convex one.
Next we demonstrate empirically that the bound shown in Theorem~\ref{thm:sgd_stab_nonconvex} is tighter than the worst-case one in Section~\ref{sec:tightness}.
Finally, we suggest application of these bounds by showcasing principled
transfer learning approaches in Section~\ref{sec:tl},
and we conclude in Section~\ref{sec:conclusions}.

\subsection{Non-convex Losses}
Our proof of a stability bound for non-convex loss functions, Theorem~\ref{thm:sgd_stab_nonconvex} (in the submission file), follows a general outline of \cite[Theorem~3.8]{hardt2016train}.
Namely,
the outputs of \ac{SGD} run on a training set $S$ and its perturbed version ${S^{(i)}}$ will not differ too much,
because by the time a perturbation is encountered, the step size has already decayed enough.
So, on the one hand, stabilization is enforced by the diminishing the step size, and on the other hand, by how much updates expand the distance between the gradients after the perturbation.
Since \cite{hardt2016train} work with uniform stability, they capture the expansiveness of post-perturbation update by the Lipschitzness of the gradient.
In combination with a recursive argument, their bound has exponential dependency on the Lipschitz constant of the gradient.
We argue that the Lipschitz continuity of the gradient can be too pessimistic in general.
Instead, we rely on a local data-driven argument: considering that we initialize \ac{SGD} at point $\boldsymbol{w}_1$, how much do updates expand the gradient under the distribution of interest?
The following crucial lemma characterizes such behavior in terms of the curvature at $\boldsymbol{w}_1$.
\begin{lemma}
\label{lem:expansiveness_nonconvex}
Assume that the loss function $f(\cdot, z)$ is $\beta$-smooth and that its Hessian is $\rho$-Lipschitz.
Then,
\begin{align}
  \left\| G_t(\boldsymbol{w}_{S,t}) - G_t(\boldsymbol{w}_{{S^{(i)}},t}) \right\|
\leq \pr{1 + \alpha_t \xi_t(S,z) } \delta_t(S,z)
\end{align}
where
\begin{align*}
&\xi_t(S,z) := \left\|\nabla^2 f(\boldsymbol{w}_1, z_t) \right\|_2
+ \frac{\rho}{2} \left\| \sum_{k=1}^{t-1} \alpha_k \nabla f(\boldsymbol{w}_{S,k}, z_k) \right\|
+ \frac{\rho}{2} \left\| \sum_{k=1}^{t-1} \alpha_k \nabla f(\boldsymbol{w}_{{S^{(i)}},k}, z_{k'}) \right\|~.
\end{align*}
Furthermore, for any $t \in [T]$,
\begin{align*}
   \E_{S,z}\br{\xi_t(S,z)} &\leq \E_{S,z}\br{\left\|\nabla^2 f(\boldsymbol{w}_1, z_t) \right\|_2}\\
  &+ 2 \rho \sqrt{ \pr{R(\boldsymbol{w}_1) - R^{\star}} c (1 + \ln(T)) }\\
  &+ \rho \sigma \pr{ \sqrt{2 c \beta} + c (1 + \ln(T)) }~.
\end{align*}
\end{lemma}
\begin{proof}
Recall that the randomness of the algorithm is realized through sampling without replacement from the uniform distribution over $[m]$.
Apart from that we will not be concerned with the randomness of the algorithm, and given the set of random variables $\{j_i\}_{i=1}^m$, for brevity we will use indexing notation $z_1, z_2, \ldots, z_m$ to indicate $z_{j_1}, z_{j_2}, \ldots, z_{j_m}$.
Next, let ${S^{(i)}} = \cbr{z_i'}_{i=1}^m$, and introduce a shorthand notation
$f_k(\boldsymbol{w}) = f(\boldsymbol{w}, z_k)$ and $f_{k'}(\boldsymbol{w}) = f(\boldsymbol{w}, z'_k)$.
We start by applying triangle inequality to get
\begin{align*}
\left\| G_t(\boldsymbol{w}_{S,t}) - G_t(\boldsymbol{w}_{{S^{(i)}},t}) \right\| \leq \|\boldsymbol{w}_{S,t} - \boldsymbol{w}_{{S^{(i)}},t}\|
+ \alpha_t \left\| \nabla f_t(\boldsymbol{w}_{S,t}) - \nabla f_t(\boldsymbol{w}_{{S^{(i)}},t}) \right\|~.
\end{align*}
In the following we will focus on the second term of r.h.s.\ above.
Given \ac{SGD} outputs $\boldsymbol{w}_{S,t}$ and $\boldsymbol{w}_{{S^{(i)}},t}$ with $t > i$, our goal here is to establish how much do gradients grow apart with every new update.
This behavior can be characterized assuming that gradient is Lipschitz continuous, however, we conduct a local analysis.
Specifically, we observe how much do updates expand gradients, given that we start at some point $\boldsymbol{w}_1$ under the data-generating distribution.
So, instead of the Lipschitz constant, expansiveness rather depends on the curvature around $\boldsymbol{w}_1$.
On the other hand, we are dealing with outputs at an arbitrary time step $t$, and therefore we first have to relate them to the initialization point $\boldsymbol{w}_1$.
We do so by using the gradient update rule and telescopic sums, and conclude that this relationship is controlled by the sum of gradient norms along the update path.
We further establish that this sum is controlled by the risk of $\boldsymbol{w}_1$ up to the noise of stochastic gradients, through stationary-point result of Lemma~\ref{lem:E_path_bounded_by_R1_and_noise}.
Thus, the proof consists of two parts: 1) Decomposition into curvature and gradients along the update path, and 2) bounding those gradients.
\paragraph{1) Decomposition.}
Introduce $\boldsymbol{\delta}_t := \boldsymbol{w}_{{S^{(i)}},t} - \boldsymbol{w}_{S,t}$.
By Taylor theorem we get that
\begin{align*}
  &\nabla f_t(\boldsymbol{w}_{S,t}) - \nabla f_t(\boldsymbol{w}_{{S^{(i)}},t}) = \nabla^2 f_t(\boldsymbol{w}_1) \boldsymbol{\delta}_t
+ \int_0^1 \Big(\nabla^2 f_t(\boldsymbol{w}_{S,t} + \tau \boldsymbol{\delta}_t) - \nabla^2 f_t(\boldsymbol{w}_1)\Big) \diff \tau \boldsymbol{\delta}_t~.
\end{align*}
Taking norm on both sides, applying triangle inequality, Cauchy-Schwartz inequality, and assuming that Hessians are $\rho$-Lipschitz we obtain
\begin{align}
  &\|\nabla f_t(\boldsymbol{w}_{S,t}) - \nabla f_t(\boldsymbol{w}_{{S^{(i)}},t})\|
\leq \rho \int_0^1 \left\| \boldsymbol{w}_{S,t} - \boldsymbol{w}_1 + \tau \boldsymbol{\delta}_t \right\| \diff \tau \|\boldsymbol{\delta}_t\|
+ \left\|\nabla^2 f_t(\boldsymbol{w}_1) \right\| \|\boldsymbol{\delta}_t\|~.\label{eq:gradient_track_int}
\end{align}
\paragraph{2) Bounding gradients.}
Using telescoping sums and \ac{SGD} update rule we get that
\begin{align*}
\boldsymbol{w}_{S,t} - \boldsymbol{w}_1 + \tau \boldsymbol{\delta}_t
&= \boldsymbol{w}_{S,t} - \boldsymbol{w}_1 + \tau \pr{\boldsymbol{w}_{{S^{(i)}},t} - \boldsymbol{w}_1 + \boldsymbol{w}_1 - \boldsymbol{w}_{S,t}} \\
&= \sum_{k=1}^{t-1} \pr{\boldsymbol{w}_{S,k+1} - \boldsymbol{w}_{S,k}}\\
&+ \tau \sum_{k=1}^{t-1} \pr{\boldsymbol{w}_{{S^{(i)}},k+1} - \boldsymbol{w}_{{S^{(i)}},k}}\\
&- \tau \sum_{k=1}^{t-1} \pr{\boldsymbol{w}_{S,k+1} - \boldsymbol{w}_{S,k}}\\
&= (\tau - 1) \sum_{k=1}^{t-1} \alpha_k \nabla f_k(\boldsymbol{w}_{S,k}) - \tau \sum_{k=1}^{t-1} \alpha_k \nabla f_{k'}(\boldsymbol{w}_{{S^{(i)}},k})~.
\end{align*}
Plugging above into the integral of~\eqref{eq:gradient_track_int} we have
\begin{align*}
&\int_0^1 \left\| \sum_{k=1}^{t-1} \alpha_k \pr{ (\tau - 1) \nabla f_k(\boldsymbol{w}_{S,k}) - \tau \nabla f_{k'}(\boldsymbol{w}_{{S^{(i)}},k}) } \right\| \diff \tau\\
  &\leq \frac{1}{2} \left\| \sum_{k=1}^{t-1} \alpha_k \nabla f_k(\boldsymbol{w}_{S,k}) \right\| + \frac{1}{2} \left\| \sum_{k=1}^{t-1} \alpha_k \nabla f_{k'}(\boldsymbol{w}_{{S^{(i)}},k}) \right\|\\
  &\leq \frac{1}{2} \sum_{k=1}^{t-1} \alpha_k \|\nabla f_k(\boldsymbol{w}_{S,k})\| + \frac{1}{2} \sum_{k=1}^{t-1} \alpha_k \|\nabla f_{k'}(\boldsymbol{w}_{{S^{(i)}},k}) \|~.
\end{align*}
Plugging this result back into~\eqref{eq:gradient_track_int} completes the proof of the first statement.
The second statement comes from Lemma~\ref{lem:E_path_bounded_by_R1_and_noise} with $\alpha_t = c/t$.
\end{proof}
Next, we need the following statement to prove our stability bound.
\begin{prop}[Bernstein-type inequality]
\label{prop:bernstein}
Let $Z$ be a zero-mean real-valued r.v., such that $|Z| \leq b$ and $\E[Z^2] \leq \sigma^2$.
Then for all $|c| \leq \frac{1}{2 b}$, we have that
$
\E\br{e^{c Z}} \leq e^{c^2 \sigma^2}~.
$
\end{prop}
\begin{proof}
Stated inequality is a consequence of a Bernstein-type inequality for moment generating functions, Theorem 2.10 in~\cite{boucheron2013concentration}.
Observe that zero-centered r.v. $Z$ bounded by $b$ satisfies Bernstein's condition, that is
\[
|\E[(Z - \E[Z])^q]| \leq \frac{q!}{2} \sigma^2 b^{k-2} \qquad \text{for all integers } q \geq 3~.
\]
This in turn satisfies condition for Bernstein-type inequality stating that
\[
\E\br{\exp\pr{c (Z - \E[Z])}} \leq \exp\pr{ \frac{c^2 \sigma^2/2}{1 - b |c|} }~.
\]
Choosing $|c| \leq \frac{1}{2 b}$ verifies the statement.
\end{proof}
Now we are ready to prove Theorem~\ref{thm:sgd_stab_nonconvex}, which bounds the $\epsilon(\mathcal{D}, \boldsymbol{w}_1)$-on-average  stability of \ac{SGD}.
\begin{proof}[Proof of Theorem~\ref{thm:sgd_stab_nonconvex}.]
For brevity denote
\[
r := \E_{S,A}\br{R(A_S)}
\] and
\[
\Delta_t(S,z) := \E_A\br{\delta_t(S,z) \ | \ \delta_{t_0}(S,z) = 0}~.
\]
By Lemma~\ref{lem:stab_decompose}, for all $t_0 \in [m]$,
\begin{align}
\E_{S,z}\E_A\br{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)}
\leq L \E_{S,z}\br{ \Delta_T(S,z) } + r \frac{t_0}{m}~. \label{eq:nonconvex_t0_tuning_bound}
\end{align}
Most of the proof is dedicated to bounding the first term in~\eqref{eq:nonconvex_t0_tuning_bound}.
We deal with this similarly as in~\cite{hardt2016train}.
Specifically, we state the bound on $\Delta_T(S,z)$ by using a recursion.
In our case, however, we also have an expectation w.r.t.\ the data, and to avoid complications with dependencies, we first unroll the recursion for the random quantities, and only then take the expectation.
At this point the proof crucially relies on the product of exponentials arising from the recursion, and all relevant random quantities end up inside of them.
We alleviate this by Proposition~\ref{prop:bernstein}.
Finally, we conclude by minimizing
~\eqref{eq:nonconvex_t0_tuning_bound}
w.r.t.\ $t_0$.
Thus we have three steps: 1) recursion, 2) bounding $\E[\exp(\cdots)]$, and 3) tuning of $t_0$.
\paragraph{1) Recursion.}
We begin by stating the bound on $\Delta_T(S,z)$ by recursion.
Thus we will first state the bound on $\Delta_{t+1}(S,z)$ in terms of $\Delta_t(S,z)$, and other relevant quantities and then unravel the recursion.
As in the convex case,
we distinguish two cases: 1) \ac{SGD} encounters the perturbed point at step $t$, that is $t = i$, and 2) the current point is the same in $S$ and ${S^{(i)}}$, so $t \neq i$.
For the first case, we will use worst-case boundedness of $G_t$, Corollary~\ref{cor:G_boundedness}, that is, $\|G_t(\boldsymbol{w}_{S,t}) - G_t(\boldsymbol{w}_{{S^{(i)}},t})\| \leq \delta_t(S,z) + 2 \alpha_t L~.$
To handle the second case we will use Lemma~\ref{lem:expansiveness_nonconvex}, namely,
\begin{equation*}
  \left\| G_t(\boldsymbol{w}_{S,t}) - G_t(\boldsymbol{w}_{{S^{(i)}},t}) \right\| \leq \pr{1 + \alpha_t \xi_t(S,z) } \delta_t(S,z)~.
\end{equation*}
In addition, as a safety measure we will also take into account that the gradient update rule is at most $(1 + \alpha_t \beta)$-expansive by Lemma~\ref{lem:worst-case-expansive}.
So we will work with the function $\psi_t(S,z) := \min\cbr{\xi_t(S,z), \beta}$ instead of $\xi_t(S,z)$.
and decompose the expectation w.r.t.\ $A$ for a step $t$.
Noting that \ac{SGD} encounters the perturbed example with probability $\frac{1}{m}$,
\begin{align}
\Delta_{t+1}(S,z) &\leq \pr{1 - \frac{1}{m}} \pr{1 + \alpha_t \psi_t(S,z) } \Delta_t(S,z)
+ \frac{1}{m} \pr{ 2 \alpha_t L + \Delta_t(S,z) }~ \nonumber\\
&= \pr{1 + \pr{1 - \frac{1}{m}} \alpha_t \psi_t(S,z) } \Delta_t(S,z)
+ \frac{2 \alpha_t L}{m} \nonumber\\
&\leq \exp\pr{ \alpha_t \psi_t(S,z) } \Delta_t(S,z)
+ \frac{2 \alpha_t L}{m}~, \label{eq:nonconvex_recursion}
\end{align}
where the last inequality follows from $1 + x \leq \exp(x)$.
This inequality is not overly loose for $x \in [0,1]$, and, in our case it becomes instrumental in handling the recursion.

Now, observe that relation $x_{t+1} \leq a_t x_t + b_t$ with $x_{t_0} = 0$ unwinds from $T$ to $t_0$ as
$ x_T \leq \sum_{t=t_0+1}^T b_t \prod_{k=t+1}^T a_k$.
Consequently, having $\Delta_{t_0}(S,z) = 0$, we unwind~\eqref{eq:nonconvex_recursion} to get
\begin{align}
  \Delta_T(S,z)
\leq &\sum_{t=t_0+1}^T \pr{ \prod_{k=t+1}^T \exp\pr{\frac{c \psi_k(S,z)}{k} } } \frac{2 c L}{m t} \nonumber \\
= &\sum_{t=t_0+1}^T \exp \pr{ c \sum_{k=t+1}^T \frac{\psi_k(S,z)}{k} } \frac{2 c L}{m t}~. \label{eq:end_up_with_exponential}
\end{align}
\paragraph{2) Bounding $\E[\exp(\cdots)]$.}
We take expectation w.r.t.\ $S$ and $z$ on both sides and
focus on the expectation of the exponential in~\eqref{eq:end_up_with_exponential}.
First, introduce $\mu_k := \E_{S,z}[\psi_k(S,z)]$, and proceed as
\begin{align}
&\E_{S,z}\br{ \exp \pr{c \sum_{k=t+1}^T \frac{\psi_k(S,z)}{k} } }
= \E_{S,z}\br{ \exp \pr{c \sum_{k=t+1}^T \frac{\psi_k(S,z) - \mu_k}{k} } }
\exp\pr{ c \sum_{k=t+1}^T \frac{\mu_k}{k} }~. \label{eq:exp_split}
\end{align}
Observe that zero-mean version of $\psi_k(S,z)$ is bounded as
\begin{align*}
  \sum_{k=t+1}^T \frac{|\psi_k(S,z) - \mu_k|}{k} \leq 2 \beta \ln(T)~,
\end{align*}
and assume the setting of $c$ as $c \leq \frac{1}{2 (2 \beta \ln(T))^2}$.
By Proposition~\ref{prop:bernstein}, we have
\begin{align*}
  \E\br{ \exp \pr{c \sum_{k=t+1}^T \frac{\psi_k(S,z) - \mu_k}{k} } }
\leq &\exp \pr{c^2 \E\br{\pr{\sum_{k=t+1}^T \frac{\psi_k(S,z) - \mu_k}{k} }^2}}\\
= &\exp \pr{\frac{c}{2} \E\br{\pr{\frac{1}{2 \beta \ln(T)} \sum_{k=t+1}^T \frac{\psi_k(S,z) - \mu_k}{k} }^2}}\\
\leq &\exp \pr{ \frac{c}{2} \E\br{\abs{\sum_{k=t+1}^T \frac{\psi_k(S,z) - \mu_k}{k} }}}\\
\leq &\exp \pr{ \frac{c}{2} \sum_{k=t+1}^T \frac{\E\br{|\psi_k(S,z) - \mu_k|} }{k} }\\
\leq &\exp \pr{ c \sum_{k=t+1}^T \frac{ \mu_k }{k} }~.
\end{align*}
Getting back to~\eqref{eq:exp_split} we conclude that
\begin{align}
\label{eq:E_exp_bound}
 \E_{S,z}\br{ \exp \pr{c \sum_{k=t+1}^T \frac{\psi_k(S,z)}{k} } } \leq \exp \pr{ c \sum_{k=t+1}^T \frac{ 2 \mu_k }{k} }~.
\end{align}
Next, we give an upper-bound on $\mu_k$, that is
$
\mu_k \leq \min\cbr{\beta, \E_{S,z}[\xi_k(S,z)]}
$.
Finally, we bound $\E_{S,z}[\xi_k(S,z)]$ using the second result of Lemma~\ref{lem:expansiveness_nonconvex}, which holds for any $k \in [T]$,
to get that $\mu_k \leq \gamma$, with $\gamma$ defined in the statement of the theorem.
\paragraph{3) Tuning of $t_0$.}
Now we turn our attention back to~\eqref{eq:end_up_with_exponential}.
Considering that we took an expectation w.r.t.\ the data, we use~\eqref{eq:E_exp_bound} and the fact that $\mu_k \leq \gamma$ to get that
\begin{align*}
\E_{S,z}[\Delta_T(S,z)] &\leq \sum_{t=t_0+1}^T \exp \pr{ 2 c \gamma \sum_{k=t+1}^T \frac{ 1 }{k} }  \frac{2 c L}{m t}\\
&\leq \sum_{t=t_0+1}^T \exp \pr{2 c \gamma \ln\pr{\frac{T}{t}} } \frac{2 c L}{m t}\\
&= \frac{2 c L}{m} \pr{T^{2 c \gamma}} \sum_{t=t_0+1}^T t^{-2 c \gamma - 1}\\
&\leq \frac{1}{2 c \gamma }\frac{2 c L}{m} \pr{\frac{T}{t_0}}^{2 c \gamma}~.
\end{align*}
Plug the above into~\eqref{eq:nonconvex_t0_tuning_bound} to get
\begin{align}
\E_{S,z}\E_A\br{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)}
\leq \frac{L^2}{ \gamma m } \pr{\frac{T}{t_0}}^{2 c \gamma}  + r \frac{t_0}{m}~. \label{eq:bound_to_minimize_wrt_t0}
\end{align}
Let $q = 2 c \gamma$. Then, setting
\[
t_0 = \pr{\frac{2 c L^2}{r}}^{\frac{1}{1 + q}} T^{\frac{q}{1 + q}}
\]
minimizes~\eqref{eq:bound_to_minimize_wrt_t0}.
Plugging $t_0$ back we get that~\eqref{eq:bound_to_minimize_wrt_t0} equals to
\[
\frac{1 + \frac{1}{q}}{m} \pr{2 c L^2}^{\frac{1}{1 + q}} (r T)^{\frac{q}{1 + q}}~.
\]
This completes the proof.
\end{proof}
\subsubsection{Optimistic Rates for Learning with Non-convex Loss Functions}
Next we will prove an optimistic bound based on Theorem~\ref{thm:sgd_stab_nonconvex}, in other words, the bound that demonstrates fast convergence rate subject to the vanishing empirical risk.
First we will need the following technical statement.
\begin{lemma}{\cite[Lemma~7.2]{cucker2007learning}}
\label{lem:poly_root_bound}
Let $c_1, c_2, \ldots, c_l > 0$ and $s > q_1 > q_2 > \ldots > q_{l-1} > 0$.
Then the equation
\[
x^s - c_1 x^{q_1} - c_2 x^{q_2} - \cdots - c_{l-1} x^{q_{l-1}} - c_l = 0
\]
has a unique positive solution $x^{\star}$. In addition,
\[
x^{\star} \leq \max\left\{(l c_1)^\frac{1}{s-q_1}, (l c_2)^\frac{1}{s-q_2}, \cdots, (l c_{l-1})^\frac{1}{s-q_{l-1}}, (l c_l)^\frac{1}{s} \right\}.
\]
\end{lemma}
Next we prove a useful technical lemma similarly as in~\cite[Lemma~7]{orabona2014simultaneous}.
\begin{lemma}
\label{lem:nonparametric_optimistic_solve}
Let $a, c > 0$ and $0 < \alpha < 1$. Then the inequality
\[
x - a x^{\alpha} - c \leq 0
\]
implies
\[
x \leq \max\cbr{ 2^{\frac{\alpha}{1 - \alpha}} a^{\frac{1}{1 - \alpha}}, \pr{2 c}^{\alpha} a } + c~.
\]
\end{lemma}
\begin{proof}
Consider a function $h(x) = x - a x^{\alpha} - c$.
Applying Lemma~\ref{lem:poly_root_bound} with $s=1$, $l=2$, $c_1=a$, $c_2=c$, and $q_1 = \alpha$ we get that $h(x) = 0$ has a unique positive solution $x^{\star}$ and
\begin{equation}
\label{eq:xstar}
x^{\star} \leq \max\cbr{ (2 a)^{\frac{1}{1 - \alpha}}, 2 c }~.
\end{equation}
Moreover, the inequality $h(x) \leq 0$ is verified for $x=0$, and $\lim_{x \rightarrow +\infty} h(x) = +\infty$, so we have that $h(x) \leq 0$ implies $x \leq x^{\star}$.
Now, using this fact and the fact that $h(x^{\star}) = 0$, we have that
\[
x \leq x^{\star} = a \pr{x^{\star}}^{\alpha} + c~,
\]
and upper-bounding $x^{\star}$ by~\eqref{eq:xstar} we finally have
\[
x \leq a \max\cbr{ (2 a)^{\frac{\alpha}{1 - \alpha}}, \pr{2 c}^{\alpha} } + c~,
\]
which completes the proof.
\end{proof}
\begin{proof}[Proof of Corollary~\ref{cor:nonconvex_optimistic}.]
Consider Theorem~\ref{thm:sgd_stab_nonconvex} and observe that it verifies condition of Lemma~\ref{lem:nonparametric_optimistic_solve} with $x = \E_{S,A}\br{R(A_S)}$, $c = \E_{S,A}\br{\wh{R}_S(A_S)}$, $\alpha = \frac{c \gamma}{1 + c \gamma}$, and
\[
a = \frac{1 + \frac{1}{c \gamma}}{m} \pr{2 c L^2}^{\frac{1}{1 + c \gamma}} T^{\frac{c \gamma}{1 + c \gamma}}~.
\]
Note that $\alpha/(1-\alpha) = c \gamma$ and $1/(1-\alpha)=1+c \gamma$.
Then, we obtain that
\begin{align*}
&\E_{S,A}\br{R(A_S) - \wh{R}_S(A_S)}\\
&\leq \max\Bigg\{
2^{c \gamma} \pr{\frac{1 + \frac{1}{c \gamma}}{m}}^{1 + c \gamma} \pr{2 c L^2} T^{c \gamma},
\pr{2 \E_{S,A}\br{\wh{R}_S(A_S)}}^{\frac{c \gamma}{1 + c \gamma}}
\pr{\frac{1 + \frac{1}{c \gamma}}{m} \pr{2 c L^2}^{\frac{1}{1 + c \gamma}} T^{\frac{c \gamma}{1 + c \gamma}}}
\Bigg\}\\
&= \max\Bigg\{
\pr{2 + \frac{2}{c \gamma}}^{1 + c \gamma} \pr{c L^2} \pr{\frac{T^{c \gamma}}{m^{1 + c \gamma}}},
\frac{1 + \frac{1}{c \gamma}}{m} \pr{2 c L^2}^{\frac{1}{1 + c \gamma}} \pr{2 \E_{S,A}\br{\wh{R}_S(A_S)} \cdot T}^{\frac{c \gamma}{1 + c \gamma}}
\Bigg\}~.
\end{align*}
This completes the proof.
\end{proof}
\section{Proofs}
\label{sec:proofs}
In this section we present proofs of all the statements.
\begin{proof}[Proof of Theorem~\ref{thm:onaverage_stab_generalization}]
Indicate by $S=\{z_i\}_{i=1}^m$ and $S'=\{z'_i\}_{i=1}^m$ independent training sets sampled i.i.d.\ from $\mathcal{D}$, and let ${S^{(i)}}=\cbr{z_1, \ldots, z_{i-1}, z_i', z_{i+1}, \ldots, z_m}$, such that $z_i' \distas{\iid} \mathcal{D}$.
We relate expected empirical risk and expected risk by
  \begin{align*}
    \E_S \E_A\br{ \wh{R}_S(A_S) }
=   &\E_S \E_A\br{ \frac{1}{m} \sum_{i=1}^m f(A_S, z_i) }\\
=   &\E_{S,S'} \E_A\br{ \frac{1}{m} \sum_{i=1}^m f(A_{{S^{(i)}}}, z'_i) }\\
=   &\E_{S,S'} \E_A\br{ \frac{1}{m} \sum_{i=1}^m f(A_S, z'_i) } - \delta\\
=   &\E_{S} \E_A\br{ R(A_S) } - \delta~,
  \end{align*}
where
\begin{align*}
  \delta = &\E_{S,S'} \E_A\br{ \frac{1}{m} \sum_{i=1}^m \pr{ f(A_S, z'_i) - f(A_{{S^{(i)}}}, z'_i) } }\\
  = &\frac{1}{m} \sum_{i=1}^m \E_{S, z'_i} \E_A\br{ f(A_S, z'_i) - f(A_{{S^{(i)}}}, z'_i) }~.
\end{align*}
Renaming $z'_i$ as $z$ and taking $\sup$ over $i$ we get that
\begin{align*}
\delta \leq \sup_{i \in [m]}\cbr{ \E_{S, z} \E_A\br{ f(A_S, z) - f(A_{{S^{(i)}}}, z) } }~.
\end{align*}
This completes the proof.
\end{proof}
\subsection{Preliminaries}
We say that the \ac{SGD} gradient update rule is an operator $G_t~:~\mathcal{H}~\mapsto~\mathcal{H}$, such that
\[
  G_t(\boldsymbol{w}) := \boldsymbol{w} - \alpha_t \nabla f(\boldsymbol{w}, z_{i_t})~,
\]
and it is also a function of the training set $S$ and a random index set $I$.
Then, $\boldsymbol{w}_{t+1} = G_t(\boldsymbol{w}_t)$, throughout $t = 1, \ldots, T$.
Recall the use of notation $\boldsymbol{w}_{S,t}$ to indicate the output of \ac{SGD} ran on a training set $S$, at step $t$,
and define
\[
\delta_t(S,z) := \|\boldsymbol{w}_{S,t} - \boldsymbol{w}_{{S^{(i)}},t}\|~.
\]
Next, we summarize a few instrumental facts about $G_t$ and few statements about the loss functions used in our proofs.
\begin{definition}[Expansiveness]
A gradient update rule is $\eta$-expansive if for all $\boldsymbol{w}, \boldsymbol{v}$,
\[
\|G_t(\boldsymbol{w}) - G_t(\boldsymbol{v})\| \leq \eta \|\boldsymbol{w} - \boldsymbol{v}\|~.
\]
\end{definition}
The following lemma characterizes expansiveness for the gradient update rule under different assumptions on $f$.
\begin{lemma}[Lemma 3.6 in \cite{hardt2016train}]
  \label{lem:worst-case-expansive}
  Assume that $f$ is $\beta$-smooth. Then, we have that:
\begin{enumerate}
  \item[1)] $G_t$ is $(1+\alpha_t \beta)$-expansive,
  \item[2)] If $f$ in addition is convex, then, for any $\alpha_t \leq \frac{2}{\beta}$, the gradient update rule $G_t$ is $1$-expansive.
\end{enumerate}
\end{lemma}
An important consequence of $\beta$-smoothness of $f$ is self-boundedness~\cite{shalev2014understanding}, which we will use on many occasions.
\begin{lemma}[Self-boundedness]
\label{lem:self_boundedness}
For $\beta$-smooth non-negative function $f$ we have that
\[
\|\nabla f(\boldsymbol{w}, z)\| \leq \sqrt{2 \beta f(\boldsymbol{w}, z)}~.
\]
\end{lemma}
Self-boundedness in turn implies the following boundedness of a gradient update rule.
\begin{cor}
\label{cor:G_boundedness}
Assume that $f$ is $\beta$-smooth and non-negative.
Then,
\[
\|\boldsymbol{w} - G_t(\boldsymbol{w})\| = \alpha_t \|\nabla f(\boldsymbol{w}, z_{j_t})\| \leq \alpha_t \min\cbr{ \sqrt{2 \beta f(\boldsymbol{w}, z_{j_t})}, L }~.
\]
\end{cor}
\begin{proof}
By Lemma~\ref{lem:self_boundedness}
  \begin{align*}
    \|\alpha_t \nabla f(\boldsymbol{w}, z_{j_t})\| \leq \alpha_t \sqrt{2 \beta f(\boldsymbol{w}, z_{j_t})}~,
  \end{align*}
and also by Lipschitzness of $f$, $\|\alpha_t \nabla f(\boldsymbol{w}, z_{j_t})\| \leq \alpha_t L$.
\end{proof}
Next we introduce a bound that relates the risk of the output at step $t$ to the risk of the initialization point $\boldsymbol{w}_1$ through the variance of the gradient.
Given an appropriate choice of step size, this bound will be crucial at stating stability bounds that depend on the risk at $\boldsymbol{w}_1$.
The proof idea is similar to the one of~\cite{ghadimi2013stochastic}.
In particular, it does not require convexity of the loss function.
\begin{lemma}
\label{lem:stationary_point}
Suppose \ac{SGD} is ran with step sizes $\alpha_1, \ldots, \alpha_{t-1} \leq \frac{1}{\beta}$ w.r.t.\ the $\beta$-smooth loss $f$.
Then we have that
  \begin{align}
  \sum_{k=1}^{t-1} \pr{\alpha_k - \frac{\alpha_k^2 \beta}{2}} \E_S\br{ \|\nabla R(\boldsymbol{w}_k)\|^2 }
\leq R(\boldsymbol{w}_1) - R(\boldsymbol{w}_t) + \frac{\beta}{2} \sum_{k=1}^{t-1} \alpha_k^2 \E_S\br{ \|\nabla f(\boldsymbol{w}_k, z_{j_k}) - \nabla R(\boldsymbol{w}_k)\|^2 }~.
\end{align}
\end{lemma}
\begin{proof}
  %
  For brevity denote $f_k(\boldsymbol{w}) \equiv f(\boldsymbol{w}, z_{j_k})$.
By $\beta$-smoothness of $R$ and recalling that the \ac{SGD} update rule $\boldsymbol{w}_{k+1} = \boldsymbol{w}_k - \alpha_k \nabla f_k(\boldsymbol{w}_k)$, we have
\begin{align*}
  R(\boldsymbol{w}_{k+1}) - R(\boldsymbol{w}_k) &\leq \nabla R(\boldsymbol{w}_k)^{\top} \pr{\boldsymbol{w}_{k+1} - \boldsymbol{w}_k} + \frac{\beta}{2} \|\boldsymbol{w}_{k+1} - \boldsymbol{w}_k\|^2 \\
  &= - \alpha_k \nabla R(\boldsymbol{w}_k)^{\top} \nabla f_k(\boldsymbol{w}_k) + \frac{\beta \alpha_k^2}{2} \|\nabla f_k(\boldsymbol{w}_k)\|^2 \\
  &= - \alpha_k \nabla R(\boldsymbol{w}_k)^{\top} \nabla f_k(\boldsymbol{w}_k) + \frac{\beta \alpha_k^2}{2} \|\nabla f_k(\boldsymbol{w}_k) - \nabla R(\boldsymbol{w}_k) + \nabla R(\boldsymbol{w}_k)\|^2 \\
  &= - \alpha_k \nabla R(\boldsymbol{w}_k)^{\top} \nabla f_k(\boldsymbol{w}_k)\\
&+ \frac{\beta \alpha_k^2}{2} \Big( \|\nabla f_k(\boldsymbol{w}_k) - \nabla R(\boldsymbol{w}_k)\|^2 + \|\nabla R(\boldsymbol{w}_k)\|^2\\
&- 2 \pr{\nabla f_k(\boldsymbol{w}_k) - \nabla R(\boldsymbol{w}_k)}^{\top} \nabla R(\boldsymbol{w}_k) \Big)\\
&=- \pr{\alpha_k + \alpha_k^2 \beta} \nabla R(\boldsymbol{w}_k)^{\top} \nabla f_k(\boldsymbol{w}_k)\\
&+ \frac{3 \alpha^2 \beta}{2} \|\nabla R(\boldsymbol{w}_k)\|^2 +
\frac{\beta \alpha_k^2}{2} \|\nabla f_k(\boldsymbol{w}_k) - \nabla R(\boldsymbol{w}_k)\|^2~.
\end{align*}
Taking expectation w.r.t. $S$ on both sides, recalling that $\E_{z_k}[\nabla f_k(\boldsymbol{w}_k)] = \nabla R(\boldsymbol{w}_k)$ and rearranging terms we get
\[
\pr{\alpha_k - \frac{\alpha^2 \beta}{2}} \E\br{ \|\nabla R(\boldsymbol{w}_k)\|^2 } \leq R(\boldsymbol{w}_k) - R(\boldsymbol{w}_{k+1}) + \frac{\beta \alpha_k^2}{2} \E\br{ \|\nabla f_k(\boldsymbol{w}_k) - \nabla R(\boldsymbol{w}_k)\|^2 }~,
\]
and summing above over $k = 1, \ldots, t-1$ we get the statement.
\end{proof}
\begin{lemma}
  \label{lem:E_path_bounded_by_R1_and_noise}
Suppose \ac{SGD} is ran with step sizes $\alpha_1, \ldots, \alpha_{t-1} \leq \frac{1}{\beta}$ on the $\beta$-smooth loss $f$.
Assume that the variance of stochastic gradients obeys
\[
\E_{S, z}\br{\left\|  \nabla f(\boldsymbol{w}_{S,k}, z) - \nabla R(\boldsymbol{w}_{S,k}) \right\|^2} \leq \sigma^2 \quad \forall k \in [T]~.
\]
Then we have that
  \begin{equation*}
    \E_S\br{\sum_{k=1}^{t-1} \alpha_k \|\nabla f(\boldsymbol{w}_{S,k}, z_k)\| }
\leq 2 \sqrt{ \pr{\sum_{k=1}^{t-1} \alpha_k} \pr{ R(\boldsymbol{w}_1) - \inf_{\boldsymbol{w} \in \mathcal{H}} R(\boldsymbol{w}) + \frac{\beta \sigma^2}{2} \sum_{k=1}^{t-1} \alpha_k^2} } + \sigma \sum_{k=1}^{t-1} \alpha_k~.
  \end{equation*}
\end{lemma}
\begin{proof}
First we perform the decomposition,
\begin{align}
 \E_S\br{ \sum_{k=1}^{t-1} \alpha_k \|\nabla f(\boldsymbol{w}_{S,k}, z_k)\| } &=
\sum_{k=1}^{t-1} \alpha_k \E_{S}\br{\left\|  \nabla R(\boldsymbol{w}_{S,k}) \right\|}
+ \sum_{k=1}^{t-1} \alpha_k \E_{S}\br{\left\|  \nabla f(\boldsymbol{w}_{S,k}, z_k) - \nabla R(\boldsymbol{w}_{S,k}) \right\|} \nonumber \\
&\leq  \sum_{k=1}^{t-1} \alpha_k \E_{S}\br{\left\|  \nabla R(\boldsymbol{w}_{S,k}) \right\|} + \sigma \sum_{k=1}^{t-1} \alpha_k~. \label{eq:path_variance_decomposition}
\end{align}
Introduce
\[
Q_t := \sum_{k=1}^{t-1} \pr{\alpha_k - \frac{\alpha_k^2 \beta}{2}}~.
\]
Now we invoke the stationary-point argument to bound the first term above as
\begin{align}
\sum_{k=1}^{t-1} \alpha_k \E_{S}\br{ \sqrt{\left\|  \nabla R(\boldsymbol{w}_k) \right\|^2} }
&\leq \sum_{k=1}^{t-1} \frac{\pr{1 - \frac{\alpha_k \beta}{2}}}{\pr{1 - \frac{\alpha_k \beta}{2}}} \cdot \alpha_k \sqrt{\E_{S}\br{ \left\|  \nabla R(\boldsymbol{w}_k) \right\|^2 } } \tag{By Jensen's inequality}\\
&\leq 2 \sum_{k=1}^{t-1} \pr{\alpha_k - \frac{\alpha_k^2 \beta}{2}} \sqrt{\E_{S}\br{ \left\|  \nabla R(\boldsymbol{w}_k) \right\|^2 } } \tag{Assuming that $\alpha_k \leq \frac{1}{\beta}$}\\
&=
\frac{2 Q_t}{Q_t} \sum_{k=1}^{t-1} \pr{\alpha_k - \frac{\alpha_k^2 \beta}{2}} \sqrt{\E_{S}\br{ \left\|  \nabla R(\boldsymbol{w}_k) \right\|^2 } } \\
&\leq 2 \sqrt{Q_t} \sqrt{ \sum_{k=1}^{t-1} \pr{\alpha_k - \frac{\alpha_k^2 \beta}{2}} \E_{S}\br{ \left\|  \nabla R(\boldsymbol{w}_k) \right\|^2 } } \tag{By Jensen's inequality}\\
    &\leq 2 \sqrt{Q_t} \sqrt{ R(\boldsymbol{w}_1) - R(\boldsymbol{w}_t) + \frac{\beta \sigma^2}{2} \sum_{k=1}^{t-1} \alpha_k^2 }~. \tag{By Lemma~\ref{lem:stationary_point}}
\end{align}

Combining this with~\eqref{eq:path_variance_decomposition} gives
\begin{align}
    \E_S\br{\sum_{k=1}^{t-1} \alpha_k \|\nabla f(\boldsymbol{w}_{S,k}, z_k)\| }
&\leq 2 \sqrt{ \pr{\sum_{k=1}^{t-1} \alpha_k} \pr{ R(\boldsymbol{w}_1) - \inf_{\boldsymbol{w} \in \mathcal{H}} R(\boldsymbol{w}) + \frac{\beta \sigma^2}{2} \sum_{k=1}^{t-1} \alpha_k^2} } + \sigma \sum_{k=1}^{t-1} \alpha_k~,
\end{align}
which completes the proof.
\end{proof}

The following lemma is similar to Lemma~3.11 of~\cite{hardt2016train}, and is instrumental in bounding the stability of \ac{SGD}.
However, we make an adjustment and state it in expectation over the data.
Note that it does not require convexity of the loss function.
\begin{lemma}
\label{lem:stab_decompose}
Assume that the loss function $f(\cdot, z) \in [0, 1]$ is $L$-Lipschitz for all $z$.
Then, for every $t_0 \in \{0, 1, 2, \ldots m\}$ we have that,
\begin{align}
  &\E_{S,z}\E_A\br{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)} \label{eq:stab_decompose_lhs}\\
&\leq L \E_{S,z}\br{\E_{A}\br{\delta_T(S,z) \ \middle| \ \delta_{t_0}(S,z) = 0}}
+ \E_{S,A}\br{R(A_S)} \frac{t_0}{m}~. \label{eq:stab_decompose_1_E_cond}
\end{align}
\end{lemma}
\begin{proof}
We proceed with elementary decomposition, Lipschitzness of $f$, and using the fact that $f$ is non-negative to have that
\begin{align}
  f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)
  &= \pr{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)} \ind{\delta_{t_0}(S,z) = 0} \label{eq:E_cond_stab_1} \\
  &+ \pr{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)} \ind{\delta_{t_0}(S,z) \neq 0} \nonumber \\
  &\leq L \delta_T(S,z) \ind{\delta_{t_0}(S,z) = 0}
  + f(\boldsymbol{w}_{S,T}, z) \ind{\delta_{t_0}(S,z) \neq 0}~. \label{eq:E_cond_stab_2}
\end{align}
Taking expectation w.r.t. algorithm randomization, we get that
\begin{align}
  \E_A\br{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)}
  &\leq L \E_A\br{\delta_T(S,z) \ind{\delta_{t_0}(S,z) = 0} }\\
  &+ \E_A\br{f(\boldsymbol{w}_{S,T}, z) \ind{\delta_{t_0}(S,z) \neq 0}}~. \label{eq:delta_neq_0_term}
\end{align}
Recall that $i \in [m]$ is the index where $S$ and $S^{(i)}$ differ, and introduce a random variable $\tau_A$ taking on the index of the first time step where \ac{SGD} uses the example $z_i$ or a replacement $z$.
Note also that $\tau_A$ does not depend on the data.
When $\tau_A > t_0$, then it must be that $\delta_{t_0}(S,z) = 0$, because updates on both $S$ and $S^{(i)}$ are identical until $t_0$.
A consequence of this is that $\ind{\delta_{t_0}(S,z) \neq 0} \leq \ind{\tau_A \leq t_0}$.
Thus the rightmost term in~\eqref{eq:delta_neq_0_term} is bounded as
\begin{align*}
  \E_A\br{f(\boldsymbol{w}_{S,T}, z) \ind{\delta_{t_0}(S,z) \neq 0}} \leq \E_A\br{f(\boldsymbol{w}_{S,T}, z) \ind{\tau_A \leq t_0} }~.
\end{align*}
Now, focus on the r.h.s.\ above.
Recall that we assume randomization by sampling from the uniform distribution over $[m]$ without replacement, and denote a realization by $\cbr{j_i}_{i=1}^m$.
Then, we can always express our randomization as permutation function $\pi_A(S) = \cbr{z_{j_i}}_{i=1}^m$.
In addition, introduce an algorithm $\text{GD} : \mathcal{Z}^m \mapsto \mathcal{H}$, which is identical to $A$, except that it passes over the training set $S$ sequentially without randomization.
That said, we have that
\begin{align*}
  \E_A\br{f(\boldsymbol{w}_{S,T}, z) \ind{\tau_A \leq t_0} }
= \E_A\br{f(\text{GD}_{\pi_A(S)}, z) \ind{\tau_A \leq t_0} }~,
\end{align*}
and taking expectation over the data,
\begin{align*}
  \E_{S,z}\br{\E_A\br{f(\boldsymbol{w}_{S,T}, z) \ind{\tau_A \leq t_0} }} = \E_A\br{\E_{S,z}\br{f(\text{GD}_{\pi_A(S)}, z) } \ind{\tau_A \leq t_0} }~.
\end{align*}
Now observe that for any realization of $A$, $\E_{S,z}\br{f(\text{GD}_{\pi_A(S)}, z) } = \E_A \E_{S,z}\br{f(A_S, z) }$ because expectation w.r.t.\ $S$ and $z$ does not change under our randomization
\footnote{Strictly speaking we could omit $\E_A[\cdot]$ and consider \emph{any} randomization by reshuffling, but we keep expectation for the sake of clarity.}.
Thus, we have that
\begin{align*}
  \E_A\br{\E_{S,z}\br{f(\text{GD}_{\pi_A(S)}, z) } \ind{\tau_A \leq t_0} } = \E_{S,A}\br{R(A_S)} \P(\tau_A \leq t_0)~.
\end{align*}
Now assuming that $\tau_A$ is uniformly distributed over $[m]$ we have that
\[
\P\pr{\tau_A \leq t_0} = \frac{t_0}{m}~.
\]
Putting this together with~\eqref{eq:E_cond_stab_1} and~\eqref{eq:E_cond_stab_2}, we finally get that
\begin{align*}
  \E_{S,z}\E_A\br{f(\boldsymbol{w}_{S,T}, z) - f(\boldsymbol{w}_{{S^{(i)}},T}, z)}
  &\leq L \E_{S,z}\br{\E_{A}\br{\delta_T(S,z) \ind{\delta_{t_0}(S,z) = 0}}}
+ \E_{S,A}\br{R(A_S)} \frac{t_0}{m}\\
&\leq L \E_{S,z}\br{\E_{A}\br{\delta_T(S,z) \ \middle| \ \delta_{t_0}(S,z) = 0}}
+ \E_{S,A}\br{R(A_S)} \frac{t_0}{m}~.
\end{align*}
This completes the proof.
\end{proof}
We spend a moment to highlight the role of conditional expectation in~\eqref{eq:stab_decompose_1_E_cond}.
Observe that we could naively bound~\eqref{eq:stab_decompose_lhs} by the Lipschitzness of $f$, but Lemma~\ref{lem:stab_decompose} follows a more careful argument.
First note that $t_0$ is a free parameter.
The expected distance in~\eqref{eq:stab_decompose_1_E_cond} between \ac{SGD} outputs $\boldsymbol{w}_{S,t}$ and $\boldsymbol{w}_{{S^{(i)}},t}$ is conditioned on the fact that at step $t_0$ outputs of \ac{SGD} are still the same.
This means that the perturbed point is encountered after $t_0$.
Then, the conditional expectation should be a decreasing function of $t_0$: the later the perturbation occurs, the smaller deviation between $\boldsymbol{w}_{S,t}$ and $\boldsymbol{w}_{{S^{(i)}},t}$ we should expect.
Later we use this fact to minimize the bound~\eqref{eq:stab_decompose_1_E_cond} over $t_0$.
\input{convex.tex}
\input{nonconvex_stab.tex}
\begin{proof}[Proof of Proposition~\ref{prop:nonconvex_transfer}]
Consider minimizing the bound given by Corollary~\ref{cor:non_convex_R1} (in the submission file) over a discrete set of source hypotheses $\cbr{\boldsymbol{w}^{\text{src}}_k}_{k=1}^K$,
\begin{align}
&\min_{k \in [K]}\epsilon(\mathcal{D}, \boldsymbol{w}^{\text{src}}_k) \nonumber\\
&\leq \min_{k \in [K]} \mathcal{O}\pr{ \frac{1 + \frac{1}{c \gamma_k}}{m}  \pr{ R(\boldsymbol{w}^{\text{src}}_k) \cdot T}^{\frac{c \gamma_k}{1 + c \gamma_k}} }~,
\label{eq:nonconvex_TL_bound}
\end{align}
and let
\begin{align*}
&\gamma_k = \mathcal{O}\pr{ \E_{z \sim \mathcal{D}}\br{ \|\nabla^2 f(\boldsymbol{w}^{\text{src}}_k, z)\|_2 } + \sqrt{R(\boldsymbol{w}^{\text{src}}_k)} }~,\\
&\widehat{\gamma}_k = \frac{1}{m} \sum_{i=1}^m \|\nabla^2 f(\boldsymbol{w}^{\text{src}}_k, z_i)\|_2 + \sqrt{\wh{R}_S(\boldsymbol{w}^{\text{src}}_k)}~.
\end{align*}
By Hoeffding inequality, with high probability, we have that
$
| \gamma_k - \widehat{\gamma}_k | \leq \mathcal{O}\pr{ \frac{1}{\sqrt[4]{m}} }
$.
Now we further upper bound~\eqref{eq:nonconvex_TL_bound}
by upper bounding $R(\boldsymbol{w}^{\text{src}}_k)$ and apply union bound to get
\begin{align*}
&\min_{k \in [K]} \epsilon(\mathcal{D}, \boldsymbol{w}^{\text{src}}_k)\\
&\leq
\min_{k \in [K]} \mathcal{O}\pr{ \pr{1 + \frac{1}{c \widehat{\gamma}^-_k}}
\wh{R}_S(\boldsymbol{w}^{\text{src}}_k)^{\frac{c \widehat{\gamma}^+_k}{1 + c \widehat{\gamma}^+_k}}
\cdot
\frac{\sqrt{\log(K)}}{m^{\frac{1}{1 + c \widehat{\gamma}^+_k}} } }~,
\end{align*}
where $\widehat{\gamma}_k^{\pm} = \widehat{\gamma}_k \pm \frac{1}{\sqrt[4]{m}}$.
This completes the proof.
\end{proof}
\section{Related Work}
\label{sec:related}
Algorithmic stability has been a topic of interest in learning theory for a long time,
however, the modern approach on the relationship between stability and generalization goes back to the milestone work of~\cite{BousquetE02}.
They analyzed several notions of stability, which fall into two categories: distribution-free and distribution-dependent ones.
The first category is usually called \emph{uniform} stability and focuses on the intrinsic stability properties of an algorithm without regard to the data-generating distribution.
Uniform stability was used to analyze many algorithms, including regularized \ac{ERM}~\cite{BousquetE02}, randomized aggregation schemes~\cite{elisseeff2005stability}, and recently \ac{SGD} by \cite{hardt2016train,london2016generalization}, and \cite{poggio2011online}.
Despite the fact that uniform stability has been shown to be sufficient to guarantee learnability,
it can be too pessimistic, resulting in worst-case rates.

In this work we are interested in the data-dependent behavior of \ac{SGD}, thus
the emphasis will fall on the distribution-dependent notion of stability, known as \emph{on-average } stability, explored throughly in~\cite{shalev2010learnability}.
The attractive quality of this less restrictive stability type is that the resulting bounds are controlled by how stable the algorithm is under the data-generating distribution.
For instance, in~\cite{BousquetE02} and~\cite{devroye1979distribution}, the on-average  stability is related to the variance of an estimator.
In \cite[Sec. 13]{shalev2014understanding}, the authors show risk bounds that depend on the expected empirical risk of a solution to the regularized \ac{ERM}.
In turn, one can exploit this fact to state improved \emph{optimistic} risk bounds, for instance, ones that exhibit \emph{fast-rate} regimes~\cite{koren2015fast,gonen2017fast},
or even to
design enhanced algorithms that minimize these bounds in a data-driven way, e.g.\
by exploiting side information as in transfer~\cite{kuzborskij2013stability,Ben-DavidU13} and metric learning~\cite{perrot2015theoretical}.
Here, we mainly focus on the later direction in the context of \ac{SGD}: how stable is \ac{SGD} under the data-generating distribution given an initialization point?
We also touch the former direction by taking advantage of our data-driven analysis and show optimistic bounds as a corollary.

We will study the on-average  stability of \ac{SGD} for both convex and non-convex loss functions.
In the convex setting, we will relate stability to the risk at the initialization point, while
previous data-driven stability arguments usually consider minimizers of convex \ac{ERM} rather than a stochastic approximation~\cite{shalev2014understanding,koren2015fast}.
Beside convex problems, our work also covers the generalization ability of \ac{SGD} on non-convex problems.
Here, we borrow techniques of~\cite{hardt2016train} and extend them to the distribution-dependent setting.
That said, while bounds of~\cite{hardt2016train} are stated in terms of worst-case quantities, ours reveal new connections to the data-dependent second-order information.
These new insights also partially justify empirical observations in deep learning about the link between the curvature and the generalization error~\cite{hochreiter1997flat,keskar2017large,chaudhuri2017entropy}.
At the same time, our work is an alternative to the theoretical studies of neural network objective functions~\cite{choromanska2015loss,kawaguchi2016deep}, as we focus on the direct connection between the generalization and the curvature.

In this light, our work is also related to non-convex optimization by \ac{SGD}.
Literature on this subject typically studies rates of convergence to the stationary points~\cite{ghadimi2013stochastic,allen2016variance,reddi2016stochastic}, and ways to avoid saddles~\cite{ge2015escaping,lee2016gradient}.
However, unlike these works, and similarly to~\cite{hardt2016train}, we are interested in the generalization ability of \ac{SGD}, and thanks to the stability approach, involvement of stationary points in our analysis is not necessary.

Finally, we propose an example application of our findings in \ac{TL}.
For instance, by controlling the stability bound in a data-driven way, one can choose an initialization that leads to improved generalization.
This is related to \ac{TL} where one transfers from pre-trained models~\cite{kuzborskij2016fast,tommasi2013learning,pentina2014pac,Ben-DavidU13}, especially popular in deep learning due to its data-demanding nature~\cite{galanti2016theoretical}.
Literature on this topic is mostly focused on the \ac{ERM} setting and PAC-bounds, while our analysis of \ac{SGD} yields such guarantees as a corollary.

\section{Main Results}
\label{sec:results}
Before presenting our main results in this section, we discuss algorithmic details and assumptions.
We will study the following variant of \ac{SGD}:
given a training set $S = \{z_i\}_{i=1}^m \distas{\iid} \mathcal{D}^m$, step sizes $\cbr{\alpha_t}_{t=1}^T$, random indices $I = \{j_t\}_{t=1}^{T}$, and an initialization point $\boldsymbol{w}_1$, perform updates
\begin{align*}
  \boldsymbol{w}_{t+1} &= \boldsymbol{w}_t - \alpha_t \nabla f(\boldsymbol{w}_t, z_{j_t})~
\end{align*}
for $T \leq m$ steps.
Moreover we will use the notation $\boldsymbol{w}_{S,t}$ to indicate the output of \ac{SGD} ran on a training set $S$, at step $t$.
We assume that the indices in $I$ are sampled from the uniform distribution over $[m]$ \emph{without} replacement, and that this is the only source of randomness for \ac{SGD}.
In practice this corresponds to permuting the training set before making a pass through it, as it is commonly done in practical applications.
We also assume that the variance of stochastic gradients obeys
\begin{equation*}
\E_{S, z}\br{\left\|  \nabla f(\boldsymbol{w}_{S,t}, z) - \nabla R(\boldsymbol{w}_{S,t}) \right\|^2} \leq \sigma^2 \quad \forall t \in [T]~.
\end{equation*}
Next, we introduce statements about the loss functions $f$ used in the following.
\begin{definition}[Lipschitz $f$]
A loss function $f$ is $L$-Lipschitz if $\|\nabla f(\boldsymbol{w}, z)\| \leq L$, $\forall \boldsymbol{w} \in \mathcal{H}$ and $\forall z \in \mathcal{Z}$.
Note that this also implies that
$
|f(\boldsymbol{w}, z) - f(\boldsymbol{v}, z)| \leq L \|\boldsymbol{w} - \boldsymbol{v}\|~.
$
\end{definition}
\begin{definition}[Smooth $f$]
A loss function is $\beta$-smooth if $\forall \boldsymbol{w}, \boldsymbol{v} \in \mathcal{H}$ and $\forall z \in \mathcal{Z}$,
$
\|\nabla f(\boldsymbol{w}, z) - \nabla f(\boldsymbol{v}, z)\| \leq \beta \|\boldsymbol{w} - \boldsymbol{v}\|~,
$
which also implies
$
f(\boldsymbol{w}, z) - f(\boldsymbol{v}, z) \leq \nabla f(\boldsymbol{v}, z)^{\top} (\boldsymbol{w} - \boldsymbol{v}) + \frac{\beta}{2} \|\boldsymbol{w} - \boldsymbol{v}\|^2~.
$
\end{definition}
\begin{definition}[Lipschitz Hessians]
A loss function $f$ has a $\rho$-Lipschitz Hessian if $\forall \boldsymbol{w}, \boldsymbol{v} \in \mathcal{H}$ and $\forall z \in \mathcal{Z}$,
$
\|\nabla^2 f(\boldsymbol{w}, z) - \nabla^2 f(\boldsymbol{v}, z)\|_2 \leq \rho \|\boldsymbol{w} - \boldsymbol{v}\|~.
$
\end{definition}
The last condition is occasionally used in analysis of \ac{SGD}~\cite{ge2015escaping} and holds whenever $f$ has a bounded third derivative.
All presented theorems assume that the loss function used by \ac{SGD} is non-negative, Lipschitz, and $\beta$-smooth.
Examples of such commonly used loss functions are the logistic/softmax losses and neural networks with sigmoid activations.
Convexity of loss functions or Lipschitzness of Hessians will only be required for some results, and we will denote it explicitly when necessary.
Proofs for all the statements in this section are given in the supplementary material.
\subsection{Convex Losses}
First, we present a new and data-dependent stability result for convex losses.
\begin{theorem}
\label{thm:sgd_stab_convex}
Assume that $f$ is convex, and that \ac{SGD}'s step sizes satisfy $\alpha_t = \frac{c}{\sqrt{t}} \leq \frac{1}{\beta},~ \forall t \in [T]$.
Then \ac{SGD} is $\epsilon(\mathcal{D}, \boldsymbol{w}_1)$-on-average  stable with
\begin{equation*}

  \epsilon(\mathcal{D}, \boldsymbol{w}_1) = \mathcal{O}\pr{ \sqrt{c\pr{R(\boldsymbol{w}_1) - R^{\star}}} \cdot \frac{\sqrt[4]{T}}{m}
  + c \sigma\frac{\sqrt{T}}{m}}~.
\end{equation*}
\end{theorem}
Under the same assumptions, taking step size of order $\mathcal{O}(1/\sqrt{t})$, \cite{hardt2016train} showed a uniform stability bound $\epsilon = \mathcal{O}(\sqrt{T/m})$.
Our bound differs since it involves a multiplicative risk at the initialization point.
Thus, our bound corroborates the intuition that whenever we start at a good location of the objective function, the algorithm is more stable and
thus generalizes better.
However, this is only the case, whenever the variance of stochastic gradient $\sigma^2$ is not too large.
In the extreme case, deterministic case, and of $\Risk_1=0$, the theorem confirms that SGD, in expectation, does not need to make any updates and is therefore perfectly stable.
On the other hand, when the variance $\sigma^2$ is large enough to make the second summand in Theorem~\ref{thm:sgd_stab_convex} dominant, the bound does not offer improvement compared to~\cite{hardt2016train}.
Note, that a result of this type cannot be obtained through the more restrictive uniform stability, precisely because such bounds on the stability must
hold even for a worst-case choice of data distribution and initialization.
In contrast, the notion of stability we employ depends on the data-generating distribution, which allowed us to introduce dependency on the risk.

Furthermore, consider that we start at arbitrary location $\boldsymbol{w}_1$:
assuming that the loss function is bounded for a concrete $\mathcal{H}$ and $\mathcal{Z}$, the rate of our bound up to a constant is no worse than that of \cite{hardt2016train}.
Finally, one can always tighten this result by taking the minimum of two bounds.
\subsection{Non-convex Losses}
Now we state a new stability result for non-convex losses.
\begin{theorem}
\label{thm:sgd_stab_nonconvex}
Assume that $f(\cdot, z) \in [0,1]$ and has a $\rho$-Lipschitz Hessian,
and that step sizes of a form $\alpha_t = \frac{c}{t}$ satisfy
$c \leq \min\cbr{\frac{1}{\beta}, \frac{1}{4 (2 \beta \ln(T))^2}}$.
Then \ac{SGD} is $\epsilon(\mathcal{D}, \boldsymbol{w}_1)$-on-average  stable with
\begin{align}
  \epsilon(\mathcal{D}, \boldsymbol{w}_1) \leq \frac{1 + \frac{1}{c \gamma}}{m} \pr{2 c L^2}^{\frac{1}{1 + c \gamma}} \pr{ \E_{S,A}\br{R(A_S)} \cdot T}^{\frac{c \gamma}{1 + c \gamma}}~,\label{eq:nonconvex}
\end{align}
where
\begin{align}
&\gamma := \mathcal{\tilde{O}}\pr{ \min\cbr{\beta,~ \E_z\br{\left\|\nabla^2 f(\boldsymbol{w}_1, z) \right\|_2}
 + \Delta_{1,\sigma^2}^{\star}  }}~, \label{eq:gamma}\\
&\Delta_{1,\sigma^2}^{\star} := \rho \pr{ c \sigma + \sqrt{c \pr{\Risk_1 - R^{\star}}} }~. \nonumber
\end{align}
\end{theorem}
In particular, $\gamma$ characterizes how the curvature at the initialization point affects stability, and hence the generalization error of \ac{SGD}.
Since $\gamma$ heavily affects the rate of convergence in~\eqref{eq:nonconvex}, and in most situations smaller $\gamma$ yields higher stability, we now look at a few cases of its behavior.
Consider a regime such that $\gamma$ is of the order $\tilde{\Theta}\pr{\E[\|\nabla^2 f(\boldsymbol{w}_1, z)\|_2] + \sqrt{\Risk_1} + \sigma}$,
or in other words, that stability is controlled by the curvature, the risk of the initialization point $\boldsymbol{w}_1$, and the variance of the stochastic gradient $\sigma^2$.
This suggests that starting from a point in a less curved region with low risk
should yield higher stability, and therefore as predicted by our theory, allow for faster generalization.
In addition, we observe that the considered stability regime offers a principled way to pre-screen a good initialization point in practice, by choosing the one that minimizes spectral norm of the Hessian and the risk.

Next, we focus on a more specific case. Suppose that we choose a step size $\alpha_t = \frac{c}{t}$ such that $\gamma = \tilde{\Theta}\pr{\E[\|\nabla^2 f(\boldsymbol{w}_1, z)\|_2]}$, yet not too small, so that the empirical risk can still be decreased.
Then, stability is dominated by the curvature around $\boldsymbol{w}_1$.
Indeed, lower generalization errors on non-convex problems, such as training deep neural networks, have been observed empirically when \ac{SGD} is actively guided~\cite{hochreiter1997flat,goodfellow2016deep,chaudhuri2017entropy} or converges to solutions with low curvature~\cite{keskar2017large}.
However, to the best of our knowledge, Theorem~\ref{thm:sgd_stab_nonconvex} is the first to establish a theoretical link between the curvature of the loss function and the generalization ability of \ac{SGD} in a data-dependent sense.

Theorem~\ref{thm:sgd_stab_nonconvex} immediately implies following statement that further reinforces the effect of the initialization point on the generalization error, assuming that $\E_S[R(A_S)] \leq \Risk_1$.
\begin{cor}
\label{cor:non_convex_R1}
Under conditions of Theorem~\ref{thm:sgd_stab_nonconvex} we have that \ac{SGD} is $\epsilon(\mathcal{D}, \boldsymbol{w}_1)$-on-average  stable with
\begin{equation}
\label{eq:non_convex_R1}
    \epsilon(\mathcal{D}, \boldsymbol{w}_1) = \mathcal{O}\pr{ \frac{1 + \frac{1}{c \gamma}}{m}  \pr{ \Risk_1 \cdot T}^{\frac{c \gamma}{1 + c \gamma}} }~.
\end{equation}
\end{cor}
We take a moment to discuss the role of the risk term in $\pr{\Risk_1 \cdot T}^{\frac{c \gamma}{1 + c \gamma}}$.
Observe that $\epsilon(\mathcal{D}, \boldsymbol{w}_1) \rightarrow 0$ as $\Risk_1 \rightarrow 0$, in other words, the generalization error approaches zero as the risk of the initialization point vanishes.
This is an intuitive behavior, however, uniform stability does not capture this due to its distribution-free nature.
Finally, we note that \cite[Theorem 3.8]{hardt2016train} showed a bound similar to~\eqref{eq:nonconvex}, however, in place of $\gamma$ their bound has a Lipschitz constant of the gradient.
The crucial difference lies in term $\gamma$ which is now not merely a Lipschitz constant, but rather depends on the data-generating distribution and initialization point of \ac{SGD}.
We compare to their bound by
considering the worst case scenario, namely, that \ac{SGD} is initialized in a point with high curvature, or altogether, that the objective function is highly curved everywhere.
Then, at least our bound is no worse than the one of~\cite{hardt2016train}, since $\gamma \leq \beta$.

Theorem~\ref{thm:sgd_stab_nonconvex} also allows us to prove an optimistic generalization bound for learning with \ac{SGD} on non-convex objectives.
\begin{cor}
Under conditions of Theorem~\ref{thm:sgd_stab_nonconvex} we have that the output of \ac{SGD} obeys
\label{cor:nonconvex_optimistic}
{\small
\begin{align*}
\E_{S,A}\br{R(A_S) - \wh{R}_S(A_S)} =
\mathcal{O}\pr{
\frac{1 + \frac{1}{c \gamma}}{m} \cdot
\max\cbr{ \pr{\E_{S,A}\br{\wh{R}_S(A_S)} \cdot T}^{\frac{c \gamma}{1 + c \gamma}},
\pr{\frac{T}{m}}^{c \gamma}
}
}~.
\end{align*}
}
\end{cor}
An important consequence of Corollary~\ref{cor:nonconvex_optimistic}, is that for a vanishing expected empirical risk, in particular for $\E_{S,A}[\wh{R}_S(A_S)] = \mathcal{O}\pr{\frac{T^{c \gamma}}{m^{1 + c \gamma}}}$, the generalization error behaves as $\mathcal{O}\pr{\frac{T^{c \gamma}}{m^{1 + c \gamma}}}$.
Considering the full pass, that is $m = \mathcal{O}(T)$,
we have an optimistic generalization error of order $\mathcal{O}\pr{1/m}$ instead of $\mathcal{O}(m^{-\frac{1}{1 + c \gamma}})$.
We note that PAC bounds with similar optimistic message (although not directly comparable), but without curvature information can also be obtained through empirical Bernstein bounds as in~\cite{maurer2009empirical}.
However, a PAC bound does not suggest a way to minimize non-convex empirical risk in general,
where, on the other hand, SGD is known to work reasonably well.
\subsubsection{Tightness of Non-convex Bounds}
\label{sec:tightness}
Next we empirically assess the tightness of our non-convex generalization bounds on real data.
In the following experiment we train a neural network with three convolutional layers interlaced with max-pooling, followed by the fully connected layer with $16$ units, on the MNIST dataset. This totals in a model with $18$K parameters.
\begin{figure}
\caption{Empirical tightness of data-dependent and uniform generalization bounds evaluated by training a convolutional neural network.}
\label{fig:mnist_optimistic_nonconvex}
\begin{center}
\includegraphics[width=7cm]{mnist_optimistic.pdf}
\end{center}
\end{figure}
Figure~\ref{fig:mnist_optimistic_nonconvex} compares our data-dependent bound~\eqref{eq:nonconvex} to the distribution-free one of~\cite[Theorem 3.8]{hardt2016train}.
As as a reference we also include an empirical estimate of the generalization error taken as an absolute difference of the validation and training average losses.
Since our bound also depends on the initialization point, we plot~\eqref{eq:nonconvex} for multiple ``warm-starts'', \text{ie.{}} with \ac{SGD} initialized from a pre-trained position.
We consider $7$ such warm-starts at every $200$ steps, and report data-dependent quantities used to compute~\eqref{eq:nonconvex} just beneath the graph.
Our first observation is that, clearly, the data-dependent bound gives tighter estimate, by roughly one order of magnitude.
Second, simulating start from a pre-trained position suggests even tighter estimates: we suspect that this is due to decreasing validation error which is used as an empirical estimate for $R(\boldsymbol{w}_1)$ which affects bound~\eqref{eq:nonconvex}.

We compute an empirical estimate of the expected Hessian spectral norm by the power iteration method using an efficient Hessian-vector multiplication method~\cite{pearlmutter1994fast}.
Since bounds depend on constants $L$, $\beta$, and $\rho$, we estimate them
by tracking maximal values of the gradient and Hessian norms throughout optimization.
We compute bounds with estimates $\widehat{L} = 78.72$, $\widehat{\beta} = 1692.28$, $\widehat{\rho} = 3823.73$, and $c = 10^{-3}$.
\subsection{Application to Transfer Learning}
\label{sec:tl}
One example application of data-dependent bounds presented before lies in \emph{\acf{TL}},
where we are interested in achieving faster generalization on a \emph{target} task by exploiting side information that originates from different but related \emph{source} tasks.
The literature on \ac{TL} explored many ways to do so, and here we will focus on the one that is most compatible with our bounds.
More formally, suppose that the \emph{target} task at hand is characterized by a joint probability distribution $\mathcal{D}$, and as before we have a training set $S \distas{\iid} {\mathcal{D}}^m$.
Some \ac{TL} approaches also assume access to the data sampled from the distributions associated with the \emph{source} tasks.
Here we follow a conservative approach -- instead of the source data,
we receive a set of \emph{source} hypotheses $\cbr{\boldsymbol{w}^{\text{src}}_k}_{k=1}^K \subset \mathcal{H}$, trained on the source tasks.
The goal of a learner is to come up with a target hypothesis, which in the optimistic scenario generalizes better by relying on source hypotheses.
In the \ac{TL} literature this is known as \ac{HTL}~\cite{kuzborskij2016fast}, that is, we transfer from the source hypotheses which act as a proxy to the source tasks and
the risk $R(\boldsymbol{w}^{\text{src}}_k)$ quantifies how much source and target tasks are related.
In the following we will consider \ac{SGD} for \ac{HTL}, where the source hypotheses act as initialization points.
First, consider learning with convex losses:
Theorem~\ref{thm:sgd_stab_convex} depends on $R(\boldsymbol{w}_1)$, thus it immediately quantifies the relatedness of source and target tasks.
So it is enough to pick the point that minimizes the stability bound to transfer from the most related source.
Then, bounding $R(\boldsymbol{w}^{\text{src}}_k)$ by $\wh{R}_S(\boldsymbol{w}^{\text{src}}_k)$ through Hoeffding bound along with union bound gives with high probability that
\[
  \min_{k \in [K]} \epsilon(\mathcal{D}, \boldsymbol{w}^{\text{src}}_k) \leq \min_{k \in [K]}  \mathcal{O}\pr{ \wh{R}_S(\boldsymbol{w}^{\text{src}}_k) + \sqrt{\frac{\log(K)}{m}} }~.
\]
Hence, the most related source is the one that simply minimizes empirical risk.
Similar conclusions where drawn in \ac{HTL} literature, albeit in the context of \ac{ERM}.
Matters are slightly more complicated in the non-convex case.
We take a similar approach, however, now we minimize stability bound~\eqref{eq:non_convex_R1}, and for the sake of simplicity assume that we make a full pass over the data, so $T=m$.
Minimizing the following empirical upper bound select the best source.
\begin{prop}
\label{prop:nonconvex_transfer}
Let
$
\widehat{\gamma}_k^{\pm} = \Theta\Big(\frac{1}{m} \sum_{i=1}^m \|\nabla^2 f(\boldsymbol{w}^{\text{src}}_k, z_i)\|_2 + \sqrt{\wh{R}_S(\boldsymbol{w}^{\text{src}}_k)}
 \pm \sqrt[4]{\log(K) / m } \Big)$.
 Then with high probability the generalization error of $\boldsymbol{w}^{\text{src}}_k$ is bounded by
\[
\min_{k \in [K]} \mathcal{O}\pr{ \pr{1 + \frac{1}{c \widehat{\gamma}_k^{-}}}
\wh{R}_S(\boldsymbol{w}^{\text{src}}_k)^{\frac{c \widehat{\gamma}^+_k}{1 + c \widehat{\gamma}^+_k}}
\cdot
\frac{\sqrt{\log(K)}}{m^{\frac{1}{1 + c \widehat{\gamma}^{+}_k}} } }~.
\]
\end{prop}
Note that $\widehat{\gamma}_k^{\pm}$ involves estimation of the spectral norm of the Hessian, which is computationally cheaper to evaluate compared to the complete Hessian matrix~\cite{pearlmutter1994fast}.
This is particularly relevant for deep learning, where computation of the Hessian matrix can be prohibitively expensive.

\section{Stability of \acl{SGD}}
\label{sec:stability}
First, we introduce definitions used in the rest of the paper.
\subsection{Definitions}
\input{def.tex}
\subsection{Uniform Stability and Generalization}
On an intuitive level, a learning algorithm is said to be \emph{stable} whenever a small perturbation in the training set does not affect its outcome too much.
Of course, there is a number of ways to formalize the perturbation and the extent of the change in the outcome,
and we will discuss some of them below.
The most important consequence of a stable algorithm is that it \emph{generalizes} from the training set to the unseen data sampled from the same distribution.
In other words, the difference between the risk $R(A_S)$ and the empirical risk $\wh{R}_S(A_S)$ of the algorithm's output is controlled by the quantity that captures how stable the algorithm is.
So, to observe good performance, or a decreasing true risk, we must have a stable algorithm \emph{and} decreasing empirical risk (training error), which usually comes by design of the algorithm.
In this work we focus on the stability of the \acf{SGD} algorithm, and thus, as a consequence, we study its generalization ability.

Recently,~\cite{hardt2016train} used a stability argument to prove generalization bounds for learning with \ac{SGD}.

Specifically, the authors extended the notion of the \emph{uniform stability} originally proposed by~\cite{BousquetE02}, to accommodate randomized algorithms.
\begin{definition}[Uniform stability]
A randomized algorithm $A$ is $\epsilon$-uniformly stable if for all datasets $S, {S^{(i)}} \in \mathcal{Z}^m$ such that $S$ and ${S^{(i)}}$ differ in the $i$-th example, we have
\[
\sup_{z \in \mathcal{Z}, i \in [m]} \cbr{ \E_{A}\br{f(A_S,z) - f(A_{{S^{(i)}}},z)} } \leq \epsilon~.
\]
\end{definition}
Since \ac{SGD} is a randomized algorithm, we have to cope with two sources of randomness: the data-generating process and the randomization of the algorithm $A$ itself, hence we have statements in expectation.
The following theorem of~\cite{hardt2016train} shows that the uniform stability implies generalization in expectation.
\begin{theorem}
\label{thm:uniform}
  Let $A$ be $\epsilon$-uniformly stable. Then,
\[
\abs{ \E_{S,A}\br{ \wh{R}_S(A_S) - R(A_S) } } \leq \epsilon~.
\]
\end{theorem}
Thus it suffices to characterize the uniform stability of an algorithm to state a generalization bound.
In particular, \cite{hardt2016train} showed generalization bounds for \ac{SGD} under different assumptions on the loss function $f$.
Despite that these results hold in expectation,
other forms of generalization bounds, such as high-probability ones, can be derived from the above~\cite{shalev2010learnability}.

Apart from~\ac{SGD}, uniform stability has been used before to prove generalization bounds for many learning algorithms~\cite{BousquetE02}.
However, these bounds typically suggest worst-case generalization rates, and rather reflect intrinsic stability properties of an algorithm.
In other words, uniform stability
is oblivious to the data-generating process and any other side information,
which might reveal scenarios where generalization occurs at a faster rate.
In turn, these insights could motivate the design of improved learning algorithms.
In the following we address some limitations of analysis through uniform stability by using a less restrictive notion of stability.
We extend the setting of~\cite{hardt2016train} by proving data-dependent stability bounds for convex and non-convex loss functions.
In addition, we also take into account the initialization point of an algorithm as a form of supplementary information,
and we dedicate special attention to its interplay with the data-generating distribution.
Finally, we discuss situations where one can explicitly control the stability of \ac{SGD} in a data-dependent way.

\section{Data-dependent Stability Bounds for \acs{SGD}}
\label{sec:data-stability}
In this section we describe a notion of data-dependent algorithmic stability,
that allows us to state generalization bounds which depend not only on the properties of the learning algorithm,
but also on the additional parameters of the algorithm.
We indicate such additional parameters by $\theta$, and therefore we denote stability as a function $\epsilon(\theta)$.
In particular, in the following we will be interested in scenarios where $\theta$ describes the data-generating distribution and the initialization point of \ac{SGD}.
\begin{definition}[On-Average  stability]
A randomized algorithm $A$ is $\epsilon(\theta)$-on-average  stable if it is true that
\[
\sup_{i \in [m]}\cbr{ \E_{A}\E_{S,z}\br{ f(A_S, z) - f(A_{{S^{(i)}}}, z)} } \leq \epsilon(\theta)~,
\]
\text{where} $S \distas{\iid} \mathcal{D}^m$ and ${S^{(i)}}$ is its copy with $i$-th example replaced by $z \distas{\iid} \mathcal{D}$.
\end{definition}
Our definition of on-average  stability resembles the notion introduced by~\cite{shalev2010learnability}.
The difference lies in the fact that we take supremum over index of replaced example.
A similar notion was also used by~\cite{BousquetE02} and later by~\cite{elisseeff2005stability} for analysis of a randomized aggregation schemes, however their definition involves absolute difference of losses.
The dependence on $\theta$ also bears similarity to recent work of~\cite{london2016generalization}, however, there, it is used in the context of uniform stability.
The following theorem shows that on-average -~stable random algorithm is guaranteed to generalize in expectation.
\begin{theorem}
\label{thm:onaverage_stab_generalization}
Let an algorithm $A$ be $\epsilon(\theta)$-on-average  stable. Then,
\[
\E_{S}\E_{A}\br{ R(A_S) - \wh{R}_S(A_S) } \leq \epsilon(\theta)~.
\]
\end{theorem}