Improving Task-free Continual Learning by
Distributionally Robust Memory Evolution

A sequence of mini-batch labeled data $({\bm{x}}_{k},y_{k},h_{k})$ sequentially arrives at each timestamp $k$ and forms a non-stationary data stream. Each data point is associated with a latent task identity $h_{k}$, where ${\bm{x}}_{k}$ denotes the mini-batch data received at timestamp $k$, $y_{k}$ is the data label associated with ${\bm{x}}_{k}$. According to (Aljundi et al., 2019b), a more general definition of task-free CL is that data distribution shift could happen at any time without explicit task splits. Our method can be directly applied to those more general cases as well. During both the training and testing time, the task identity $h_{k}$ is not available to the learner. At the same time, a small memory buffer ${\mathcal{M}}$ is maintained, and replay the data in ${\mathcal{M}}$ when learning the new task to avoid forgetting the previously learned knowledge. The memory buffer ${\mathcal{M}}$ is updated by reservoir sampling (RS), similar to (Riemer et al., 2019). The goal is to learn a model $f({\bm{x}},{\bm{\theta}})$ to perform well on all the tasks seen until timestamp $k$. Standard memory replay for CL (Chaudhry et al., 2019b) is to optimize an objective under a known probability distribution $\mu_{0}$. Formally speaking, CL with standard memory replay can be expressed as:

where ${\bm{\theta}}$ are model parameters and ${\bm{x}}$ is the raw memory buffer data with probability measure (density) $\mu_{0}$, i.e., $\forall{\bm{x}}\in{\mathcal{M}},{\bm{x}}\sim\mu_{0}$. ${\mathcal{L}}({\bm{\theta}},{\bm{x}},y)$ is the loss function associated with the data $({\bm{x}},y)$. In the following, we temporally omit the term ${\mathcal{L}}({\bm{\theta}},{\bm{x}}_{k},y_{k})$ due to the fact that $({\bm{x}}_{k},y_{k})$ is the mini-batch data arrived at timestamp $k$ and does not depend on the memory data distribution.

Distributionally Robust Optimization (DRO) (Rahimian & Mehrotra, 2019) is a systematic and elegant framework to model the decision-making with ambiguity in the underlying probability distribution. For task-free CL, different from the standard memory-replay methods in Section 3.1, implicitly assuming $\mu_{0}$ is known. In contrast, our proposed DRO framework takes that the underlying actual probability distribution of memory data $\mu$ is unknown and lies in an ambiguity set of probability distributions. Modeling the memory data uncertainty is particularly useful when the memory buffer is small compared to the whole dataset since the memory has limited coverage to approximate the stationary distribution of all examples seen so far, illustrated in Figure 2 (a). Thus, there is significant uncertainty in modeling the multi-task learning scenarios (the performance upper bound) with only a small memory buffer. The proposed DRO framework optimizes the worst-case performance in the ambiguity set of probability distributions since we cannot access the distribution of all the previous data. Therefore, it helps the model generalize to previous tasks since it can potentially narrow the gap between the memory data distribution and the distribution of all the previous data, illustrated in Figure 2 (b). As a by-product of this optimization framework, it also helps learn features robust to data distribution perturbations. On the other hand, the memory buffer is updated slowly during CL (e.g., by reservoir sampling), and standard memory-replay repeatedly trains the memory buffer, and the CL model can easily overfit the memory buffer, as illustrated in Figure 1 (a). Thus, the memory buffer could become less effective for mitigating forgetting. By optimizing the worst-case evolved memory data distribution at each iteration, our proposed DRO can also alleviate the memory overfitting problem by transforming the memory data to make them more difficult to memorize. This is illustrated in Figure 1 (b) and (c). Mathematically speaking, the proposed DRO for task-free CL can be expressed as:

where the inner $\sup$ optimization is to gradually make the memory data distribution increasingly harder to be memorized. ${\mathcal{P}}$ in Eq. (3) denotes the ambiguity set of probability measures (distributions or densities) for the memory data distribution to characterize its uncertainty. One common choice to define ${\mathcal{P}}$ is through Kullback-Leibler (KL) divergence. ${\mathcal{D}}(\mu_{0}||\pi)$ denotes the KL divergence between probability measure $\mu_{0}$ and $\pi$, where $\pi$ is the target worst-case evolved memory data distribution, i.e., the probability distribution $\pi$ that achieves $\sup_{\mu\in{\mathcal{P}}}\mathbb{E}_{\mu}{\mathcal{L}}({\bm{\theta}},{\bm{x}},y)$. $\epsilon$ is a constant threshold to characterize the closeness between $\mu_{0}$ and $\pi$ to ensure the worst-case evolved memory data distribution $\pi$ does not deviate from the raw memory data distribution $\mu_{0}$ too much. Eq. (4) constrains the gradient dot product between the worst-case evolved memory data distribution and raw memory data distribution, i.e., $\nabla_{{\bm{\theta}}}{\mathcal{L}}({\bm{\theta}},{\bm{x}},y)\cdot\nabla_{{\bm{\theta}}}{\mathcal{L}}({\bm{\theta}},{\bm{x}}^{\prime},y)$, to avoid the evolved memory data deviate from the raw memory data too much. Intuitively, if the gradient dot product is positive, it means the evolved memory data has a similar update direction compared to the raw data. $\lambda$ is a threshold to determine the constraint magnitude.

To solve the optimization, i.e., Eq. (2-4), the worst-case optimization that involves the $\sup$ optimization within the KL-divergence ball is generally computationally intractable since it involves the optimization over infinitely many probability distributions. To address this problem, by Lagrange duality (Boyd & Vandenberghe, 2004), we convert Eq. (2-4) into the following unconstrained optimization problem (detailed derivations are put in Appendix B):

Where $\gamma$ and $\beta$ control the magnitude of regularization. Since the KL-divergence term ${\mathcal{D}}(\mu||\pi)$ is implicitly handled by gradient flow (discussed in the following sections), we set $\gamma=1$ for simplicity throughout this paper. The gradient of the third term (gradient dot product term) can be efficiently approximated by finite difference in practice. The optimization Eq. (5) is still computationally hard to solve because the inner $\sup$ optimization is over probability measure space, which is an infinite-dimensional function space. We name Eq. (5) as task-free DRO.

To make it tractable to solve the task-free DRO Eq. (5), we formulate it from a new continuous dynamics perspective. This new perspective brings significant advantages over directly solving Eq. (5): 1) the optimization of Eq. (5) over probability measure space can be converted into a continuous probability distribution evolution, equivalent as a WGF, enabling gradient-based solutions in probability measure space (function space); 2) we can efficiently solve the WGF by various methods, which provide potential derivations of many different memory evolution methods. We then propose a novel solution by decomposing the task-free DRO Eq. (5) into a gradient flow system. Specifically, we solve the inner $\sup$ optimization problem for memory evolution with WGF in Wasserstein space of probability measures and solve the outer optimization problem for the model parameters with gradient flow in Euclidean space. We convert the task-free DRO into a gradient flow system that alternately updates the memory evolution and model parameters.

Given a memory buffer at timestamp $k$, the raw memory data is ${\mathcal{M}}=\{({\bm{x}}^{1}_{0},y^{1}),({\bm{x}}^{2}_{0},y^{2}),\cdots,({\bm{x}}^{N}_{0},y^{N})\}$, where $N$ is the memory buffer size. We perform a similar memory evolution procedure at each CL timestamp and omit the CL timestamp $k$ for notation clarity. We denote ${\bm{x}}_{t}^{i}$ as the $i^{th}$ datapoint in the evolved memory buffer after $t$ evolution steps. The raw memory data is assumed to be i.i.d. sampled from the random variable $X_{0}$, i.e., $\{{\bm{x}}_{0}^{1},{\bm{x}}_{0}^{2},\cdots,{\bm{x}}_{0}^{N}\}\sim X_{0}$. $X_{0}$ follows the probability distribution $\mu_{0}$, i.e., $X_{0}~{}\sim\mu_{0}$. At evolution time $t$, the memory data ${\mathcal{M}}$ is evolved as random variable $X_{t}$, i.e., $\{{\bm{x}}_{t}^{1},{\bm{x}}_{t}^{2},\cdots,{\bm{x}}_{t}^{N}\}\sim X_{t}$ and probability distribution of random variable $X_{t}$ follows the probability measure $\mu_{t}$, i.e., $X_{t}~{}\sim\mu_{t}$. The empirical measure on the evolved memory buffer at time $t$ is defined as $\hat{\mu}_{t}=\frac{1}{N}\sum_{i=1}^{i=N}\delta({\bm{x}}^{i}_{t})$ and $\delta$ is the Dirac measure. We model the memory evolution process as continuous WGF in probability measure space, i.e., use $(\mu_{t})_{t\geq 0}$ to model the probability distribution evolution of memory data. The evolving $(\mu_{t})_{t\geq 0}$ will in turn determine the memory evolution process $(X_{t})_{t\geq 0}$ in Euclidean space.

