AdaLip: An Adaptive Learning Rate Method per Layer for Stochastic Optimization

Before describing our contributions, we will present a generic framework, which can represent all adaptive optimization methods. In the sequel, we will adopt standard notation and relevant mathematical techniques from the literature [21, 23]. We will use \(\mathcal {F} \in \Re ^d\) to denote the set of feasible points for the weights \(w_t\). We assume that the set \( \mathcal {F}\) has bounded diameter \(D_{\infty } \in \Re \), if \(\Vert x - y \Vert \le D_{\infty }\), for all \(x,y \in \mathcal {F}\). In the feasible set \(\mathcal {F}\), there is another assumption that \(\Vert \nabla f_t(x) \Vert _{\infty }\) is bounded \(\forall \, t \in [0, 1, ..., T]\). This generic framework is portrayed in Algorithm 1.

Table 1 summarizes how various algorithms fit with the generic framework. The main differences in these algorithms lie in the selection of the “averaging" functions \(\phi \) and \(\psi \), through which the parameters \(m_t\) and \(V_t\) are updated. Through this abstraction, the differences among the various optimizers become more apparent.

Let \(f: \Re ^d \rightarrow \Re \) be a function with a smooth gradient:Consider an optimization problem, where the goal is to find the set of parameters w that minimize the loss function f (i.e. Eq. 1), through Gradient Descent (Eq. 2).

The proof for that can be found in the Appendix A. Computing the optimal learning rate, however, requires prior knowledge of the Lipschitz constant of the loss function’s gradient, which is generally not the case.

In practise, it is difficult to find an overall good learning rate because that can change from dataset to dataset and from model to model. Also, because the majority of deep learning models are trained with variants of SGD, meaning noisy gradients, calculating the exact value of L is not easily feasible. If, however, the Lipschitz constant could be learned during training and its approximate value is accurate, then one could adapt the learning rate towards its optimal value as training progresses. Up till now this has remained an open research question [18].

Considering that the update steps are quite small for each iteration, we can calculate the constant L of a small subspace, the one that the optimizer explores, assuming that it is L-smooth. This means that the information of the subspace of the loss function needed was computed during a single forward pass. A multiple forward pass scheme could be approached by computing the gradients of different nearby directions, but it costs extra training time making it extremely infeasible for large scale datasets and deep models. Below, we present the approximation of L in a stochastic environment. It can be derived from Eq. 4, if \((w_1, w_2)\) is substituted by \((w_t, w_{t-1})\):From Eq. 2:The analysis so far has been based around Gradient Descent. In a realistic scenario the stochastic version will take its place. This translates to the gradients being calculated from a subset of the dataset (i.e. a batch) and, as a result, from a subspace of the loss space. These gradients, \(g_t\), are noisy versions of the total gradient but their expected value is equal to the total:By renaming \(\nabla f(w_t)\) as \(g_t\) and using Eq. 5, the above inequality transforms to this:where \(\alpha ^*\) is the optimal learning rate and \(\alpha _t\) is the initial learning rate.

Equation 6 offers a way of approximating the optimal value of the learning rate \(\alpha ^*\). Ideally, we would want to set the learning rate at each to the value indicated. This practice, however, isn’t recommended, as this equation exhibits a high degree of variance from iteration to iteration, which would result in very unstable training.

One issue that arises is due to the denominator, which is the norm of the difference of the gradients of two iterations. This can be seen as the magnitude of change of the direction the optimizer chooses. The problem happens especially near local minima, where the updates are smaller and close to each other. In this situation the denominator might reach values close to zero, causing the learning rate to explode. A workaround is to add a small positive term, \(c_t\), to the denominator:This creates an artificial low bound to the denominator, which can change over the course of training. In the following sections, the impact of the hyperparameter \(c_t\) to the training process, will be explored in more detail.

The instability of the algorithm, however, isn’t solely attributed to the denominator; the stochastic nature of the algorithm also plays its part. Equation 7 helps with approximating the optimal learning rate of the loss’ subspace, visible through each batch \(x_t\). Our goal is to approximate the optimal learning rate of the underlying loss function, though. To achieve this goal, the moving average of \(\alpha ^*\) for each batch is computed:with \(\gamma \in (0,1)\) being the coefficient of the moving average. In closed form this can be written as:with \(S_0 = 1\). Thus, the approximation \(A_t\) of the optimal learning rate \(\alpha ^*\) is calculated as a product of \(\alpha _t\) and the above exponential moving average:

