A Neural Dirichlet Process Mixture Model for Task-Free Continual Learning

Parisi et al. (2019) classify CL approaches into three branches: regularization (Kirkpatrick et al., 2017; Aljundi et al., 2018), replay (Shin et al., 2017) and expansion (Aljundi et al., 2017; Rusu et al., 2016; Yoon et al., 2018) methods. Regularization and replay approaches fix the model architecture before training and prevent catastrophic forgetting by regularizing the change of a specific set of weights or replaying previously learned data. Hybrids of replay and regularization also exist, such as Gradient Episodic Memory (GEM) (Lopez-Paz & Ranzato, 2017; Chaudhry et al., 2019a). On the other hand, methods based on expansion add new network components to learn new data. Conceptually, such direction has the following advantages compared to the first two: (i) catastrophic forgetting can be eliminated since new information is not overwritten on pre-existing components and (ii) the model capacity is determined adaptively depending on the data.

Task-Free Continual Learning. All the works mentioned above heavily rely on explicit task definition. However, in real-world scenarios, task definition is rarely given at training time. Moreover, the data domain may gradually shift without any clear task boundary. Despite its importance, task-free CL has been largely understudied; to the best of our knowledge, there are only a few works (Aljundi et al., 2019a; b; Rao et al., 2019), each of which is respectively based on regularization, replay, and a hybrid of replay and expansion. Specifically, Aljundi et al. (2019a) extend MAS (Aljundi et al., 2018) by adding heuristics to determine when to update the importance weights with no task definition. In their following work (Aljundi et al., 2019b), they improve the memory management algorithm of GEM (Lopez-Paz & Ranzato, 2017) such that the memory elements are carefully selected to minimize catastrophic forgetting. While focused on unsupervised learning, Rao et al. (2019) is a parallel work that shares several similarities with our method, e.g., model expansion and short-term memory. However, due to their model architecture, expansion is not enough to stop catastrophic forgetting; consequently, generative replay plays a crucial role in Rao et al. (2019). As such, it can be categorized as a hybrid of replay and expansion. More detailed comparison between our method and Rao et al. (2019) is deferred to Appendix M.

We briefly review the Dirichlet process mixture (DPM) model (Antoniak, 1974; Ferguson, 1983), and a variational method to approximate the posterior of DPM models in an online setting: Sequential Variational Approximation (SVA) (Lin, 2013). For a more detailed review, refer to Appendix A.

Dirichlet Process Mixture (DPM). The DPM model is often applied to clustering problems where the number of clusters is not known in advance. The generative process of a DPM model is

where $x_{n}$ is the $n$-th data, and $\theta_{n}$ is the $n$-th latent variable sampled from $G$, which itself is a distribution sampled from a Dirichlet process (DP). The DP is parameterized by a concentration parameter $\alpha$ and a base distribution $G_{0}$. The expected number of clusters is proportional to $\alpha$, and $G_{0}$ is the marginal distribution of $\theta$ when $G$ is marginalized out. Since $G$ is discrete with probability 1 (Teh, 2010), same values can be sampled multiple times for $\theta$. If $\theta_{n}=\theta_{m}$, the two data points $x_{n}$ and $x_{m}$ belong to the same cluster. An alternative formulation uses the variable $z_{n}$ that indicates to which cluster the $n$-th data belongs such that $\theta_{n}=\phi_{z_{n}}$ where $\phi_{k}$ is the parameter of the $k$-th cluster. In the context of this paper, $\phi_{k}$ refers to the parameters of the $k$-th expert.

Approximation of the Posterior of DPM Models. Since the exact inference of the posterior of DPM models is infeasible, approximate inference methods are applied. Among many approximation methods, we adopt the Sequential Variational Approximation (SVA) (Lin, 2013). While the data is given one by one, SVA sequentially determines $\rho_{n}$ and $\nu_{k}$, which are the variational approximation for the distribution of $z_{n}$ and $\phi_{k}$ respectively. Since $\rho_{n}$ satisfies $\sum_{k}\rho_{n,k}=1$ and $\rho_{n,k}>=0$, $\rho_{n,k}$ can be interpreted as the probability of $n$-th data belonging to $k$-th cluster and is often called responsibility. $\rho_{n+1}$ and $\nu^{(n+1)}$ at step $n+1$ are computed as:

