On Tiny Episodic Memories in Continual Learning

In plain words

A model that learns tasks one after another tends to forget the earlier ones. This paper looks at a simple fix: keep a very small notebook of examples from past tasks, about one per class, and practise on a few of them every time the model learns something new. The model still sees each new example only once.

Several newer methods try to protect old knowledge with more complicated machinery. The authors find that this plain "mix in some old examples" recipe already does better than those methods in their tests. What surprised them is that training again and again on such a tiny notebook does not ruin performance on old tasks. Their explanation is that the data of the new task acts as a regulariser, which keeps the repeated practice from overfitting. How the notebook is filled matters once it is this small, and the paper compares several filling rules.

TL;DR

• Experience Replay (ER): each SGD step uses the current task's mini-batch plus a same-size mini-batch sampled from an episodic memory alg. 1.
• In single-pass experiments on four split benchmarks, ER with one example per class beats fine-tuning, EWC, A-GEM and MER on average accuracy tab. 3–tab. 6. These are author-reported numbers.
• When the memory is tiny, class-balanced writers (ring buffer, k-means, mean of features) beat reservoir sampling; reservoir is better for larger memories §5.4.
• Cost is close to fine-tuning: one extra memory batch per step tab. 2. The only public code covers the MNIST and CIFAR ER runs, with reservoir writing.

Problem

Continual learning (CL) asks a learner to acquire tasks from a stream while keeping earlier ones. Catastrophic forgetting makes this hard §1. The setting here is: each task is fully supervised, each example is seen once in the stream, and the learner may store only a small episodic memory of past examples §1, §3.1. Each example carries a task id that selects a task-specific classifier head §5.2.

Regularisation-based methods (such as EWC) store a measure of parameter importance, whose best-case memory cost is the number of network parameters §2. GEM and A-GEM use episodic memory as a constraint on the gradient §2, §4, which the paper finds costlier than ER §5.4. The paper argues that directly training on the memory is simpler and more efficient §4.

Method

Experience Replay starts from the fine-tuning baseline, which trains on the new task starting from the previous task's parameters. ER makes two changes §4, alg. 1:

1. The episodic memory is updated at every step, after the gradient update alg. 1, line 8.
2. The gradient batch is doubled: the current mini-batch is stacked with a mini-batch drawn at random from memory alg. 1, lines 6–7.

Memory is used as training data, not as a gradient constraint §4. Both halves of the stacked batch are used in the same step, and the paper takes the average of the two gradients in its comparison with A-GEM §4.

# Experience replay for continual learning  [alg. 1]
# Input:  task streams, memory size, batch size K, learning rate
# Output: parameters theta, episodic memory

def experience_replay(tasks, mem_size, batch_size, lr):
    # Start with an empty episodic memory  [alg. 1, line 2]
    memory = []
    # Count of stream examples seen so far  [alg. 1, line 3]
    num_seen = 0
    for task in tasks:
        # Mini-batches of the current task, each used once  [alg. 1, line 5]
        for batch in task.minibatches(batch_size):
            # Draw K examples from memory  [alg. 1, line 6]
            memory_batch = sample(memory, batch_size)
            # One SGD step on the stacked batches  [alg. 1, line 7]
            theta = sgd_step(theta, batch + memory_batch, lr)
            # Write the current batch into memory  [alg. 1, line 8; §4]
            memory = update_memory(memory, mem_size, task, num_seen, batch)
            num_seen = num_seen + batch_size
    return theta, memory

The paper compares four memory-writing rules. They all receive the current mini-batch and decide which examples to keep. The ring buffer, k-means and MoF also key their storage by task and class alg. 3–5.

Reservoir sampling §4, alg. 2. Each stream example is kept with probability mem_size / n, where n is the number of examples seen so far. The memory is filled first; after that, a random slot is overwritten when a random index falls inside the memory.

# Reservoir writing  [alg. 2; §4]
# An example is kept with probability mem_size / num_seen.
def update_reservoir(memory, mem_size, num_seen, batch):
    for example in batch:
        if len(memory) < mem_size:
            memory.append(example)
        else:
            # Uniform index over the examples seen so far  [alg. 2]
            index = random_int(0, num_seen)
            if index < mem_size:
                memory[index] = example
        num_seen = num_seen + 1
    return memory