Below, we define and present the gradient flow in Wasserstein space. Let ${\mathcal{P}}_{2}({\mathbb{R}}^{d})$ denote the space of probability measures on ${\mathbb{R}}^{d}$ with finite second-order moments. Each element $\mu\in{\mathcal{P}}_{2}({\mathbb{R}}^{d})$ is a probability measure represented by its density function $\mu:{\mathbb{R}}^{d}\to{\mathbb{R}}$ with respect to Lebesgue measure $dx$. The Wasserstein distance between two probability measures $\mu_{1},\mu_{2}\in{\mathcal{P}}_{2}({\mathbb{R}}^{d})$ is defined as:

where $\prod(\mu_{1},\mu_{2})=\{\omega|\omega(A\times{\mathbb{R}}^{d})=\mu_{1}(A),\omega({\mathbb{R}}^{d}\times B)=\mu_{2}(B)\}$. $\omega$ is the joint probability measure with marginal measure of $\mu_{1}$ and $\mu_{2}$ respectively. Thus, ${\mathbb{W}}^{2}=({\mathcal{P}}_{2}({\mathbb{R}}^{d}),W_{2})$ forms a metric space.

where $:=$ means defined as, and $div({\bm{r}}):=\sum_{i=1}^{d}\partial_{{\bm{z}}^{i}}{\bm{r}}^{i}({\bm{z}})$ is the divergence operator of a vector-valued function ${\bm{r}}:{\mathbb{R}}^{d}\to{\mathbb{R}}^{d}$, where ${\bm{z}}^{i}$ and ${\bm{r}}^{i}$ are the $i$ th element of ${\bm{z}}$ and ${\bm{r}}$; $\nabla$ is the gradient of a scalar-valued function. $\nabla_{W_{2}}F(\mu_{t}):=div(\mu_{t}\nabla\frac{\delta F}{\delta\mu}(\mu_{t}))$ is the Wasserstein gradient of functional $F$ at $\mu_{t}$, where $\frac{\delta F}{\delta\mu}(\mu_{t})$ is the first variation of $F$ at $\mu_{t}$. The first variation of a functional in function space is analogous to the first-order gradient of a function in Euclidean distance. We put more detailed definition in Appendix D. Intuitively, the WGF describes that the probability measure $\mu_{t}$ follows the steepest curve of the functional $F(\mu)$ in Wasserstein space of probability measures (function space) to gradually move towards the target probability measure.

We define the energy functional $F(\mu)$ for memory evolution as the following:

By defining such energy functional $F(\mu)$, the Eq. (5) can be equivalently solved by the following gradient flow system Eq. (9, 10), we name it as Dynamic DRO.

where Eq. (9) solves the inner $\sup$ problem in Eq. (5) with WGF in Wasserstein space and Eq. (10) solves the outer minimization problem in Eq. (5) for parameter update with gradient flow in Euclidean space. In the following, we focus on solving Eq. (9) and Eq. (10).

