Step-size Optimization for Continual Learning — part 1/2

Cleaned-text code points [0, 15889). Bibliography omitted here; source.md includes it.

# Step-size Optimization for Continual Learning

###### Abstract

In continual learning, a learner has to keep learning from the data over its whole life time. A key issue is to decide what knowledge to keep and what knowledge to let go. In a neural network, this can be implemented by using a step-size vector to scale how much gradient samples change network weights. Common algorithms, like RMSProp and Adam, use heuristics, specifically normalization, to adapt this step-size vector. In this paper, we show that those heuristics ignore the effect of their adaptation on the overall objective function, for example by moving the step-size vector away from better step-size vectors. On the other hand, stochastic meta-gradient descent algorithms, like IDBD ( Sutton,, 1992 ) , explicitly optimize the step-size vector with respect to the overall objective function. On simple problems, we show that IDBD is able to consistently improve step-size vectors, where RMSProp and Adam do not. We explain the differences between the two approaches and their respective limitations. We conclude by suggesting that combining both approaches could be a promising future direction to improve the performance of neural networks in continual learning.

## 1 The Role of Step-size in Continual Learning

Continual learning is a setting where learning needs to always adapt to the latest data to learn new things or track moving targets. Continual learning contrasts with other problem settings where the goal is to converge to some fixed point. A key problem in continual learning is to able to learn from data what needs to be maintained, for example to avoid catastrophic forgetting ( French,, 1999 ) , and what needs to be updated to continue to track the objective function, for example because of limited capacity ( Sutton et al.,, 2007 ) .

A Stochastic Gradient Descent (SGD) method updates a set of weights 𝒘 t \bm{w}_{t} by incrementally accumulating the product of a step-size 1 1 1 We prefer to use step-size to learning-rate because we think it is more accurate. Indeed, a step-size parameter does not indicate the rate at which a system learns. A large step-size may or may not be correlated with a high rate of learning. scalar parameter η \eta and a sample gradient estimate ∇ J t ​ ( 𝒘 t ) ^ \widehat{\nabla J_{t}(\bm{w}_{t})} given an objective function J t ​ ( 𝒘 t ) J_{t}(\bm{w}_{t}) : 𝒘 t + 1 ← 𝒘 t − η ​ ∇ J t ​ ( 𝒘 t ) ^ \bm{w}_{t+1}\leftarrow\bm{w}_{t}-\eta\widehat{\nabla J_{t}(\bm{w}_{t})} We name this method Classic SGD in the rest of this paper. In Classic SGD, the step-size parameter η \eta is a key parameter to determine how much the weights of the learner are updated given the latest sample. For example, in a non-continual learning setting, this step-size parameter is often scheduled to converge to 0, forcing the changes to the parameters to be smaller over time until further adaptation becomes impossible. Using the same scalar step-size for all the weights is limited because it does not differentiate across dimensions.

Other conventional step-size adaptation methods include RMSProp ( Hinton et al.,, 2012 ) and Adam ( Kingma and Ba,, 2015 ) . RMSProp and Adam both normalize the gradients before updating the weights. Additionally, Adam also uses momentum to smooth the gradients. In practice, the weight update is composed of a component wise product of two parts: first, a step-size vector 𝜶 t + 1 \bm{\alpha}_{t+1} slowly changed and the gradient estimate ∇ J t ​ ( 𝒘 t ) ^ \widehat{\nabla J_{t}(\bm{w}_{t})} : 𝒘 t + 1 ← 𝒘 t − 𝜶 t + 1 ⋅ ∇ J t ​ ( 𝒘 t ) ^ \bm{w}_{t+1}\leftarrow\bm{w}_{t}-\bm{\alpha}_{t+1}\cdot\widehat{\nabla J_{t}(\bm{w}_{t})} (1) Both RMSProp and Adam adapt the step-size vector 𝜶 t \bm{\alpha}_{t} using a heuristic, normalising with an estimate of the gradient magnitude, irrespective of the effect of such changes on the objective function. To illustrate this, we used a simplification of a problem introduced in Sutton, (1992) , that is, a non-stationary 2-dimensional linear regression problem. On the first dimension, the optimal weight w 1 w_{1} is equals to 0 and constant. On the second dimension, the optimal weight w 2 w_{2} is non-stationary. Every 20 steps, w 2 w_{2} flips between -1 and 1 with a probability of 0.5. Features for both dimensions are independently sampled from a constant normal distribution. One would expect that a good method to optimize a step-size vector ( α 1 , α 2 ) (\alpha_{1},\alpha_{2}) would learn a low step-size α 1 \alpha_{1} for the first dimension to ignore noise, and learn a higher step-size α 2 \alpha_{2} to track the flipping weight.