Ring buffer §4, alg. 3. For each task, every class gets its own FIFO stack. The stacks have equal size mem_size / C, where C is the total number of classes across all tasks §4. Old entries fall out as new ones arrive, so samples from old tasks do not change once the stack has been filled.

# Ring buffer per task and class  [alg. 3; §4]
# Each stack holds mem_size / C examples, C = total number of classes [§4].
def update_ring(stacks, task, batch):
    for example, label in batch:
        # Append; a full FIFO drops its oldest example  [alg. 3]
        stacks[task][label].append(example)
    return stacks

k-means §4, alg. 4. For each class, online k-means is run on the features before the last classification layer. The memory keeps, for each centroid, the input example closest to that centroid.

# Online k-means per (task, class); keep the example nearest each centroid
# phi(x): features before the last classification layer  [§4]
def update_kmeans(slots, centroids, counts, task, batch, phi):
    for example, label in batch:
        feature = phi(example)
        # Nearest centroid j for this class  [alg. 4]
        j = nearest_centroid(feature, centroids[task][label])
        counts[task][label][j] += 1
        # Move the centroid toward the feature by 1 / count  [alg. 4]
        centroid = centroids[task][label][j]
        centroid = centroid + (feature - centroid) / counts[task][label][j]
        centroids[task][label][j] = centroid
        # Keep the example if closer than the stored one in slot j  [alg. 4]
        d = norm(feature - centroid)
        if d < slots[task][label][j].distance:
            slots[task][label][j] = (example, d)
    return slots

Mean of features (MoF) §4, alg. 5. For each class, a running mean of the features is kept. The memory holds the examples closest to that mean. It uses the same balancing guarantee as the ring buffer and k-means §4.

# Mean of features per (task, class); keep the nearest examples  [alg. 5; §4]
def update_mof(heaps, means, task, batch, phi, alpha):
    for example, label in batch:
        feature = phi(example)
        # Exponential moving average; alpha is not stated  [alg. 5]
        mean = alpha * means[task][label] + (1 - alpha) * feature
        means[task][label] = mean
        d = norm(feature - mean)
        heap = heaps[task][label]
        # Max-heap by distance: evict the farthest stored example  [alg. 5]
        if heap.find_max() > d:
            heap.delete_max()
            heap.insert(example, d)
    return heaps

Hybrid writer §5.4, fig. 2. The paper proposes a hybrid that starts with reservoir sampling and switches to the ring-buffer rule once any class seen so far has only one example left in memory. At the switch, it keeps min(n, |M| / K) randomly chosen examples per class, where n is the number of stored examples of that class and K is the number of classes observed so far §5.4. The switch also lets over-represented classes shrink as new classes arrive §5.4.

Analysis of why repeated training does not overfit §5.5, fig. 3. The authors use two tasks from MNIST Rotations with 10 memory slots, one per class of task 1. They report that training on the memory alone overfits, that training on the new task alone can help or hurt depending on the relatedness of the two tasks, and that training on both gives better accuracy on task 1 than the baseline in every case they tried. They read the new-task data as a data-dependent regulariser on the old task §5.5.

Hyperparameters

Values are from the paper unless marked as code-only. "Not stated" means the paper does not give the value.

Results

Author-reported. Setup: four split benchmarks, single pass per task, average accuracy after the last evaluation task §3.2, §5.1. The numbers below are the paper's own tables.

Metrics §3.2. Let a(i, j) be accuracy on test task j after training on task i. Average accuracy at task T and forgetting are:

A_{T} = \frac{1}{T} \sum_{j=1}^{T} a_{T, j} eq. 1

f_j^i = \max_{l \in \{1,\dots,i-1\}} a_{l, j} - a_{i, j} eq. 2

F_T = \frac{1}{T-1} \sum_{j=1}^{T-1} f_j^T eq. 3

Average accuracy (%) with one example per class tab. 3–tab. 6. The ± values are as printed; the paper does not define them in the captions.