We propose three different methods for efficiently solving the Dynamic DRO in Eq.  (9)-Eq. (10). The first solution is Langevin dynamics with a diffusion process to evolve the memory data distribution; we name this method WGF-LD. Next, different from WGF-LD, which relies on randomness to transform the memory data, we kernelize the WGF in Eq. (9) by solving the WGF in reproducing kernel Hilbert space (RKHS). It deterministically transforms the memory data; we name this method WGF-SVGD. Furthermore, we generalize the above WGF and improve their flexibility to incorporate prior knowledge or geometry information. One instantiation of this general WGF uses Hamiltonian dynamics; we name it WGF-HMC. More novel memory evolution methods are worth exploring in future work.

Before introducing the proposed solutions, we first present some preliminaries for solving the WGF. By calculus of variation (Miersemann, 2012), the first variation for the KL divergence and $V(\mu)$ are (Ambrosio et al., 2008):

CIFAR10, following (Aljundi et al., 2019a), we split the CIFAR-10 dataset into 5 disjoint tasks with the same training, validation, and test sets.

MiniImagenet, following (Aljundi et al., 2019a), we split MiniImagenet (Vinyals et al., 2016) dataset with 100 image classes into 20 disjoint tasks. Each task has 5 classes.

CIFAR-100 (Krizhevsky, 2009), contains 100 image classes. We also split it into 20 disjoint tasks.

Baselines. We performed comprehensive experiments by comparing to the following strong baselines, including Experience Replay (ER) (Chaudhry et al., 2019b), Maximally Interfering Retrieval (MIR) (Aljundi et al., 2019a), AGEM (Chaudhry et al., 2019a), Gradient-Based Sample Selection (GSS-Greedy) (Aljundi et al., 2019c) and GMED (Jin et al., 2021). Furthermore, following (Jin et al., 2021), we also compare data augmentation, such as random rotations, scaling, and horizontal flipping applied to memory buffer data in ER and name this baseline as ER_aug. Following (Aljundi et al., 2019a), we also compare (1) fine-tuning, which trains on each latent task sequentially when new batches of each task arrive without any forgetting mitigation mechanism; (2) iid online: which trains the model with a single-pass through the iid sampled data on the same set of samples; (3) iid offline: (upper-bound) which trains the model with multiple passes through the iid sampled data. We train the model with 5 epochs for this baseline. Detailed descriptions of the compared baselines are placed in Appendix A.

In addition, our proposed methods are orthogonal to existing memory-replay-based CL methods. Thus, they are versatile, and we can seamlessly combine the proposed methods with them. For illustration, we combine our proposed methods with ER, MIR, and GMED to show the effectiveness. We name the combination methods ER+WGF-LD, ER+WGF-SVGD, ER+WGF-HMC, MIR+WGF-LD, GMED+WGF-LD, etc. Combining with other memory-replay-based methods is straightforward.

Implementation Details. We use the Resnet-18 as (Aljundi et al., 2019a). The number of evolution steps $T$ is set to be $5$, the evolution rate $\alpha$ is set to be $0.01$ for CIFAR10, $0.05$ for CIFAR100 and $0.001$ for MiniImagenet, $\beta=0.003$ and momentum $\tau=0.1$. We set the memory buffer size to be 500 for CIFAR-10 dataset, $5K$ for CIFAR-100 and $10K$ for MiniImagenet. All other hyperparameters are the same as (Aljundi et al., 2019a). All reported results in our experiments are the average accuracy of 10 runs with standard deviation.

We compare the proposed methods to various task-free CL baselines. Table 1 shows the effectiveness of combining the proposed method with existing memory-replay methods, e.g., ER, MIR and GMED. We can observe that our method outperforms these strong baselines. In particular, for ER and ER + WGF methods, our method outperforms baselines by $4.5\%$, $1.4\%$, and $2.8\%$ on CIFAR10, CIFAR-100, and MiniImageNet, respectively. For MIR and MIR + WGF methods, our method outperforms baselines by $3.8\%$, $1.6\%$, and $2.1\%$ on CIFAR10, CIFAR-100, and MiniImageNet, respectively. For GMED and GMED + WGF methods, our method outperforms baselines by $3.6\%$, $0.9\%$, and $1.4\%$ on CIFAR10, CIFAR-100, and MiniImageNet, respectively. Our methods outperform baselines because they dynamically evolve the memory data distribution and sufficiently explore the input space in a principled way. The proposed methods generate more diverse memory data, and the evolved memory becomes more difficult for the CL model to memorize than baseline methods. In addition, WGF-SVGD generally performs better than WGF-SGLD and WGF-HMC. We believe this is because, in RKHS, the evolved memory data is encouraged to be far apart by the kernel repulsive forces; the evolved memory data may better represent the distribution of all the previous data.