In practice, SVA adds a new component only when $\rho_{K+1}$ is greater than a certain threshold $\epsilon$. If $G_{0}$ and $p(x_{i}|\phi)$ are not a conjugate pair, stochastic gradient descent (SGD) is used to find the MAP estimation $\hat{\phi}$ with a learning rate of $\lambda$ instead of calculating the whole distribution $\nu_{k}$:

DPM for Discriminative Tasks. DPM can be extended to discriminative tasks where each data point is an input-output pair $(x,y)$, and the goal is to learn the conditional distribution $p(y|x)$. To use DPM, which is a generative model, for discriminative tasks, we first learn the joint distribution $p(x,y)$ and induce the conditional distribution from it: $p(y|x)=p(x,y)/\int_{y}p(x,y)$. The joint distribution modeled by each component can be decomposed as $p(x,y|z)=p(y|x,z)p(x|z)$ (Rasmussen & Ghahramani, 2002; Shahbaba & Neal, 2009).

DPM in Related Fields. Recent works of Nagabandi et al. (2019) and Jerfel et al. (2019) exploit the DPM framework to add new components without supervision in the meta-learning context. Nagabandi et al. (2019) apply DPM to the model-based reinforcement learning to predict the next state from a given state-action pair. When a new task appears, they add a component under the DPM framework to handle predictions in the new task. Jerfel et al. (2019) apply DPM to online meta-learning. Extending MAML (Finn et al., 2017), they assume that similar tasks can be grouped into a super-task in which the parameter initialization is shared among tasks. DPM is exploited to find the super-tasks and the parameter initialization for each super-task. Therefore, it can be regarded as a meta-level CL method. These works, however, lack generative components, which are often essential to infer the responsible component at test time, as will be described in the next section. As a consequence, it is not straightforward to extend their algorithms to other CL settings beyond model-based RL or meta-learning. In contrast, our method implements a DPM model that is applicable to general task-free CL.

We can formulate a CL scenario as a stream of data involving different tasks $\mathcal{D}_{1},\mathcal{D}_{2},...$ where each task $\mathcal{D}_{k}$ is a set of data sampled from a (possibly) distinct distribution $p(x|z=k)$. If $K$ tasks are given so far, the overall distribution is expressed as the mixture distribution:

where $p(z=k)$ can be approximated by $N_{k}/N$ where $N_{k}=|\mathcal{D}_{k}|$ and $N=\sum_{k}N_{k}$. The goal of CL is to learn the mixture distribution in an online manner. Regularization and replay methods directly model the approximate distribution $p(x;\phi)$ parameterized by a single component $\phi$ and update it to fit the overall distribution $p(x)$. When updating $\phi$, however, they do not have full access to all the previous data, and thus the information of previous tasks is at risk of being lost as more tasks are learned. Another way to solve CL is to use a mixture model: approximating each $p(x|z=k)$ with $p(x;\phi_{k})$. If we learn a new task distribution $p(x|z=K+1)$ with new parameter $\phi_{K+1}$ and leave the existing parameters intact, we can preserve the knowledge of the previous tasks. The expansion-based CL methods follow this idea.

Similarly, in the discriminative task, the goal of CL is to model the overall conditional distribution, which is a mixture of task-wise conditional distribution $p(y|x,z=k)$:

Prior expansion methods use expert networks each of which models a task-wise conditional distribution $p(y|x;\phi_{k})$¹¹1 The models with multiple output heads sharing the same base network (Rusu et al., 2016; Yoon et al., 2018) can also fall into this category as the expert correspond to each subnetwork connected to an output head. . However, a new problem arises in expansion methods: choosing the right expert given $x$, i.e., $p(z|x)$ in Eq.(6). Existing methods assume that explicit task descriptor $z$ is given, which is generally not true in human-like learning scenarios. That is, we need a gating mechanism that can infer $p(z|x)$ only from $x$ (i.e., which expert should process $x$). With the gating, the model prediction naturally reduces to the sum of expert outputs weighted by the gate values, which is the mixture of experts (MoE) (Jacobs et al., 1991) formulation: $p(y|x)\approx\sum_{k}p(y|x;\phi_{k})p(z=k|x)$.