Forgetting at one example per class tab. 1. Lower is better.

Learning time on the evaluation stream, seconds tab. 2. The paper uses ER with the ring buffer as the representative ER variant. The text gives no hardware.

Trends the paper reports §5.4, fig. 1, fig. 2:

• The ER variants outperform fine-tuning, EWC, A-GEM and MER, including at the smallest memory §5.4, fig. 1.
• Average accuracy increases with memory size and does not saturate at 13 examples per class §5.4.
• Reservoir is best for larger memories, and the class-balanced writers are best when memory is tiny §5.4, tab. 4.
• The hybrid writer is at least as good as the others across the whole stream in the one experiment shown (Split CIFAR, 85 slots, 3 runs). Its numbers appear only in a figure, not in a table fig. 2.
• ER with the memory has near-perfect accuracy on stored examples; A-GEM does not reach that even after 1000 samples §5.5, tab. 7.

Paper vs code

Official code: github.com/facebookresearch/agem, linked from the paper's abstract footnote. Checked at commit 45421499483b28935491251e9e821c55e8b3c089. The repository's LICENSE is the MIT licence. The repository holds the MNIST, CIFAR, CUB and AWA training scripts; it has no miniImageNet script. Checked: conv_split_cifar.py, fc_permute_mnist.py, conv_split_cub.py, model/model.py, utils/utils.py and replicate_results_er.sh. The list below compares the paper with the code.

• Mini-batch sizes differ between defaults and the paper. The paper uses 10 for both halves §5.2. The CIFAR script's defaults are 16 and 256 conv_split_cifar.py. The ER replicate script passes 10 and 10 replicate_results_er.sh, which matches the paper.
• Memory budget differs between MNIST and CIFAR. The CIFAR script sets the memory to mem_size × 100 slots, with 100 classes in total, regardless of how many tasks have been seen conv_split_cifar.py. The MNIST script sets mem_size × num_tasks × 10 slots, so its budget grows with the number of tasks fc_permute_mnist.py. The paper's wording is "per class" §4, §5.4.
• ER writes only with reservoir sampling in the released code. ER's write step calls the reservoir update [code: conv_split_cifar.py#L483-L486, fc_permute_mnist.py#L352-L354]. The ring-buffer writer appears only in the A-GEM branch conv_split_cifar.py. No k-means code was found in the training scripts or utils. MoF (is_herding) appears only in the CUB and AWA scripts, not in the ER path conv_split_cub.py. The hybrid scripts have a --set-hybrid flag that I did not trace; it is not on the ER path. The ER-ring, ER-k-means, ER-MoF and hybrid results cannot be reproduced from this repository's ER code.
• No ER code for CUB or miniImageNet. The CUB script's model list has no ER entry conv_split_cub.py. The repository has no miniImageNet script.
• Learning rates. The MNIST script's default learning rate is 1e-3 fc_permute_mnist.py. The replicate script passes 0.1, which matches Table 8 for Permuted MNIST and CIFAR tab. 8 replicate_results_er.sh.
• Permuted MNIST task count differs. The paper says 23 tasks in total §5.1. The MNIST script sets 20 tasks in total and uses 20 − 3 = 17 for evaluation fc_permute_mnist.py.
• Training examples per MNIST task are not stated in the paper. The script defaults to 1000 fc_permute_mnist.py, and the replicate script passes 1000 replicate_results_er.sh.
• Loss over stacked batches. The ER loss sums each example's cross-entropy over the classes of its own task, then divides by the size of the stacked batch model.py. The paper's "average of the two gradients" matches this only when the two halves are the same size, as they are here (10 and 10).
• Single pass means one pass over the stream, not one pass over the memory. The script runs ceil(n / batch) steps per task conv_split_cifar.py. Memory examples are re-used across many steps.
• Optimiser. Plain SGD with a constant learning rate in the ER runs conv_split_cifar.py. This matches the paper's SGD §5.2.
• Task heads. The paper describes a task-specific head §5.2. The code uses one shared output layer with per-task logit masks in training conv_split_cifar.py. Test-time masking is also used conv_split_cifar.py.
• Architecture details. The paper says "reduced ResNet18" §5.2. The code's RESNET-S uses 3×3 kernels with widths 20, 20, 40, 80, 160 model.py. The paper does not state these widths.
• No code for the analyses. No script for the MNIST Rotations analysis was found (grep for "rotation" in the Python files finds nothing).