Figure 1 shows the loss landscape as the average squared error computed by running for 1,000,000 steps using linear regression as the objective function J t ​ ( 𝒘 t ) J_{t}(\bm{w}_{t}) , given the step-size vector at a point ( α 1 , α 2 ) (\alpha_{1},\alpha_{2}) , and using the update of Equation 1 . Thus, Figure 1 is a representation of the step-size space with respect to the loss, and not the weight space as commonly depicted. In other words, each point represents how well regression is able to track the target given the fixed step-size vector at that point. Weights are initialized to 0. We can see that the pair of optimal step-sizes is α 1 = 0 \alpha_{1}=0 and α 2 ≈ 0.33 \alpha_{2}\approx 0.33 , as indicated by the diamond at the bottom of the plots.

Figure 1 -left shows the trajectory of the step-size vector 𝜶 t \bm{\alpha}_{t} in Equation 1 when updated by RMSProp using three different values for the step-size parameter. On this problem, we observe that RMSProp was not able to learn the best step-size vector to decrease the overall loss. Perhaps more surprisingly, in the two lower trajectories, the step-size vector ends up at a worse position from the optimal step-size vector compared to where it started from. Also note how RMSProp mostly moves both step-sizes α 1 \alpha_{1} and α 2 \alpha_{2} similarly, on a diagonal, and is not really able to distinguish between the different properties—a constant weight and a flipping weight—of the two dimensions. Although not shown in the figure, Adam behaves similarly to RMSProp on this problem.

Figure 1 -right shows the trajectory of the step-size vector when updated by the Incremental-Delta-Bar-Delta (IDBD) algorithm ( Sutton,, 1992 ) . IDBD is an online stochastic meta-gradient descent algorithm explicitly optimizing the step-size vector with respect to an objective function. More specifically, the IDBD algorithm learns a step-size vector by accumulating online estimates of the gradient of the objective function with respect to the step-sizes (IDBD is presented in detail later). Unlike RMSProp, IDBD is able to consistently update the step-size vector in the direction that is closer to the optimal point at the bottom of the landscape. Note that IDBD, RMSProp, and Adam all share the same compute and memory complexity.

We now explain this result in more details, highlighting the difference between step-size normalization algorithms (RMSProp, Adam) and step-size optimization algorithms (IDBD).

## 2 Setting

We consider an online learning setting where the learner observes a possibly non-stationary sequence J 1 , J 2 , … J_{1},J_{2},\ldots of loss functions. The learner starts with some initial weight vector 𝒘 1 \bm{w}_{1} at time t = 1 t=1 . At time t = 1 , 2 , … t=1,2,\ldots , upon observing a new loss function J t ​ ( ⋅ ) J_{t}(\cdot) , the learner incurs a sample of the loss J t ​ ( 𝒘 t ) J_{t}(\bm{w}_{t}) and then updates its weight vector 𝒘 t \bm{w}_{t} to 𝒘 t + 1 \bm{w}_{t+1} via some learning algorithm. The goal of the learner is to minimize the average lifetime loss.

Given a gradient sample ∇ J t ​ ( 𝒘 t ) \nabla J_{t}(\bm{w}_{t}) , a learner can use different algorithms to update its weight vector 𝒘 t \bm{w}_{t} . Perhaps the simplest method is the Classic SGD algorithm introduced in Section 1 and presented below. Gradient samples ∇ J t ​ ( 𝒘 t ) \nabla J_{t}(\bm{w}_{t}) are multiplied by a fixed scalar step-size parameter η \eta before being added to the current weight vector 𝒘 t \bm{w}_{t} .

[noitemsep,nolistsep]

η \eta : step-size parameter for the weight update

