Making Alchemy More Scientific (V): Fine-Tuning the Learning Rate Based on Gradients
In the previous four articles, we explored a series of convergence results for SGD, moving from bounded domains to unbounded domains, and from average-loss convergence to endpoint-loss convergence. Some readers might feel that, no matter how we phrase it, it's still all about SGD — surely these must be "ancient" results by now? Not at all! The core identity that the fourth post, Making Alchemy More Scientific (IV): New Identity, New Learning Rate, relies on comes from as recently as 2023; the conclusion in the third post, Making Alchemy More Scientific (III): Endpoint-Loss Convergence of SGD, is only slightly older, dating from 2020.
Also in the fourth post, we derived the practically common "linear decay" learning-rate schedule, showing that this chain of theoretical derivations is not mere armchair theorizing but can genuinely guide practice. Next, we will discuss finer-grained gradient-based learning-rate strategies. This will help us understand the principles behind learning-rate scheduling, and it also forms the basis of various adaptive learning-rate optimizers.
The Original Starting Point
If we carefully revisit the earlier proofs, we'll notice that the starting point of this whole series of results is an utterly unremarkable identity:
\begin{equation}\begin{aligned} \Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\varphi}\Vert^2=&\, \Vert\boldsymbol{\theta}_t - \eta_t \boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)- \boldsymbol{\varphi}\Vert^2 \\ =&\, \Vert\boldsymbol{\theta}_t - \boldsymbol{\varphi}\Vert^2 - 2\eta_t (\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t) + \eta_t^2\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 \end{aligned}\label{eq:begin}\end{equation}more
I call it "unremarkable" because it is so elementary that any decent high-school student could understand and prove it. And yet, this simple identity underpins an entire series of results about stochastic optimization — which really makes one appreciate how the greatest truths are often the simplest, and how returning to basics can be so powerful.
To understand this identity, we should note two kinds of "arbitrariness." First, $\boldsymbol{\varphi}$ is arbitrary — in most cases we take it to be the theoretical optimum $\boldsymbol{\theta}^*$, in order to obtain the most valuable results, but this doesn't change the fact that $\boldsymbol{\varphi}$ itself is arbitrary. We already know that one of the key steps in the derivations of Making Alchemy More Scientific (III): Endpoint-Loss Convergence of SGD and Making Alchemy More Scientific (IV): New Identity, New Learning Rate was substituting $\boldsymbol{\varphi}=\boldsymbol{\theta}_{T-k}$ and $\boldsymbol{\varphi}=\boldsymbol{\theta}_k$.
Second, $\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$ is arbitrary too, which might seem a bit surprising at first glance, since previous articles always took $\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$ to be the gradient of the loss function, $\nabla_{\boldsymbol{\theta}_t}L(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$. But in fact, the only role that setting it to the gradient plays is to combine with the convexity of $L$ to obtain
\begin{equation}L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi}) \leq (\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\end{equation}
which links the identity to the loss function. If we don't need this link, or have some other way of establishing it, then this choice is not necessary. This arbitrariness is crucial: it allows us to construct more complex update rules, and it is also the key that lets us generalize these results to non-SGD optimizers later on.
The Classical Result
Rearranging equation $\eqref{eq:begin}$ slightly gives
\begin{equation}2\eta_t (\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)= \Vert\boldsymbol{\theta}_t - \boldsymbol{\varphi}\Vert^2 - \Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\varphi}\Vert^2 + \eta_t^2\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 \label{eq:begin-2}\end{equation}
The "old-school" approach from here is to first divide both sides by $2\eta_t$, and then sum over $t$ from $1\sim T$. The advantage of this is that the left-hand side can be directly related to the loss via convexity; the disadvantage is that it requires assumptions of a bounded domain and a non-increasing learning rate in order to bound the right-hand side conveniently. So far, we have only used the "old-school" method in the first post, Making Alchemy More Scientific (I): Average-Loss Convergence of SGD, where the final result was
\begin{equation}\frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - \frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}^*)\leq \frac{R^2}{2T\eta_T} + \frac{1}{2T}\sum_{t=1}^T\eta_t\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2\label{leq:old-avg}\end{equation}
This time we explicitly kept $\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert$ rather than assuming $\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert\leq G$ for further simplification. Since we assumed the learning rate is non-increasing, we have
\begin{equation}\frac{R^2}{2T\eta_T} + \frac{1}{2T}\sum_{t=1}^T\eta_t\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2\geq \frac{R^2}{2T\eta_T} + \frac{\eta_T}{2T}\sum_{t=1}^T\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 \geq \frac{R}{T}\sqrt{\sum_{t=1}^T\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2}\label{leq:old-avg-optimal}\end{equation}
Let $V_t = \sum_{k=1}^t\Vert\boldsymbol{g}(\boldsymbol{x}_k,\boldsymbol{\theta}_k)\Vert^2$; then equality holds under the condition $\eta_1 = \eta_2 = \cdots = \eta_T = R/\sqrt{V_T}$. That is, the right-hand side of inequality $\eqref{leq:old-avg}$ attains its minimum under a constant learning rate $R/\sqrt{V_T}$, which is the fastest-converging learning rate. However, this result violates causality: at time $t$ we cannot "know in advance" the gradient norms at future time steps. One way to fix this while respecting causality is
\begin{equation}\eta_t = \frac{R}{\sqrt{V_t}} = \frac{R}{\sqrt{\sum_{k=1}^t\Vert\boldsymbol{g}(\boldsymbol{x}_k,\boldsymbol{\theta}_k)\Vert^2}}\label{eq:adagrad-mini}\end{equation}
After this modification, we need to re-prove
\begin{equation}\sum_{t=1}^T\eta_t\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 = R\sum_{t=1}^T\frac{V_t - V_{t-1}}{\sqrt{V_t}}\leq 2R\sum_{t=1}^T\frac{V_t - V_{t-1}}{\sqrt{V_t} + \sqrt{V_{t-1}}} = 2R\sum_{t=1}^T (\sqrt{V_t} - \sqrt{V_{t-1}}) = 2R\sqrt{V_T}\end{equation}
Substituting back into inequality $\eqref{leq:old-avg}$ gives
\begin{equation}\frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - \frac{1}{T}\sum_{t=1}^T L(\boldsymbol{x}_t,\boldsymbol{\theta}^*)\leq \frac{3R}{2T}\sqrt{\sum_{t=1}^T\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2}\end{equation}
The result is only about 50% larger than the ideal $\eqref{leq:old-avg-optimal}$, which is already quite good — and crucially, it doesn't violate causality, making it a practically feasible learning-rate strategy. It is also the prototype of the AdaGrad optimizer: if we apply the form of equation $\eqref{eq:adagrad-mini}$ element-wise to each component, we get the standard version of AdaGrad, which we will return to discuss later.
Being Careful with Expectations
Starting from the second post, Making Alchemy More Scientific (II): Generalizing the Results to Unbounded Domains, we adopted the "new-school" approach, which directly sums both sides of equation $\eqref{eq:begin-2}$:
\begin{equation}\begin{aligned} \sum_{t=1}^T 2\eta_t (\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t) =&\, \Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2 - \Vert\boldsymbol{\theta}_{T+1} - \boldsymbol{\varphi}\Vert^2 + \sum_{t=1}^T \eta_t^2\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 \\ \leq &\, \Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2 + \sum_{t=1}^T \eta_t^2\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2 \end{aligned}\end{equation}
Clearly, the advantage of the "new-school" approach is that it doesn't assume monotonicity of the learning rate, nor does it require a bounded domain — the $\Vert\boldsymbol{\theta}_t - \boldsymbol{\varphi}\Vert^2$ term on the right-hand side is naturally eliminated. The cost is that the left-hand summation acquires an extra weight $\eta_t$. This weight is not a big problem for SGD, but it is nearly fatal for adaptive learning-rate optimizers. So even though the new-school method has many elegant features, it almost never generalizes to adaptive learning-rate optimizers — in those cases we are still forced to rely on the old-school method.
Without straying too far, let's return to analyzing SGD. Applying convexity to the left-hand side of the equation above gives
\begin{equation}\sum_{t=1}^T \eta_t [L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})]\leq \frac{\Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2}{2} + \frac{1}{2}\sum_{t=1}^T \eta_t^2\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2\end{equation}
In the first three articles, the next step was to take the expectation $\mathbb{E}$ of both sides. But here we need to be very careful! Taking the expectation is over the entire $\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_T$, and in the first three articles we always assumed the learning rate $\eta_t$ was independent of the data, so $\mathbb{E}[\eta_t [L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})]] = \eta_t\mathbb{E} [L(\boldsymbol{x}_t,\boldsymbol{\theta}_t) - L(\boldsymbol{x}_t,\boldsymbol{\varphi})] = \eta_t\mathbb{E} [L(\boldsymbol{\theta}_t) - L(\boldsymbol{\varphi})]$ held, which let us relate the left-hand side to the target loss. In this article, however, we want to consider gradient-dependent learning rates, so $\eta_t$ cannot simply be factored out.
One relatively simple remedy is to assume that $\eta_t$ depends at most on $\boldsymbol{x}_1,\boldsymbol{x}_2,\cdots,\boldsymbol{x}_{t-1}$, in which case the expectation of $\boldsymbol{x}_t$ can be computed separately, e.g. $\mathbb{E}[\eta_t L(\boldsymbol{x}_t,\boldsymbol{\theta}_t)] = \mathbb{E}[\eta_t \mathbb{E}_{\boldsymbol{x}_t}[L(\boldsymbol{x}_t,\boldsymbol{\theta}_t)]] = \mathbb{E}[\eta_t L(\boldsymbol{\theta}_t)]$. But that feels somewhat restrictive, so instead we will simply assume that $\eta_t$ is independent of the data. Does this mean we can no longer implement gradient-based learning-rate adjustment? Not at all — it just means we can only implement adjustment based on the expected gradient $G_t = \sqrt{\mathbb{E}[\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2]}$, i.e.,
\begin{equation}\sum_{t=1}^T \eta_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\varphi})]\leq \frac{\mathbb{E}[\Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert^2]}{2} + \frac{1}{2}\sum_{t=1}^T \eta_t^2\underbrace{\mathbb{E}[\Vert\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)\Vert^2]}_{G_t^2}\label{leq:E-L}\end{equation}
Still the Old Story
With inequality $\eqref{leq:E-L}$ in hand, the subsequent derivations are fairly routine. We will generalize, one by one, the conclusions of the second, third, and fourth posts, and see what new insight they offer for learning-rate adjustment. The second post, Making Alchemy More Scientific (II): Generalizing the Results to Unbounded Domains, has two conclusions. The first assumes the learning rate is non-increasing; replacing $\eta_t$ uniformly on the left-hand side with $\eta_T$ and substituting $\boldsymbol{\varphi}=\boldsymbol{\theta}^*$ gives
\begin{equation}\frac{1}{T}\sum_{t=1}^T \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)] \leq \frac{R^2}{2T\eta_T} + \frac{1}{2T\eta_T}\sum_{t=1}^T \eta_t^2 G_t^2\end{equation}
where $R = \Vert\boldsymbol{\theta}_1 - \boldsymbol{\theta}^*\Vert$. Since the learning rate is non-increasing, it's easy to show the right-hand side also attains its minimum under a constant learning rate, in which case
\begin{equation}\eta_t = \frac{R}{ \sqrt{\sum_{k=1}^T G_k^2}}\end{equation}
Of course, this result also violates causality, so we must modify it to
\begin{equation}\eta_t = \frac{R}{\sqrt{\sum_{k=1}^t G_k^2}} = \frac{R}{ \sqrt{\sum_{k=1}^t \mathbb{E}[\Vert\boldsymbol{g}(\boldsymbol{x}_k,\boldsymbol{\theta}_k)\Vert^2]}}\end{equation}
But the $\mathbb{E}$ in the denominator is also unrealistic, since it would require repeatedly running many trials to compute the average. We simply drop $\mathbb{E}$ as an approximation to the above expression — or equivalently, assume the variance of $\Vert\boldsymbol{g}(\boldsymbol{x}_k,\boldsymbol{\theta}_k)\Vert^2$ is very small, so that a single sample is accurate enough. We will default to this operation for all subsequent results. In short, after removing the expectation $\mathbb{E}$, the result is essentially the same as $\eqref{eq:adagrad-mini}$, offering no new insight for now.
Inverse Proportion to the Gradient
The second conclusion of Making Alchemy More Scientific (II): Generalizing the Results to Unbounded Domains is a weighted-average form of inequality:
\begin{equation}\frac{\sum_{t=1}^T \eta_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)]}{\sum_{t=1}^T \eta_t}\leq \frac{R^2}{2\sum_{t=1}^T \eta_t} + \frac{\sum_{t=1}^T \eta_t^2 G_t^2}{2\sum_{t=1}^T \eta_t}\label{leq:avg-weighted}\end{equation}
Using the Cauchy–Schwarz inequality gives
\begin{equation}\sum_{t=1}^T \eta_t^2 G_t^2 = \frac{\sum_{t=1}^T (\eta_t G_t)^2 \sum_{t=1}^T G_t^{-2}}{\sum_{t=1}^T G_t^{-2}} \geq \frac{\left(\sum_{t=1}^T \eta_t G_t G_t^{-1}\right)^2}{\sum_{t=1}^T G_t^{-2}} = \frac{\left(\sum_{t=1}^T \eta_t\right)^2}{\sum_{t=1}^T G_t^{-2}}\end{equation}
Equality holds under the condition $\eta_t G_t \propto G_t^{-1}$, i.e., $\eta_t \propto G_t^{-2}$. Substituting this into the right-hand side of $\eqref{leq:avg-weighted}$, we get
\begin{equation}\frac{R^2}{2\sum_{t=1}^T \eta_t} + \frac{\sum_{t=1}^T \eta_t^2 G_t^2}{2\sum_{t=1}^T \eta_t} \geq \frac{R^2}{2\sum_{t=1}^T \eta_t} + \frac{\sum_{t=1}^T \eta_t}{2\sum_{t=1}^T G_t^{-2}}\geq \frac{R}{\sqrt{\sum_{t=1}^T G_t^{-2}}}\end{equation}
The overall condition for equality is $\eta_t = R G_t^{-2}/\sqrt{Q_T}$, where $Q_t = \sum_{k=1}^t G_k^{-2}$. The most notable feature of this result is that it tells us the learning rate should be inversely proportional to the squared gradient norm. This can be used to explain the necessity of Warmup: at the beginning of training, gradient norms tend to be large, and later they decrease and remain approximately constant for a long stretch. Modifying this into a causality-respecting form gives
\begin{equation}\eta_t = \frac{R G_t^{-2}}{\sqrt{Q_t}} = \frac{R G_t^{-2}}{\sqrt{\sum_{k=1}^t G_k^{-2}}}\end{equation}
After this modification, we again need to re-prove:
\begin{gather}\sum_{t=1}^T \eta_t = \sum_{t=1}^T \frac{R(Q_t - Q_{t-1})}{\sqrt{Q_t}} \geq \sum_{t=1}^T \frac{R(Q_t - Q_{t-1})}{\sqrt{Q_t} + \sqrt{Q_{t-1}}} = \sum_{t=1}^T R(\sqrt{Q_t} - \sqrt{Q_{t-1}}) = R \sqrt{Q_T} \\ \sum_{t=1}^T \eta_t^2 G_t^2 = \sum_{t=1}^T \frac{R^2(Q_t - Q_{t-1})}{Q_t} = R^2 + \sum_{t=2}^T \frac{R^2(Q_t - Q_{t-1})}{Q_t} \leq R^2 + R^2\sum_{t=2}^T \ln \frac{Q_t}{Q_{t-1}}= R^2+ R^2\ln \frac{Q_T}{Q_1}\end{gather}
Substituting these two results into $\eqref{leq:avg-weighted}$ gives
\begin{equation}\frac{\sum_{t=1}^T \eta_t \mathbb{E}[L(\boldsymbol{\theta}_t) - L(\boldsymbol{\theta}^*)]}{\sum_{t=1}^T \eta_t}\leq \frac{R}{\sqrt{Q_T}}\left(1 + \frac{1}{2}\ln \frac{Q_T}{Q_1}\right)\end{equation}
The key difference from the optimum is an extra logarithmic-growth factor, $\ln (Q_T/Q_1)$ — a common phenomenon that arises when switching from a static to a dynamic learning rate.
Focusing on the Present
Starting from the third post, Making Alchemy More Scientific (III): Endpoint-Loss Convergence of SGD, our conclusions began to address endpoint-loss convergence. Generalizing the core conclusion of that article to dynamic gradients yields
\begin{equation}\mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] \leq \frac{R^2}{2T\eta_T} + \frac{1}{2\eta_T}\sum_{t=1}^{T}\frac{\eta_t^2 G_t^2}{\max(1,\,T-t)}\label{leq:last-1}\end{equation}
Since this result also requires the learning rate to be non-increasing, it's easy to show the minimum of the right-hand side is $R\sqrt{V_T/T}$, attained under a constant learning rate:
\begin{equation}\eta_t = \frac{R}{\sqrt{T V_T}},\qquad V_t = \sum_{k=1}^t\frac{G_k^2}{\max(1,\,t-k)}\end{equation}
This learning rate is likewise quite interesting: it has a denominator similar to that of the "mini" AdaGrad learning rate $\eqref{eq:adagrad-mini}$ — both $V_t$ take the form of a sum of squared gradients. The difference is that here the gradient at time $k$ is weighted by $1/\max(1,\,t-k)$, meaning greater emphasis is placed on the current gradient. This is a new feature that emerges from moving from average-loss to endpoint-loss convergence, and it is already very close to the way RMSProp, Adam, and similar optimizers update their second moment via an EMA.
Intuitively, the causality-respecting version would be $\eta_t = R/\sqrt{t V_t}$, but this is not quite precise enough — the correct version should be $\eta_t = R/\sqrt{T V_t}$. Substituting into the right-hand side of equation $\eqref{leq:last-1}$ gives
\begin{equation}\begin{aligned} \frac{R^2}{2T\eta_T} + \frac{1}{2\eta_T}\sum_{t=1}^{T}\frac{\eta_t^2 G_t^2}{\max(1,\,T-t)} =&\, \frac{R}{2}\sqrt{\frac{V_T}{T}}\left(1 + \sum_{t=1}^{T}\frac{G_t^2/V_t}{\max(1,\,T-t)}\right) \\ \leq&\, \frac{R}{2}\sqrt{\frac{V_T}{T}}\left(1 + \sum_{t=1}^{T}\frac{1}{\max(1,\,T-t)}\right) \\ \leq&\, \frac{R}{2}\sqrt{\frac{V_T}{T}} (3 + \ln T)\\ \end{aligned}\end{equation}
where we used $G_t^2\leq V_t$. Compared to the optimum, there is again an extra logarithmic-growth factor, $\ln T$.
The Grand Synthesis
In Making Alchemy More Scientific (IV): New Identity, New Learning Rate, we obtained the strongest endpoint-loss convergence result so far. Generalizing it to dynamic gradients gives the "grand synthesis" of equations $\eqref{leq:avg-weighted}$ and $\eqref{leq:last-1}$:
\begin{equation}\mathbb{E}[L(\boldsymbol{\theta}_T) - L(\boldsymbol{\theta}^*)] \leq \frac{R^2}{2\eta_{1:T}} + \frac{1}{2}\sum_{t=1}^T\frac{\eta_t^2 G_t^2}{\eta_{\min(t+1, T):T}}\label{leq:last-2}\end{equation}
This result looks simple but is actually quite complex — at the very least, we cannot directly see under what learning-rate pattern the right-hand side attains its minimum. However, in the previous article the author offered an approach: approximate the problem continuously and solve it using the calculus of variations. Let's try that again here: set $S_t = \eta_{\min(t+1, T):T}$, so that for $t < T - 1$ we have $\eta_t = S_{t-1} - S_t\approx -\dot{S}_t$. Uniformly approximating $-\dot{S}_t$ by $\eta_t$, and approximating the sum by an integral, the right-hand side of the above equation becomes approximately (substituting $S_t = W_t^2$)
\begin{equation}\frac{R^2}{2S_0} + \frac{1}{2}\int_0^T \frac{\dot{S}_t^2 G_t^2}{S_t}dt = \frac{R^2}{2W_0^2} + 2\int_0^T \dot{W}_t^2 G_t^2 dt \label{eq:int-approx}\end{equation}
By the definition $W_T=0$, and fixing $W_0$, the integral part becomes a variational problem with fixed boundaries. The Euler-Lagrange equation gives $\frac{d}{dt}(\dot{W}_t G_t^2)=0$, i.e., $\dot{W}_t \propto G_t^{-2}$. Integrating both sides and combining with $W_T=0$ gives $W_t = W_0\int_t^T G_s^{-2} ds/\int_0^T G_s^{-2} ds$; substituting back into $\eqref{eq:int-approx}$ gives $R^2/2W_0^2 + 2 W_0^2 / \int_0^T G_s^{-2} ds$, whose minimum is attained at $2W_0^2 = R(\int_0^T G_s^{-2} ds)^{1/2}$. Finally, using the approximation $\eta_t\approx -\dot{S}_t=-2W_t\dot{W}_t$, we ultimately obtain
\begin{equation}\eta_t \approx \frac{R G_t^{-2} \int_t^T G_s^{-2}ds}{(\int_0^T G_s^{-2}ds)^{3/2}}\end{equation}
If we restore the discretization, we might guess that the optimal learning rate is roughly of the form
\begin{equation}\eta_t = \frac{R G_t^{-2} (Q_T - Q_t)}{Q_T^{3/2}} = \frac{R G_t^{-2}}{\sqrt{Q_T}} (1 - Q_t/Q_T)\label{eq:last-2-opt-lr}\end{equation}
where $Q_t$ is exactly the quantity $\sum_{k=1}^t G_k^{-2}$ defined earlier. As it turns out, this "guess" is indeed the correct answer! However, verifying it by substituting it back into equation $\eqref{leq:last-2}$ turns out to be extremely tedious, because the denominator of equation $\eqref{leq:last-2}$ is $\eta_{t+1:T}$ rather than $\eta_{t:T}$, so there's no guarantee that $\eta_t / \eta_{t+1:T}$ is bounded, which makes the bounding process especially difficult. The author tried for a week without success. We won't push further here — we will instead prove it in the next post using a more elegant construction.
For now, let's appreciate equation $\eqref{eq:last-2-opt-lr}$. Notice that $R G_t^{-2}/\sqrt{Q_T}$ is exactly the optimal learning rate for $\eqref{leq:avg-weighted}$, characterized by being proportional to $G_t^{-2}$ — as noted earlier, this can be used to explain Warmup. Equation $\eqref{eq:last-2-opt-lr}$ additionally multiplies by a factor $1 - Q_t/Q_T$, which is monotonically decreasing and, under suitable assumptions, decays linearly — so this term can explain Decay. Thus, equation $\eqref{eq:last-2-opt-lr}$ corresponds precisely to the classical "Warmup-Decay" learning-rate strategy.
What's even more interesting: the optimal learning rate from the previous section, $\eqref{leq:last-1}$, told us to focus more on the current gradient, but the result in equation $\eqref{eq:last-2-opt-lr}$ is even more extreme — it depends only on the current and future gradients, with no dependence on historical gradients whatsoever. This is undoubtedly a violation of causality. How can we modify it to respect causality? $\sqrt{Q_T}$ could be replaced with $\sqrt{Q_t}$; as for $Q_T$ within $Q_t/Q_T$, it can be written as $(Q_T/T)\times T$, and then we could consider replacing $Q_T/T$ with $Q_t/t$, giving
\begin{equation}\eta_t = \frac{R G_t^{-2}}{\sqrt{Q_t}} (1 - t/T)\end{equation}
This looks quite reasonable, but actually proving its validity by substituting it back into equation $\eqref{leq:last-2}$ is not an easy task either.
Summary
Starting with this article, we have begun considering gradient-based learning-rate scheduling. This helps us understand the underlying principles behind strategies such as Warmup and Decay, and also offers a useful reference point for building various kinds of adaptive learning-rate optimizers.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.