Continual Learning via Manifold Expansion Replay

In the replay strategy, the number of samples collected for replay is limited. This poses a challenge in recalling the knowledge of the entire task from these samples. If a classification loss is used when training replay samples, the model can only learn to classify these samples rather than previous tasks. As a result, the model may still suffer from catastrophic forgetting as the number of tasks increases. For a specific task, we assume that each sample represents a piece of meta-knowledge. The model learns the entire knowledge manifold by learning from these samples. We now consider the first problem encountered: how to integrate this meta-knowledge when learning new tasks to form a more comprehensive knowledge manifold. To do this, we need to measure the distance between two knowledge manifolds. Intuitively, the distance between the new and old knowledge manifolds should be small because previously acquired meta-knowledge does not change. If we view the knowledge manifold as a distribution of meta-knowledge, one possible choice for a distance metric is KL divergence, which is commonly used to measure the distance between two distributions. However, KL divergence cannot be used as a strict distance function because it is asymmetric and cannot provide distance information when two distributions do not overlap. Our method, MaER, uses Wasserstein distance to measure the distance between two knowledge manifolds.

In mathematics, the Wasserstein distance is a distance function defined between probability distributions on a given metric space $(M,\rho)$, where $\rho(x,y)$ is a distance function for two instances $x$ and $y$ in the set $M$. Formally, the $p$-th Wasserstein distance between two probability measures $\mu$ and $\nu$ on $M$ with $p$-moment is defined as:

where $\Gamma(\mu,\nu)$ represents the set of all coupling of $\mu$ and $\nu$. Wasserstein distance has some good properties. In contrast to KL divergence, Wasserstein distance is a symmetric metric and provides information even if the distributions do not overlap.

Now we describe how MaER uses the Wasserstein distance to facilitate meta-knowledge fusion in continual learning. When learning task $i$, we have a teacher model $f_{t}$ that has learned the previous $i-1$ tasks and a student model $f_{s}$ responsible for learning task $i$. For a sample $x$, we use $\Phi_{s}(x;\theta)$ to denote the knowledge representation of $f_{s}$ for that sample, where $\Phi_{s}$ denotes the feature extractor for $f_{s}$. Our objective is for the student model $f_{s}$ to effectively learn task $i$. To achieve this, we train $f_{s}$ on samples from task $i$ using the cross-entropy loss function for classification. Concurrently, it is imperative that $f_{s}$ retains the knowledge acquired from previous tasks. During training, we replay samples from the memory buffer $\mathcal{M}$, and in addition to learning to classify these samples accurately, we aim to minimize the Wasserstein distance between the knowledge representation learned by $f_{s}$ and that of the teacher model $f_{t}$. Taking these considerations into account, the new loss function is defined as:

where $\ell$ is the cross-entropy function in classification tasks. For ease of computation, we use the Wasserstein distance with p = 2 and assume that meta-knowledge is equally important. In MaER, the calculation of $W_{2}$ can be simplified to the following form:

where $d$ is the distance function in metric space $M$, which can be Euclidean distance or geodesic distance.

We now turn to the second problem: how to sample such that a small number of meta-knowledge can represent the knowledge manifold of the entire task as much as possible. We assume that the meta-knowledge is uniformly distributed over this knowledge manifold. We can only sample a small fraction of meta-knowledge to represent the whole manifold. Intuitively, our sampling method needs to account for two aspects: (1) sample as uniformly as possible to maintain the geometric properties, (2) the sampled meta-knowledge should span the entire knowledge manifold as much as possible, avoiding bias towards new knowledge and preventing forgetting during learning. To address this problem, MaER employs a greedy strategy that incrementally enlarges the diameter of the manifold during the sampling process. First, the sampling process is stochastic. When a new sample arrives, if $\mathcal{M}$ is already full, then each sample in $\mathcal{M}$ has an equal chance of being replaced. Stochastic sampling preserves the consistency of the data distribution in $\mathcal{M}$ and the original distribution. Second, MaER’s sampling strategy guarantees that samples that can augment the diameter of the manifold are always collected. For other samples, MaER will collect them with a certain probability. The criterion for determining whether a sample can augment the diameter of the manifold is essential for understanding MaER’s sampling strategy. We begin by introducing the concepts of the centroid and diameter of a manifold. The Fréchet mean is a natural generalization of the centroid and can be applied to any manifold. For a metric space $\mathcal{X}=(X,d,\mu)$, its Fréchet mean is defined as:

With the manifold centroid $\mathcal{C}$, we can estimate the diameter of the manifold in a simple and efficient way. We define the diameter as the largest distance from the centroid $\mathcal{C}$ to the sample $x$. Having defined these concepts, we can describe the sampling process for MaER, as shown in Algorithm 1.

Here, we begin by describing the continual learning benchmarks, implementation details, and compared methods.

Benchmarks. We conducted experiments on several continual learning datasets: Permuted MNIST, Rotated MNIST, Split CIFAR10, Split CIFAR100 and Split TinyImageNet. Permuted MNIST is derived from applying a random permutation to the pixels. Rotated MNIST is derived from rotating the image at a random angle. Split CIFAR10, Split CIFAR100, and Split TinyImageNet are derived from splitting CIFAR10, CIFAR100, and miniImageNet, respectively, such that the classes in the different tasks are disjoint. Split CIFAR10 consists of 5 tasks, while the other datasets consist of 20 tasks each.

Implementation details. In the network architecture, we utilized a three-layer MLP with 256 neurons in the hidden layer for MNIST and a standard resnet18 for CIFAR10, CIFAR100 and TinyImageNet. We optimize the parameters during training using SGD. The learning rate is set to 0.01 for MNIST and 0.003 for others. The batch size is set to 16 for all experiments. We trained the model for 10 epochs on MNIST, 5 epochs on Split CIFAR100 and 20 epochs on others.

Compared methods. We compare MaER with several baseline methods. The rehearsal-based baselines include iCaRL [14], ER [30], ER-RING [31], ER-RESERVOIR [31], A-GEM [32], HAL [33], DER [34], and SNCL [35], with DER and SNCL being the strongest baselines. We also compare with other strategy-based methods, including SI [10], EWC [11], oEWC [36], LwF [19], and PNN [12]. Additionally, we compare with two non-continual learning methods, JOINT and SGD, as upper and lower bounds.

Table I shows the results of our method using a tiny buffer on Split CIFAR10, Split TinyImagNet, Permuted MNIST, and Rotated MNIST. On Split CIFAR10 and Rotated MNIST, our method achieved competitive results, differing from state-of-the-art replay methods by $0.5\%\sim 0.73\%$. On Permuted MNIST and Split TinyImageNet, our method achieves state-of-the-art performance, significantly outperforming baseline methods. On Split TinyImageNet, MaER surpassed SNCL by $1.80\%\sim 3.33\%$ in average accuracy when taking different buffer sizes. On Permuted P-MNIST, this gap was even more pronounced, with our method leading SNCL by $3.81\%$ in average accuracy when taking different buffer sizes. On Split CIFAR10 and Rotated MNIST, most baselines achieved considerable results. However, on Permuted MNIST and Split TinyImageNet, the gap between different methods became more pronounced. These results indicate that Permuted MNIST and Split TinyImageNet are more challenging when using a tiny buffer. MaER can work well in challenging scenarios and achieves state-of-the-art in average accuracy. On Split CIFAR100, we also compared MaER with more replay strategies. In Table II, our method surpassed all baselines when taking different buffer sizes, leading the baseline by $0.9\%\sim 7.64\%$ in average accuracy. As the buffer size increased, MaER was able to perform correspondingly better. Table III shows the ablation study results of different components of MaER. Here, $\mathcal{L}_{WD}$. represents the Wasserstein distance loss. Compared to the naive replay, both the replay strategy and memory management in MaER achieve certain performance improvements.