Some "Model Training Strategies" Derived from the Ideas Behind the Amos Optimizer
If we think of training a model as "alchemy," then the "furnace" is clearly the optimizer. It's said that AdamW is currently the fastest scheme for training neural networks — I haven't compared them one by one myself, so I can't say for sure how true this is, but it's certainly the case that most pretraining nowadays uses AdamW or its variant LAMB. However, just as having a furnace doesn't guarantee you'll produce good elixirs, even once we've settled on AdamW, there are still many questions without definitive answers, such as:
1. How should the learning rate adapt to different initializations and parameterizations?
2. How should the weight decay rate be tuned?
3. What schedule should the learning rate follow?
4. Can we reduce the optimizer's GPU memory footprint?
Although in practice we can mostly just borrow the parameters and strategies that others have already tuned, the lack of a systematic tuning guide always leaves us feeling a bit unsure of ourselves when "training." In this post, based on the ideas from Google's recently proposed Amos optimizer, I'll give some reference results.
Background Recap
The Amos optimizer comes from Google's recent paper Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale, which provides fairly complete derivations for the questions above, and confirms their effectiveness through experiments. However, the derivations in the original paper are really not easy to read — the notation and approximations are quite arbitrary in places, giving an overall impression of "clutter." Fortunately the underlying ideas of Amos are not too complicated, so we can borrow them here.
Before starting the derivation, let's first review what the existing solutions to the questions above look like.
First, regarding the first question, many readers may not fully understand what "initialization" and "parameterization" mean — these are simply two different ways of setting up model weights. A common example is a $n\times n$ matrix, typically initialized with "mean 0, variance $1/n$"; for details see my earlier posts Understanding Model Parameter Initialization from a Geometric Perspective and A Brief Discussion on Transformer Initialization, Parameterization, and Normalization. From "variance $1/n$" we can already see that different parameters have different scales (or orders of magnitude); if we update all parameters with the same learning rate, the update magnitude for each parameter ends up being different. In my view, a rather elegant solution to this problem is the LAMB optimizer, where the norm of each update directly depends on the norm of the parameter itself, and the learning rate is only used to describe the relative size of the update.
As for the weight decay rate, at least in the pretraining domain, from what I've observed, people just stick with the earliest choice of 0.01, and I haven't seen much work on tuning this parameter. As for the learning rate schedule, everyone knows the learning rate should be gradually decayed to zero, but there isn't much theoretical guidance on exactly what decay schedule to use — most results are just experimentally derived summaries. Finally, on the question of saving GPU memory, the classic work here is the AdaFactor optimizer, which I discussed earlier in A Brief Analysis of the AdaFactor Optimizer (with Open-Source Implementation). There are basically two main approaches to reducing an optimizer's memory footprint: one is dropping momentum, the other is performing low-rank decomposition of the second moment. Amos essentially follows both of these approaches as well.
Problem Setup
This post is mainly concerned with the first three questions above, hoping to derive some "plug-and-play" results. First, let's write the optimizer's update rule in shorthand form as:
\begin{equation}\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \alpha_t \boldsymbol{u}_t\end{equation}
Here $\boldsymbol{\theta}_t, \boldsymbol{\theta}_{t+1}$ respectively represent the parameter values at time $t,t+1$, $\boldsymbol{u}_t$ represents the update vector at time $t$ (which depends on the task and data), and the scalar $\alpha_t > 0$ (with every element of the vector greater than 0) represents the learning rate at time $t$.
Starting with AdamW, mainstream optimizers have tended to separate the weight decay term out from $\boldsymbol{u}_t$, i.e.
\begin{equation}\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - (\alpha_t \boldsymbol{u}_t + \rho_t\boldsymbol{\theta}_t)\end{equation}
where $\rho_t > 0$ is the weight decay rate. The main task of this post is to try to resolve how $\alpha_t$ and $\rho_t$ should be set.
Weight Decay
We know that whether it's weight decay or L2 regularization, it is in itself unrelated to the training objective — it's just an auxiliary term whose purpose is to improve the model's generalization ability. Since it's auxiliary, a basic requirement is that it shouldn't "steal the show." To this end, let's introduce a constraint:
\begin{equation}\mathcal{O}(\alpha_t^2) = \mathcal{O}(\rho_t)\end{equation}
That is, throughout the entire update process, the update contributed by weight decay should always be one order higher than the update related to the training objective; since $\alpha_t,\rho_t$ is basically always less than 1, "higher order" here means "smaller."
Suppose the optimization endpoint of the parameters is $\boldsymbol{\theta}^*$, and let $\boldsymbol{\varepsilon}_t = \boldsymbol{\theta}_t - \boldsymbol{\theta}^*$. According to the update rule we can get
\begin{equation}\begin{aligned} \Vert\boldsymbol{\varepsilon}_{t+1}\Vert^2 =&\, \Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\theta}^*\Vert^2 \\ =&\, \Vert\boldsymbol{\theta}_t - (\alpha_t \boldsymbol{u}_t + \rho_t\boldsymbol{\theta}_t) - \boldsymbol{\theta}^*\Vert^2 \\ \approx&\, \Vert\boldsymbol{\varepsilon}_t\Vert^2 - 2 \alpha_t \boldsymbol{u}_t \cdot \boldsymbol{\varepsilon}_t + \left(\alpha_t^2 \Vert\boldsymbol{u}_t\Vert^2 - 2 \rho_t \boldsymbol{\theta}_t \cdot \boldsymbol{\varepsilon}_t\right) \end{aligned}\label{eq:base-approx}\end{equation}
The final approximation retains only terms of order no higher than $\mathcal{O}(\alpha_t^2)$.
Clearly, $\Vert\boldsymbol{\varepsilon}_t\Vert$ is the distance between the current result and the endpoint, and naturally we want it to be as small as possible — so we naturally hope that every update step reduces this distance, i.e. $\Vert\boldsymbol{\varepsilon}_{t+1}\Vert < \Vert\boldsymbol{\varepsilon}_t\Vert$. Now looking at equation $\eqref{eq:base-approx}$, $- 2 \alpha_t \boldsymbol{u}_t \cdot \boldsymbol{\varepsilon}_t$ can be either positive or negative — if it's negative, that helps achieve $\Vert\boldsymbol{\varepsilon}_{t+1}\Vert < \Vert\boldsymbol{\varepsilon}_t\Vert$, but $\alpha_t^2 \Vert\boldsymbol{u}_t\Vert^2$ is necessarily positive, which works against achieving $\Vert\boldsymbol{\varepsilon}_{t+1}\Vert < \Vert\boldsymbol{\varepsilon}_t\Vert$. However, once weight decay is introduced, there's an extra term $- 2 \rho_t \boldsymbol{\theta}_t \cdot \boldsymbol{\varepsilon}_t$; if this term can cancel out the negative effect of $\alpha_t^2 \Vert\boldsymbol{u}_t\Vert^2$, then introducing weight decay would not only enhance generalization but also help with convergence.
Feasibility Analysis
So next, we need to examine the feasibility of
\begin{equation}\alpha_t^2 \Vert\boldsymbol{u}_t\Vert^2 = 2 \rho_t \boldsymbol{\theta}_t \cdot \boldsymbol{\varepsilon}_t\label{eq:base-cond}\end{equation}
By "feasibility" I mean whether $\boldsymbol{\theta}_t \cdot \boldsymbol{\varepsilon}_t$ can be greater than 0 — only if it is greater than 0 is it possible for the two sides to be equal. Using the definition of $\boldsymbol{\varepsilon}_t$ we get $\boldsymbol{\theta}_t = \boldsymbol{\varepsilon}_t + \boldsymbol{\theta}^*$, so
\begin{equation}\boldsymbol{\theta}_t \cdot \boldsymbol{\varepsilon}_t = (\boldsymbol{\varepsilon}_t + \boldsymbol{\theta}^*) \cdot \boldsymbol{\varepsilon}_t = \Vert \boldsymbol{\varepsilon}_t\Vert^2 + \boldsymbol{\theta}^* \cdot \boldsymbol{\varepsilon}_t\end{equation}
Note that $\boldsymbol{\theta}^*$ is our target — a fixed point — while $\boldsymbol{\varepsilon}_t$ is the difference vector between the current time step and the target; generally speaking these two have no necessary correlation, so we can approximately treat them as two random vectors in a high-dimensional space. According to The Angle Distribution Between Two Random Vectors in n-Dimensional Space, we know that two random vectors in high-dimensional space are almost always orthogonal, so $\boldsymbol{\theta}^* \cdot \boldsymbol{\varepsilon}_t\approx 0$, i.e. $\boldsymbol{\theta}_t \cdot \boldsymbol{\varepsilon}_t \approx \Vert \boldsymbol{\varepsilon}_t\Vert^2$. Of course, if we want to be more careful, we can introduce a parameter $q$:
\begin{equation}\boldsymbol{\theta}_t \cdot \boldsymbol{\varepsilon}_t \approx q\Vert \boldsymbol{\varepsilon}_t\Vert^2\end{equation}
Then equation $\eqref{eq:base-cond}$ becomes
\begin{equation}\alpha_t^2 \Vert\boldsymbol{u}_t\Vert^2 \approx 2 \rho_t q\Vert \boldsymbol{\varepsilon}_t\Vert^2\label{eq:base-cond-approx}\end{equation}
Both sides are greater than 0, so equation $\eqref{eq:base-cond}$ could indeed hold.
Asymptotic Estimation
If equation $\eqref{eq:base-cond}$ holds, then equation $\eqref{eq:base-approx}$ simplifies to \begin{equation}\Vert\boldsymbol{\varepsilon}_{t+1}\Vert^2 \approx \Vert\boldsymbol{\varepsilon}_t\Vert^2 - 2 \alpha_t \boldsymbol{u}_t \cdot \boldsymbol{\varepsilon}_t = \Vert\boldsymbol{\varepsilon}_t\Vert^2 - 2 \alpha_t \Vert\boldsymbol{u}_t\Vert \Vert\boldsymbol{\varepsilon}_t\Vert \cos(\boldsymbol{u}_t, \boldsymbol{\varepsilon}_t)\end{equation}
We said that $\boldsymbol{u}_t$ represents the task-related update, and on average it should necessarily be beneficial to the task (otherwise the original optimizer would be flawed), so on average we should have $\cos(\boldsymbol{u}_t, \boldsymbol{\varepsilon}_t) > 0$. Let's further assume there exists a $p > 0$ such that $\cos(\boldsymbol{u}_t, \boldsymbol{\varepsilon}_t)\sim p$, so we have
\begin{equation}\Vert\boldsymbol{\varepsilon}_{t+1}\Vert^2 \approx \Vert\boldsymbol{\varepsilon}_t\Vert^2 - 2 \alpha_t p\Vert\boldsymbol{u}_t\Vert \Vert\boldsymbol{\varepsilon}_t\Vert\end{equation}
Using the approximation $\eqref{eq:base-cond-approx}$ we get $\alpha_t \Vert\boldsymbol{u}_t \Vert \Vert \boldsymbol{\varepsilon}_t\Vert \approx \sqrt{2 \rho_t q}\Vert \boldsymbol{\varepsilon}_t\Vert^2$. Substituting into the equation above gives
\begin{equation}\Vert\boldsymbol{\varepsilon}_{t+1}\Vert^2 \approx \Vert\boldsymbol{\varepsilon}_t\Vert^2(1 - 2 p\sqrt{2 \rho_t q})\approx \Vert\boldsymbol{\varepsilon}_t\Vert^2\exp(- 2 p\sqrt{2 \rho_t q})\end{equation}
Recursing this step by step, we get
\begin{equation}\Vert\boldsymbol{\varepsilon}_t\Vert^2 \approx\Vert\boldsymbol{\varepsilon}_0\Vert^2\exp\left(- 2 \sum_{i=1}^{t-1}p\sqrt{2 \rho_i q}\right)\label{eq:varepsilon-t}\end{equation}
We can see that the exponent on the right-hand side is necessarily monotonically decreasing — it's a decay function. Now let's revisit the approximation $\eqref{eq:base-cond-approx}$: it has two parameters, $\alpha_t$ and $\rho_t$, to be tuned, but only one (approximate) equation. To make $\alpha_t$ and $\rho_t$ decay at the same rate, we set $2\rho_t q \approx \lambda^2 \Vert\boldsymbol{\varepsilon}_t\Vert^2$, which gives us
\begin{equation}\begin{aligned}\alpha_t \approx \frac{\lambda\Vert\boldsymbol{\varepsilon}_t\Vert^2}{\Vert\boldsymbol{u}_t\Vert} \approx&\, \frac{\lambda\Vert\boldsymbol{\varepsilon}_0\Vert^2}{\Vert\boldsymbol{u}_t\Vert} \exp\left(- 2 \sum_{i=1}^{t-1}p\sqrt{2 \rho_i q}\right) \\ \rho_t \approx \frac{\lambda^2\Vert\boldsymbol{\varepsilon}_t\Vert^2}{2q} \approx&\, \frac{\lambda^2\Vert\boldsymbol{\varepsilon}_0\Vert^2}{2q} \exp\left(- 2 \sum_{i=1}^{t-1}p\sqrt{2 \rho_i q}\right) \end{aligned}\label{eq:alpha-rho}\end{equation}
This is the variation pattern for $\alpha_t,\rho_t$ derived in this post. Of course, now that we have the pattern, there are still four parameters $\lambda,\Vert\boldsymbol{\varepsilon}_0\Vert,p,q$ to determine; $q$ is relatively simple — just setting $q=1$ is fine — but even so there are still three parameters left to pin down.
Scale Estimation
By definition, $\Vert\boldsymbol{\varepsilon}_0\Vert = \Vert\boldsymbol{\theta}_0 - \boldsymbol{\theta}^*\Vert$ is the distance between the initial parameters and the target parameters, which can be understood as the scale of parameter variation. There are several different cases.
The first is when the parameter is a matrix-multiplication kernel, such as the kernel matrix of a fully-connected layer or a convolutional layer. These are typically initialized as "mean 0, variance $\sigma^2$" random initialization ($\sigma$ depending on the shape); this way, if $\boldsymbol{\theta}\in\mathbb{R}^k$, then we can estimate that $\Vert\boldsymbol{\theta}_0\Vert^2\approx k\sigma^2$. Additionally, this class of parameters has a property: under reasonable initialization, the mean and variance of the parameters won't change much after training completes — at least the order of magnitude stays consistent — so we can also take $\Vert\boldsymbol{\theta}^*\Vert^2\approx k\sigma^2$, and since the initialization is random, $\boldsymbol{\theta}_0 \cdot \boldsymbol{\theta}^*\approx 0$, and hence
\begin{equation}\Vert\boldsymbol{\varepsilon}_0\Vert^2 = \Vert\boldsymbol{\theta}_0 - \boldsymbol{\theta}^*\Vert^2 = \Vert\boldsymbol{\theta}_0\Vert^2 + \Vert\boldsymbol{\theta}^*\Vert^2 - 2\boldsymbol{\theta}_0 \cdot \boldsymbol{\theta}^* \approx 2k\sigma^2\end{equation}
The second case is additive bias terms, such as the bias vector of a fully-connected or convolutional layer, and the $\boldsymbol{\beta}$ vector of a Normalization layer. These parameters are typically "all-zero initialized," so $\Vert\boldsymbol{\varepsilon}_0\Vert^2 = \Vert\boldsymbol{\theta}^*\Vert^2$. If we predict from experience that the trained model's bias terms will be around $\pm\sigma$, then we can also estimate $\Vert\boldsymbol{\theta}^*\Vert^2\approx k\sigma^2$; the original Amos paper takes $\sigma=0.5$. Finally, there's the $\boldsymbol{\gamma}$ vector of the Normalization layer, which is typically "all-one initialized," and after training remains around 1 as well; assuming the error is $\pm\sigma$, we can estimate $\Vert\boldsymbol{\theta}^*\Vert^2\approx k\sigma^2$. Here $k$ refers to the vector dimension.
We can see that all of these $\Vert\boldsymbol{\varepsilon}_0\Vert^2$ results share a common form — they can all be written as $k\sigma^2$, where $\sigma$ is our prediction for the scale of parameter variation. For multiplicative matrices, $\sigma$ can simply be taken as the standard deviation of the initialization; for additive biases or the $\boldsymbol{\gamma}$ vector, we can simply take $\sigma=0.5$, or handle other special parameters with special treatment as needed.
Separating Out the Scale
Now let's look at the full update quantity. According to equation $\eqref{eq:alpha-rho}$, we have
\begin{equation}\alpha_t \boldsymbol{u}_t \approx \lambda\Vert\boldsymbol{\varepsilon}_0\Vert^2 \times \frac{\boldsymbol{u}_t}{\Vert\boldsymbol{u}_t\Vert} \times \exp\left(- 2 \sum_{i=1}^{t-1}p\sqrt{2 \rho_i q}\right)\end{equation}
where $\frac{\boldsymbol{u}_t}{\Vert\boldsymbol{u}_t\Vert}$ is a unit vector controlling the direction of the update, and $\exp$ is a decay term which we can set aside for now; so the magnitude of the update is controlled by $\lambda\Vert\boldsymbol{\varepsilon}_0\Vert^2$.
Going back to the first question at the start of this post — "How should the learning rate adapt to different initializations and parameterizations?" — the intuitive idea is clearly that parameters with a larger variation scale should get larger updates per step, or more simply, the update should be proportional to the variation scale. And we've just estimated the variation scale, which can be described by $\Vert\boldsymbol{\varepsilon}_0\Vert$, so we posit that $\lambda\Vert\boldsymbol{\varepsilon}_0\Vert^2=\alpha_0 \Vert\boldsymbol{\varepsilon}_0\Vert$ should hold, where $\alpha_0$ is the global initial learning rate. Solving this in reverse gives $\lambda=\alpha_0/\Vert\boldsymbol{\varepsilon}_0\Vert$; substituting into equation $\eqref{eq:alpha-rho}$ gives
\begin{equation}\alpha_t \approx \frac{\alpha_0\Vert\boldsymbol{\varepsilon}_0\Vert}{\Vert\boldsymbol{u}_t\Vert} \exp\left(- 2 \sum_{i=1}^{t-1}p\sqrt{2 \rho_i q}\right),\quad \rho_t \approx \frac{\alpha_0^2}{2q} \exp\left(- 2 \sum_{i=1}^{t-1}p\sqrt{2 \rho_i q}\right)\label{eq:alpha-rho-2}\end{equation}
Here $\alpha_0$ represents the relative update magnitude at each step (the global learning rate) — there's not much room for further derivation here; generally taking it around $10^{-3}$ is fine, or up to $10^{-2}$ for simpler tasks. $\Vert\boldsymbol{\varepsilon}_0\Vert$ was already estimated in the previous section, roughly $\sqrt{k}\sigma$; $\sigma$ represents the average variation scale of the parameter, which differs across different parameters — it's precisely through this term that we explicitly separate out the parameter scale, thereby achieving the effect of adapting to parameter scale (the update is proportional to $\sigma$). In particular, if we replace $\Vert\boldsymbol{\varepsilon}_0\Vert$ in the above equation with $\Vert\boldsymbol{\theta}_t\Vert$, we get exactly the LAMB optimizer. This also shows why, if the initialization mean of $\boldsymbol{\theta}$ is not 0 (as with the $\boldsymbol{\gamma}$ vector), replacing $\Vert\boldsymbol{\theta}_t\Vert$ with $\Vert\boldsymbol{\varepsilon}_0\Vert$ would be problematic — which is exactly why LAMB's approach is to simply leave the update for such parameters untransformed (i.e., keep the original update rule).
Analytic Approximation
Actually, the results so far are already suitable for implementation, except that the parameter $p$ is hard to tune. To further understand how the parameter $p$ affects the decay function, we can go further and derive an analytic approximation for $\rho_t$!
Multiplying both sides of $\rho_t$ in equation $\eqref{eq:alpha-rho-2}$ by $2q$, and then taking the square root of both sides, gives
\begin{equation}\sqrt{2q\rho_t} \approx \alpha_0 \exp\left(-\sum_{i=1}^{t-1}p\sqrt{2 \rho_i q}\right)\end{equation}
Denote the sum in the exponent $\sum\limits_{i=1}^{t-1}p\sqrt{2 \rho_i q}$ as $S_t$; then the equation above corresponds to the difference equation
\begin{equation}\frac{S_t - S_{t-1}}{p} \approx \alpha_0 \exp\left(- S_{t-1}\right) \quad \Rightarrow \quad S_{t+1} - S_t \approx \alpha_0 p\exp\left(- S_t\right)\end{equation}
The decay function is then $\exp\left(-2S_t\right)$. To find an asymptotic approximation, we replace the difference with a derivative (see Perturbation Methods for Difference Equations), giving
\begin{equation}\frac{dS_t}{dt} \approx \alpha_0 p \exp\left(- S_t\right)\end{equation}
This is a simple differential equation, which can be solved (combined with $S_0=0$) to give
\begin{equation}\exp\left(-2S_t\right) \approx \frac{1}{(\alpha_0 p t + 1)^2}\end{equation}
This is the explicit solution for the decay function, showing that the hyperparameter should decay as the inverse square of the step count. Substituting this into equation $\eqref{eq:alpha-rho-2}$, the full result is
\begin{equation}\alpha_t \approx \frac{\alpha_0\Vert\boldsymbol{\varepsilon}_0\Vert}{\Vert\boldsymbol{u}_t\Vert} \frac{1}{(\alpha_0 p t + 1)^2},\quad \rho_t \approx \frac{\alpha_0^2}{2q} \frac{1}{(\alpha_0 p t + 1)^2}\label{eq:alpha-rho-3}\end{equation}
This explicit solution not only makes implementation more convenient, it also makes the meaning of $p$ clearer. For example, if we want the learning rate to drop to half its original value after $T$ steps, then we need $(\alpha_0 p T + 1)^2=2$, from which we can solve
\begin{equation}\alpha_0 p = \frac{\sqrt{2}-1}{T}\end{equation}
As for what $T$ should actually be, that depends on the task difficulty and the amount of data, and there isn't much room for further derivation there either.
Dynamic Convergence
The discussion above assumes there exists a constant $p > 0$ such that $\cos(\boldsymbol{u}_t, \boldsymbol{\varepsilon}_t)\sim p$ — which can be understood as the model converging at a fixed rate. This is hard to satisfy in practice; it's more common for the convergence rate to slow relatively as training progresses further. To account for this, let's further assume that $p$ is a function $p_t$ of the step count $t$. With this change, the earlier derivation largely still holds, except that the corresponding constant $p$ needs to be replaced by the subscripted version $p_i$:
\begin{equation}\sqrt{2\rho_t q} \approx \alpha_0 \exp\left(- \sum_{i=1}^{t-1}p_i\sqrt{2 \rho_i q}\right)\end{equation}
Repeating the derivation from the previous section, we get
\begin{equation}\frac{S_t - S_{t-1}}{p_t} \approx \alpha_0 \exp\left(- S_{t-1}\right) \quad \Rightarrow \quad S_{t+1} - S_t \approx \alpha_0 p_t\exp\left(- S_t\right)\end{equation}
The approximate differential equation is then
\begin{equation}\frac{dS_t}{dt} \approx \alpha_0 p_t \exp\left(- S_t\right)\end{equation}
and the result of integration is
\begin{equation}\exp\left(-S_t\right) \approx \frac{1}{\alpha_0 \int_0^t p_{\tau} d\tau + 1}\end{equation}
But now there's an extra $p_t$ to determine. To reduce the cost of tuning, let's assume that the rate at which convergence slows down matches the rate at which $\Vert\boldsymbol{\varepsilon}_t\Vert$ decays; and according to equation $\eqref{eq:varepsilon-t}$, the decay function for $\Vert\boldsymbol{\varepsilon}_t\Vert$ is $\exp\left(-S_t\right)$, so we set $p_t = p_0\exp\left(-S_t\right)$, and substituting into the equation above gives
\begin{equation}\exp\left(-S_t\right) \approx \frac{1}{\alpha_0 p_0 \int_0^t \exp\left(-S_{\tau}\right) d\tau + 1}\end{equation}
This is essentially just a simple differential equation, which is easily solved to give
\begin{equation}\exp\left(-2S_t\right) \approx \frac{1}{2\alpha_0 p_0 t + 1}\end{equation}
Substituting this into equation $\eqref{eq:alpha-rho-2}$, we get
\begin{equation}\alpha_t \approx \frac{\alpha_0\Vert\boldsymbol{\varepsilon}_0\Vert}{\Vert\boldsymbol{u}_t\Vert} \frac{1}{2\alpha_0 p_0 t + 1},\quad \rho_t \approx \frac{\alpha_0^2}{2q} \frac{1}{2\alpha_0 p_0 t + 1}\label{eq:alpha-rho-4}\end{equation}
Looking purely at the decay schedule, this is exactly "inverse time decay," which is also one of the common learning rate decay schedules. Theoretically speaking, this result rests on more reasonable assumptions than the earlier equation $\eqref{eq:alpha-rho-3}$.
Summary
This post borrowed the ideas from the Amos optimizer to derive some results regarding the learning rate and weight decay rate, $\eqref{eq:alpha-rho-3}$ and $\eqref{eq:alpha-rho-4}$, which can be applied to existing optimizers in a plug-and-play fashion, and can to some extent simplify the difficulty of tuning hyperparameters.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.