The above lemma proves that the moving average of Eq. 8 is bounded and, as a result, the learning rate \(A_t\) is also bounded. This will be helpful for the next theorem to prove the convergence of a SGD optimizer with such learning rate. For the proof of convergence, the term of regret will be used. Regret (R) is the sum of all previous differences between the current network’s prediction \(f_t(w_t)\) and the best possible prediction \(f_t(w^*)\), obtained from the optimal set of weights \(w^*\). The goal of the proof is to show that the regret averaged over time reaches zero. The regret can be written in the following form:

The proofs for both Lemma 2 and Theorem 1 can be found in the Appendix A.

Usually in neural networks, parameters within the same layer share similar properties. From layer to layer, though, they might differ. An occasion where this is important is the case of vanishing gradients [28], where the early layers of the network won’t be trained adequately. While this problem has been addressed successfully with non-saturating activation functions [29], better weight initialization strategies [30] and transformation layers [28], we can see, in practice, none of the above solve the issue directly.

In order to have a better insight of the training process and how layers differ from each other, an experimental scheme was performed in order to observe the magnitude of the gradients and the magnitude of the updates in 3 different datasets comparing two optimizers, SGD and Adam with steady learning rates. To measure that we computed the mean of the absolute values of the gradients and the updates of each layer of the network for every iteration. The results are shown in Fig. 1. The networks that were used are shown in Appendix B. The blue lines correspond to the earlier iterations of the training and as the colormap reaches the red color, it is the indication for the last iterations. The layers displayed are only Convolutional and Dense layers. We can see that the left two columns display the networks that were trained with SGD. As expected, the magnitude of the updates is proportional to the magnitude of the gradients and their shapes are quite similar. On the other hand, the networks trained by the Adam optimizer showed faster convergence and achieved higher accuracy. Analysing the two right columns of the Fig. 1, we can assess that the magnitude of the updates is not proportional to the magnitude of the gradients and does not mimic the shape of the gradients. This is an insight on what makes Adam better than SGD in this specific scenario. Adam is able to construct updates closer to the optimal ones independently of the magnitude of the gradient. SGD usually suffers from that, meaning that some layers in the middle might display low gradients, hindering the path to the optimal updates. This shows the usefulness of an adaptive learning rate per layer that will help any optimizer, like SGD or Adam, scale the gradients and, in consequence, the updates in a way that it will make it easier to reach the optimal weights.

All the weights and gradients that were mentioned in the previous equations of Sect. 3 were the full vectors containing every parameter of the network. However, in a realistic setting it is difficult to construct one large vector of this size (usually would have millions of dimensions) and perform calculations, such as multiplications, computing norms, etc. This is the second reason why it is preferable to compute the updates of the network for every layer separately. This led to the motivation of constructing an optimizer that employs a different adaptive learning rate per layer. In the next section, the full algorithm will be displayed and how it fits with other optimizers.

Algorithms 2, 3 and 4 show the three new optimizers that are created. The first is the vanilla version of AdaLip, where the algorithm adapts the learning rates of a SGD optimizer. As discussed previously, the algorithm estimates the Lipschitz constant for each batch and, through a moving average, approximates the optimal learning rate \(A_t\) (Eq. 9). This is done for each layer independently for the reasons mentioned in Sect. 3.4. It is important to note that the norms displayed in the Algorithms 2, 3 and 4 are computed using only the gradients of the weights that are in the same layer with the weight that is being updated in that iteration.

The AdaLip methodology can be applied on top of existing optimizers. One such instance is AdaLip + RMSProp, which will be referred to as RMSLip and is described in Algorithm 3. Again, the optimal learning rate is estimated and is applied for each layer separately. The difference, here, is that each parameter is scaled by the moving average of the squares of its past gradients, as dictated by the RMSProp update rule.

Finally, we’ll examine the combination of AdaLip and Adam, which will be referred to as AdamLip. The update rule of Adam is left unchanged (i.e. the moving averages of the gradients \(m_t\) and their squares \(\nu _t\) and their bias-correction), however in this version, the constant learning rate is substituted with \(A_t\) (Eq. 9).

This Section will explain how the hyperparameter values for AdaLip were chosen and initialized.

