Efficient continual learning at the edge with progressive segmented training

Terminology. The continual learning problem can be formulated as follows: the machine learning system is continuously exposed to a stream of labeled input data X¹, X², ..., where ${X}^{y}=\left\{{x}_{1}^{\,y},\dots ,{x}_{{\,n}_{y}}^{\,y}\right\}$ correspond to all examples of class $y\in \mathbb{N}$. When the new task $\left\{{X}^{s},\dots ,{X}^{t}\right\}$ comes in, the data of old tasks $\left\{{X}^{1},\dots ,{X}^{s-1}\right\}$ are no longer available, except for a small amount of previously seen data stored in the memory set $\mathcal{P}=\left({P}_{1},\dots ,{P}_{s-1}\right)$.

For deep neural networks such as VGG-Net [32] and ResNet [34], the network parameter Θ usually consists of feature extractor $\varphi :\mathcal{X}\to {\mathbb{R}}^{d}$ and classification weight vectors $w\in {\mathbb{R}}^{d}$. The network keeps updating its parameter Θ according to the previously seen data $\mathcal{X}$, in order to predict labels ${\mathcal{Y}}^{\ast }$ with its output $\mathcal{Y}={w}^{\top }\varphi (\mathcal{X})$. During the network training with data corresponding to classes $\left\{{X}^{1},\dots ,{X}^{s-1}\right\}$, our target is to minimize the loss function $\mathcal{L}(\mathcal{Y};{\mathcal{X}}_{s-1};{\Theta})$ of this (s − 1)-class classifier. Similarly, with the introduction of a new task with classes $\left\{{X}^{s},\dots ,{X}^{t}\right\}$, the target now is to minimize $\mathcal{L}(\mathcal{Y};{\mathcal{X}}_{t};{\Theta})$ of this t-class classifier.

Training routine. Every time when a new task is available, PST calls a training routine (figure 1 and algorithm 1) to update the parameter Θ to Θ', and the memory set $\mathcal{P}$ to ${\mathcal{P}}^{\prime }$, according to the current training data $\left\{{X}^{s},\dots ,{X}^{t}\right\}$ and a small amount of previously seen data (memory set) $\mathcal{P}$. The training routine consists of three major components: (1) memory-assisted training and balancing, (2) significance sampling, and (3) model segmentation, as described in the following subsections.

Algorithm 1. PST training routine.

Input: $\left\{{X}^{s},\dots ,{X}^{t}\right\}$//Current task data in per-class sets

Require Θ = (Θfixed; Θfree)//Current network parameters, Θfree is trainable

Require $\mathcal{P}=({P}_{1},\dots ,{P}_{s-1})$//Memory sample sets from previous data

1:  Memory-assisted training and balancing: ${{\Theta}}_{\text{free}}\to {{\Theta}}_{\text{free}}^{\prime }$//${{\Theta}}^{\prime }=({{\Theta}}_{\text{fixed}};{{\Theta}}_{\text{free}}^{\prime })$

2:  Significance sampling: identify Θimportant in ${{\Theta}}_{\text{free}}^{\prime }$//Θ' = (Θfixed; Θimportant; Θsecondary)

3:  Model segmentation: ${{\Theta}}_{\text{important}}\to {{\Theta}}_{\text{important}}^{\prime }$//${{\Theta}}^{\prime }=({{\Theta}}_{\text{fixed}};{{\Theta}}_{\text{important}}^{\prime };{{\Theta}}_{\text{secondary}})$

4:  $({{\Theta}}_{\text{fixed}};{{\Theta}}_{\text{important}}^{\prime })\to {{\Theta}}_{\text{fixed}}^{\prime }$

5:  ${{\Theta}}_{\text{secondary}}\to {{\Theta}}_{\text{free}}^{\prime }$

Output: ${{\Theta}}^{\prime }=({{\Theta}}_{\text{fixed}}^{\prime };{{\Theta}}_{\text{free}}^{\prime })$//Updated network parameters

Output: ${\mathcal{P}}^{\prime }=({P}_{1},\dots ,{P}_{t})$//Updated memory set

