Efficient Data-Parallel Continual Learning with Asynchronous Distributed Rehearsal Buffers

TL;DR

This work tackles catastrophic forgetting in class-incremental continual learning under streaming data by introducing a distributed rehearsal buffer that integrates with data-parallel training on HPC hardware. It distributes per-GPU buffers across processes, uses asynchronous management and RDMA-enabled distributed sampling to create diverse, augmented mini-batches of size $b+r$, and overlaps buffer operations with training to minimize overhead. Empirical results on 128 GPUs across multiple models and a four-task ImageNet-derived sequence show that rehearsal achieves top-5 accuracy close to the train-from-scratch upper bound (e.g., 80.55% vs 91%) while maintaining runtime near the incremental lower bound, demonstrating strong scalability and practicality. The approach offers a scalable, efficient pathway for continual learning in environments where data arrive continuously and retraining from scratch is infeasible.

Abstract

Deep learning has emerged as a powerful method for extracting valuable information from large volumes of data. However, when new training data arrives continuously (i.e., is not fully available from the beginning), incremental training suffers from catastrophic forgetting (i.e., new patterns are reinforced at the expense of previously acquired knowledge). Training from scratch each time new training data becomes available would result in extremely long training times and massive data accumulation. Rehearsal-based continual learning has shown promise for addressing the catastrophic forgetting challenge, but research to date has not addressed performance and scalability. To fill this gap, we propose an approach based on a distributed rehearsal buffer that efficiently complements data-parallel training on multiple GPUs, allowing us to achieve short runtime and scalability while retaining high accuracy. It leverages a set of buffers (local to each GPU) and uses several asynchronous techniques for updating these local buffers in an embarrassingly parallel fashion, all while handling the communication overheads necessary to augment input mini-batches (groups of training samples fed to the model) using unbiased, global sampling. In this paper we explore the benefits of this approach for classification models. We run extensive experiments on up to 128 GPUs of the ThetaGPU supercomputer to compare our approach with baselines representative of training-from-scratch (the upper bound in terms of accuracy) and incremental training (the lower bound). Results show that rehearsal-based continual learning achieves a top-5 classification accuracy close to the upper bound, while simultaneously exhibiting a runtime close to the lower bound.

Figures (7)

• Figure 1: For every process $n$, a rehearsal buffer $\mathcal{B}_n$ contains representatives from the classes seen so far. The distributed rehearsal buffer $\mathcal{B}$ contains representatives from the $K$ classes.
• Figure 2: For a given process $n$, $c$ candidates from the incoming mini-batch are sampled and used to populate $\mathcal{B}_n$. If the buffer for class $i$ is full, representatives from $R^i_n$ are replaced at random. The figure depicts the rehearsal buffer $\mathcal{B}_n$ state for two subsequent iterations for $c = 2$.
• Figure 3: On a given process $n$, every incoming mini-batch is augmented with $r$ representatives sampled randomly and without replacement from the distributed rehearsal buffer $\mathcal{B}$. Here, $r = 2$ on two subsequent iterations. Sampling from $\mathcal{B}$ introduces communication between the $N$ distributed processes.
• Figure 4: Asynchronous updates of the rehearsal buffers and global augmentations: $r$ representatives sampled globally beginning with the previous iteration are used by the training loop to assemble an augmented mini-batch on each process $n$. Meanwhile, the distributed rehearsal buffer extracts candidates from the current mini-batch to update each $B_n$ locally, then collects the next $r$ representatives using global sampling.
• Figure 5: Top-5 accuracy for ResNet-50, 16 GPUs, ImageNet (4 tasks).