However, it is not possible to use a single gate network as in Shazeer et al. (2017) to model $p(z|x)$ in CL; since the gate network is a classifier that finds the correct expert for a given data, training it in an online setting causes catastrophic forgetting. Thus, one possible solution to replace a gating network is to couple each expert $k$ with a generative model that represents $p(x|z=k)$ as in Rasmussen & Ghahramani (2002) and Shahbaba & Neal (2009). As a result, we can build a gating mechanism without catastrophic forgetting as

where $p(z=k)\approx N_{k}/N$. We also differentiate the notation for the parameters of discriminative models for classification and generative models for gating by the superscript $D$ and $G$.

If we know the true assignment of $z$, which is the case of task-based CL, we can independently train a discriminative model (i.e., $p(y|x;\phi_{k}^{D})$) and a generative model (i.e., $p(x;\phi_{k}^{G})$) for each task $k$. In task-free CL, however, $z$ is unknown, so the model needs to infer the posterior $p(z|x,y)$. Even worse, the total number of experts is unknown beforehand. Therefore, we propose to employ a Bayesian nonparametric framework, specifically the Dirichlet process mixture (DPM) model, which can fit a mixture distribution with no prefixed number of components. We use SVA described in section 2.2 to approximate the posterior in an online setting. Although SVA is originally designed for the generative tasks, it is easily applicable to discriminative tasks by making each component $k$ to model $p(x,y|z)=p(y|x,z)p(x|z)$.

The proposed approach for task-free CL, named Continual Neural Dirichlet Process Mixture (CN-DPM) model, consists of a set of experts, each of which is associated with a discriminative model (classifier) and a generative model (density estimator). More specifically, the classifier models $p(y|x,z=k)$, for which we can adopt any classifier or regressor using deep neural networks, while the density estimator describes the marginal likelihood $p(x|z=k)$, for which we can use any explicit density model such as VAEs (Kingma & Welling, 2014) and PixelRNN (Oord et al., 2016). We respectively denote the classifier and the density estimator of expert $k$ as $p(y|x;\phi_{k}^{D})$ and $p(x;\phi_{k}^{G})$, where $\phi_{k}^{D}$ and $\phi_{k}^{G}$ are the parameters of the models. Finally, the prediction $p(y|x)$ can be obtained from Eq.(7) by plugging in the output of the classifier and the density estimator. Note that the number of experts is not prefixed but expanded via the DPM framework. Figure 1 illustrates the overall training and inference process of our model.

Training. We assume that samples sequentially arrive one at a time during training. For a new sample, we first decide whether the sample should be assigned to an existing expert or a new expert should be created for it. Suppose that samples up to $(x_{n},y_{n})$ are sequentially processed and $K$ experts are already created when a new sample $(x_{n+1},y_{n+1})$ arrives. We compute the responsibility $\rho_{n+1,k}$ as follows:

where $G_{0}$ is a distribution corresponding to the weight initialization. If $\operatorname*{arg\,max}_{k}\rho_{n+1,k}\neq K+1$, the sample is assigned to the existing experts proportional to $\rho_{n+1,k}$, and the parameters of the experts are updated with the new sample by Eq.(4) such that $\hat{\phi}_{k}$ is the MAP approximation given the data assigned up to the current time step. Otherwise, we create a new expert.

Short-Term Memory. However, it is not a good idea to create a new expert immediately and initialize it to be the MAP estimation given $x_{n+1}$. Since both the classifier and density estimator of an expert are neural networks, training the new expert with only a single example leads to severe overfitting. To mitigate this issue, we employ short-term memory (STM) to collect sufficient data before creating a new expert. When a data point is classified as new, we store it to the STM. Once the STM reaches its maximum capacity $M$, we stop the data inflow for a while and train a new expert with the data in the STM for multiple epochs until convergence. We call this procedure sleep phase. After sleep, the STM is emptied, and the newly trained expert is added to the expert pool. During the subsequent wake phase, the expert is learned from the data assigned to it. This STM trick assumes that the data in the STM belong to the same expert. We empirically find that this assumption is acceptable in many CL settings where adjacent data are highly correlated. The overall training procedure is described in Algorithm 1. Note that we use $\rho_{n,0}$ instead of $\rho_{n,K+1}$ in the algorithm for brevity.

