Multiband VAE: Latent Space Alignment for Knowledge Consolidation in Continual Learning

1 Introduction

Recent advances in generative models [Goodfellow et al., 2014; Kingma and Welling, 2014] led to their unprecedented proliferation across many real-life applications. This includes high energy physics experiments at Large Hadron Collider (LHC) at CERN, where they are employed to speed up the process of particles collisions simulations [Paganini et al., 2018; Deja et al., 2020; Kansal et al., 2021].

Those applications are possible, thanks to the main objective of generative methods, which is the modelling of complex data manifolds with simpler distributions. Unfortunately, this goal remains difficult to deliver in real-life situations where training data is presented to the model in separate portions, e.g., from consecutive periods of data gathering at CERN. The distributions of data within these portions often vary significantly, hence updating model with new examples leads to catastrophic forgetting of previous knowledge. In generative modelling this is observed through limited distribution of generated examples.

Generative continual learning methods aim to address these challenges usually in one of three ways: through regularization (e.g. [Nguyen et al., 2018]), adjustment of the structure of a network to the next task (e.g. [Rao et al., 2019]), or rehearsal of previously seen samples when training with new data (e.g. [Rebuffi et al., 2017]). Nevertheless, practical applications of those methods are yet limited, so most of them focus on the artificial class-incremental (CI) training scenario. In this approach, consecutive portions of data (tasks) contain disjoint classes and share almost no similarity. While this is the most difficult scenario for discriminative models, we argue that the assumption of classes separation greatly simplifies the problem in generative modelling where task index might be used without reducing the method's generality (detailed discussion in the appendix).

Moreover, the assumption of task independence in CI scenario reduces the complexity of continual learning [Ke et al., 2021]. Therefore, in this work, we postulate to investigate the adaptation of generative continual learning methods to the ever-changing data distribution. While, for the CI scenario, we expect no forgetting of previous knowledge, in other scenarios, where model is retrained with additional partially similar data, we should aim for performance improvement. This can be observed through forward knowledge transfer - higher performance on a new task, thanks to already incorporated knowledge, and backward knowledge transfer - better generations from previous tasks, when retrained on additional similar examples [Lopez-Paz and Ranzato, 2017].

Therefore, to simulate real-life conditions, we prepare a set of diversified continual learning scenarios with data splits following Dirichlet distribution, inspired by a similar approach in federated learning [Hsu et al., 2019]. Our experiments indicate that this is indeed a more challenging setup for the majority of recent state-of-the-art continual generative models, which lack sufficient knowledge sharing between tasks.

To mitigate this problem, we propose a Multiband VAE. The core idea behind our method is to split the process of model retraining into two steps: (1) a local encoding of data from the new task into a new model's latent space and (2) a global rearrangement and consolidation of new and previous data. In particular, we propose to align local data representations from consecutive tasks through the additional neural network. In reference to the way how radio spectrum frequencies are allocated, we name data representations from different tasks bands. As in telecommunication, our goal is to limit interference between bands. However, we train our model to align parts that represent the same or similar data. To support knowledge consolidation between different bands, we additionally propose a controlled forgetting mechanism that enables the substitution of degraded reconstructions of past samples with new data from the current task.

The main contributions of this work are:

• A novel method for generative continual learning of Variational Autoencoder that counteracts catastrophic forgetting while being able to align even partially similar tasks at the same time.
• A simple method for controlled forgetting of past examples whenever a new similar data is presented.
• A novel knowledge consolidation training scenario that underlines limitations of recent state-of-the-art methods.

2 RelatedWorks

Most of the works incorporating generative models in continual learning relate to generative rehearsal. In this technique, the base model is trained with a mixture of new data examples from the current task and recreation of previous samples generated by a generative model. This idea was first introduced by [Shin et al., 2017], with Generative Adversarial Networks (GAN) trained with the self rehearsal method so-called Generative Replay (GR). [Lesort et al., 2019] overview different generative models trained with the GR method. Our Multiband VAE is a direct extension to this technique.

Continual learning of generative models [Nguyen et al., 2018] adapt regularization-based methods such as Elastic Weight Consolidation (EWC) [Kirkpatrick et al., 2017], and Synaptic Intelligence (SI) [Zenke et al., 2017] to the continual learning in generative models regularizing the adjustment of the most significant weights. The authors also introduce Variational Continual Learning (VCL), with adjustments in parts of the model architecture for each task.

