Understanding the unstable convergence of gradient descent
April 03, 2022
Abstract
Most existing analyses of (stochastic) gradient descent rely on the condition that for \(L\)-smooth costs, the step size is less than \(2/L\). However, many works have observed that in machine learning applications step sizes often do not fulfill this condition, yet (stochastic) gradient descent still converges, albeit in an unstable manner. We investigate this unstable convergence phenomenon from first principles, and discuss key causes behind it. We also identify its main characteristics, and how they interrelate based on both theory and experiments, offering a principled view toward understanding the phenomenon.
1 Introduction↩︎
Gradient descent (GD) runs the iteration \[\begin{align} {\boldsymbol{\theta}}^{t+1} = {\boldsymbol{\theta}}^{t} - \eta \nabla f({\boldsymbol{\theta}}^{t})\,, \end{align}\] seeking to optimize a cost function \(f\). It also provides a conceptual foundation for stochastic gradient descent (SGD), one of the key algorithms in modern machine learning. A vast body of literature that analyzes (S)GD assumes that the cost \(f\) is \(L\)-smooth (we say \(f\) is \(L\)-smooth if \(\|\nabla f({\boldsymbol{\theta}}^{})-\nabla f({\boldsymbol{\theta}}^{'})\|\le L\|{\boldsymbol{\theta}}^{ }-{\boldsymbol{\theta}}^{'}\|\) for all \({\boldsymbol{\theta}}^{},{\boldsymbol{\theta}}^{'}\)), and subsequently exploits the associated “descent lemma”: \[\begin{align} \label{descent} f({\boldsymbol{\theta}}^{t+1}) \leq f({\boldsymbol{\theta}}^{t}) -\eta \bigl(1 - L\frac{\eta}{2}\bigr) \left\|{\nabla f({\boldsymbol{\theta}}^{t})} \right\|^2\,. \end{align}\tag{1}\] To ensure descent via inequality 1 , the condition \[\begin{align} \label{cond:step} L < \frac{2}{\eta}, \end{align}\tag{2}\] is imposed. This condition ensure that GD decreases the cost \(f\) at each iteration. Whenever condition 2 holds, we call it the “stable” regime in this paper.
When the cost is quadratic, condition 2 is in fact necessary for stablility: if \(\eta > \frac{2}{L}\), then GD diverges (see [fact:1]). This observation carries over to most convex optimization settings and also neural networks when using the neural tangent kernel approximations [1]–[3]. Thus, it is reasonable to assume condition 2 for those analyses. However, for general nonconvex costs, it is not clear whether the stable regime condition 2 is required or even reasonable.
Figure 1: Example of unstable convergence for training CIFAR-10 with GD. We follow the experimental setup of [4]; see [ex:cifar] for details. We use a ReLU network. Here, condition 2 fails, but the training loss still (non-monotonically) decreases in the long run..
Recently, it has been observed that GD on neural networks often violates condition 2 . More specifically, [4] observe that when we run GD to train a neural network, the condition 2 fails, but contrary to the common wisdom from convex optimization, the training loss still (non-monotonically) decreases in the long run. See Figure 1 for an example of this phenomenon. We call this phenomenon “unstable” convergence.
Unfortunately, very little is known about unstable convergence. The causes and implications of this phenomenon have not been explored in the literature. More importantly, the main features as well as the scope of this phenomenon have not been discussed. Characterizing the main features is important because it not only furnishes better understanding of this bizarre phenomenon, but also lays a foundation for future theoretical studies; especially, the main characteristics of this phenomenon will help build a more practical theory of the neural network optimization.
Contributions. In light of the above motivation, the main contributions of this paper are as follows:
We discuss the main causes driving the unstable convergence phenomenon (3).
We identify the main features that characterize unstable convergence in terms of how loss, iterates, and sharpness1 evolve with GD updates. Moreover, we investigate and clarify the relations between them. Our characterizations demonstrate that the features of unstable convergence are in stark contrast with those of traditional stable convergence, suggesting that their optimization mechanisms are significantly different.
In particular, the main features considered in this work provide alternative ways to identify unstable convergence in practice.
2 provides a more technical overview of our main findings, along with their interpretations.
1.1 Related Work↩︎
Under various contexts, several recents works have observed the unstable convergence phenomenon in training neural networks with (S)GD [5]–[9]. We refer readers to the related work section of [4] for greater context.
The unstable convergence phenomenon is first formally identified by [4], and in their paper it is named edge of stability. More specifically, they observe a more refined version of the unstable convergence: when training a neural network with GD, the sharpness at the iterate goes beyond the threshold \(\frac{\eta}{2}\), and often saturates right at (or above) the threshold. In 5, we will explore the relations between our main features and their observed phenomenon.
Concurrent works. Recently, [10] also investigate the causes of unstable convergence based on their empirical observations. Their main observation is that unstable convergence might be due to the landscape of loss function where the loss grows slower than a quadratic near the local minima. As we will see in 3.2, their main finding is consistent with our explanation. They also demonstrate through examples that such “sub-quadratic” growth near the minima is caused by the heterogeneity of data; see their Section 6 for details.
Another work by [11] identifies a setting in which one can prove the unstable convergence phenomenon theoretically. More specifically, they show that the normalized gradient descent dynamics of form \({\boldsymbol{\theta}}^{t+1} = {\boldsymbol{\theta}}^{t} - \eta \nabla f({\boldsymbol{\theta}}^{t})/\left\|{\nabla f({\boldsymbol{\theta}}^{t})} \right\|\) can provably exhibit the unstable convergence phenomenon near the minima under some suitable assumptions; see their Section 4 for details.
2 Unstable GD Can’t Reach Stationary Points↩︎
In this section, we build intuitions about what the unstable regime \(\eta >\frac{2}{L}\) suggests. First, note that the fixed points \({\boldsymbol{\theta}}^{\infty}\) of the GD dynamics \({\boldsymbol{\theta}}^{t+1} = {\boldsymbol{\theta}}^{t} - \eta \nabla f({\boldsymbol{\theta}}^{t})\) are the stationary points, i.e., points such that \(\nabla f({\boldsymbol{\theta}}^{\infty})=0\). Hence, the GD dynamics will eventually approach one of the stationary points, and in order to understand the unstable regime, we first need to understand the behavior of the dynamics near the stationary points whose sharpness is greater than \(\frac{2}{\eta}\).
As a warm up, we first consider the simplest setting of quadratic costs where the sharpness is constant globally. We begin with the following well-known fact.
On a quadratic cost \(f({\boldsymbol{\theta}}^{})= \frac{1}{2}{\boldsymbol{\theta}}^{\top} P {\boldsymbol{\theta}}^{} +\mathbf{q}^\top {\boldsymbol{\theta}}^{} +r\), GD will diverge if any eigenvalue of \(P\) exceeds the threshold \(2/\eta\). For convex quadratics, this condition is “iff.”0◻
Below we quickly illustrate this fact through an example.
Example 1. Consider optimizing a quadratic cost \(f(\theta_1,\theta_2) = 20 \theta_1^2+ \theta_2^2\). Note that in this case \(L=40\). Let us run GD on this cost with \(\eta = 2/39,\;2/40,\;2/41\).
As shown in the above plots, GD converges to the optimum if \(\eta <2/L\) and it diverges if \(\eta >2/L\). 0◻
Due to the above fact, one can build the following intuition: when a stationary point has sharpness greater than \(\frac{2}{\eta}\), the GD dynamics cannot converge to the stationary point.
We first formalize this intuition. In particular, we show that GD cannot converge to a stationary point that has sharpness greater than \(2/\eta\). We make the following assumptions about the spectrum of Hessian at stationary points and the non-degeneracy of the GD dynamics.
Assumption 1. Let \(F({\boldsymbol{\theta}}^{})= {\boldsymbol{\theta}}^{}- \eta \nabla f({\boldsymbol{\theta}}^{})\) and assume that for any subset \(S\) of measure zero, \(F^{-1}(S)\) is of measure zero. Moreover, each stationary point \({\boldsymbol{p}}\) satisfies \(\frac{1}{\eta} \not\in \lambda (\nabla^2 f({\boldsymbol{p}}))\), where \(\lambda (\nabla^2 f({\boldsymbol{p}}))\) denotes the set of eigenvalues of the Hessian of \(f\) at \({\boldsymbol{p}}\).
Theorem 1. For a given subset \(\mathcal{X}\) of the domain of parameter \({\boldsymbol{\theta}}^{}\), assume that \(f\) is \(C^2\) in \(\mathcal{X}\). Suppose that for each stationary point \({\boldsymbol{p}}\in \mathcal{X}\), it holds that either \(\lambda_{\min}(\nabla^2 f({\boldsymbol{p}}))<0\) or \(\lambda _{\max}(\nabla^2 f({\boldsymbol{p}})) > \frac{2}{\eta}\). Then under 1, there is a measure-zero subset \(\mathcal{N}\) s.t. for all initializations \({\boldsymbol{\theta}}^{0}\in \mathcal{X}\setminus \mathcal{N}\), the GD dynamics \({\boldsymbol{\theta}}^{t+1} = {\boldsymbol{\theta}}^{t} - \eta \nabla f({\boldsymbol{\theta}}^{t})\) do not converge to any of the stationary points in \(\mathcal{X}\).
Proof of 1. The proof is inspired by those of [12], [13]. First, we recall Stable Manifold Theorem [14].
Lemma 1. Let \({\boldsymbol{p}}\) be a fixed point for the \(C^r\) local diffeomorphism \(h: U\to \mathbb{R}^n\) where \(U\) is an open neighborhood of \({\boldsymbol{p}}\) in \(\mathbb{R}^n\) and \(r\geq 1\). Let \(E^{\sf s} \oplus E^{\sf c} \oplus E^{\sf u}\) be the invariant spliiting of \(\mathbb{R}^n\) into the generalized eigenspaces of \(Dh({\boldsymbol{p}})\) corresponding to eigenvalues of absolute value less than one, equal to one, and greater than one. To the \(Dh({\boldsymbol{p}})\)-invariant subspace \(E^{\sf s} \oplus E^{\sf c}\) there is an associated local \(h\)-invariant \(C^r\) embedded disc \(W^{\sf s\oplus c}_{\sf loc}\) of dimension \(\dim(E^{\sf s} \oplus E^{\sf c})\) and ball \(B\) around \({\boldsymbol{p}}\) such that \(h(W^{\sf s\oplus c}_{\sf loc})\cap B \subset W^{\sf s\oplus c}_{\sf loc}\). Moreover, if \(h^n(\mathbf{x})\in B\) for all \(n\geq 0\), then \(\mathbf{x}\in W^{\sf s\oplus c}_{\sf loc}\).
To apply 1, we first show that the GD dynamics \(F\) is a local diffeomorphism at each stationary point \({\boldsymbol{p}}\) satisfying \(\frac{1}{\eta}\notin \lambda(\nabla^2 f({\boldsymbol{p}}))\). This follows from the inverse function theorem: (i) Note that \(F\) is a \(C^1\) vector field since \(f\) is \(C^2\), (ii) the Jacobian of \(F\) is equal to \(DF({\boldsymbol{p}}) = I-\eta \nabla^2f({\boldsymbol{p}})\), and since \(\frac{1}{\eta}\notin \lambda(\nabla^2 f({\boldsymbol{p}}))\), the Jacobian is invertible. Hence, by inverse function theorem, we conclude that \(F\) is a local diffeomorphism around \({\boldsymbol{p}}\).
Hence for each stationary point \({\boldsymbol{p}}\) satisfying \(1/\eta\notin \lambda(\nabla^2 f({\boldsymbol{p}}))\), we can apply 1 at \({\boldsymbol{p}}\). Let \(B_{\boldsymbol{p}}\) be the open ball due to 1. Let \(\mathcal{S}\) be the set of stationary points. Consider the following open cover \[\begin{align} \label{cover} \mathcal{C}:=\bigcup_{\substack{{\boldsymbol{p}}: \text{ stationary point}\\ \frac{1}{\eta}\notin \lambda(\nabla^2 f({\boldsymbol{p}}))}} B_{\boldsymbol{p}}\,. \end{align}\tag{3}\] Then from 1, it follows that \(\mathcal{S}\subset \mathcal{C}\) and hence \(\mathcal{C}\) is an open cover of \(\mathcal{S}\). Thus, Lindelöf’s lemma guarantees that there exists a countable subcover of \(\mathcal{C}\), i.e., there exist \({\boldsymbol{p}}_1,{\boldsymbol{p}}_2,\dots\) s.t. \(\mathcal{C}=\cup_{i=1}^\infty B_{{\boldsymbol{p}}_i}\). If GD converges to a stationary point \({\boldsymbol{p}}\), there must exist \(t_0\) and \(i\) such that \(F^{t}({\boldsymbol{p}}_0)\in B_{{\boldsymbol{p}}_i}\) for all \(t\geq t_0\). From 1, we conclude that \(F^{t}({\boldsymbol{\theta}}^{}_0)\in W^{\sf s\oplus c}_{\sf loc}({\boldsymbol{p}}_i)\). In other words, we have \({\boldsymbol{\theta}}^{}_0 \in F^{-t}(W^{\sf s\oplus c}_{\sf loc}({\boldsymbol{p}}_i))\) for all \(t\geq t_0\). Hence the set of initial points in \(\mathcal{X}\) for which GD converges to a stationary point is a subset of \[\begin{align} \mathcal{N}:= \bigcup_{i=1}^\infty \bigcup_{t=0}^\infty F^{-t}(W^{\sf s\oplus c}_{\sf loc}({\boldsymbol{p}}_i)) \,. \end{align}\] Now from the assumption that either \(\lambda_{\min}(\nabla^2 f({\boldsymbol{p}}))<0\) or \(\lambda _{\max}(\nabla^2 f({\boldsymbol{p}})) > \frac{2}{\eta}\), it follows that \(I-\eta \nabla^2 f({\boldsymbol{p}})\) has an eigenvalue whose absolute value is greater than \(1\). Hence, for each stationary point \({\boldsymbol{p}}\), \(\dim( E^{\sf u})\geq 1\). This implies that each \(W^{\sf s\oplus c}_{\sf loc}({\boldsymbol{p}})\) has measure zero, and from the assumption that \(F^{-1}(\mathcal{X})\) is of measure zero if \(\mathcal{X}\) is of measure zero, it holds that each \(F^{-t}(W^{\sf s\oplus c}_{\sf loc}({\boldsymbol{p}}_i))\) is of measure zero. Thus, being a countable union of measure zero sets, \(\mathcal{N}\) is measure zero. It follows that for initialization \({\boldsymbol{\theta}}^{}_0 \in \mathcal{X}\setminus\mathcal{N}\), the GD dynamics never converge to a stationary point in \(\mathcal{X}\). ◻
Note that 1 applies to the case when stationary points are not isolated. Moreover, the condition that every stationary point \({\boldsymbol{p}}\) satisfies \(\frac{1}{\eta} \not\in \lambda (\nabla^2 f({\boldsymbol{p}}))\) can be relaxed to the condition that the open cover \(\mathcal{C}\) in 3 covers the entire set of stationary points \(\mathcal{S}\).
The main takeaway of 1 is that for almost all initializations, GD cannot converge to the stationary point whose sharpness is larger than \(2/\eta\) even when there is only a single eigenvector whose eigenvalue exceeds the threshold \(2/\eta\). Having this intuition, we next discuss how “convergence” could happen under the unstable regime.
3 How Can Unstable GD “Converge”?↩︎
In the previous section, we saw that when stationary points have large sharpness relative to the step size, GD cannot converge to those stationary points. However, as we saw in 1, GD can still “converge” under the unstable regime; GD still manages to (non-monotonically) decrease the training loss in the long run. In this section, we understand this bizzare co-occurrence. We first discuss some possible causes for the unstable regime.
3.1 What Causes the Unstable Regime↩︎
One possible cause for the unstable regime is that the landscape has only “trivial” stationary points; we will formalize the meaning of “trivial” shortly. This situation turns out to be quite common for neural networks as illustrated by the following result.
Assume the loss of neural network parametrized \({\boldsymbol{\theta}}^{}\) contains a weight decay term as follows, \[\begin{align} \ell({\boldsymbol{\theta}}^{}) = \frac{1}{n}\sum_{i=1}^n f(\boldsymbol{x}_i, {\boldsymbol{\theta}}^{}) + \gamma \|{\boldsymbol{\theta}}^{}\|^2_2. \end{align}\] If we partition the network parameter \({\boldsymbol{\theta}}^{} = [\boldsymbol{\xi}; \boldsymbol{\zeta}]\) such that a subset of the network parameters \(\boldsymbol{\zeta}\) is positive homogeneous, i.e. for any input data \(x_i\) and positive number \(c > 0\) \[\begin{align} f(\boldsymbol{x}_i, [\boldsymbol{\xi}, \boldsymbol{\zeta}]) = f(x_i, [\boldsymbol{\xi}, c\boldsymbol{\zeta}]), \end{align}\] Then the loss \(\ell({\boldsymbol{\theta}}^{})\) has no stationary point if \(\boldsymbol{\zeta} \neq 0\).
Proof of [prop:trivial]. This statement follows by a simple observation that from positive homogeneity, \[\left \langle \nabla_{\boldsymbol{\zeta}} f(\boldsymbol{x}_i, [\boldsymbol{\xi}, \boldsymbol{\zeta}]), \boldsymbol{\zeta} \right\rangle = 0.\] Therefore, if \(\nabla_{\boldsymbol{\zeta}} \gamma \|{\boldsymbol{\theta}}^{}\|^2_2 \neq 0\), we have \[\nabla_{\boldsymbol{\zeta}} f(\boldsymbol{x}_i, [\boldsymbol{\xi}, \boldsymbol{\zeta}]) + \nabla_{\boldsymbol{\zeta}} \gamma \|{\boldsymbol{\theta}}^{}\|^2_2 \neq 0,\] which concludes the proof. ◻
Notice that the positive homogeneity parameters exist in many networks such as ResNet or Transformer when normalization layers exist (\(\boldsymbol{a}_{L+1} = \boldsymbol{a}_L / \|\boldsymbol{a}_L\|\), where \(\boldsymbol{a}_L\) denote the input to layer \(L\), and \(\| \cdot \|\) denotes a norm of choice).
Note that in practical settings, [prop:trivial] also suggests that there could be lack of flat minima near the GD trajectory. In the above example of ResNet or Transformer, the networks often add a small \(\epsilon\) term to the normalization \(\boldsymbol{a}_{L+1} = \boldsymbol{a}_L / (\epsilon + \|\boldsymbol{a}_L\|)\) to avoid the loss being undefined at \(\boldsymbol{a}_L = 0\). However, the stationary points only exist when \(\|\boldsymbol{a}_L\| \approx \epsilon\), in which case the sharpness of the stationary point is very large (on the order of \(\sim 1/{\epsilon}\)).
In fact, it has been extensively observed in the literature that the sharpness around GD with practical stepsize choices often goes beyond the threshold \(2/\eta\). This claim is verified through a comprehensive set of experiments and called progressive sharpening in [4]; we refer readers to their Section 3.1 for details. For instance, the sharpness curve in 1 shows this phenomenon. Moreover, a similar phenomenon was observed in [5], and they speculated that the density of sharp minima is much larger than the density of flat minima in the neural network landscape. See their Section 4.1 for details.
We summarize our discussion regarding the causes of unstable regime as follows.
For practical stepisze choices, lack of (flat) non-trivial stationary points near the GD trajectory can cause GD to enter the unstable regime.
3.2 Causes for Convergence↩︎
As we discussed in [fact:1], for quadratic costs (or more generally for most convex costs), GD being in the unstable regime implies that GD will diverge entirely. However, as demonstrated by [4] through a comprehensive set of experiments, in neural network training, this situation no longer holds. In this section, we discuss how in the unstable regime “convergence” could happen through examples. As a warm-up, let us revisit the quadratic cost considered in [fact:1], but this time with some modifications.
Example 2 (“Flattened” quadratic cost). For the same quadratic cost as in [fact:1], we chose the same diverging step size \(\eta = 2/39 > 2/L\), but this time we change the cost a bit by applying \(\tanh(\cdot )\) on top of the quadratic cost. More formally, we consider the cost \(\tanh(20\cdot \theta_1^2 + \theta_2^2)\). Due to the fact \(\tanh\approx x\) near zero, this transformation wouldn’t change the geometry near the global minimum. We run GD on the modified cost, and the result looks as follows (we include the result for the original quadratic cost on the left for comparison):
As one can see from the above plot, for the transformed cost, GD does not diverge in the unstable regime. 0◻
The above toy example illustrates that indeed for nonconvex costs, being in the unstable regime does not necessarily imply complete divergence. For the above example, this was possible because of \(\tanh(\cdot)\), which ‘flattens’’ out the landscape of the quadratic cost away from the minimum.
More formally, let us denote the GD dynamics by \(F({\boldsymbol{\theta}}^{}):= {\boldsymbol{\theta}}^{}- \eta \nabla f({\boldsymbol{\theta}}^{})\). Then the role of \(\tanh(\cdot)\) in the above example is that it creates a compact subset near the minimum that is forward-invariant: we say \(S\) is forward-invariant with respect to the dynamics \(F\) if \(F(S)\subseteq S\). Because the gradient of \(\tan(\text{quadratic})\) vanishes as the point gets farther away from the minimum, there exists a forward-invariant compact subset \(\mathcal{X}\) near the minimum.
In a very recent concurrent work by [10], this phenomenon is discussed in a more principled manner using the subquadratic growth property. More specifically, they observed that for practical neural network settings, the loss landscape near the minima exhibits growth that is slower that quadratic, in which case the GD dynamics do not diverge entirely even in the unstable regime. See their Section 4 for details.
We demonstrate this point for neural network examples. We first consider the simplest neural network example, namely a single hidden neuron network.
Example 3 (Single neuron networks). We consider a trivial task of fitting the data \((\textcolor{black}{1},\textcolor{black}{0})\) with a single hidden neuron neural network. Formally, we consider two types of networks:
linear network: \(f(\theta_1,\theta_2) = (\theta_1\cdot (\textcolor{black}{1}\cdot\theta_2)-\textcolor{black}{0})^2\).
\(\tanh\) network: \(f(\theta_1,\theta_2) = (\theta_1\cdot \tanh(\textcolor{black}{1}\cdot\theta_2)-\textcolor{black}{0})^2\).
We initialize both networks at \({\boldsymbol{\theta}}^{0} =(13, 0.01)\) choose step size \(\eta=2/150\) to train them.
As one can see from the above plots, for a linear network, the iterate quickly diverges, while for the \(\tanh\) network, the iterate does not diverge and converges to a minimum (whose sharpness is indeed approximately equal to \(2/\eta\)). 0◻
3 illustrates that the use of activation function like \(\tanh\) can create a compact forward-invariant subset near the minima, which helps GD not diverge in the unstable regime. In fact, the above example suggests that GD indeed exhibits some convergence behaviour where while being in the unstable regime, GD travels along the valley of minima until it finds a flat enough minimum where it can stabilize.
We now consider more practical neural network examples inspired by the settings considered in [4].
For this example, we follow the setting of the main experiment [4] in their Section 3. Specifically, we use (full-batch) GD to train a neural network on \(5,000\) examples from CIFAR-10 with the CrossEntropy loss, and the network is a fully-connected architecture with two hidden layers of width \(200\). Under this common setting, we consider three types of networks: (i) linear network without activations; (ii) \(\tanh\) activations; (ii) ReLU activations. We choose the step size \(\eta=2/30\) and the results are as follows:
As one can see from the above plot, GD converges for the networks with activation functions, while GD diverges without activation functions. 0◻
We summarize our discussion regarding the causes for convergence as follows.
Ingredients of neural networks such as activation functions create a compact forward-invariant set near the minima, which helps GD (non-monotonically) converge in the unstable regime.
In this section, we have discussed the causes of unstable convergence and explain how the intuitions differ from those of conventional convex optimization. We next move on to study the main characteristics of unstable convergence. For instance, we observe that under the unstable convergence phenomenon, the loss is very non-monotonic. Can we understand the behavior of loss in a more principled way?
4 Characteristics of the Unstable Convergence↩︎
In this section, we aim to quantify unstable convergence through several quantities that can be computed during the training. In particular, we will characterize the unstable convergence in terms of the loss behavior and the iterate behavior. We will later demonstrate that the two different behaviors are interconnected with each other.
4.1 Characteristics in Loss Behavior↩︎
We first investigate what happens to the loss under unstable convergence. As a warm-up, we first consider the loss behavior under stable convergence.
4.1.1 Warm-up: The Stable Regime↩︎
Recall from the descent lemma 1 that when GD is in the stable regime, then we have \(f({\boldsymbol{\theta}}^{t+1})-f({\boldsymbol{\theta}}^{t}) \leq - c \eta \left\|{\nabla f({\boldsymbol{\theta}}^{t})} \right\|^2\) for some constant \(c>0\). Putting it differently, we have \[\begin{align} \frac{f({\boldsymbol{\theta}}^{t+1})-f({\boldsymbol{\theta}}^{t})}{\eta \left\|{\nabla f({\boldsymbol{\theta}}^{t})} \right\|^2} \leq -\mathrm{const.} \end{align}\] Let us give the ratio on the LHS a name for convenience:
Definition 1 (Relative progress ratio). We define \[\begin{align} {\sf RP}({\boldsymbol{\theta}}^{}):=\frac{f({\boldsymbol{\theta}}^{}-\eta \nabla f({\boldsymbol{\theta}}^{}) )-f({\boldsymbol{\theta}}^{})}{\eta \left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2}\,. \end{align}\]
Let us revisit [ex:cifar] and verify that for smaller step sizes the relative progress ratio is indeed a negative number.
We use the same setting as [ex:cifar], which follows the setting of the main experiment in [4]. For activations, we choose \(\tanh\) following [4]. We choose much smaller step sizes so that GD is in the stable regime. We plot the loss and the relative progress ratio until the training accuracy hits \(95\%\).
From the above plots, one can see that the relative progress ratio stays negative for all iterations. Moreover, there is no non-monotonic behavior in the loss curve. 0◻
Given the result above, one might wonder why the relative progress saturates around \(-1\). Although we do not have a clear explanation, we suspect that this happens because the trajectory of GD quickly converges to a single direction. We will quickly revisit this later this section. See [rmk:stable].
4.1.2 Relative Progress Ratio under Unstable Convergence↩︎
Given that relative progress is strictly negative number in the stable regime, we now investigate how relative progress ratio behaves in the case of unstable convergence.
We use the same setting as [exp:stable]. This time we choose step sizes larger so that GD operates in the unstable convergence regime. We plot the loss and the relative progress ratio until the training accuracy hits \(95\%\).
The above experiment shows that in the the unstable regime, the relative progress ratio saturates around \(0\) unlike the stable regime. 0◻
One curious aspect of the above results is that the optimization seems to get faster as we choose larger step sizes. This is in fact one of the main observations in [4], suggesting that the unstable convergence is preferred in practice for its faster optimization. However, that does not mean one can increase the step size too large. For example, in the above experiment, we observe that the training loss diverges for step size \(\eta = 2/10\).
Based on [exp:unstable], we raise the following question:
Q. why does \({\sf RP}({\boldsymbol{\theta}}^{t})\) oscillate around \(0\) under unstable convergence?
We begin with explaining why \({\sf RP}({\boldsymbol{\theta}}^{t})\) cannot stay above \(0\). Since the loss is converging in a long term, it cannot be that \({\sf RP}({\boldsymbol{\theta}}^{t})>0\) for many iterations; otherwise, the loss will keep increasing, contradicting the convergence.
More curious part is the fact that \({\sf RP}({\boldsymbol{\theta}}^{t})\) cannot stay below zero, which directly contrasts with the stable regime. To understand this phenomenon, we begin with some intuition.
We have seen that when GD encounters sharp minima, it oscillates near the minima because it cannot stabilize to the minima (due to 1). In other words, the loss change \(f({\boldsymbol{\theta}}^{t+1})-f({\boldsymbol{\theta}}^{t})\) would be much smaller compared to \(\left\|{\eta^2\nabla f({\boldsymbol{\theta}}^{t})} \right\|^2\) the square of the distance that GD travels. Hence, intuitively, one might expect that relative progress cannot be too negative under the unstable convergence regime. We would like to formalize this intuition next.
4.2 Characteristics in Iterates Movement↩︎
To that end, let us formally define what it means for GD to oscillate. More generally, consider the situation where \({\boldsymbol{\theta}}^{}\) is updated by moving along the vector \(-\mathbf{v}\). Then this update is oscillatory if the directional derivative at the updated parameter \({\boldsymbol{\theta}}^{}-\mathbf{v}{}\) is nearly negative of that at \({\boldsymbol{\theta}}^{}\), i.e., \[\begin{align} \left\langle \mathbf{v},~\nabla f({\boldsymbol{\theta}}^{} - \mathbf{v}) \right\rangle \approx -\left\langle \mathbf{v},~\nabla f({\boldsymbol{\theta}}^{}) \right\rangle\,. \end{align}\] Inspired by this, we consider the following definition.
Definition 2 (Directional smoothness). For an update vector \(\mathbf{v}\), we define \[\begin{align} L({\boldsymbol{\theta}}^{};\mathbf{v}):=\frac{1}{\left\|{\mathbf{v}} \right\|^2}\left\langle \mathbf{v},~ \nabla f({\boldsymbol{\theta}}^{})- \nabla f\big({\boldsymbol{\theta}}^{} - \mathbf{v}\big) \right\rangle. \end{align}\]
Now coming back to the gradient descent where the update vector is \(\mathbf{v}= \eta \nabla f({\boldsymbol{\theta}}^{})\), we have \[\begin{align} L({\boldsymbol{\theta}}^{};\eta \nabla f({\boldsymbol{\theta}}^{})) = \frac{\left\langle \nabla f({\boldsymbol{\theta}}^{}),~ \nabla f({\boldsymbol{\theta}}^{}) - \nabla f\big({\boldsymbol{\theta}}^{} - \eta \nabla f({\boldsymbol{\theta}}^{}) \big) \right\rangle}{\eta \left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2}. \end{align}\] When GD is exhibiting oscillatory behaviour, we would have \[\begin{align} \left\langle \nabla f({\boldsymbol{\theta}}^{}),~\nabla f({\boldsymbol{\theta}}^{} - \mathbf{v}) \right\rangle \approx -\left\langle \nabla f({\boldsymbol{\theta}}^{}),~\nabla f({\boldsymbol{\theta}}^{}) \right\rangle\,, \end{align}\] in which case, it holds that \[\begin{align} \label{eq:oscil} L({\boldsymbol{\theta}}^{};\eta \nabla f({\boldsymbol{\theta}}^{})) \approx \frac{2}{\eta}\quad \text{(when GD iterates oscillate)}. \end{align}\tag{4}\]
For intuition, let us quickly verify 4 for quadratic costs.
Example 4 (Quadratics). Consider a quadratic loss function \(f({\boldsymbol{\theta}}^{}) = {\boldsymbol{\theta}}^{\top} P{\boldsymbol{\theta}}^{}\) with \(P \succeq 0\). Then, the GD update reads \({\boldsymbol{\theta}}^{t+1} = ( I-\eta P) {\boldsymbol{\theta}}^{t}\). For an eigenvector/eigenvalue pair \((\mathbf{q},\lambda )\) of \(P\), the quantity \(\left\langle \mathbf{q}_{\max} ,~{\boldsymbol{\theta}}^{t} \right\rangle\) evolves as \[\begin{align} \left\langle \mathbf{q}_{\max} ,~{\boldsymbol{\theta}}^{t} \right\rangle &= \mathbf{q}^\top (I-\eta P){\boldsymbol{\theta}}^{t-1} = (1-\eta \lambda) \left\langle \mathbf{q}_{\max} ,~{\boldsymbol{\theta}}^{t-1} \right\rangle\\ &= (1-\eta \lambda)^t \left\langle \mathbf{q}_{\max} ,~{\boldsymbol{\theta}}^{0} \right\rangle\,. \end{align}\] This implies that if \(\lambda <2/\eta\), then \(\mathbf{q}^\top {\boldsymbol{\theta}}^{t}\to 0\). Hence, if \(\eta = 2/\lambda_{\max}^{} (P)\), then after sufficiently large iterations \(t\), we have \({\boldsymbol{\theta}}^{t} \approx (-1)^t \left\langle \mathbf{q}_{\max} ,~{\boldsymbol{\theta}}^{0} \right\rangle \mathbf{q}_{\max}\), in which case \(L({\boldsymbol{\theta}}^{};\eta \nabla f({\boldsymbol{\theta}}^{})) \approx \frac{2}{\eta}\). 0◻
Given the above view on “oscillating” iterates, we now measure directional smoothness under unstable convergence.
Under the same setting as Experiments [exp:stable] and [exp:unstable], we measure the value \(L({\boldsymbol{\theta}}^{t};\eta \nabla f({\boldsymbol{\theta}}^{t}))\) at each iteration.
Indeed, one can see that for the unstable regime \(L({\boldsymbol{\theta}}^{t};\eta \nabla f({\boldsymbol{\theta}}^{t}))\) saturates around \(2/\eta\), indicating that GD is exhibiting an oscillating behavior. 0◻
[exp:dir] verifies that GD is indeed showing an oscillating behavior under unstable convergence. We remark that a similar conclusion is made in [6] as well as the recent concurrent works [10], [11]. Now coming back to our original question: can we show a formal relation between the directional smoothness and the relative progress ratio?
4.3 Relation between Relative Progress Ratio and Directionl Smoothness↩︎
2 formalizes our intuition that under the oscillating behavior of GD, \({\sf RP}({\boldsymbol{\theta}}^{t})\) cannot stay below zero.
Theorem 2. The following identity holds: \[\begin{align} {\sf RP}({\boldsymbol{\theta}}^{}) = -1+\frac{\eta}{2}\cdot 2\int_{0}^1 \tau\cdot L({\boldsymbol{\theta}}^{};\eta \tau \nabla f({\boldsymbol{\theta}}^{})) \;\mathrm{d}\tau \,. \end{align}\]
Proof. See 7. ◻
2 implies that if the weighted average of \(L({\boldsymbol{\theta}}^{};\eta \tau \nabla f({\boldsymbol{\theta}}^{}))\) is close to \(2/\eta\), namely \[\begin{align} 2\int_{0}^1 \tau\cdot L({\boldsymbol{\theta}}^{};\eta \tau \nabla f({\boldsymbol{\theta}}^{})) \;\mathrm{d}\tau \approx \frac{2}{\eta}\,, \end{align}\] then \({\sf RP}({\boldsymbol{\theta}}^{})\) is indeed approximately equal to zero. This formally justifies that when GD shows an oscillating behavior, \({\sf RP}({\boldsymbol{\theta}}^{t})\) cannot stay below zero.
In our last experiment of this subsection, we verify that the above weighted average is approximately equal to the single value \(L({\boldsymbol{\theta}}^{};\eta \nabla f({\boldsymbol{\theta}}^{}))\), building a stronger relation between the directional smoothness and the relative progress ratio.
In the same setting as [exp:unstable], we choose step size \(\eta = 2/60\) and in every \(5\) iterations, we compute the following values: \[\begin{align} L({\boldsymbol{\theta}}^{t};\eta \tau \nabla f({\boldsymbol{\theta}}^{t}))\quad \text{for}~\tau\in \{0.01,0.02,\dots, 1\}. \end{align}\] In the plot below, we report the mean of \(L({\boldsymbol{\theta}}^{t};\eta \tau \nabla f({\boldsymbol{\theta}}^{t}))\) among \(\tau\in \{0.01,0.02,\dots, 1\}\) together with the shades which indicate the standard deviations.
This experiment verifies that \(L({\boldsymbol{\theta}}^{t};\eta \tau \nabla f({\boldsymbol{\theta}}^{t}))\) does vary too much across \(\tau \in [0,1]\). Hence, the single value \(L({\boldsymbol{\theta}}^{};\eta \nabla f({\boldsymbol{\theta}}^{}))\) well represents the weighted average in 2. 0◻
Hence, [exp:const] justifies the relation \[\begin{align} \label{rel:rp} \boxed{ {\sf RP}({\boldsymbol{\theta}}^{}) \approx -1+\frac{\eta}{2}\cdot L({\boldsymbol{\theta}}^{};\eta \nabla f({\boldsymbol{\theta}}^{}))\,,} \end{align}\tag{5}\] which precisely explains how the oscillatory behavior of GD results in a small relative progress ratio.
Interestingly, the validity of equation 5 and [exp:const] suggests that even though the gradient Lipschitzness is not a good assumption for neural networks, some form of Hessian Lipschitzness is valid along the GD trajectory.
We summarize the finding in this section as follows.
Under the unstable convergence regime, \({\sf RP}({\boldsymbol{\theta}}^{t})\) oscillates near \(0\) for the following two reasons:
\({\sf RP}({\boldsymbol{\theta}}^{t})\) can’t stay above \(0\) because otherwise the loss would not decrease in the long run.
\({\sf RP}({\boldsymbol{\theta}}^{t})\) can’t stay below \(0\) due to the oscillating behavior of GD iterates. This is formalized via 5 .
Given 5 , one can have a better explanation for [rmk:stable-1] regarding why \({\sf RP}({\boldsymbol{\theta}}^{t})\) saturates around \(-1\). In the second result of [exp:dir], the directional smoothness remains very small in the stable regime. Based on 5 , this implies that \({\sf RP}({\boldsymbol{\theta}}^{t})\) is close to \(-1\), which was indeed the case in [exp:stable].
4.3.1 Additional experiments↩︎
In this subsection, we verify the relation 5 for other experimental settings.
Under the same setting as [exp:stable] (the setting of the main experiments in [4]), this time we choose ReLU as activation functions.
In the next set of experiments, we put \(2\) more hidden layers of width \(200\) (total \(4\) hidden layers of width \(200\)).
The results are largely similar to those for \(\tanh\) activations, and the relation 5 holds for all cases. 0◻
4.4 Implications for Sharpness↩︎
In the previous subsection, we saw that one characteristics of unstable convergence is that \({\sf RP}({\boldsymbol{\theta}}^{t})\) saturates around zero. In this section, we investigate implications of this characteristics for sharpness. In particular, we discuss some relations to a curious phenomenon called edge of stability (EoS) recently observed in [4]. The gist of their observation is that for GD on neural networks often satisfies the following properties:
\(\lambda_{\max}^{}(\nabla^2 f({\boldsymbol{\theta}}^{t}))> 2/\eta\) for most iterates.
In fact, in many cases \(\lambda_{\max}^{}(\nabla^2 f({\boldsymbol{\theta}}^{t}))\) saturates right at (or slightly above) \(2/\eta\).
To that end, we begin with the following consequence of 2.
Corollary 1. Let \(L_{t}\) be the maximum sharpness along the line segment between the iterates \({\boldsymbol{\theta}}^{t}\) and \({\boldsymbol{\theta}}^{t+1}\), i.e., \(L_{t}:=\sup_{{\boldsymbol{\theta}}^{} \in \overline{{\boldsymbol{\theta}}^{t}{\boldsymbol{\theta}}^{t+1}}}\{\lambda_{\max}^{}(\nabla^2 f({\boldsymbol{\theta}}^{})) \}\). Then, the following inequality holds: \[\begin{align} \frac{2}{\eta }\cdot ({\sf RP}({\boldsymbol{\theta}}^{}) +1) \leq L_{t} \,. \end{align}\]
Proof. It follows from the fact that for each \(\tau \in [0,1]\) \(L({\boldsymbol{\theta}}^{};\eta \tau \nabla f({\boldsymbol{\theta}}^{})) \leq \sup\{ \lambda_{\max}^{}(\nabla^2 f({\boldsymbol{\theta}}^{}))~:~{\boldsymbol{\theta}}^{}\) lies on the line segment between \({\boldsymbol{\theta}}^{t}\) and \({\boldsymbol{\theta}}^{t}-\eta\tau \nabla f({\boldsymbol{\theta}}^{})\}\). Clearly, the right hand side is upper bounded by \(L_{t}\). ◻
1 implies that when \({\sf RP}({\boldsymbol{\theta}}^{t})\) oscillates around zero, then \(L_{t}\) has to be frequently above the threshold \(2/\eta\). One can actually refine this statement to understand the part A of EoS, given our results so far. In light of [exp:const], if \(L({\boldsymbol{\theta}}^{t};\eta \tau \nabla f({\boldsymbol{\theta}}^{t}))\) does vary much across \(\tau \in [0,1]\), then one can actually write \[\begin{align} \frac{2}{\eta } ({\sf RP}({\boldsymbol{\theta}}^{t}) +1) &\approx \lim_{\tau \to 0 } L({\boldsymbol{\theta}}^{t};\eta \tau \nabla f({\boldsymbol{\theta}}^{t}))\\ &\overset{\clubsuit}{\leq} \lambda_{\max}^{}(\nabla^2 f({\boldsymbol{\theta}}^{t})) \,, \end{align}\] which is the part A of EoS.
Moreover, let us for a moment additionally assume that \(\nabla f({\boldsymbol{\theta}}^{t})\) is approximately parallel to the largest eigenvector of the Hessian \(\nabla^2 f({\boldsymbol{\theta}}^{t})\). This might look stringent at first glance, but given the calculations in 4, this assumption is true for unstable GD on a quadratic cost. Also, recently, this behavior is theoretically proven for the normalized gradient descent dynamics [11]. Under this assumption, one can further deduce that the inequality \((\clubsuit)\) holds with approximate equality, and the part B of EoS would hold in that case.
5 Relative Progress for SGD↩︎
In this section, we extend our discussion to the stochastic gradient descent (SGD): \[\begin{align} {\boldsymbol{\theta}}^{t+1} = {\boldsymbol{\theta}}^{t} - \eta g({\boldsymbol{\theta}}^{t}),~~\text{where }\mathbb{E}[g({\boldsymbol{\theta}}^{t})] = \nabla f({\boldsymbol{\theta}}^{t})\,. \end{align}\] For the case of SGD, there is one obvious challenge. With SGD, the training loss does not decrease monotonically since SGD is a random algorithm. Hence, it is not clear how to precisely define what it means for SGD to be in the unstable regime. On the other hand, inspired by our discussion in 4.1, a more transparent way to define the unstable regime for SGD is via the relative progress ratio. In particular, we consider the following extension.
Definition 3 (Expected relative progress ratio). \[\begin{align} \mathbb{E}[{\sf RP}({\boldsymbol{\theta}}^{})]:=\frac{\mathbb{E}f({\boldsymbol{\theta}}^{}-\eta g({\boldsymbol{\theta}}^{}) )-f({\boldsymbol{\theta}}^{})}{\eta \left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2}\,. \end{align}\]
Under the same setting as [exp:relu], this time we train the network with SGD with minibatch size of \(32\) and step size \(\eta=2/100\). We compute the full-batch loss and the expected relative progress ratio at the end of each epoch.
Note that \(\mathbb{E}[{\sf RP}({\boldsymbol{\theta}}^{t})]\) does not stay below zero. Based on our discussion in 4.1, this suggests that SGD is in the unstable regime. 0◻
One very surprising aspect of the above results is that \(\mathbb{E}[{\sf RP}({\boldsymbol{\theta}}^{t})]\) is not negative for a majority of iterations. This is rather counter-intuitive given that in the loss plot SGD decreases the loss in the long run. On the other hand, we note that this counter-intuitive phenomenon is also observed by [4] in a comprehensive set of experiments. In particular, they mention “what may be more surprising is that SGD is not even decreasing the training loss in expectation.” See [4] and [4] for details.
We now establish a relation analogous to 5 for SGD. Similarly to 2, one can prove the following: \[\begin{align} &\mathbb{E}[{\sf RP}({\boldsymbol{\theta}}^{})] =-1 + \frac{\eta}{2} \cdot 2 \int_{0}^1 \tau \mathbb{E}_g\left[{\textstyle \frac{\left\|{g({\boldsymbol{\theta}}^{})} \right\|^2 \cdot L({\boldsymbol{\theta}}^{};\eta \tau g({\boldsymbol{\theta}}^{}))}{\left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2} }\right] \; \mathrm{d}\tau \,. \end{align}\] See 7 for derivation. Hence, the analogous relation to 5 is: \[\begin{align} \label{rel:sgd} \mathbb{E}[ {\sf RP}({\boldsymbol{\theta}}^{})] \approx -1+\frac{\eta}{2}\cdot \mathbb{E}\left[\textstyle \frac{\left\|{g({\boldsymbol{\theta}}^{})} \right\|^2}{\left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2} L({\boldsymbol{\theta}}^{};\eta g({\boldsymbol{\theta}}^{}))\right]. \end{align}\tag{6}\]
Under the same setting as [exp:sgd95relu], we compute \(\mathbb{E}[ {\sf RP}({\boldsymbol{\theta}}^{})]\) and the RHS of 6 at the end of each epoch and compare those values. We choose step sizes \(\eta = 2/50,\;2/100,\;2/150\).
Note that the LHS and RHS of 6 are very similar in all results, verifying the relation 6 . 0◻
We repeat [exp:sgd95verify] with \(\tanh\) activations.
The results are similar to those for ReLU networks.0◻
6 Discussion and Future Directions↩︎
This work demonstrated characteristics of unstable convergence that are in stark contrast with those of stable convergence. Consequently, this work leads to several interesting future directions.
\(\blacksquare\)
Better optimizer? The characteristics based on the directional smoothness suggests that the adjacent iterates have nearly opposite gradients in the unstable regime. This obviates the efficacy of many efficient methods (e.g. variance reduced methods) which are designed based on the intuition that the adjacent iterates have similar gradients. Hence, it would be interesting to design an efficient optimizer that respects the new characteristics.
Faster training under unstable convergence? As discussed in [rmk:fast], another striking feature of unstable convergence is that the training loss seems to converge faster. One interesting question is whether one can elucidate this faster optimization by further exploring our characterization of relative progress.
What assumptions are valid for neural network optimization? It is clear that unstable convergence cannot be reasoned with the widespread condition of \(\eta < \frac{2}{L}\). Then what assumptions would be valid for neural networks? As discussed in [rmk:hessian95lip], our [exp:const] suggests that although the gradient Lipschitzness is not a good assumption for neural networks, some form of Hessian Lipschitzness might be a valid one.
Unstable regime for adaptive methods? Our characterizations are limited to constant step size (S)GD, and it is not clear how these characterizations carry over to adaptive methods such as Adam and RMSProp. Investigating adaptive methods will help us understand how they differ from (S)GD for neural network optmization.
Acknowledgements↩︎
Kwangjun Ahn thanks Pourya Habib Zadeh for various discussions and his help with experiments during the initial stage of this work. Kwangjun Ahn also thanks Stephen J. Wright and Sinho Chewi for helpful discussions while visiting the Simons Institute for the Theory of Computing. Kwangjun Ahn also thanks Daniel Soudry for fruitful discussions regarding 1.
Kwangjun Ahn and Suvrit Sra acknowledge support from an NSF CAREER grant (1846088), and NSF CCF-2112665 (TILOS AI Research Institute). Kwangjun Ahn also acknowledges support from Kwanjeong Educational Foundation. Jingzhao Zhang acknowledges support from IIIS young scholar fellowship.
7 Proof of 2↩︎
Using the fact \[\begin{align} \frac{\mathrm{d}}{\mathrm{d}\tau} \left[ f({\boldsymbol{\theta}}^{} - \eta \tau \nabla f({\boldsymbol{\theta}}^{} )\right] = \left\langle -\eta \nabla f({\boldsymbol{\theta}}^{}),~\nabla f\big({\boldsymbol{\theta}}^{} - \eta \tau \nabla f({\boldsymbol{\theta}}^{})\big) \right\rangle\,, \end{align}\] we obtain \[\begin{align} f\big({\boldsymbol{\theta}}^{} -\eta \nabla f({\boldsymbol{\theta}}^{})\big) - f({\boldsymbol{\theta}}^{}) &= -\eta \int_{0}^1 \left\langle \nabla f({\boldsymbol{\theta}}^{}),~ \nabla f\big({\boldsymbol{\theta}}^{} - \eta \tau \nabla f({\boldsymbol{\theta}}^{}) \big) \right\rangle \mathrm{d}\tau\\ &=-\eta \left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2-\eta \int_{0}^1 \left\langle \nabla f({\boldsymbol{\theta}}^{}),~ \nabla f\big({\boldsymbol{\theta}}^{} - \eta \tau \nabla f({\boldsymbol{\theta}}^{}) \big)- \nabla f({\boldsymbol{\theta}}^{}) \right\rangle \mathrm{d}\tau\,. \end{align}\] Hence, after rearranging, we get \[\begin{align} \frac{\eta}{2} \cdot 2\int_{0}^1 \tau \cdot L({\boldsymbol{\theta}}^{};\eta \tau \nabla f({\boldsymbol{\theta}}^{})) \;\mathrm{d}\tau - 1= \frac{f\big({\boldsymbol{\theta}}^{} -\eta \nabla f({\boldsymbol{\theta}}^{})\big) - f({\boldsymbol{\theta}}^{})}{\eta \cdot \left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2 } = {\sf RP}({\boldsymbol{\theta}}^{})\,, \end{align}\] which is precisely the relation in 2.
7.0.0.1 Derivation for the SGD case.
For the SGD case, the derivation is similar. \[\begin{align} & f\big({\boldsymbol{\theta}}^{} -\eta g({\boldsymbol{\theta}}^{})\big) - f({\boldsymbol{\theta}}^{}) \\ &\quad \overset{(a)}{=} -\eta \int_{0}^1 \left\langle g({\boldsymbol{\theta}}^{}),~ \nabla f\big({\boldsymbol{\theta}}^{} - \eta \tau g({\boldsymbol{\theta}}^{}) \big) \right\rangle \mathrm{d}\tau \\ &\quad = -\eta \left\langle g({\boldsymbol{\theta}}^{}),~\nabla f({\boldsymbol{\theta}}^{}) \right\rangle -\eta \int_{0}^1 \left\langle g({\boldsymbol{\theta}}^{}),~ \nabla f\big({\boldsymbol{\theta}}^{} - \eta \tau g({\boldsymbol{\theta}}^{}) \big)-\nabla f({\boldsymbol{\theta}}^{}) \right\rangle \mathrm{d}\tau \end{align}\] where \((a)\) is due to the fact \[\begin{align} \frac{\mathrm{d}}{\mathrm{d}\tau} \left[ f({\boldsymbol{\theta}}^{} - \eta \tau g({\boldsymbol{\theta}}^{} )\right] = \left\langle -\eta g({\boldsymbol{\theta}}^{}),~\nabla f\big({\boldsymbol{\theta}}^{} - \eta \tau g({\boldsymbol{\theta}}^{})\big) \right\rangle\,. \end{align}\] After taking expectation over the randomness in the stochastic gradient, we obtain \[\begin{align} & \mathbb{E}f\big({\boldsymbol{\theta}}^{} -\eta g({\boldsymbol{\theta}}^{})\big) - f({\boldsymbol{\theta}}^{}) = -\eta \left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2 -\eta \int_{0}^1 \mathbb{E}\left\langle g({\boldsymbol{\theta}}^{}),~ \nabla f\big({\boldsymbol{\theta}}^{} - \eta \tau g({\boldsymbol{\theta}}^{}) \big)-\nabla f({\boldsymbol{\theta}}^{}) \right\rangle \mathrm{d}\tau \end{align}\] Hence, after rearranging we obtain the desired equation: \[\begin{align} & \frac{\mathbb{E}f\big({\boldsymbol{\theta}}^{} -\eta g({\boldsymbol{\theta}}^{})\big) - f({\boldsymbol{\theta}}^{})}{\eta \left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2} = -1 + \frac{\eta}{2} \cdot 2 \int_{0}^1 \tau\cdot \mathbb{E}_g\left[ \frac{\left\|{g({\boldsymbol{\theta}}^{})} \right\|^2 \cdot L({\boldsymbol{\theta}}^{};\eta \tau g({\boldsymbol{\theta}}^{}))}{\left\|{\nabla f({\boldsymbol{\theta}}^{})} \right\|^2} \right] \; \mathrm{d}\tau \,. \end{align}\] This completes the derivation.