Hinton et al., (2012) proposed the SGD algorithm named RMSProp, described below, which normalizes the gradient samples before they are accumulated in the weight vector. It does so with an additional vector 𝒈 t \bm{g}_{t} which maintains a component-wise moving average of the square of the gradient, ( ∇ J t ​ ( 𝒘 t ) ) 2 \left(\nabla J_{t}(\bm{w}_{t})\right)^{2} . RMSProp introduces a new scalar step-size parameter γ g \gamma_{g} that sets how fast the average 𝒈 t \bm{g}_{t} tracks the square of the gradient. The average 𝒈 t \bm{g}_{t} is then used to normalise the update applied to the weight vector by dividing each component of the gradient by the square root 𝒈 t \sqrt{\bm{g}_{t}} of the average. A small constant ϵ \epsilon is added to 𝒈 t \bm{g}_{t} for stability. Thus, the step-size vector 𝜶 t \bm{\alpha}_{t} for RMSProp becomes η 𝒈 t + ϵ \frac{\eta}{\sqrt{\bm{g}_{t}+\epsilon}} .

[noitemsep,nolistsep]

η \eta : step-size parameter for the weight update

γ g \gamma_{g} : step-size parameter for normalization

ϵ \epsilon : constant for numerical stability

The Adam algorithm ( Kingma and Ba,, 2015 ) , described below, adds two ideas to RMSProp. The first idea is momentum , where Adam replaces the gradient ∇ J t ​ ( 𝒘 t ) \nabla J_{t}(\bm{w}_{t}) in the weight update with a tracking average of the gradient denoted 𝒎 t \bm{m}_{t} . This tracking average is updated at every step given an additional step-size parameter denoted γ m \gamma_{m} . Replacing the gradient with its average can be seen as a form of gradient smoothing. Setting γ m = 1 \gamma_{m}=1 disables momentum. The average of the gradient 𝒎 t \bm{m}_{t} and the squared gradient 𝒈 t \bm{g}_{t} are both initialized to zero. The second idea is to correct for the bias in the estimates of 𝒈 t \bm{g}_{t} and 𝒎 t \bm{m}_{t} because of that initialization to zero. Adam adjusts 𝒈 t \bm{g}_{t} and 𝒎 t \bm{m}_{t} to 𝒎 ^ t \bm{\hat{m}}_{t} and 𝒈 ^ t \bm{\hat{g}}_{t} respectively and uses these unbiased estimates in the weight update. The equivalent step-size vector 𝜶 t \bm{\alpha}_{t} for Adam is η 𝒈 ^ 𝒕 + ϵ \frac{\eta}{\sqrt{\bm{\hat{g}_{t}}+\epsilon}} .

[noitemsep,nolistsep]

η \eta : step-size parameter for the weight update

γ m \gamma_{m} : step-size parameter for gradient smoothing

γ g \gamma_{g} : step-size parameter for normalization

ϵ \epsilon : constant for numerical stability

## 3 Limitations of Step-size Normalization

This section highlights the limitations of step-size normalization on two simple learning problems, namely weight-flipping and rate-tracking . For both problems, we examine a simple setting with a linear least mean squared regression loss function. For t = 1 , 2 , … t=1,2,\ldots , the losses J 1 , J 2 , … J_{1},J_{2},\ldots are of the form J t ​ ( 𝒘 t ) = ( 𝒙 t ⊺ ​ 𝒘 t − y t ∗ ) 2 J_{t}(\bm{w}_{t})=(\bm{x}_{t}^{\intercal}\bm{w}_{t}-y^{*}_{t})^{2} , where 𝒙 t ∈ ℝ d \bm{x}_{t}\in\mathbb{R}^{d} is the feature vector and y t ∗ ∈ ℝ y^{*}_{t}\in\mathbb{R} is the label at time t t . In such setting, ∇ J t ​ ( 𝒘 t ) = δ t ​ 𝒙 t \nabla J_{t}(\bm{w}_{t})=\delta_{t}\bm{x}_{t} where δ t = 𝒙 t ⊺ ​ 𝒘 t − y t ∗ \delta_{t}=\bm{x}_{t}^{\intercal}\bm{w}_{t}-y^{*}_{t} .

The weight-flipping problem.