Figure 1 illustrates PST training routine for task T_i and task T_i+1. In figure 1(a), which is the moment that task T_i comes in, the network consists of two portions: parameters Θ_fixed (gray blocks) are fixed for previous tasks, and parameters Θ_free (light blue blocks) are trainable for current and future tasks. We allow Θ_free to be updated for task T_i , with Θ_fixed included in the feedforward path. To mitigate the parameter bias toward the new task, a memory set is used to assist the training. The memory set is sampled uniformly and randomly from all the classes in previous tasks, which is a simple yet highly efficient approach, as explained in the RWalk work [4]. For example, if the memory budget is K and s − 1 classes have been learned in previous tasks, then the memory set stores $\frac{K}{s-1}$ images for each class. We mix samples from this memory set with equal samples per class from the current task, i.e. K samples of the memory and $\frac{K}{s-1}\times (t-s+1)$ samples from the current task, and provide them to the network: (i) for a few epochs at the beginning of the training; (ii) periodically (e.g. every 3 epochs) during training; (iii) for a few epochs at the end of the training to fine-tune classification layer (i–iii are noted in figure 3).

In comparison to most related works that adopt the single-stage optimization technique, the proposed three-step optimization strategy performs much better. One of the primary reasons behind catastrophic forgetting is knowledge drift in both feature extraction and classification layers. The three-pronged strategy helps minimize this drift in the following ways: step (i) provides a well-balanced initialization; step (ii) reviews previous data and thus, consolidates previously learned knowledge for the entire network; step (iii) corrects bias by balancing classification layers, which is simple yet efficient as compared to [27] that utilizes an extra bias correction layer after the classifier. After memory-assisted training and balancing, the network parameters are updated from Θ = (Θ_fixed; Θ_free) to ${{\Theta}}^{\prime }=({{\Theta}}_{\text{fixed}};{{\Theta}}_{\text{free}}^{\prime })$, as stated in algorithm 1 line 1.

After the network has learned task T_i , PST samples crucial learning units for the current task: for feature extraction layers (i.e. convolutional layers), PST samples important filters; for fully-connected layers, PST samples important neurons. The definitions of filter and neuron are as follows: the lth convolutional layer can be formulated as: the output of this layer ${\mathcal{Y}}_{l}={\mathcal{X}}_{l}\,\ast \,{{\Theta}}_{l}$, where ${{\Theta}}_{l}\in {\mathbb{R}}^{{O}_{l}\times {I}_{l}\times K\times K}$. The set of weights that generates the oth output feature map is denoted as a filter ${{\Theta}}_{l}^{o}$, where ${{\Theta}}_{l}^{o}\in {\mathbb{R}}^{{I}_{l}\times K\times K}$. The lth fully-connected layer can be represented by: ${\mathcal{Y}}_{l}={\mathcal{X}}_{l}\cdot {{\Theta}}_{l}$, where ${{\Theta}}_{l}\in {\mathbb{R}}^{{O}_{l}\times {I}_{l}}$. The set of weights ${{\Theta}}_{l}^{t}$ that connected to the tth class can be denoted as a neuron, where ${{\Theta}}_{l}^{t}\in {\mathbb{R}}^{1\times {I}_{l}}$.

The filter/neuron sampling is based on an importance score that is adopted in PST to measure the effect of a single filter/neuron on the loss function, i.e. the importance of each filter/neuron. The importance score is developed from the Taylor expansion of the loss function. Previously, Molchanov et al [35] applied it on pruning secondary parameters. The importance score represents the difference between the loss with and without each filter/neuron. In other words, if the removal of a filter/neuron leads to relatively small accuracy degradation, this unit is recognized as an unimportant unit, and vice versa. Thus, the objective function to obtain the filter with the highest importance score is formulated as:

$$\begin{equation}\stackrel{\mathrm{arg min}}{{{\Theta}}_{l}^{o}}\quad \vert {\Delta}\mathcal{L}({{\Theta}}_{l}^{o})\vert {\Leftrightarrow}\stackrel{\mathrm{arg min}}{{{\Theta}}_{l}^{o}}\,\vert \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})-\mathcal{L}(\mathcal{Y};\mathcal{X};{{\Theta}}_{l}^{o}=\mathbf{0})\vert .\end{equation} \tag{ 1 }$$

Using the first-order of Taylor expansion of $\vert \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})-\mathcal{L}(\mathcal{Y};\mathcal{X};{{\Theta}}_{l}^{o}=\mathbf{0})\vert$ at ${{\Theta}}_{l}^{o}=\mathbf{0}$, we get:

$$\begin{align}\hfill \vert {\Delta}\mathcal{L}({{\Theta}}_{l}^{o})\vert & \simeq \vert \frac{\partial \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})}{\partial {{\Theta}}_{l}^{o}}{{\Theta}}_{l}^{o}\vert ={\sum }_{i=0}^{{I}_{l}}{\sum }_{m=0}^{K}{\sum }_{n=0}^{K}\vert \frac{\partial \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})}{\partial {{\Theta}}_{l}^{o,i,m,n}}{{\Theta}}_{l}^{o,i,m,n}\vert \hfill \end{align} \tag{ 2 }$$

where $\frac{\partial \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})}{\partial {{\Theta}}_{l}^{o,i,m,n}}$ is the gradient of the loss function with respect to parameter ${{\Theta}}_{l}^{o,i,m,n}$.

Similarly, the saliency score of a neuron is derived as:

$$\begin{equation}\vert {\Delta}\mathcal{L}({{\Theta}}_{l}^{t})\vert \simeq \vert \frac{\partial \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})}{\partial {{\Theta}}_{l}^{t}}{{\Theta}}_{l}^{t}\vert ={\sum }_{i=0}^{{I}_{l}}\vert \frac{\partial \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})}{\partial {{\Theta}}_{l}^{t,i}}{{\Theta}}_{l}^{t,i}\vert \end{equation} \tag{ 3 }$$

where $\frac{\partial \mathcal{L}(\mathcal{Y};\mathcal{X};{\Theta})}{\partial {{\Theta}}_{l}^{t,i}}$ is the gradient of the loss with respect to parameter ${{\Theta}}_{l}^{t,i}$.

Based on the importance score, we sort the learning units layer by layer and identify the top β units (dark blue blocks in figure 1(b)). In the following model segmentation step, we deal with the location of important parameters, rather than the value of these parameters, which will be explained in the next subsection. β is an empirical hyper-parameter that should be approximately proportional to the complexity of the current task. For example, when incrementally learning 10 classes of CIFAR-100 at a time, β can be 10%; when learning 20 classes per task, β can be 20%.

Due to the nature of continual learning, the total number of tasks is not known beforehand, so the network can be reserved with a larger capacity in order to freeze enough knowledge for previous tasks and leave enough space for future tasks. Once the continual learning is complete, one can leverage model compression approaches [36–40] to compress the model size. It is also worth mentioning that significance sampling is only performed once after each task so that the computation cost of this step is minimized.

After important units are sampled according to the importance score, current network parameter Θ' = (Θ_fixed; Θ_important; Θ_secondary), where Θ_fixed are the frozen parameters for all the previous tasks, Θ_important are important parameters for the current task, and Θ_secondary are unimportant parameters for the current task, as stated in algorithm 1 line 2. Our ideal target is to reinforce Θ_important in a way such that their contribution to the current task is as crucial as possible. Previously, Liu et al [41] observed that the sampled network architecture itself (rather than the selected parameters) is more indispensable to the learning efficacy. Inspired by this conclusion, we keep the Θ_fixed and Θ_secondary intact, randomly initialize Θ_important and retrain them with current training data assisted by a memory set to obtain ${{\Theta}}_{\text{important}}^{\prime }$. This step reinforces the contribution of Θ_important to the learning, as proved by our experimental results demonstrated in figure 2 and table 2. After model segmentation, ${{\Theta}}_{\text{important}}^{\prime }$ along with the aforementioned Θ_fixed will be kept frozen in future tasks, and Θ_secondary will be used to learn new knowledge.

Zoom In Zoom Out Reset image size

We divide CIFAR-10 into 2 tasks (5 classes each) and analyze the PST training routine step by step in this subsection. Figure 3 presents the learning curve for training 2 tasks (5 classes each) in CIFAR-10.

Zoom In Zoom Out Reset image size

From epoch 0 to epoch 180, T₁ is trained and reaches baseline accuracy. The weight distribution after training T₁ is present in figure 2(a). At epoch 180, we sample the top 50% (since there are two tasks in total) important parameters and retrain them with the secondary parameters untouched (epoch 180 to epoch 300), which is the model segmentation step. The weight distribution after this step is shown in figure 2(d). It is worth mentioning that previous works, such as PackNet [20] and Piggyback [19], prune the secondary parameters and thus, distort the weight distribution (figure 2(b)). At epoch 300, task T₂ appears and updates the parameters. At the same time, the acquired knowledge of T₁ is disturbed by T₂ updating, leading to an accuracy degradation on T₁ (see the green curve at epoch 300). From epoch 300 to the end is the step of T₂ training, during which the memory data is injected following step (i)–(iii) to balance.

