Layerwise Optimization by Gradient Decomposition for Continual Learning

Continual learning has been a long-standing research problem in the field of machine learning [32, 45, 36, 47]. Generally speaking, there are three different types of scenarios for continual learning: (1) class-incremental scenario, where the number of class labels keeps growing but no explicit task boundary is defined under test; (2) task-incremental scenario, where the boundaries among tasks are assumed known and the information about the task under test is given. (3) data-incremental scenario, where the set of class labels are identical for all tasks. Existing work of above three scenarios mainly falls into two categories: Regularization-based and Memory-based approaches.

Stochastic Gradient Descent is the most widely used optimization techniques for training DNNs [3, 31, 2]. However, it applied the same hyper-parameters to update all parameters in different layers, which may not be optimal for loss minimization. Therefore, layerwise adaptive optimization algorithms were proposed [10, 21]. RMSProp [41] altered the learning rate of each layer by dividing the square root of its exponential moving average. LARS [54] let the layerwise learning rate be proportional to the ratio of the norm of the weights to the norm of the gradients. Both layerwise adaptive optimizers solved the variation of update frequencies for different layers and thus outperformed SGD in various large-scale benchmarks. Layerwise gradient update strategy is also applied in meta-learning for rapid loss convergence on transfer learning [13, 14, 48] and few-shot learning [49, 28]. [12] proposed WarpGrad which layerwisely meta-learned to warp task loss surfaces across the joint task-parameter distribution to facilitate gradient descent. MT-Net [27] enabled the meta-learner to learn on each layer’s activation space, a subspace that the task-specific learner performed gradient descent on. Our method differs from the above methods in two aspects. Our layerwise gradient update strategy aims at preserving old knowledge. Instead, RMSProp, LARS aimed to learn new tasks more efficiently, and WarpGrad, MT-Net aimed to transfer more related information in source data to target samples, which have different targets with our method. Secondly, RMSProp and LARS tried to handle the different update frequency between parameters in different layers but our method aims at handling large magnitude variation of parameters in different layers.

Biologically, continual learning is inspired by human learning. We learn the new tasks by applying the knowledge previously learned from the related tasks. In this way, there exists common knowledge shared among the old tasks. Given the losses of old task $\{\mathcal{L}^{\text{old}}(\theta_{t},\mathcal{M}_{i})\}_{i=1}^{t-1}$, we assume each loss has two components. One is the shared loss, which is the same for any previous task. Minimizing the shared loss would improve the overall performance across all previous tasks. The other is task-specific component, reflecting the task-specific knowledge of every old task. For $\mathcal{L}^{old}(\theta_{t},\mathcal{M}_{i})$ about the $i$-th episodic memory, it can be written as:

where $\mathcal{L}_{\text{shared}}^{old}(\theta^{t},\mathcal{M}_{t-1}^{cor})$ is the shared loss term, which is the same for different old tasks, and $\mathcal{R}^{old}_{\text{specific}}(\theta_{t},\mathcal{M}_{i})$ is the residual for the task $\mathcal{T}_{i}$.

where the equality holds since partial derivatives are linear.

Our algorithm aims to find a gradient $w$ that reduces the loss of the new task but does not increase the losses of any memory tasks. Inspired by GEM, if we assume the shared loss function is locally linear, the change of loss can be diagnosed by the sign of the inner product between its corresponding gradient $\bar{g}$ and the update $w$. The positive, zero, or negative inner product of $w$ and $\bar{g}$ indicates that the shared loss will decrease, preserve or increase respectively if we update the network by $-\eta w$, where $\eta$ is a small positive value. Similar consideration applies to the task-specific residuals by replacing $\bar{g}$ with $\hat{g}_{i}$.

Based on the above observations, our algorithm requires that the gradient update $w$ should be close to $g$ by minimizing $\|w-g\|_{2}^{2}$ and will not increase the shared loss by constraining $\bar{g}^{\top}w\geq 0$, That is:

However, when considering the task-specific gradients $\hat{g}_{i}$, we find that the only solution to $\hat{g}_{i}^{\top}w\geq 0,\forall i<t$ is $\hat{g}_{i}^{\top}w=0,\forall i<t$. According to Equ. (5), for any gradient update $w$, we have $\sum_{i=1}^{t-1}\hat{g}_{i}^{\top}w=0$ because $\sum_{i=1}^{t-1}\hat{g}_{i}=\mathbf{0}$. It means that unless $\hat{g}_{i}^{\top}w=0,\forall i<t$, any other $w$ would result in $\hat{g}_{i}^{\top}w>0$ for some $i$ while $\hat{g}_{i}^{\top}w<0$ for some other $i$, i.e., the losses of some tasks will increase but the losses of others will decrease. To better preserve the knowledge from each old task, the only choice is to require the gradient update $w$ orthogonal to the column space of the task-specific gradients, i.e.,

For the ease of optimization, $\hat{G}^{\top}w=\mathbf{0}$ can be transformed according to the following lemma:

Based on the above analysis, the gradient update $w$ can be obtained by solving the following optimization problem:

After transformation, we can solve the optimization problem in Equ. (10) by Karush-Kuhn-Tucker condition [4]. The solution is as follows:

where $P=\mathbf{I}-BB^{\top}$, which is positive semidefinite. The schematic representation of gradient optimization is present in Figure 1(b).

Traditional gradient-based CL algorithms, e.g. GEM, concatenate gradients from all layers together and the optimization is based on this concatenated gradient. The loss decrease of episodic memory $\mathcal{M}_{i}$ after updating parameters $\theta_{t}$ by $-\eta w$ is

where $w$ is obtained by Equ. (11). Considering $\frac{\partial\mathcal{L}^{\text{old}}(\theta_{t},\mathcal{M}_{i})}{\partial\theta_{t}}$ is the summation of the shared gradient $\bar{g}$ and its task-specific gradient $\hat{g}$, the loss change of episodic memory $\mathcal{M}_{i}$ will be negative if $\bar{g}Pg\geq 0$ holds but will be zero if $\bar{g}Pg<0$ after replacing $w$ with Equ (11). Specifically,

We notice that the condition $\bar{g}^{\top}Pg$ heavily depends on the elements with large magnitude in concatenated gradients as inner product is used. Since we have observed that the magnitude of gradients have a large variation among different layers from Figure 3, $\Delta\mathcal{L}_{1}^{\text{old}}(\theta_{t},\mathcal{M}_{i})$ will be dominated by some layers where the magnitude of new gradient $g$, shared gradient $\bar{g}$ and task-specific gradients $\hat{G}$ is large.

For our proposed layerwise gradient update strategy, we replace concatenated $g,P$ in Equ.(11) with layer-specific gradients $g^{(l)},P^{(l)}$, where $l$ is layer index. Suppose the network is updated layerise, i.e., $\theta_{t}^{(l)}\leftarrow\theta_{t}^{(l)}-\eta w^{(l)}$, the loss change of episodic memory $\mathcal{M}_{I}$ is $-\eta\sum_{l=1}^{L}\bar{g}^{(l)^{\top}}P^{(l)^{\top}}g^{(l)^{\top}}$, where $L$ is the number of layers. For layer $l$ that satisfies $\bar{g}^{(l)^{\top}}P^{(l)^{\top}}g^{(l)}\geq 0$, its contribution to $\Delta\mathcal{L}_{2}^{\text{old}}$ is $-\eta\bar{g}^{(l)^{\top}}P^{(l)^{\top}}g^{(l)}\leq 0$. For layer $l$ that satisfies $\bar{g}^{(l)^{\top}}P^{(l)^{\top}}g^{(l)}<0$, its contribution to $\Delta\mathcal{L}_{2}^{\text{old}}$ is zero. Therefore, the total loss change of episodic memory $\mathcal{M}_{i}$, $\Delta\mathcal{L}_{2}^{\text{old}}$, can be presented as

