Exclusive Supermask Subnetwork Training for Continual Learning (2023)

Abstract: Continual Learning (CL) methods focus on accumulating knowledge over time while avoiding catastrophic forgetting. Recently, Wortsman et al. (2020) proposed a CL method, SupSup, which uses a randomly initialized, fixed base network (model) and finds a supermask for each new task that selectively keeps or removes each weight to produce a subnetwork. They prevent forgetting as the network weights are not being updated. Although there is no forgetting, the performance of SupSup is sub-optimal because fixed weights restrict its representational power. Furthermore, there is no accumulation or transfer of knowledge inside the model when new tasks are learned. Hence, we propose ExSSNeT (Exclusive Supermask SubNetwork Training), that performs exclusive and non-overlapping subnetwork weight training. This avoids conflicting updates to the shared weights by subsequent tasks to improve performance while still preventing forgetting. Furthermore, we propose a novel KNN-based Knowledge Transfer (KKT) module that utilizes previously acquired knowledge to learn new tasks better and faster. We demonstrate that ExSSNeT outperforms strong previous methods on both NLP and Vision domains while preventing forgetting. Moreover, ExSSNeT is particularly advantageous for sparse masks that activate 2-10% of the model parameters, resulting in an average improvement of 8.3% over SupSup. Furthermore, ExSSNeT scales to a large number of tasks (100).