From Theorem 1 the optimal \(c_t\) has been calculated in order to have a guarantee of convergence. In practice, the experiments were run with a constant \(c_t = c = 10^{-8}\), which is close to the theoretical \(c_t\). This will be discussed in more detail in Sect. 6.2.

For the initial learning rate \(\alpha _0\) we selected many values in order to test its robustness and it will be further analyzed in Sect. 5. The initial value of the moving average \(S_0\) is set to one, because then the overall learning rate will start at the initial learning rate \(\alpha _0\).

Another important hyperparameter to tune is the coefficient \(\gamma \) of the moving average. Its range is (0, 1) and the proposed value is 0.8, which was found empirically. Lower values tend to affect the learning rate to change drastically from iteration to iteration. This could be beneficial to overcome bad local minima that the the network can get stuck, but it also makes the algorithm unstable. On the other hand, higher values will smooth out the overall learning rate containing any possible spikes.

First, AdaLip was tested on MNIST, a dataset of small grayscale images of handwritten digits, with a size of \(28 \times 28\). It consists of 60000 training images and 10000 for testing. The images were normalized to the range of [0, 1] and no further preprocessing or augmentation was added.

The SGD-based optimizers were examined for 8 different learning rates for this task (i.e. 0.005, 0.01, 0.05, 0.1, 0.5, 0.7, 1.0, 1.3), Adam-based ones were examined for 4 (i.e. 0.0005, 0.001, 0.005, 0.01) and RMSProp optimizers for 6 (i.e. 0.0003, 0.0005, 0.0007, 0.001, 0.005, 0.01).

AdaLip in conjunction with every optimizer seems to perform better in both max and mean performance. Also, it has lower std meaning that the results are more stable and fluctuate less with the change of learning rate. Regarding the Epochs of convergence we see a slight improvement in mean and median of the Epochs of conversion. Considering that MNIST is an easy to train dataset and it converges extremely fast, there is not any room for big improvements to achieve. For a closer look, the mean curve of all learning rates for SGD and Adam based optimizers is displayed in Fig. 2. The shading represents the variance from different learning rate experiments. The AdaLip versions seem to converge faster and reach higher accuracies.

CIFAR10 was the second dataset to test our methodology. CIFAR10 consists of colored images with \(32 \times 32\) size. The training set has 50000 images and the test set 10000, while they are divided in 10 classes. The images were normalized as in MNIST between [0, 1] and no further preprocessing was added.

In this task, the SGD-based optimizers were examined for 10 learning rates (i.e. 0.005, 0.01, 0.05, 0.07, 0.1, 0.2, 0.3, 0.5, 0.7, 1.0), Adam-based ones for 7 learning rates (i.e. 0.0002, 0.0003, 0.0005, 0.0007, 0.001, 0.005, 0.01) and RMSProp-based ones for 5 (i.e. 0.0003, 0.0005, 0.0007, 0.001, 0.005). The same procedure was followed here, which yielded similar results with MNIST. AdamLip and RMSLip outperformed their counterparts on mean and max scores. AdaLip scored 0.5% lower in max performance but the mean and median results were quite an improvement compared to SGD. Regarding the std scores RMSLip shows great stability while the rest display similar behavior. As far as the Epochs of convergence analysis, Adam and SGD are slight ahead (by 1 epoch or less on average) but RMSLip is faster than its counterpart. The mean curve of all learning rates for SGD and Adam based optimizers is plotted in Fig. 3. It is clear that the AdaLip versions converge faster and to higher accuracy.

The last dataset used for testing the proposed framework was CIFAR100. It has the same images as CIFAR10, but, instead of 10 classes, the dataset is divided in 100. This makes it a lot more difficult to train a neural network and to achieve a good generalization score. To help with the generalization, image augmentation was performed after the normalization. The augmentation techniques were random rotations, flips and width/height shifts that help with the better training of the network.

Here we examined 6 learning rates (i.e. 0.005, 0.01, 0.05, 0.1, 0.5, 0.7) for SGD-based optimizers, 4 learning rates (i.e. 0.0005, 0.001, 0.005, 0.01) for Adam-based ones and 6 (i.e. 0.0003, 0.0005, 0.0007, 0.001, 0.005, 0.01) for RMSProp-based ones. Again, in all max and mean scores the AdaLip-based algorithms show an improvement in the overall training procedure. AdaLip has while having a lower std. Regarding the Epochs of convergence AdaLip shows great improvement over SGD. AdamLip converges by 2 epochs faster on average than Adam, while RMSLip displays a faster overall training progress. Figure 3 shows the mean curves, like in the previous datesets.