In HyperCL, [von Oswald et al., 2019] propose entirely different approach, where a hypernetwork generates the weights of the continually trained model. This yields stateof-the-art results in discriminative models task-incremental training but is also applicable to the generative models. In order to differentiate tasks, [Rao et al., 2019] propose CURL that learns task-specific representation and deals with task ambiguity by performing task inference within the generative model. This approach directly addresses the problem of forgetting by maintaining a buffer for original instances of poorly-approximated samples and expanding the model with a new component whenever the buffer is filled. In BooVae, [Egorov et al., 2021] propose an approach for continual learning of VAE with an additive aggregated posterior expansion. Several works train GANs in the continual learning scenarios either with memory replay [Wu et al., 2018], with the extension to VAEGAN in Lifelong-VAEGAN by [Ye and Bors, 2020].

Continual learning with disentanglement In VASE by [Achille et al., 2018], authors propose a method for continual learning of shared disentangled data representation. While encoding images with a standard VAE, VASE also seeks shared generative factors. A similar concept of mixed-type latent space was introduced in LifelongVAE [Ramapuram et al., 2020], where it is composed of discrete and continuous values. In this work we also use a disentanglement method with binary latent space.

3 Method

In this section, we introduce Multiband VAE - a method for consolidating knowledge in a continually learned generative model. We propose to split generative replay training into two parts: (1) a local training that allows us to build a new data representations band in the latent space of VAE, and (2) global training where we attach a newly trained band to the already trained global model. As a part of the global training, we propose a controlled forgetting mechanism where we replace selected reconstructions from previous tasks with currently available data.

3.1 Knowledge Acquisition - Local Training

In the local training, we learn a new data representations band by training a VAE using only currently available data.

Let xij denote the j-th sample of i-th task. Then, for given sample xij , and latent variable λij we use a decoder pθ, which is trained to maximize posterior probability p(xij |λij ). To get the latent variable λij , we use encoder qϕ parametrized with weights vector ϕ that approximates probability q(λij |xij ).

To simplify the notation, let us focus on specific task i and drop the index. As in standard VAE, we follow optimization introduced by [Kingma and Welling, 2014] that maximizes the variational lower bound of log likelihood:

where θ and ϕ are weights of encoder and decoder respectively. In the first task, this is the only part of the training, after which local decoder is remembered as a global one. In other cases we drop local decoder.

3.2 Shared Knowledge Consolidation

In the second - global part of the training, we align the newly trained band with already encoded knowledge. The simplest method to circumvent interference between bands is to partition the latent space of VAE and place new data representation in a separate area of latent space. However, such an approach limits information sharing across separate tasks and hinders forward and backward knowledge transfer. Therefore, in Multiband VAE we propose to align different latent spaces through an additional neural network that we call translator. Translator maps individual latent spaces which are conditioned with task id into the common global one where examples are stored independently of their source task, as presented in Fig 2.

To that end, we define a translator network tρ(λi , i) that learns a common alignment of separate latent spaces λi conditioned with task id i to a single latent variable Z , where all examples are represented independently of their source task. Finally, we propose a global decoder pω(x|Z) that based on distribution approximated with latent variables Z learns to approximate original data distribution x.

To counteract forgetting, when training translator and global decoder we use auto-rehearsal as in standard generative replay, with a copy of the translator and decoder frozen at the beginning of the task. As training pairs, we use combination of original images x with their encodings from local encoder λ, and for previous tasks, random values λ with generations x reconstructed with a frozen translator and global decoder. Fig. 3 presents the overview of this procedure.

We start translator training with a frozen global decoder, to find the best fitting part of latent space Z for a new band of data without disturbing previous generations. For that end we minimize the reconstruction loss:

where k is the number of all tasks.

Then, we optimize parameters of translator and global decoder jointly, minimizing the reconstruction error between outputs from the global decoder and training examples

To generate new example t with Multiband VAE, we randomly sample task id i ∼ U({1, . . . , k}), where k is the number of all tasks and latent representation λt ∼ N(⃗0, I). These values are mapped with translator network to latent variable zt, which is the input to global decoder to generate xt. Therefore, translator and global decoder are the only models that are stored in-between tasks.

3.3 Controlled Forgetting