In this section, we evaluate the robustness of the CL model to adversarial perturbed examples. Given a classifier $f({\bm{x}},{\bm{\theta}})$, for an image ${\bm{x}}$, the goal is to find another example ${\bm{x}}^{\prime}$ that is close enough to ${\bm{x}}$ measured by some distance function $D({\bm{x}}.{\bm{x}}^{\prime})\leq\epsilon$ such that the classifier classifies it into another different class, i.e. $f({\bm{x}},{\bm{\theta}})\neq f({\bm{x}}^{\prime},{\bm{\theta}})$. In this paper, we focus on the most commonly used $\ell_{\infty}$ and $\ell_{2}$ norm as the distance function.

Our memory evolution methods optimize on the worst-case evolved memory data distribution during CL; the model would be thus naturally robust to adversarial perturbations. For $\ell_{\infty}$ norm attack, we evaluate the robustness under the strong PGD $\ell_{\infty}$ attack (Madry et al., 2018), which constructs adversarial examples with projected gradient descent and $\ell_{\infty}$ norm constraints. The adversarial perturbation magnitude ranges from $[1/255,2/255,\cdots,10/255]$ with 20 steps attack and stepsize of $\frac{2}{255}$ on CIFAR-100 and Mini-ImageNet. For $\ell_{2}$ norm attack, we adopt the strong Carlini $\&$ Wagner attack (Carlini & Wagner, 2017). For illustration, we evaluate the model robustness to adversarial examples after training with ER, GMED, ER + WGF-LD, ER + WGF-SVGD, and ER + WGF-HMC, respectively. Figure 5 shows the PGD $\ell_{\infty}$ attack results on CIFAR-100 and Mini-ImageNet. Our WGF-HMC and WGF-LD memory evolution significantly outperforms naive ER baseline by $4\%-12\%$ depending on the perturbation magnitude. In addition, WGF-HMC and WGF-LD perform more robust than WGF-SVGD. We believe this is because the randomness introduced in WGF-HMC and WGF-LD can better explore the input space and thus can generate harder evolved memory data. WGF-SVGD smooths the function gradient by the kernel function, thus may generate less hard examples. Table 2 shows the Carlini $\&$ Wagner attack results. We can see that under the strong Carlini $\&$ Wagner attack, ER baseline accuracy becomes zero, and our methods still outperform baselines ranging from $6.1\%,3.0\%,3.1\%$ on CIFAR-10, CIFAR-100, Mini-ImageNet, respectively. Both results demonstrate the robustness of our proposed methods to adversarial examples.

Effects of different memory size. To investigate the effect of different smaller memory sizes on the model performance, we evaluate the effects on Mini-ImageNet with memory sizes of 3000, 5000, and 10000. We evaluate the effects on CIFAR-100 with the memory sizes of 2000, 3000, and 5000. We show the results in Table 3. In most cases, our WGF memory evolution substantially outperforms the baselines with different memory buffer sizes.

Effect of number of evolution steps. To investigate the effect of different evolution steps, we compare 3, 5, and 7 evolution steps, respectively. Table 4 shows the performance variation of different number of evolution steps. We can find that the performance improves slightly with an increasing number of evolution steps. For efficiency and sufficiently exploring the input space to evolve harder memory examples, we choose the evolution step of 5.

Hyperparameter sensitivity. Due to space limitation, we put hyperparameter sensitivity analysis, including the regularizer weight $\beta$ and evolution rate $\alpha$, in Appendix C.

Computation cost. We compare the proposed ER + WGF to ER to evaluate its running efficiency. Table 5 shows the efficiency comparison results. We set the simple baseline ER with a running time unit of 1. Our method has 3-4 times the computational cost compared to a simple ER baseline. In future work, we will improve its computational efficiency.