This post is a “companion piece” to our ICML submission, covered in the previous post. I wrote this up in a much more detailed technical report as well, for fun.
If you read our ICML paper, you’ll find yourself transported on an epic journey through a new paradigm of thinking of functions, not as just convex, smooth, strongly convex, but of *nonuniformly* convex, smooth, strongly convex etc. and their prevalence in logistic regression, softmax policies, and a specialized 2-layer neural network. (By the way, we do not impose convexity all the time, but just use it some of the time for improved rates).
This post is a review of the concepts in the paper, but from a point of view of an oversimplified version of the problem — both in dimension and in discretization. This was an exercise I had done to help me build intuition during the initial part of the project. The methods here are a lot simpler, a lot less groundbreaking, but also a lot easier to follow. In general, what I realized after doing this exercise, is that continuous-time analysis is great at producing intuition, but there’s no conversion to proofs over discrete — not for upper, nor lower bounds.
Logistic regression loss landscape
First, I want to revisit the loss landscape for a 2-parameter logistic regression model, when data is separable, and when it is not (below)
When the data is separable, the optimization landscape for logistic regression has no attainable minimum. Instead, it has what we call a descent cone, e.g. a direction for $\theta$ to move where $f(\theta)$ continues to decrease, forever. This must seem very strange because $f(\theta)$ is always decreasing and bounded below by 0, so in the strictest sense we do converge to 0, but it happens at $\theta\to \infty$, e.g. there is no reasonable definition of $\theta^*$. On the other hand, when data is not separable, the optimization landscape has a clear and finite, attained $\theta^*$.
The other important phenomenon to note is the curvature of the function in different places. In both separable and nonseparable cases, in both functions, the curvature is basically nil when $\|\theta\|$ is large, but there is a region where the Hessian of $\theta$ is quite large. This corresponds to the point where the classifier is most uncertain about all points. When the data is separable, this is not a good solution — you can find a separating hyperplane and then just become more confident about it! But when the data is not separable, this is likely the best solution.
What does this mean in practice? Well, it means that when you implement gradient descent over either problem, you don’t observe the similar phenomenon as if it were over linear regression, or something more resembling a pure quadratic (especially if ridge regularization is used). First, if the data is separable, obtaining a high-precision solution is really hard, because the second half of training is just eking your way down the descent cone, and just seems to take forever. (In practice, for simple log reg problems, this isn’t really an issue, since early stopping at low generalization error is usually possible. However, there’s reason to believe that this won’t work for DNNs, where the margin-maximizing part of training is also really important for updating the weights in the earlier layers.) Second, outside of this high-uncertainty region, the function is 1. super simple but 2. not at all curvy. Note that this is exactly the set of functions that reach the worst-case rates for gradient descent, and is also not the right set of functions for something like momentum acceleration methods.
Well that was a lot of griping. What are you going to do about it?
The recent series of papers [1-6] and ours, discussing this idea of nonuniform smoothness is a very nice toolset for handling this problem. Specifically, as we saw in the previous post, it can be shown that
- log. reg. and multiclass log. reg. is nonuniform smooth and nonuniform PL when the data is separable, and
- (not shown but proof by picture?) when the data is not separable, both problems exhibit uniform smooth behaviors far from the optimum, and switch to a more quadratic mode close to the optimum.
This suggests that in fact, a different set of methods, optimized for nonuniform smooth and nonuniform PL, is better suited for these types of problems. What exactly are such problems? So glad you asked!
- Normalized gradient descent. This method is best suited for the nonuniform set of problems, especially when the solution is not attained.
- Clipped GD. (Not us, but see [1, 5] and many others) This method specifically transitions from NGD to GD, as one gets close to a 0-gradient solution, which perfectly fits both separable and nonseparable regimes.
- RMSProp and Adam. This actually closely resembles NGD, with a bit of a lag. Moreover, the often-neglected $\epsilon$ term in the denominator does exactly the “transition to GD” thing that Clip-GD does, in the close-to-0-gradient regime.
And with that, the stage is set…
A 1-D mental picture of nonuniformity
To recap, the following are common conditions, applied to convex $f:\mathbb R\to\mathbb R$ (1-D functions): for all $\theta$,
\begin{align*}
f”(\theta) &\le L,
&&
&& \text{(Uniform smoothness)} \\[6pt]
f”(\theta) &\ge \mu,
&&
&& \text{(Strong convexity)} \\[6pt]
f”(\theta) &\le L_0 + L_1 |f'(\theta)|,
&&
&& \text{(Non-uniform smoothness, [1])} \\[6pt]
f”(\theta) &\le L_0 + L_1 (f(\theta)-f^*),
&&
&& \text{(Non-uniform smoothness, [2])} \\[6pt]
|f'(\theta)| &\ge \mu (f(\theta)-f^*)^\tau,
&& \tau \in (0,1]
&& \text{(Non-uniform Lojasiewicz [2])}
\end{align*}
If the first two conditions are true, and if additionally $f'(\theta^*) = 0$, then
\[
\frac{1}{2\mu}\|\nabla f(\theta)\|^2 \overset{\text{strong convexity}}{\geq} f(\theta)-f^* \overset{\text{uniform smooth}}{\geq} \frac{1}{2L}\|\nabla f(\theta)\|^2
\]
Using standard tricks, we see that this is equivalent to the following when $f$ is twice-differentiable:
\[
\mu \leq f”(\theta) \leq L.\tag{Uniform-1}
\]
And, it also follows that
\[
\frac{(f'(\theta))^2}{2L} \leq f(\theta)-f^* \leq \frac{(f'(\theta))^2}{2\mu}. \tag{Uniform-2}
\]
Now if we merge the second set of conditions, and kind of make an argument that the sign of $f'(\theta)$ does not change during the training trajectory (true in the continuous analysis case) then we can kind of summarize it as
\[
\mu |f'(\theta)| \leq f”(\theta) \leq L|f'(\theta)|\tag{Nonuniform-1}
\]
and
\[
\frac{|f'(\theta) |}{L} \leq f(\theta)-f^* \leq\frac{|f'(\theta)|}{\mu}.\tag{Nonuniform-2}
\]
Quick and dirty convergence rates
So now is the coolest part of continuous time analysis: we can do really fast and intuition-building convergence rates using actually only the back of an envelope. Specifically, for a method of the form
\begin{equation}
\theta^{(k+1)} = \theta^{(k)} + \eta \mathcal A(\theta^{(k)}, \theta^{(k-1)},…)
\label{eq:method-discretization}
\end{equation}
a continuous time analysis focuses on the specific ODE
$$\dot \theta(t) = \mathcal A(\{\theta(\tau)\}_{\tau \leq t})$$
which is the limiting behavior as $\eta\to 0$. Then, for $f:\mathbb R\to \mathbb R$ (1-D functions) we may characterize the trajectory as
\[
\dot \theta = -f'(\theta) \quad \mathrm{(GD~flow),} \qquad\qquad \dot \theta = 1 \quad \mathrm{(NGD~Flow)}
\]
Now consider a function $f$ which is non-uniform smooth and non-uniform PL, and thus satisfies
\begin{equation}
-\mu f'(\theta) \leq f”(\theta) \leq – L f'(\theta).
\end{equation}
Then defining the trajectory $g(t) = f'(\theta(t))$, then $\dot g = f”(\theta) \dot \theta$, and
\begin{align*}
\dot g &\in -[\mu,L]\, g^2
&\Longrightarrow&
& \frac{1}{L t + g_0} \;\le\; g(t) \;\le\; \frac{1}{\mu t + g_0}
&& \text{(GD flow)} \\[6pt]
\dot g &\in -[\mu,L]\, g
&\Longrightarrow&
& g_0 e^{-Lt} \;\le\; g(t) \;\le\; g_0 e^{-\mu t}
&& \text{(NGD flow).}
\end{align*}
Under a few additional reasonable assumptions, we can show that $f(\theta(t))-f^*$ decays at the same rate, and thus in only a few lines have shown a clear separation between the performance of NGD and GD for nonuniform smooth and nonuniform PL functions; specifically, that NGD is fundamentally superior.
Also, using very similar techniques, we can show that over uniformly smooth and strongly convex functions,
\begin{equation}
\mu \leq f”(\theta) \leq L.
\label{eq:assp:uniform}
\end{equation}
then using very similar techniques, we see that
\begin{align*}
\dot g &\in -[\mu,L]\, g
&\Longrightarrow&
& g_0 e^{-Lt} \;\le\; g(t) \;\le\; g_0 e^{-\mu t}
&& \text{(GD flow)} \\[6pt]
\dot g &\in -[\mu,L]
&\Longrightarrow&
& g_0 (1 – Lt)\;\le\; g(t) \;\le\; g_0 (1-\mu t)
&& \text{(NGD flow).}
\end{align*}
Again, we see clear superior behavior of NGD flow, which suggests finite-time convergence, over the linear convergence of GD flow. (Admittedly, this second result is less exciting; while the first three bounds are also reflected in the method convergence rates, this last result does not manifest in practical convergence rates, as discretization errors dominate for NGD near attained minima. So, the method separation for uniformly smooth and strongly convex functions is only true for continuous trajectories.)
Key point: Adam and RMSProp are actually NGD-like in disguise
This is more or less what we are able to show in the previous paper, but I wanted to reemphasize this point for our toy problem. In the continuous world, we can show this by demonstrating that for both methods, in the limit, $\dot \theta = c$ a constant that depends on the hyperparameter of the method, but is not time-dependent, decaying or increasing! This was simply done by writing the Adam and RMSProp methods as 2nd order ODEs
\[
\dot \theta = – \frac{ f'(\theta)}{\sqrt{v}+\epsilon} , \qquad
\dot v = \eta^{-1}(1-\beta) (f'(\theta)^2 – v), \tag{RMSFlow}
\]
\[
\dot \theta = – \frac{ m \tau }{\sqrt{v}}, \qquad \dot m = b_1 (f'(\theta)-m), \qquad \dot v = 2b_2(f'(\theta)^2-v)
\tag{Adam Flow}
\]
and then characterizing their stationary points. Don’t get me wrong, the arguments were super messy and required a lot of back and forth checking to make sure there wasn’t any corner cases. But all in all, it seemed to really corroborate some of the results we are seeing in literature!
So, tl;dr, what did we learn?
- Problems of interest in machine learning are not that similar to quadratics, and are better described by non-uniform smooth and non-uniform PL bounds
- When considering the class of non-uniform smooth and PL (hereafter referred to as NUS) problems, normalized gradient descent (NGD) can fundamentally converge at an order of magnitude faster than gradient descent
- In practice, this can be used to understand the superiority of Adam and RMSProp, which can be shown asymptotically to behave like NGD, using continuous flow type analysis.
References
[1] Zhang, Jin, Fang, Wang. Improved Analysis of Clipping Algorithms for Nonconvex Optimization. Neurips, 2020.
[2] Chen, Z., Zhou, Y., Liang, Y., and Lu, Z. Generalized smooth nonconvex optimization is as efficient as smooth nonconvex optimization. ICML 2023
[3] Wang, B., Zhang, H., Meng, Q., Sun, R., Ma, Z. M., and Chen, W. On the convergence of Adam under non-uniform smoothness: Separability from SGDM and beyond.
[4] Vankov, D., Rodomanov, A., Nedich, A., Sankar, L., and Stich, S. U. Optimizing (L0,L1)-smooth functions by gradient methods.
[5] Gorbunov, E., Tupitsa, N., Choudhury, S., Aliev, A., Richtarik, P., Horvath, S., and Takac, M. Methods for convex (L0, L1)-smooth optimization: Clipping, acceleration, and adaptivity
[6] Vaswani, S. and Harikandeh, R. B. Armijo line-search can make (stochastic) gradient descent provably faster. ICML, 2025