After T₂ training, we again plot the weight distribution for the pruning-based approach (in figure 2(c)) and PST approach (in figure 2(e)). It is observed that the pruning approach fails to preserve the prior knowledge, as the weight distribution after learning T₂ shifts far away from the previous one. In contrast, PST well preserves prior knowledge (i.e. similar weight distribution after learning T₁ and after learning T₂). Compared to the baseline accuracy, pruning-based approaches forget 31% on overall accuracy while segmentation-based PST only forgets 5%.

On the CIFAR-100 dataset, our experimental results show that: (1) in overall accuracy, PST outperforms most of the previous work [1, 2, 4, 10, 16, 18] and is on par with iCaRL [8]. (2) With model segmentation, PST successfully preserves prior knowledge. (3) PST reduces more than 24× computation cost in edge computing, as compared to classic regularization approaches.

Accuracy for incrementally learning multi-classes. We compare PST with state-of-the-art approaches that reported single-head accuracy: MAS [16], EWC [1], RWalk [4], SI [2], LwF.MC [18], DMC [10], iCaRL.MC [8] and two baselines: fixed representation, finetuning. Fixed representation denotes the method that fixes the feature extraction layers for the previously learned tasks and only trains classification layers for new tasks. Finetuning denotes the method that the network trained on previous tasks is directly fine-tuned by new tasks, without strategies to prevent catastrophic forgetting. LwF.MC denotes the method that uses LwF [18] but is evaluated with multi-class single-head classification. iCaRL.MC denotes the method uses iCaRL but replaces their nearest-mean-of-exemplar [8] classifier with a regular output classifier for a fair comparison with PST. The results of MAS, EWC, RWalk, SI, and DMC are from [10], which is implemented with the official code ⁴ . The results of fix representation, finetune, LwF.MC and iCaRL are from [8]. We adopt the same memory size for a fair comparison between the baselines and PST.

The single-head overall accuracy when incrementally learning 20 tasks (5 classes per task), 10 tasks (10 classes per task), 5 tasks (20 classes per task), and 2 tasks (50 classes per task) are reported in figure 4. Among 9 different approaches, PST achieves the best accuracy on the 2-task scenario and the second best accuracy on the other scenarios. Compared to finetuning, PST largely prevents the model from catastrophic forgetting. Although PST achieves lower accuracy than iCaRL in some cases, PST is more than 24× efficient in the computation cost, as shown in figure 6. This efficiency is benefiting from model segmentation: iCaRL has to update the entire network parameters for every new observation, but PST only requires the update of partial network parameters, as the parameters related to previous tasks are frozen.

Zoom In Zoom Out Reset image size

Accuracy of the first task. Figure 5(a) compares the single-head accuracy on the first task T₁ in PST with several previous approaches that reported T₁ accuracy in their papers. PST achieves the best single-head accuracy on T₁ among all the approaches, i.e. the least forgetting. Moreover, when T₁ data is evaluated in a multi-head classification setting, as shown in figure 5(b), PST is stable and always on par with the baseline (the model that is only trained on T₁, so without forgetting). This phenomenon demonstrates that PST effectively preserves the knowledge related to T₁ through model segmentation. Without these strategies, it is difficult to maintain the previously acquired knowledge. For example, GEM [7] reported unstable multi-head T₁ accuracy, because the parameters gradually drift away from T₁ knowledge after a long period of learning new tasks.

Zoom In Zoom Out Reset image size

In a more realistic situation, continual learning may not be used to train a model from scratch at the edge. Instead, we will have a model which is well trained in the cloud and once deployed, might only be required to learn a few new classes in an online manner on the edge devices. In this section, we developed experiments to show that PST benefits continual training at the edge from the perspective of accuracy and computation cost.