where $\{W^{+}\}$ is the set of layer index whose condition $\bar{g}^{(l)^{\top}}P^{(l)^{\top}}g^{(l)}\geq 0$ is satisfied. We notice that since the gradient magnitude will not have a large variation within one layer, $\Delta\mathcal{L}_{2}^{\text{old}}$ calculated layerwise is not dominated by a few elements with large magnitudes in $g^{(l)}$ and $P^{(l)}$.

Comparing Equ. (12) with Equ. (13), the layerwise method has two advantages: first, instead of heavily depending on the gradients with large magnitude in the concatenating method, the layerwise method treats each layer equally by letting each layer determine its own condition.

Second, as each layer determines its own condition, $\bar{g}^{(l)^{\top}}P^{(l)^{\top}}g^{(l)}\geq 0$ holds for some layers even though $\bar{g}^{\top}P^{\top}g<0$, which may reduce the loss more efficiently.

Since GEM and A-GEM are two closest algorithms to our method, we would make more explicit comparisons with them to clarify our innovations.

To investigate the contribution of new components in our proposed method, we incrementally evaluate each of them on MNIST Permutation, Split CIFAR10, Split CIFAR100 and Split tinyImagenet. Denote SCC as the Shared Component Constraint, TSCC as Task-Specific Compoent Constraint, LGU as Layerwise Gradient Update, and PCA as Principal Compoent Analysis. We choose a single predictor fine-tuned across all tasks as baseline. Seven variants are then constructed on top of baseline: (a) baseline; (b) baseline+SCC; (c) baseline+SCC+LGU; (d) baseline+SCC+TSCC; (e) baseline+SCC+TSCC+PCA; (f) baseline+SCC+LGU+TSCC; (g) baseline+SCC+LGU+TSCC+PCA. Table 1 presents the performance of all variants. All reported results are tested on the task-incremental setting.

Figure 4 shows that as K increases, ACC increases initially but decreases afterwards, while BWT keeps increasing. The continued increase of BWT is expected, as more knowledge from the old tasks is preserved as K increases. However, ACC decreases after K $=5$, indicating that although the performance on the old tasks benefits from a large K, the ability to learn new task would be degraded if K is too large.

Our main results except MNIST are tested on the online task-incremental setting while MNIST is tested on the data-incremental setting. We compare our proposed method with several existing and state-of-the-art methods, including Single [40], Independent¹¹1For fair comparison, we follow the setting in GEM. In GEM, the channels of each feature map are divided by the number of tasks. As such, the model size in Independent method equals to the model size in other methods. [30], Multimodel [30], EWC [22], iCARL [38], GEM [30], A-GEM [7] and S-GEM [7]. For fair comparison, the size of each episodic memory is also 256 for GEM, A-GEM and S-GEM. The results are reported in Tab. 3. Our method achieves state-of-the-art results on all four datasets in both one epoch and three epochs settings.

Different from task-incremental continual learning setting mainly discussed in this paper, class-incremental continual learning setting do not provide task identity in the test data [19, 52]. In such setting, task boundaries naturally formed in the continual training set ought to be broken because we need classify classes from different continual tasks. In order to break the task boundary, we first stack all episodic memories as one replay buffer. When training, we randomly split the replay buffer into $N$ sub-replay buffers. The gradient decomposition is implemented on such $N$ sub-replay buffers instead of episodic memories. We assess our method on Split CIFAR100 dataset for consistency with other approaches. We employ Resnet34 as our backbone and train the network for 200 epochs for each continual step. The initial learning rate is 0.1 and decays when the epoch equals 80, 160. The size of replay buffer is restricted to be 2000 images and the memory update strategy is the same as that in  [38, 5] for fair comparison.

We compare our approach with other competing methods, including LWF [29], iCarl [38], DR [17], End2End [5], LUCR²²2LUCR is short for Learning a Unified Classifier via Rebalancing [18], GD [26], MUC [11] and TPL [51]. We fix $N=5$ when implementing the experiments. The comparisons with other competing methods are presented in Table 4. As illustrated in Table 4, our method outperforms GEM by $3\%$ which shows the consistency improvement with results in task-incremental scenario. We also achieve the state-of-the-art results which are the same MUC [11].