Inference. At test time, we infer $p(y|x)$ from the collaboration of the learned experts as in Eq.(7).

Techniques for Practicality. Naively adding a new expert has two major problems: (i) the number of parameters grows unnecessarily large as the experts redundantly learn common features and (ii) there is no positive transfer of knowledge between experts. Therefore, we propose a simple method to share parameters between experts. When creating a new expert, we add lateral connections to the features of the previous experts similar to Rusu et al. (2016). To prevent catastrophic forgetting in the existing experts, we block the gradient from the new expert. In this way, we can greatly reduce the number of parameters while allowing positive knowledge transfer. More techniques such as sparse regularization in Yoon et al. (2018) can be employed to reduce redundant parameters further. As they are orthogonal to our approach, we do not use such techniques in our experiments. Another effective technique that we use in the classification experiments is adding a temperature parameter to the classifier. Since the range of $\log p(x|z)$ is far broader than $\log p(y|x,z)$, the classifier has almost no effect without proper scaling. Thus, we can increase overall accuracy by adjusting the relative importance of images and labels. We also introduce an algorithm to prune redundant experts in Appendix D, and discuss further practical issues of CN-DPM in Appendix B.

A CL scenario defines a sequence of tasks where the data distribution for each task is assumed to be different from others. Below we describe the task-free CL scenarios used in the experiments. At both train and test time, the model cannot access the task information. Unless stated otherwise, each task is presented for a single epoch (i.e., a completely online setting) with a batch size of 10.

Split-MNIST (Zenke et al., 2017). The MNIST dataset (LeCun et al., 1998) is split into five tasks, each containing approximately 12K images of two classes, namely (0/1, 2/3, 4/5, 6/7, 8/9). We conduct both classification and generation experiments in this scenario.

MNIST-SVHN (Shin et al., 2017). It is a two-stage scenario where the first consists of MNIST, and the second contains SVHN (Netzer et al., 2011). This scenario is different from Split-MNIST; in Split-MNIST, new classes are introduced when transitioning into a new task, whereas the two stages in MNIST-SVHN share the same set of class labels and have different input domains.

Split-CIFAR10 and Split-CIFAR100. In Split-CIFAR10, we split CIFAR10 (Krizhevsky & Hinton, 2009) into five tasks in the same manner as Split-MNIST. For Split-CIFAR100, we build 20 tasks, each containing five classes according to the pre-defined superclasses in CIFAR100. The training sets of CIFAR10 and CIFAR100 consist of 50K examples each. Note that most of the previous works (Rebuffi et al., 2017; Zenke et al., 2017; Lopez-Paz & Ranzato, 2017; Aljundi et al., 2019c; Chaudhry et al., 2019a), except Maltoni & Lomonaco (2019), use task information at test time in Split-CIFAR100 experiments. They assign distinct output heads for each task and utilize the task identity to choose the responsible output head at both training and test time. Knowing the right output head, however, the task reduces to 5-way classification. Therefore, our setting is far more difficult than the prior works since the model has to perform 100-way classification only from the given input.

All the following baselines use the same base network that will be discussed in section 4.3.

iid-offline and iid-online. iid-offline shows the maximum performance achieved by combining standard training techniques such as data augmentation, learning rate decay, multiple iterations (up to 100 epochs), and larger batch size. iid-online is the model trained with the same number of epoch and batch size with other CL baselines.

Fine-tune. As a popular baseline in the previous works, the base model is naively trained as data enters.

Reservoir. As Chaudhry et al. (2019b) show that simple experience replay (ER) can outperform most CL methods, we test the ER with reservoir sampling as a strong baseline. Reservoir sampling randomly chooses a fixed number of samples with a uniform probability from an indefinitely long stream of data, and thus, it is suitable for managing the replay memory in task-free CL. At each training step, the model is trained using a mini-batch from the data stream and another one of the same sizes from the memory.