The focus of this paper is mainly revolving around the training performance of the optimizers, but, the generalization performance is extremely important as well. In order to measure the testing scores of each optimizer, we computed the mean of the maximum test accuracy of each run of the aforementioned experimental process. Meaning that for each optimizer and for each learning rate a mean testing accuracy is extracted for the three benchmark datasets. Table 5 displays the best test accuracies of the optimizers among all the starting learning rates for each case. We can see that in most cases the scores are close, but AdaLip’s version display an improvement. Specifically, AdamLip’s performance is slightly better than Adam in all 3 datasets. RMSProp seems to achieve better generalization in CIFAR10 than RMSLip, but on the the other datasets the improvement is smaller. Lastly, regarding SGD and AdaLip, AdaLip seems to perform better in all 3 datasets. Overall, the test scores show a slight advantage for the AdaLip-based optimizers.

As previously mentioned in Sect. 3.4, the norms of the weights of the layers follow a U-like shape. Specifically, the first and last layers tend to exhibit larger norms. In Figs. 5, 6 and 7 we can see that in most situations the algorithm chose a different learning rate for different layers, which seems to validate our initial intuition. It is important to note that the AdaLip and AdamLip optimizers show a wider variance over time compared to RMSLip. Another observation is that despite the norms of the weights have a fixed tendency the learning rates per layer do not. This means that the learning rate successfully adapts to the needs of each architecture and dataset uniquely.

One of the main discussion points is the selection of \(c_t\) and how it stands between the two most important features of an optimizer, the theoretical guarantee of convergence and the overall performance. As mentioned in the Related Work section one of the most used schedules for a guarantee in convergence is dividing the initial learning rate by the square root of the number of the iterations (see Eq. 3). However, in practice this is not preferred. The decay applied by the square root is too great for any real application that needs thousands of iterations to reach a good performance. With that many iterations the learning rate will become extremely small at an early stage of the training resulting in converging to a bad local minimum. On the other hand, a more steady learning rate seems to provide better results. This can be seen in Fig. 8 where constant learning rate achieves better convergence than the theoretical one.

Recent studies [8, 9, 12] show that an oscillating learning rate performs in some cases better than a strictly decreasing one, because it helps the optimizer overcome local minima.

Similarly to the case of \(c_t\), the theoretical value has the same performance as a constant value, except for some cases where constant \(c_t\) has a slight edge over the theoretical one. This can be seen in Fig. 9, where all runs with the same learning rate are really close, but constant \(c_t\) shows a small improvement. This trade-off between the guarantee of convergence and better results is quite important and can be rephrased as a trade-off of stability versus peak performance. One of the reasons that this behavior exists lies in the nature of SGD. In order to prove that SGD’s (or its variants) Regret converges to 0 over time, various assumptions are being made about the loss landscape. In practice, though, the landscape may not be ideal and converging to the first local minimum might not be satisfactory [35, 36]. Naturally, escaping saddle points and local minima has been a focal point in optimization research [11, 37, 38].

In order to find the optimal learning rate we used Lemma 1, which is based on Gradient Descent. However, we have to mention the reason that the optimal learning rate based on SGD was not used. Various experiments showed that the equation derived from SGD leads to a quite unstable learning progress. The times when the network was trained normally, the convergence speed was notably slower to the point that no amount of fine-tuning could have made it competitive. In the Appendix in Lemma 6 there is the full formula for the optimal learning rate based on SGD. It would be very interesting for future work to see if there is a way to apply this efficiently.

Another important observation about the norms of the layers arises with the use of the Batch Normalization. As can be seen from Fig. 10, the magnitude of the BatchNormalization weights does not always follow the pattern shown earlier in Fig. 1. This means that the gamma and beta weights of the BatchNormalization layers display a different behaviour compared with the rest of the weights. Although Batchnormalization helps the training to smooth the gradients of the layers [28], it creates weights that suffer from the same problem. This is the reason that AdaLip is able to construct a suitable learning rate for all these parameters. This is enforced by the fact that in practice, with networks that apply Batch Normalization (Sec. 5.3), AdaLip seems to perform better than other optimizers.