The problem was first introduced in Sutton, (1992) and is a larger version of the problem introduced in Section 1 . It defines a vector 𝒙 t \bm{x}_{t} of 20 inputs where each input is sampled independently according to a normal distribution with zero mean and unit variance. The first 15 components of the target weight vector 𝒘 ∗ t \bm{w^{*}}_{t} are all zeros. The last 5 components are either +1 or -1. To make it a continual learning problem, one of the non-zero weights is selected every 20 samples and flipped from +1 to -1 or vice-versa. Finally, the target prediction y ∗ y^{*} is defined as the linear combination of the weight vector and the input y ∗ = 𝒘 t ∗ ⊺ ​ 𝒙 t y^{*}={\bm{w}^{*}_{t}}^{\intercal}{\bm{x}_{t}} .

In this problem, a performance close to the optimal performance can be obtained by learning a different but constant step-size per component; that is, a learning algorithm able to learn a low step-size for the first 15 constant weights and a high step-size for the 5 flipping weights will perform better than an algorithm using the same step-size for all weights.

The asymptotic performance of Classic SGD, RMSProp, and Adam on the weight-flipping problem as a function of the step-size parameter is shown in Figure 2 . We treated RMSProp as a special case of Adam where γ m = 1 \gamma_{m}=1 because we considered that the problem is run for long enough to clear any bias due to initialisation. We did a sweep over γ g \gamma_{g} and η \eta . The performance for the best values of γ g \gamma_{g} and η \eta are reported in the figure. As an additional baseline, we ran Oracle SGD where the weights and step-sizes for the first 15 constant components are set to an optimal value of 0 and ran a sweep over the step-size for the remaining non-constant last 5 components. Consequently, Oracle SGD shows the best performance possible with a constant vector step-sizes in this problem setting. We report the error averaged across all steps after running for 200,000 steps.

On this problem, Adam/RMSProp performed slightly worse than Classic SGD and much worse than Oracle SGD. Their poor performance can be explained by looking at the algorithms. Adam/RMSProp maintains an average of δ t ​ 𝒙 t 2 \delta_{t}\bm{x}_{t}^{2} ; however, because all components of 𝒙 t \bm{x}_{t} are from the same normal distribution, they have the same variance. Moreover, because the error δ t \delta_{t} is global, all components end up with the same normalised step-size estimate η 𝒈 t + ϵ \frac{\eta}{\sqrt{\bm{g}_{t}+\epsilon}} . Consequently, on such tracking problem, Adam/RMSProp is not able to learn different step-sizes for each component. We conclude that normalization, as done in Adam/RMSProp, is not enough to differentiate between weights that should be fixed and weights that should change.

The 1D noisy rate-tracking problem.

In the previous problem, learning constant step-size parameters lead to an optimal solution. Generally, in a continual learning setting, the optimal step-sizes may not be constant and may need to be adapted over time. The noisy-tracking problem illustrates such a case. This problem is still a linear regression problem of a single dimension where the feature x t x_{t} equals 1 1 , at all times t t .

The learner aims to predict the target y y that changes at every step as: y t \displaystyle y_{t} = z t + 𝒩 ⁡ ( 0 , 1 ) \displaystyle=z_{t}+\mathcal{N}(0,1) z t + 1 \displaystyle z_{t+1} = z t + 𝒩 ⁡ ( 0 , σ t 2 ) \displaystyle=z_{t}+\mathcal{N}(0,\sigma_{t}^{2}) where σ t \sigma_{t} is sampled every 50,000 steps from a uniform distribution from 0 to 3. Sutton, (1981) showed that the optimal step-size α t ∗ \alpha^{*}_{t} for this problem can be computed as: α t ∗ = − σ t 2 + σ t 4 + 4 ​ σ t 2 2 . \alpha^{*}_{t}=\frac{-\sigma_{t}^{2}+\sqrt{\sigma_{t}^{4}+4\sigma_{t}^{2}}}{2}. (2) Figure 3 shows that Classic SGD and RMSProp are not able to track the optimal step-size α t ∗ \alpha^{*}_{t} on the 1D noisy rate-tracking problem. Indeed, because x t x_{t} is always equals to 1 1 , RMSProp will decrease the step-size every time the magnitude of the error increases. However, the opposite is needed for this problem. The increase in the magnitude of the error is due to a faster rate of drift in the target, which requires a larger step-size. Figure 3 shows that RMSProp does the opposite of what needs to be done: increases the step-size when it should be decreased and vice-versa.