Differentiable Hebbian Plasticity for Continual Learning
Abstract
Catastrophic forgetting poses a grand challenge for continual learning systems, which prevents neural networks from protecting old knowledge while learning new tasks sequentially.
We propose a Differentiable Hebbian Plasticity (DHP) Softmax layer which adds a fast learning plastic component to the slow weights of the softmax output layer.
The DHP Softmax behaves as a compressed episodic memory that reactivates existing memory traces, while creating new ones.
We demonstrate the ﬂexibility of our model by combining it with existing well-known consolidation methods to prevent catastrophic forgetting.
We evaluate our approach on the Permuted MNIST and Split
MNIST benchmarks, and introduce Imbalanced
Permuted MNIST — a dataset that combines the challenges of class imbalance and concept drift.
Our model requires no additional hyperparameters and outperforms comparable baselines by reducing forgetting.
1. Introduction
A key aspect of human intelligence is the ability to continually adapt and learn in dynamic environments, a characteristic which is challenging to embed into artiﬁcial intelligence.
Recent advances in machine learning (ML) have shown tremendous improvements in various problems, by learning to solve one complex task very well, through extensive training on large datasets with millions of training examples or more.
Most of the ML models that we use during deployment assume that the real-world is stationary, where in fact it is non-stationary and the distribution of acquired data changes over time.
Therefore, after learning is complete, and these models are ﬁne-tuned with new data, performance degrades with respect to the original data.
This phenomenon known as catastrophic forgetting or catastrophic interference (McCloskey & Cohen, 1989; French, 1999) serves to be a crucial problem for deep neural networks (DNNs) that are tasked with continual learning (Ring, 1994) or lifelong learning (Thrun & Mitchell, 1995).
In this learning paradigm, the goal is to adapt and learn consecutive tasks without forgetting how to perform previously learned tasks.
Some of the real-world applications that typically require this kind of learning include perception for autonomous vehicles, recommender systems, fraud detection, etc.
In most supervised learning methods, DNN architectures require independent and identically distributed (iid) samples from a stationary training distribution.
However, for ML systems that require continual learning in the real-world, the iid assumption is easily violated when: (1) There is concept drift or class imbalance in the training data distribution.
(2) Data representing all scenarios in which the learner is expected to perform are not initially available.
In such situations, DNNs face the “stability-plasticity dilemma” (Carpenter & Grossberg, 1987; Abraham & Robins, 2005).
This presents a continual learning challenge for models that need to balance plasticity (integrate new knowledge) and stability (preserve existing knowledge).
Two major theories have been proposed to explain a human’s ability to perform continual learning.
The ﬁrst theory is inspired by synaptic consolidation in the mammalian neocortex (Benna & Fusi, 2016) where a subset of synapses are rendered less plastic and therefore preserved for a longer timescale.
The second theory is the complementary learning systems (CLS) theory (McClelland et al., 1995; O’Reilly et al., 2014; Kumaran et al., 2016), which suggests that humans extract high-level structural information and store it in a different brain area while retaining episodic memories.
Here, we extend the work on differentiable plasticity (Miconi, 2016; Miconi et al., 2018) to a continual learning setting and develop a model that is capable of adapting quickly to changing environments as well as consolidating previous knowledge by selectively adjusting the plasticity of synapses.
We modify the traditional softmax layer and propose to augment the slow weights with a set of plastic weights implemented using Differentiable Hebbian Plasticity (DHP).
The model’s slow weights learn deep representations of data and the fast weights implemented with DHP learn to quickly “auto-associate” the class labels to representations.
We also demonstrate the ﬂexibility of our model by combining it with recent task-speciﬁc synaptic consolidation based methods to overcoming catastrophic forgetting such as elastic weight consolidation (Kirkpatrick et al., 2017;
Schwarz et al., 2018), synaptic intelligence (Zenke et al., 2017) and memory aware synapses (Aljundi et al., 2018).
Our model uniﬁes core concepts from Hebbian plasticity, synaptic consolidation and CLS theory to enable rapid adaptation to new unseen data, while consolidating synapses and leveraging compressed episodic memories to remember previous knowledge and mitigate catastrophic forgetting.
2. Relevant Work
Plastic Neural Networks: One of the major theories that have been proposed to explain a human’s ability to learn continually is Hebbian learning (Hebb, 1949), which suggests that learning and memory are attributed to weight plasticity, that is, the modiﬁcation of the strength of existing synapses according to variants of Hebb’s rule (Paulsen & Sejnowski, 2000; Song et al., 2000; Oja, 2008).
Recent approaches in the meta-learning literature have shown that we can incorporate fast weights into a neural network (Munkhdalai & Trischler, 2018; Rae et al., 2018).
Munkhdalai & Trischler (2018) augmented fully-connected (FC) layers preceding the softmax with a matrix of fast weights.
Here, the fast weights were implemented with non-trainable Hebbian learning-based associative memory.
Rae et al.
(2018) proposed a softmax layer that can improve learning of rare classes by interpolating between Hebbian updates and stochastic gradient descent (SGD) updates on the output layer using an arbitrarily engineered scheduling scheme.
Miconi et al.
(2018) proposed differentiable plasticity, which uses SGD to optimize the plasticity of each synaptic connection composed of a slow weight and a plastic (fast) weight.
Although this approach served to be a powerful new method for training neural networks, it was mainly demonstrated on RNNs for solving simple tasks.
Overcoming Catastrophic Forgetting: This work leverages two biologically inspired strategies to overcome the catastrophic forgetting problem: 1) Task-speciﬁc
Synaptic
Consolidation — Protecting old knowledge by dynamically adjusting the synaptic strengths to consolidate and retain memories.
2) CLS Theory — A dual memory system where, structural knowledge is acquired through slow learning via the neocortex and rapid learning via the hippocampus.
There have been several notable works inspired by taskspeciﬁc synaptic consolidation for overcoming catastrophic forgetting (Kirkpatrick et al., 2017; Zenke et al., 2017;
Aljundi et al., 2018).
All of these approaches propose a method to estimate the importance of each parameter or synapse, Ω k , where the least plastic synapses can retain memories for a long time and the more plastic synapses are considered less important.
The Ω k and network parameters θ k are updated online or after learning task T n .
Therefore, when learning new task T n , a regularizer is added to the original loss function L n ( θ ) , so that we dynamically adjust the plasticity w.r.t.
Ω k and prevent any changes to the important parameters of previously learned tasks:
[Equation]
where θ n − 1 k are the learned network parameters after training on the previous n − 1 tasks and λ is a hyperparameter for the regularizer to control the amount of forgetting.
In Elastic Weight Consolidation (EWC), Kirkpatrick et al.
(2017) use the diagonal values of an approximated Fisher information matrix for Ω k , and it is computed ofﬂine after training on a task is completed.
Schwarz et al.
(2018) proposed an online variant of EWC to improve scalability by ensuring the computational cost of the regularization term does not grow with the number of tasks.
Zenke et al.
(2017) proposed an online method called Synaptic Intelligence (SI) for computing the parameter importance where, Ω k is the cumulative change in individual synapses over the entire training trajectory on a given task.
Memory Aware
Synapses (MAS) from Aljundi et al.
(2018) measures Ω k by the sensitivity of the learned function to a perturbation in the parameters and use the cumulative change in individual synapses on the squared L2-norm of the penultimate layer.
There have been numerous approaches based on CLS principles involving pseudo-rehersal (Robins, 1995; Ans et al., 2004;
Atkinson et al., 2018), episodic replay (Lopez-Paz & Ranzato, 2017; Li & Hoiem, 2018) and generative replay (Shin et al., 2017; Wu et al., 2018)
However, in our work, we are primarily interested in neuroplasticity techniques inspired from CLS theory for representing memories.
Hinton & Plaut (1987) showed how each synaptic connection can be composed of a ﬁxed weight where slow learning stores long-term knowledge and a fast weight for temporary associative memory.
Recent research in this vein has included replacing soft attention mechanism with fast weights in RNNs (Ba et al., 2016), the Hebbian Softmax layer (Rae et al., 2018), augmenting the FC layer with a fast weights matrix (Munkhdalai & Trischler, 2018),
differentiable plasticity (Miconi et al., 2018) and neuromodulated differentiable plasticity (Miconi et al., 2019).
However, all of these methods were focused on rapid learning on simple tasks or meta-learning over a distribution of tasks.
Furthermore, they did not examine learning a large number of new tasks while, alleviating catastrophic forgetting in continual learning.
3. Model
In our model, each synaptic connection in the softmax layer has two weights: 1)
The slow weights, θ ∈ R m × d , where m is the number of units in the ﬁnal hidden layer.
A
Hebbian plastic component of the same cardinality as the slow weights, composed of the plasticity coefﬁcient, α , and the Hebbian trace, Hebb .
The α is a scaling parameter for adjusting the magnitude of the Hebb .
Hebb accumulates the mean activations of the penultimate layer for each target label in the mini-batch { y 1: B } of size B which are denoted by ˜ h ∈ R 1 × m (refer to Algorithm 1).
Given the activation of each neuron in h at the pre-synaptic connection i , the unnormalized log probabilities z at the post-synaptic connection j can be more formally computed using Eq. 2.
Then, the softmax function is applied on z to obtain the desired logits
ˆ y thus, ˆ y = softmax( z ) .
The η parameter in Eq.is a “learning rate” that learns how quickly to acquire new experiences into the plastic component.
The η parameter also acts as a decay term to prevent instability caused by a positive feedback loop in the Hebbian traces.
[Equation]
The network parameters α i,j , η
and θ i,j are optimized by gradient descent as the model is trained sequentially on different tasks in the continual learning setup.
Hebb is initialized to zero only at the start of learning the ﬁrst task
T 1 and is automatically updated based on Algorithm 1 in the forward pass during training.
Speciﬁcally, the Hebbian update for the active class c in y 1: B is computed on line 6.
This Hebbian update 1 s (cid:80) Bb =
1 h
[ y b = c ] is analogous to another formulaic description of the Hebbian learning update rule w i,j = 1 N (cid:80)
Nk
=1 a ki a kj (Hebb, 1949), where w i,j is the change in weight at connection
i, j and a ki , a kj denote the activation levels of neurons i and j , respectively, for the k th input.
Therefore, in our model, w = ˜ h
the Hebbian weight update, a i = h the hidden activations of the last hidden layer, a j = y the active target class in y 1: B and
N = s the number of inputs for the corresponding class in y 1: B (see Algorithm 1).
Across the model’s lifetime, we only update Hebb during training and during test time, we use the most recent Hebbian traces to make predictions.
The plastic component learns rapidly and performs sparse parameter updates to quickly store memory traces for each recent experience without interference from other similar recent experiences.
Furthermore, the hidden activations corresponding to the same active class are accumulated into one vector ˜ h , thus forming a compressed episodic memory in the Hebb to reﬂect individual episodic memory traces.
This method improves learning of rare classes and speeds up binding of class labels to deep representations of the data.
Hebb (Hebbian trace) do 4: s
(cid:80)
Bb =
[ y b = c ] /
*/ 5: if then
Hebb : , [Equation] ← [Equation]
[Equation]
Hebb
[Equation]
8: end if 9: end for 10: z ← ( θ [Equation] α Hebb [Equation]
/*Compute softmax pre-activations.*/
Updated Loss:
Following the existing work for overcoming catastrophic forgetting such as EWC, Online EWC, SI and
MAS (see Eq. 1), we regularize the loss L n ( θ, α, η ) and update the synaptic importance parameters of the network in an online manner.
We rewrite Eq.
1 to obtain Eq. 4 and show that the network parameters θ i,j are the weights of the connections between pre- and post-synaptic activity, as seen in Eq.
[Equation]
We adapt these existing consolidation approaches to our model and only compute the synaptic importance parameters on the slow weights of the network.
The plastic part of our model can alleviate catastrophic forgetting of learned classes by optimizing the plasticity of the synaptic connections.
4. Experiments
We tested our continual learning approach on the Permuted
MNIST, Imbalanced Permuted MNIST and Split MNIST benchmarks.
We evaluated the methods based on the average classiﬁcation accuracy on all previously learned tasks.
To establish a baseline for comparison of well-known synaptic consolidation methods, we trained neural networks with
Online EWC, SI and MAS, respectively, on all tasks in a sequential manner.
In the Permuted MNIST and Imbalanced
Permuted benchmarks we trained a multi-layered perceptron (MLP) network on a sequence of 10 tasks using plain SGD.
Detailed descriptions of the hyperparameters and training setups for all benchmarks can be found in Appendix A.
Permuted MNIST:
In this benchmark, all of the MNIST pixels are permuted differently for each task with a ﬁxed random permutation.
Although the output domain is constant, the input distribution changes between tasks thus, there exists a concept drift.
Figure 1 shows the average test accuracy as new tasks are learned.
The network with
DHP Softmax alone showed signiﬁcant improvement in its ability to alleviate catastrophic forgetting across all tasks compared to the baseline ﬁnetuned vanilla MLP network we refer to as Finetune in Figure 1.
Then we compared the performance with and without DHP Softmax using the synaptic consolidation methods.
We ﬁnd our DHP Softmax with synaptic consolidation maintains a higher test accuracy after T 10 tasks than without DHP Softmax for all variants.
[Figure]
Imbalanced Permuted MNIST:
This benchmark is identical to the Permuted MNIST benchmark but, now each task is an imbalanced distribution.
The statistics of the class distribution in each task are presented in Appendix A.2, Table 1.
Figure 2 shows the average test accuracy as new tasks are learned.
We see that DHP Softmax achieves 80.85 % after learning 10 tasks, thus providing signiﬁcant improvement over the standard neural network baseline of 76.4 % .
The signiﬁcance of the compressed episodic memory mechanism in the Hebbian traces is more apparent in this benchmark because the plastic component allows rare classes that are encountered infrequently to be remembered for a longer period of time.
We ﬁnd that DHP Softmax with MAS achieves 88.8 % ; outperforming all other methods and across all tasks.
Split MNIST:
A sequence of T n =1:5 tasks are generated by splitting the original MNIST training dataset into binary classiﬁcation problems (0/1, 2/3, 4/5, 6/7, 8/9), making the output spaces disjoint between tasks.
Similar to Zenke et al.
(2017), we trained a multi-headed MLP network on a sequence of 5 tasks.
We compute the cross entropy loss at the softmax output layer only for the digits present in the current task, T n .
We observe that DHP Softmax provides a 4.7 % improvement on test performance compared to a ﬁnetuned MLP network (Figure 3).
Also, combining
DHP Softmax with task-speciﬁc consolidation consistently improves performance across all tasks T n =1:5 .
[Figure] [Figure]
5. Discussion and conclusion
We have shown that the problem of catastrophic forgetting in continual learning environments can be alleviated by adding compressed episodic memory in the softmax layer through
DHP.
DHP Softmax alone showed noticeable improvement across all benchmarks when compared to a neural network with a traditional softmax layer.
We demonstrated the ﬂexibility of our model where, in addition to DHP Softmax, we can regularize the slow weights using EWC, SI or MAS to improve a model’s ability to alleviate catastrophic forgetting.
The approach where we combine DHP Softmax and MAS consistently leads to overall superior results compared to other baseline methods on several benchmarks.
This gives a strong indication that Hebbian plasticity enables neural networks to learn continually and remember distant memories, thus reducing catastrophic forgetting when learning from sequential datasets in dynamic environments.
A. Appendix
The model used for the Permuted MNIST and Imbalanced Permuted MNIST benchmarks is a multilayered perceptron (MLP) network with two hidden layers consisting of 400 units each with ReLU nonlinearities, and a cross-entropy loss L ( θ ) .
We train the network on a sequence of 10 tasks T n =1:10 with mini-batches of size 64 and optimized using plain SGD with a ﬁxed learning rate of 0.01.
We train for atleast 10 epochs and perform early-stopping once the validation error does not improve for 5 epochs.
If the validation error increases for more than 5 epochs, then we terminated the training on the task
T n , reset the network weights and Hebbian traces to the values that had the lowest validation error, and proceeded to the next task.
For all of the benchmarks we tested on, the η of the plastic component was set to be a small value of 0.001 and we want to emphasize that we spent little to no efforts on tuning this parameter.
Also, when training the ﬁrst task T n =1 , the synaptic importance parameter, Ω i,j in Eq. 4, was set to 0 for all of the task-speciﬁc consolidation methods that we tested on except for SI.
This is because SI is the only method we evaluated that estimates Ω i,j while training, whereas Online EWC and
MAS compute Ω i,j after learning a task.
A.1. Permuted MNIST
For the Permuted MNIST experiments shown in Figure 1, the regularization hyperparameter λ for each of the task-speciﬁc consolidation methods is λ = 100 for Online EWC (Schwarz et al., 2018), λ = 0 .
1 for SI (Zenke et al., 2017) and λ = 0 .for
MAS (Aljundi et al., 2018).
In SI, the damping parameter, ξ , was set to 0.1.
To ﬁnd the best hyperparameter combination for each of these synaptic consolidation methods, we performed a grid search using a task sequence determined by a single seed.
The hyperparameters of the consolidation methods (i.e. Online EWC, SI and MAS) remain the same with and without
DHP Softmax, and the plastic components are not regularized.
A.2. Imbalanced Permuted MNIST
For each task in the Imbalanced Permuted MNIST problem, we artiﬁcially removed training samples from each class in the original MNIST dataset (LeCun et al., 2001) based on some random probability.
The distribution of classes in each dataset corresponding to tasks T n =1:10 is given in Table
[Table] [Table]
For the Imbalanced Permuted MNIST experiments shown in Figure 2, the regularization hyperparameter λ for each of the task-speciﬁc consolidation methods is λ = 400 for Online EWC (Schwarz et al., 2018), λ = 1 . 0 for SI (Zenke et al., 2017) and λ = 0 .
1 for MAS
(Aljundi et al., 2018).
In SI, the damping parameter, ξ , was set to 0.1.
Similar to the Permuted MNIST benchmark, to ﬁnd the best hyperparameter combination for each of these synaptic consolidation methods, we performed a grid search using a task sequence determined by a single seed.
Across all experiments, we maintained the the same random probabilities detemined by a single seed to artiﬁcially remove training samples from each class.
The hyperparameters of the synaptic consolidation methods (i.e. Online EWC, SI and MAS) remain the same with and without DHP Softmax, and the plastic components are not regularized.
A.3. Split MNIST
We split the original MNIST dataset (LeCun et al., 2001) into a sequence of 5 binary classiﬁcation tasks: T 1 = { 0/1 } , T 2 = { 2/3 } , T 3 = { 4/5 } , T 4 = { 6/7 } and T 5 = { 8/9 } .
Similar to network used by Zenke et al.
(2017), we use a MLP network with two hidden layers of 256 ReLU nonlinearities each, and a cross-entropy loss.
A multi-headed approach was used to avoid interference between digits at the softmax output layer due to changes in the label distribution.
We compute the cross-entropy loss, L ( θ ) , at the softmax output layer for the digits present in the current task, T n .
We train the network on a sequence of 5 tasks T n =1:5 with mini-batches of size 64 and optimized using plain SGD with a ﬁxed learning rate of 0.
for 10 epochs.
For the Split MNIST experiments shown in Figure 3, the regularization hyperparameter λ for each of the task-speciﬁc consolidation methods is λ = 400 for Online EWC (Schwarz et al., 2018), λ = 1 . 0 for SI (Zenke et al., 2017) and λ = 1 . 5 for
MAS (Aljundi et al., 2018).
In SI, the damping parameter, ξ , was set to 0.001.
To ﬁnd the best hyperparameter combination for each of these synaptic consolidation methods, we performed a grid search using the 5 task binary classiﬁcation sequence (0/1, 2/3, 4/5, 6/7, 8/9).
The hyperparameters of the consolidation methods (i.e. Online EWC, SI and MAS) remain the same with and without DHP Softmax, and the plastic components are not regularized.