In a real-life scenario, it is common to encounter similar data examples in many tasks. In such a case, we would like our continuously trained model to refresh the memory of examples instead of combining vague, distorted memories with new instances. Therefore, we propose a mechanism for controlled forgetting of past reconstructions during the translator and global decoder joint training. To that end, when creating new training pairs, we compare representations of previous data reconstructions generated as new targets with representations of data samples from the current task in the common latent space Z. If these representations are similar enough, we substitute previous data reconstruction with the current data sample as presented in Fig. 4.

More specifically, when training on task i, we first create a subset Z i = tρ(qϕ(xi ), i) with representations of all currently available data in joint latent space Z. Now, for each data sample xlj generated as a rehearsal target from previous task l < i and random variable λlj , we compare its latent representation zj = tρ(λlj , j) with all elements of set Z i

If sim(zj) ≥ γ we substitute target sampled reconstruction xlj with respective original image fromxi . Intuitively, γ controls how much do we want to forget from task to task, with γ = 0.9 being a default value for which we observe a stable performance across all benchmarks.

4 Experiments

To visualize the difference between Generative Replay and Multiband VAE, in Fig. 5 we present a toy-example with the MNIST dataset limited to 3 tasks with data examples from 3 classes. When presented with data from a new distribution (different class in task 2), our method places a new band of data in a separate part of a common latent space Z. On the other hand, the standard generative replay model learns to transform some of the previous data examples into currently available samples before it can distinguish them, even with additional conditioning on task identity. At the same time, when presented data with partially same classes as in task 3, our translator is able to properly align bands of data representations so that similar data examples (in this case ones) are located in the same area of latent space Z independently of the source task, without interfering with zeros and twos.

4.1 Evaluation Setup

For fair comparison, in all evaluated methods we use a Variational Autoencoder architecture similar to the one introduced by [Nguyen et al., 2018], with nine dense layers. However, our Multiband VAE is not restricted to any particular architecture, so we also include experiments with a convolutional version. The exact architecture and training hyperparameters are enlisted in the appendix and code repository 1 . We do not condition our generative model with class identity since it greatly simplifies the problem of knowledge consolidation and applies to all evaluated methods. However, similarly to [Ramapuram et al., 2020], we use additional binary latent space trained with Gumbel softmax [Jang et al., 2016].

1 https://github.com/KamilDeja/multiband vae

4.2 Evaluation

To assess the quality of our method, we conduct a series of experiments on benchmarks commonly used in continual learning (MNIST, Omniglot [Lake et al., 2015]) and generative modeling - FashionMNIST [Xiao et al., 2017]. Since the performance of VAE on diverse datasets like CIFAR is limited, in order to evaluate how our method scales to more complex data, we include tests on CelebA [Liu et al., 2015]. For each dataset, we prepare a set of training scenarios designed to evaluate various aspects of continual learning. This is the only time we access data classes, since our solution is fully unsupervised.

To assess whether the model suffers from catastrophic forgetting, we run class incremental scenarios introduced by [Van de Ven and Tolias, 2019]. However, CI simplifies the problem of learning data distribution in the generative model's latent space since the identity of the task conditions final generations. Therefore, we also introduce more complex data splits with no assumption of independent task distributions. To that end, we split examples from the same classes into tasks, according to the probability q ∼ Dir(αp) sampled from the Dirichlet distribution, where p is a prior class distribution over all classes, and α is a concentration parameter that controls similarity of the tasks, as presented in Fig. 6. In particular, we exploit the Dirichlet α = 1 scenario, where the model has to learn the differences between tasks while consolidating representations for already known classes. In such a scenario we expect forward and backward knowledge transfer between tasks.

To measure the quality of generations from different methods, we use the Fr´echet Inception Distance (FID) [Heusel et al., 2017]. As proposed by [Bi´nkowski et al., 2018], for simpler datasets, we calculate FID based on the LeNet classifier pre-trained on the whole target dataset. Additionally, we report the precision and recall of the distributions as proposed by [Sajjadi et al., 2018]. As authors indicate, those metrics disentangle FID score into two aspects: the quality of generated results (Precision) and their diversity (Recall).

For each experiment, we report the FID, Precision, and Recall averaged over the final scores for each task separately. For methods that do not condition generations on the task index (CuRL and LifelongVAE), we calculate measures in comparison to the whole test set. The results of our experiments are presented in Tab. 1 and Tab. 2, where we show scores averaged over three runs with different random seeds.