## introduction artificial intelligence aims to develop agents that can learn to accomplish a set of tasks. continual learning (cl) @xcite is crucial for this, but when a model is sequentially trained on different tasks with different data distributions, it can lose its ability to perform well on previous tasks, a phenomenon is known as catastrophic forgetting (cf) @xcite @xcite . this is caused by the lack of access to data from previous tasks, as well as conflicting updates to shared model parameters when sequentially learning multiple tasks, which is called parameter interference @xcite . recently, some cl methods avoid parameter interference by taking inspiration from the lottery ticket hypothesis @xcite and supermasks @xcite to exploit the expressive power of sparse subnetworks. given that we have a combinatorial number of sparse subnetworks inside a network, @xcite noted that even within randomly weighted neural networks, there exist certain subnetworks known as supermasks that achieve good performance. a supermask is a sparse binary mask that selectively keeps or removes each connection in a fixed and randomly initialized network to produce a subnetwork with good performance on a given task. we call this the subnetwork as supermask subnetwork that is shown in figure 1 , highlighted in red weights. building upon this idea, @xcite proposed a cl method, supsup, which initializes a network with fixed and random weights and then learns a different supermask for each new task. this allows them to prevent catastrophic forgetting (cf) as there is no parameter interference (because the model weights are fixed). although supsup @xcite prevents cf, there are some problems with using supermasks for cl: (1) fixed random model weights in supsup limits the supermask subnetwork's representational power resulting in sub-optimal performance. (2) when learning a task, there is no mechanism for transferring learned knowledge from previous tasks to better learn the current task. moreover, the model is not accumulating knowledge over time as the weights are not being updated. 0) . for task 1, we first learn a supermask m1 (the corresponding subnetwork is marked by red color, column 2 row 1) and then train the weight corresponding to m1 resulting in weights w (1) (bold red lines, column 1 row 2). for task 2, we learn the mask m2 over fixed weights w (1) . if mask m2 weights overlap with m1 (marked by bold dashed green lines in column 3 row 1), then only the non-overlapping weights (solid green lines) of the task 2 subnetwork are updated (as shown by bold and solid green lines column 3 row 2). these already trained weights (bold lines) are not updated by any subsequent task. finally, for task 3, we learn the mask m3 (blue lines) and update the solid blue weights. to overcome the aforementioned issues, we propose our method, exssnet (exclusive supermask subnetwork training), pronounced as 'excess-net', which first learns a mask for a task and then selectively trains a subset of weights from the supermask subnetwork. we train the weights of this subnetwork via exclusion that avoids updating parameters from the current subnetwork that have already been updated by any of the previous tasks. in figure 1 , we demonstrate exssnet that also helps us to prevent forgetting. training the supermask subnetwork's weights increases its representational power and allows exssnet to encode task-specific knowledge inside the subnetwork (see figure 2 ). this solves the first problem and allows exssnet to perform comparably to a fully trained network on individual tasks; and when learning multiple tasks, the exclusive subnetwork training improves the performance of each task while still preventing forgetting (see figure 3 ). to address the second problem of knowledge transfer, we propose a k-nearest neighbors-based knowledge transfer (kkt) module that is able to utilize relevant information from the previously learned tasks to improve performance on new tasks while learning them faster. our kkt module uses knn classification to select a subnetwork from the previously learned tasks that has better than random predictive power for the current task and use it as a starting point to learn the new tasks. next, we show our method's advantage by experimenting with both natural language and vision tasks. for natural language, we evaluate on webnlp classification tasks (de masson d'autume et al., 2019) and glue benchmark tasks @xcite , whereas, for vision, we evaluate on splitmnist @xcite , splitcifar100 @xcite , and splittiny-imagenet @xcite datasets. we show that for both language and vision domains, exssnet outperforms multiple strong and recent continual learning methods based on replay, regularization, distillation, and parameter isolation. for the vision domain, exssnet outperforms the strongest baseline by 4.8% and 1.4% on splitci-far and splittinyimagenet datasets respectively, while surpassing multitask model and bridging the gap to training individual models for each task. in addition, for glue datasets, exssnet is 2% better than the strongest baseline methods and surpasses the performance of multitask learning that uses all the data at once. moreover, exssnet obtains an average improvement of 8.3% over supsup for sparse masks with 2 -10% of the model parameters and scales to a large number of tasks (100). furthermore, exssnet with the kkt module learns new tasks in as few as 30 epochs compared to 100 epochs without it, while achieving 3.2% higher accuracy on the splitcifar100 dataset. in summary, our contributions are listed below: ## motivation using sparsity for cl is an effective technique to learn multiple tasks, i.e., by encoding them in different subnetworks inside a single model. supsup @xcite is an instantiation of this that initializes the network weights randomly and then learns a separate supermask for each task (see figure 7 ). they prevent cf because the weights of the network are fixed and never updated. however, this is a crucial problem as discussed below. problem 1 -sub-optimal performance of supermask: although fixed network weights in supsup prevent cf, this also restricts the representational capacity, leading to worse performance compared to a fully trained network. in figure 2 , we report the test accuracy with respect to the fraction of network parameters selected by the mask, i.e., the mask density for an underlying resnet18 model on a single 100-way classification on ci-far100 dataset. the fully trained resnet18 model (dashed green line) achieves an accuracy of 63.9%. similar to @xcite , we observe that the performance of supsup (yellow dashed line) is at least 8.3% worse compared to a fully trained model. as a possible partial remedy, we propose a simple solution, ssnet (supermask subnetwork training), that first finds a subnetwork for a task and then trains the subnetwork's weights. this increases the representational capacity of the subnetwork because there are more trainable parameters. for a single task, the test accuracy of ssnet is better than supsup for all mask densities and matches the performance of the fully trained model beyond a density threshold. but as shown below, when learning multiple tasks sequentially, ssnet gives rise to parameter interference that results in cf. problem 2 -parameter interference due to subnetwork weight training for multiple tasks: next, we demonstrate that when learning multiple tasks sequentially, ssnet can still lead to cf. in figure 3 , we report the average test accuracy versus the fraction of overlapping parameters between the masks of different tasks, i.e., the sparse overlap (see equation 2 ) for five different 20-way classification tasks from splitcifar100 dataset with resnet18 model. we observe that ssnet outperforms supsup for lower sparse overlap but as the sparse overlap increases, the performance declines because the supermask subnetworks for different tasks have more overlapping (common) weights (bold dashed lines in figure 1 ). this leads to higher parameter interference resulting in increased forgetting which suppresses the gain from subnetwork weight training. our final proposal, exssnet, resolves both of these problems by selectively training a subset of the weights in the supermask subnetwork to prevent parameter interference. when learning multiple tasks, this prevents cf, resulting in strictly better performance than supsup (figure 3 ) while having the representational power to match bridge the gap with fully trained models (figure 2 ). ## method as shown in figure 1 , when learning a new task t i , exssnet follows three steps: (1) we learn a supermask m i for the task; (2) we use all the previous tasks' masks m 1 , . . . , m i-1 to create a free parameter mask m f ree i , that finds the parameters selected by the mask m i that were not selected by any of the previous masks; (3) we update the weights corresponding to the mask m f ree i as this avoids parameter interference. now, we formally describe all the step of our method exssnet (exclusive supermask subnetwork training) for a multi-layer perceptron (mlp). notation: during training, we can treat each layer l of an mlp network separately. an intermediate layer l has n l nodes denoted by v (l) = {v 1 , . . . , v n l }. for a node v in layer l, let i v denote its input and z @xmath0, where σ(.) is the activation function. given this notation, i v can be written as ## conclusion we introduced a novel continual learning method, exssnet (exclusive supermask subnetwork training), that delivers enhanced performance by utilizing exclusive, non-overlapping subnetwork weight training, overcoming the representational limitations of the prior supsup method. through the avoidance of conflicting weight updates, exssnet not only improves performance but also eliminates forgetting, striking a delicate balance. moreover, the inclusion of the knowledge transfer (kkt) module propels the learning process, utilizing previously acquired knowledge to expedite and enhance the learning of new tasks. the efficacy of exssnet is substantiated by its superior performance in both nlp and vision domains, its particular proficiency for sparse masks, and its scalability up to a hundred tasks.