No other discrepancies found.

Pitfalls

• Stacked batch when memory is short. The paper does not say what to do at the very start, when memory is empty. The code uses only the examples stored so far, so the first step uses only the current batch conv_split_cifar.py.
• Memory budget semantics. "One example per class" can mean one per class in total (CIFAR code), or one per class per task (MNIST code, which grows with the number of tasks). Pick the convention the reported numbers use, and check it when comparing memory sizes across datasets.
• Reservoir index convention. The paper's pseudocode draws a random index from a range whose bounds are not clearly stated. The code draws from an exclusive upper bound equal to the number of examples seen before the current one, so an example replaces a slot with probability mem_size / N utils.py. Implement that rule, or state your own.
• k-means details are missing. The paper does not state the number of centroids (presumably one slot per centroid per class), how centroids are initialised, or whether the stored distance is recomputed. The algorithm compares a new distance computed with the current features against a distance stored earlier alg. 4. The features change during training, so these two distances come from different model states. This is a choice the implementer must make.
• MoF details are missing. The decay α is not stated alg. 5. The heap size per class is also not given in the text; presumably it is the per-class slot count.
• Ring buffer needs the total class count. The stack size is mem_size / C, with C the total number of classes across all tasks §4. An implementer must know C in advance.
• Hybrid numbers are in a figure. The hybrid writer's switch rule is described in prose §5.4. Its results appear only in fig. 2, with 3 runs.
• Gain claim in the abstract has no named baseline. The abstract says the gain at one example per class is 7–17%, but does not name the baseline. From the one-example table tab. 3–tab. 6, ER with the ring buffer beats fine-tuning by 16.7 points on Permuted MNIST and EWC by 7.1 points on the same dataset. These are computed here, not stated by the paper.
• Gain numbers in the text match a specific variant. The text says ER beats A-GEM by 1.7 points on CIFAR at one example per class §5.4. That matches ER with k-means (56.6 vs 54.9), not ER with the ring buffer (56.2 vs 54.9).
• Memory-size limit text versus table. The text says the self-imposed limit is 13 examples per class §5.4, but the Permuted MNIST table goes to 15 tab. 3.
• Permuted MNIST task count. The text says 23 tasks §5.1. The code runs 20 in total fc_permute_mnist.py. Check which count your comparison needs.
• Units in the rotation table. In the MNIST Rotations table tab. 7, the ER memory column reads "1", while the A-GEM column reads percentages (for example 73.3). The text reads the ER entry as perfect accuracy on memory. The paper does not state the unit, so treat the ER entry as a fraction (1.0 = 100%) until the source confirms.
• ± definition. The table captions do not say what ± is. Figure 1 reports averages over 5 runs. The run count for the tables is not stated.
• Hardware for timings. The learning-time table tab. 2 gives no hardware. Compare ratios within one table, not absolute times across papers.
• Cross-validation. Hyperparameters are chosen on the 3-task cross-validation stream and then applied to the evaluation stream §5.1. In the code's evaluation mode, train and validation splits are concatenated for training conv_split_cifar.py. The paper does not state this.
• Pretrained CUB. CUB uses an ImageNet-pretrained ResNet18 §5.2, so the CUB comparison starts from pretrained weights. The other three benchmarks do not.

Lineage

• Builds on: A-GEM, which this paper uses as a main baseline and whose memory-as-constraint approach it replaces with direct training on memory §4. GEM is the original memory-as-constraint method §2.
• Compared against: fine-tuning, EWC, A-GEM and MER §5.3.
• Writing rules borrowed from earlier work: reservoir sampling (Vitter, cited in the paper §4), the ring buffer as in GEM §4, and MoF as in iCaRL §4.
• Superseded by: not stated in this wiki.