Gradient-Based Sample Selection (GSS). Aljundi et al. (2019b) propose a sampling method called GSS that diversifies the gradients of the samples in the replay memory. Since it is designed to work in task-free settings, we report the scores in their paper for comparison.

Split-MNIST. Following Hsu et al. (2018), we use a simple two-hidden-layer MLP classifier with ReLU activation as the base model for classification. The dimension of each layer is 400. For generation experiments, we use VAE, whose encoder and decoder have the same hidden layer configuration with the classifier. Each expert in CN-DPM has a similar classifier and VAE with smaller hidden dimensions. The first expert starts with 64 hidden units per layer and adds 16 units when a new expert is added. For classification, we adjust hyperparameter $\alpha$ such that five experts are created. For generation, we set $\alpha$ to produce 12 experts since more experts produce a better score. We set the memory size in both Reservoir and CN-DPM to 500 for classification and 1000 for generation.

MNIST-SVHN and Split-CIFAR10/100. We use ResNet-18 (He et al., 2016) as the base model. In CN-DPM, we use a 10-layer ResNet for the classifier and a CNN-based VAE. The encoder and the decoder of VAE have two CONV layers and two FC layers. We set $\alpha$ such that 2, 5, and 20 experts are created for each scenario. The memory sizes in Reservoir, GSS, and CN-DPM are set to 500 for MNIST-SVHN and 1000 for Split-CIFAR10/100. More details can be found in Appendix C.

All reported numbers in our experiments are the average of 10 runs. Table 1 and 3 show our main experimental results. In every setting, CN-DPM outperforms the baselines by significant margins with reasonable parameter usage. Table 3 and Figure 2 shows the results of Split-CIFAR10 experiments. Since Aljundi et al. (2019b) test GSS using only 10K examples of CIFAR10, which is 1/5 of the whole train set, we follow their setting (denoted by 0.2 Epoch) for a fair comparison. We also test a Split-CIFAR10 variant where each task is presented for 10 epochs. The accuracy and the training graph of GSS are excerpted from the original paper, where the accuracy is the average of three runs, and the graph is from one of the runs. In Figure 2, the bold line represents the average of 10 runs (except GSS, which is a single run), and the faint lines are the individual runs. Surprisingly, Reservoir even surpasses the accuracy of GSS and proves to be a simple but powerful CL method.

One interesting observation in Table 3 is that the performance of Reservoir degrades as each task is extended up to 10 epochs. This is due to the nature of replay methods; since the same samples are replayed repeatedly as representatives of the previous tasks, the model tends to be overfitted to the replay memory as training continues. This degradation is more severe when the memory size is small, as presented in Appendix I. Our CN-DPM, on the other hand, uses the memory to buffer recent examples temporarily, so there is no such overfitting problem. This is also confirmed by the CN-DPM’s accuracy consistently increasing as learning progresses.

In addition, CN-DPM is particularly strong compared to other baselines when the number of tasks increases. For example, Reservoir, which performs reasonably well in other tasks, scores poorly in Split-CIFAR100, which involves 20 tasks and 100 classes. Even with the large replay memory of size 1000, the Reservoir suffers from the shortage of memory (e.g., only 50 slots per task). In contrast, CN-DPM’s accuracy is more than double of Reservoir and comparable to that of iid-online.

Table 3 analyzes the accuracy of CN-DPM in Split-CIFAR10/100. We assess the performance and forgetting of individual components. At the end of each task, we measure the test accuracy of the responsible classifier and report the average of such task-wise classifier accuracies as Classifier (init). We report the average of the task-wise accuracies after learning all tasks as Classifier (final). With little difference between the two scores, we confirm that forgetting barely occurs in the classifiers. In addition, we report the gating accuracy measured after training as Gating (VAEs), which is the accuracy of the task identification performed jointly by the VAEs. The relatively low gating accuracy suggests that CN-DPM has much room for improvement through better density estimates.

Overall, CN-DPM does not suffer from catastrophic forgetting, which is a major problem in regularization and replay methods. As a trade-off, however, choosing the right expert arises as another problem in CN-DPM. Nonetheless, the results show that this new direction is especially promising when the number of tasks is very large.