To compare different continual-learning generative methods in a real-life scenario we also use real data from detector responses in the LHC experiment. Calorimeter response simulation is one of the most profound applications of generative models where those techniques are already employed in practice [Paganini et al., 2018]. In our studies, we use a dataset of real simulations from Zero Degree Calorimeter in the ALICE experiment at CERN introduced by [Deja et al., 2020], where a model is to learn outputs of 44×44 resolution energy depositions in calorimeter. Following [Deja et al., 2020], instead of using FID, for evaluation, we benefit from the nature of the data and compare the distribution of real and generated channels - the sum of selected pixels that well describe the physical properties of simulated output. We report the Wasserstein distance between original and generated channels distribution to measure generations' quality. We prepare a continual learning scenario for this dataset by splitting examples according to their input energy, simulating changing conditions in the collider. In practice, such split lead to continuous change in output shapes with partial overlapping between tasks - similarly to what we can observe with Dirichlet based splits on standard benchmarks (see appendix for more details and visualisations).

As presented in Tab. 1, our model outperforms comparable methods in terms of quality of generated samples. Results of comparison on the Omniglot dataset with 20 splits (Tab. 2) indicate that for almost all of related methods, training with the data splits according to the Dirichlet α = 1 distribution poses a greater challenge than the class incremental scenario. However, our Multiband VAE can precisely consolidate knowledge from such complex setups, while still preventing forgetting in CI scenario. This is only comparable to CURL that achieves this goal through additional model expansion. Experiments on more complex joint datasets, where examples are introduced from one dataset after another, indicate the superiority of Multiband VAE over similar approaches. In the real-life CERN scenario, our model also clearly outperforms other solutions. In Fig. 7 we present how generations quality for this dataset changes in standard generative replay and Multiband VAE, showing both forward and backward knowledge transfer in Multiband VAE.

Finally, we evaluate our model with a more complex dataset - CelebA with over 200 000 images of celebrity faces in 64x64 resolution. Based on annotated features, we split the dataset into 10 classes based on the hair color/cover (blonde, black, brown, hat, bald or gray). In Tab. 3 we show the results of experiments with this dataset split in class incremental and Dirichlet scenarios. For class incremental scenario, Multiband VAE learns to separate bands of examples from different tasks with disjoint distributions, while results improve if in training scenario model is presented with more similar examples. In the latter case, with Dirichlet α = 100 splits, our model reaches the quality of the upper bound, which is a standard Variational Autoencoder trained with full access to all examples in the stationary training.

Ablation study The main contribution of this work is a multiband training procedure, yet we also introduce several mechanisms that improve knowledge consolidation. Tab. 4

shows how those components contribute to the final score.

4.3 Memory Requirements and Complexity

The memory requirements of Multiband VAE are constant and equal to the size of the VAE with an additional translator, which is a small neural model with 2 fully connected layers. When training on the new task, our method requires additional temporary memory for the local model freed when finished. This is contrary to similar methods (HyperCL, VCL, CURL) which have additional constant or growing memory requirements. Computational complexity of our method is the same as for methods based on generative rehearsal (VCL, LifelongVAE, Lifelong-VAEGAN). In experiments, we use the same number of epochs for all methods, while for Multiband VAE we split them between local and global training.

5 Conclusion

In this work, we propose a new method for unsupervised continual learning of generative models. We observe that the currently employed class-incremental scenario simplifies the continual learning of generative models. Therefore, we propose a novel, more realistic scenario, with which we experimentally highlight the limitations of state-of-the-art methods. Finally, we introduce a new method for continual learning of generative models based on the constant consolidation of VAE's latent space. To our knowledge, this is the first work that experimentally shows that with continually growing data with even partially similar distribution, we can observe both forward and backward performance improvement. Our experiments on various benchmarks and with real-life data show the superiority of Multiband VAE over related methods, with upper-bound performance in some training scenarios.

Acknowledgments

This research was funded by Foundation for Polish Science (grant no POIR.04.04.00-00-14DE/18-00 carried out within the Team-Net program co-financed by the European Union under the European Regional Development Fund) and National Science Centre, Poland (grant no 2018/31/N/ST6/02374 and 2020/39/B/ST6/01511).