In table 1, we test such a system where the base model is pre-trained (similar to training on the cloud) with 10, 30, 50, 70, or 90 classes of CIFAR-100 as task T₁, while the new task T₂ that consists of 10 disjoint classes has to be learned at the edge continually. The number of trainable parameters for T₂ remains the same across these 5 experiments. As shown in table 1, if large amounts of data have been well trained in the cloud and stored in the segmented PST model, the training of incremental data at the edge causes marginal forgetting (e.g. 0.08) of the acquired knowledge.

Moreover, we estimate the computation cost during training, i.e. the number of floating point operations (FLOPs), required by PST and regularization approaches such as iCaRL [8] and EWC [1], as shown in figure 6. Computation cost is a critical overhead when deploying deep neural networks on edge devices [36, 37, 39, 44]. Edge learning prefers algorithms with low computation cost rather than that with higher one. Training at the edge includes three paths [45, 46], i.e. (1) forward path, (2) backward path, and (3) weight update path. As more and more tasks come in, the trainable parameters become fewer and fewer in PST, i.e., the weight update path gradually requires fewer operations, but regularization methods require a constant number of operations at all times, as the model is not segmented. Thus, given the model is pre-trained in the cloud with a large amount of data and loaded at the edge, PST reduces the FLOPs in the weight update path by more than 24×, and by more than 1.5× in the complete path (including all three paths), as compared to the regularization methods such as iCaRL [8]. Especially, the weight update path usually consumes 2× latency than the other two paths so that PST can largely speed up the training. Benefiting from segmentation, PST outperforms other continual learning schemes in computation efficiency.

Zoom In Zoom Out Reset image size

In this section, we demonstrate online CIFAR-10 CNN learning on an FPGA-based 16 bit fixed-point training accelerator [45, 47] on Intel Stratix-10 MX FPGA [48]. The CNN training hardware is flexible to support forward pass (FP), backward pass (BP) and weight update (WU) phases of training.

Figure 7 presents the overall FPGA system setup [49] to train CNNs using PST algorithm. For simplicity, the CNN structure used here is 16C3-16C3-MP-32C3-32C3-MP-64C3-64C3-MP-FC, where 'NCk' refers to the convolution layer with 'N' output feature maps and a kernel size of 'k', 'MP' refers to max-pooling layer and 'FC' refers to fully-connected layer. First, as shown in figure 7(a), a large amount of weights is pre-trained and selected with 9 classes from CIFAR-10 dataset. The pre-trained model and a binary mask representing the frozen weights are fed to the RTL generator. The RTL generator generates the customized training accelerator based on the pre-trained model structure and generates HBM2 memory initialization files to load the model parameters, as shown in figure 7(b). The frozen weights stored in HBM2 are used by the FPGA training accelerator to perform inference on pre-trained classes.

Zoom In Zoom Out Reset image size

The model is then exposed to a new, unlearned class from the CIFAR-10 dataset, and updated accordingly in real-time on the FPGA, as shown in figure 7(c). The entire system is demonstrated on Intel Stratix-10 MX FPGA board (figure 7(d)). Benefiting from the model segmentation, the online training of new observations requires much less computation cost and lower latency, as compared to the traditional continual learning scheme that updates the entire network. As shown in figure 7(e), the breakdown graph shows that the PST scheme saves 4.2× latency per image in the weight update (WU) phase as compared to traditional algorithms.

We remove each component from PST and repeat the experiments performed in figure 4. The overall accuracy change after the last task is reported in table 2. Replacing significance sampling with random sampling leads to model Hybrid 1; removing the model segmentation step (no reinforcement on Θ_important) leads to model Hybrid 2; removing the memory-assisted balancing leads to model Hybrid 3. The results of hybrid models prove that each component in PST is contributing to the overall performance. Removing significance sampling or model segmentation leads to a significant accuracy drop, while removing memory leads to a small accuracy drop. It shows that significance sampling and model segmentation are indispensable steps for PST, and memory-assisted balancing is supplementary.

For PST, the accuracy gap between single-head and multi-head of T₁ could be caused by the imbalance between old and new knowledge (the network is biased to new knowledge than old knowledge since old data are no longer used to train the network). Memory-assisted balancing in PST alleviates this obstacle but cannot completely prevent it. Indeed, there has hitherto been no approach to prevent this knowledge asymmetry. With more data saved from previous tasks, forgetting is reduced. But such a trend gradually saturates, as shown in figure 8.

Zoom In Zoom Out Reset image size