Watching a neural network forget

This notebook accompanies my post Continual learning: what it is, and what I've learned working with it. I wrote it because I think forgetting is something you should watch happen at least once, on a model you trained yourself, rather than take from a survey table. Everything here runs end to end in about two minutes on a laptop, no GPU.

The claim under test: train a network on a sequence of tasks, one after another, and it overwrites what it learned earlier. McCloskey and Cohen named this catastrophic interference back in 1989, and I want to see it, then compare three standard responses side by side. Naive fine-tuning, which is just carrying on training and hoping. Elastic Weight Consolidation (Kirkpatrick et al., 2017), which penalises moving the weights that earlier tasks relied on. And experience replay, which keeps a small buffer of old examples and mixes them into every new batch. As the ceiling to measure everyone against, I also train a model jointly on all data seen so far, which is what you would do if storage and access were unlimited.

My setup, chosen for speed and clarity rather than glory: scikit-learn's digits dataset (1,797 images of handwritten digits, 8 by 8 pixels), split into five tasks of two classes each, {0,1} through {8,9}. The model is a small multilayer perceptron with a single 10-way output head, and at test time it is not told which task an image came from. That makes this the class-incremental scenario in van de Ven and Tolias's taxonomy (2019), deliberately: it is the hardest of the three standard scenarios, the closest to real deployments, and, as you will see, the one where a famous method falls over.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.manual_seed(0)
np.random.seed(0)

digits = load_digits()
X = (digits.data / 16.0).astype(np.float32)   # pixel values in [0, 1]
y = digits.target

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=0)

TASKS = [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9)]

def task_split(Xa, ya, classes):
    mask = np.isin(ya, classes)
    return torch.tensor(Xa[mask]), torch.tensor(ya[mask])

train_tasks = [task_split(X_train, y_train, c) for c in TASKS]
test_tasks  = [task_split(X_test,  y_test,  c) for c in TASKS]

print("Training samples per task:", [len(t[0]) for t in train_tasks])
print("Test samples per task:    ", [len(t[0]) for t in test_tasks])

Training samples per task: [251, 252, 254, 252, 248]
Test samples per task:     [109, 108, 109, 108, 106]

Model and training loop

Nothing exotic here, and that is the point: a 64-100-100-10 MLP trained with Adam and cross-entropy, deliberately too boring to blame. train() accepts an optional penalty term, which is the hook EWC will use later. The one habit I want to model properly, because I preach it in the post: after each task I record accuracy on the test split of every task, giving an accuracy matrix $R$ where $R_{ij}$ is accuracy on task $j$ after finishing task $i$. Nearly every question I ever ask of a continual learning run gets answered by this matrix (the framing comes from Lopez-Paz and Ranzato, 2017), so I build it in from the start rather than bolting it on when something looks odd.

def make_model():
    return nn.Sequential(
        nn.Linear(64, 100), nn.ReLU(),
        nn.Linear(100, 100), nn.ReLU(),
        nn.Linear(100, 10))

def train(model, Xa, ya, epochs=30, lr=1e-3, batch=32, penalty=None):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        perm = torch.randperm(len(Xa))
        for i in range(0, len(Xa), batch):
            idx = perm[i:i + batch]
            loss = F.cross_entropy(model(Xa[idx]), ya[idx])
            if penalty is not None:
                loss = loss + penalty(model)
            opt.zero_grad()
            loss.backward()
            opt.step()

@torch.no_grad()
def accuracy(model, Xa, ya):
    return (model(Xa).argmax(1) == ya).float().mean().item()

def eval_all(model):
    return [accuracy(model, Xt, yt) for Xt, yt in test_tasks]

Strategy 1: naive fine-tuning

Train on task 1, then keep training the same network on task 2, and so on. No memory of the past, no protection for old weights. This is what happens by default when nobody on a team is thinking about forgetting, which is why I always run it first: it is less a strategy than a control group for wishful thinking.

model = make_model()
R_naive = []
for X_c, y_c in train_tasks:
    train(model, X_c, y_c)
    R_naive.append(eval_all(model))
R_naive = np.array(R_naive)

print("Accuracy on each task after the final one (naive):")
print(np.round(R_naive[-1], 3))

Accuracy on each task after the final one (naive):
[0. 0. 0. 0. 1.]

I have run this cell dozens of times now and the row of zeros still lands. Near-perfect accuracy on the most recent pair of digits, essentially nothing on every pair before it, and not a slow fade either; one task after the data disappears, it is gone. The gradients from the new task point wherever the new loss surface wants them to, and nothing in the objective speaks for the past.

Strategy 2: Elastic Weight Consolidation

EWC is the method I most wanted to work when I first met it, because the idea is lovely: estimate, per weight, how much the earlier tasks care about it, using the diagonal of the Fisher information matrix, then hold the important weights near the values they had when the earlier task finished:

$$\mathcal{L}(\theta) = \mathcal{L}_{\text{new}}(\theta) + \sum_{t} \frac{\lambda}{2} \sum_i F_i^{(t)} \, (\theta_i - \theta_i^{(t)*})^2$$

where $F^{(t)}$ is the Fisher diagonal for task $t$ and $\theta^{(t)*}$ the weights at the end of that task. (If you like Bayesian readings, this is a Laplace approximation: the old task's posterior recycled as the new task's prior. I say more about that in the post.)

def fisher_diagonal(model, Xa, ya):
    fisher = [torch.zeros_like(p) for p in model.parameters()]
    for i in range(len(Xa)):
        model.zero_grad()
        log_p = F.log_softmax(model(Xa[i:i + 1]), dim=1)[0, ya[i]]
        log_p.backward()
        for f, p in zip(fisher, model.parameters()):
            f += p.grad.detach() ** 2
    model.zero_grad()
    return [f / len(Xa) for f in fisher]

class EWC:
    def __init__(self, lam=5000.0):
        self.lam = lam
        self.snapshots = []   # one (fisher, weights_star) pair per finished task

    def add_task(self, model, Xa, ya):
        star = [p.detach().clone() for p in model.parameters()]
        self.snapshots.append((fisher_diagonal(model, Xa, ya), star))

    def penalty(self, model):
        if not self.snapshots:
            return torch.tensor(0.0)
        total = 0.0
        for fisher, star in self.snapshots:
            for f, p, ps in zip(fisher, model.parameters(), star):
                total = total + (f * (p - ps) ** 2).sum()
        return self.lam / 2 * total

model = make_model()
ewc = EWC(lam=5000.0)
R_ewc = []
for X_c, y_c in train_tasks:
    train(model, X_c, y_c, penalty=ewc.penalty)
    ewc.add_task(model, X_c, y_c)
    R_ewc.append(eval_all(model))
R_ewc = np.array(R_ewc)

print("Accuracy on each task after the final one (EWC):")
print(np.round(R_ewc[-1], 3))

Accuracy on each task after the final one (EWC):
[0.    0.    0.    0.    0.934]

And it gives me essentially nothing: on average, level with doing nothing at all. Before you suspect my implementation, as I did the first time, try raising lam by an order of magnitude in either direction; the picture does not change, and that is the interesting part. Protecting old weights cannot teach the output layer to separate classes it never saw together. No batch ever contained both a 0 and a 7, so no penalty on weight movement can create that decision boundary. van de Ven and Tolias (2019) document this failure systematically, and I regard it as one of the most useful negative results in the field: it is why task-incremental benchmark wins can evaporate the moment the task label is taken away. When I read a new regularisation method now, the first thing I look for is which scenario its results table lives in.

Strategy 3: experience replay

The crudest idea available: keep a small buffer of examples from previous tasks, here 20 per class, so at most 200 images or about 16% of the training set, and mix it into the training data for each new task. I resisted how well this works for longer than I should have.

BUFFER_PER_CLASS = 20

model = make_model()
buffer_X, buffer_y = [], []
R_replay = []
for (X_c, y_c), classes in zip(train_tasks, TASKS):
    if buffer_X:
        X_mix = torch.cat([X_c] + buffer_X)
        y_mix = torch.cat([y_c] + buffer_y)
    else:
        X_mix, y_mix = X_c, y_c
    train(model, X_mix, y_mix)
    for c in classes:
        idx = (y_c == c).nonzero().flatten()[:BUFFER_PER_CLASS]
        buffer_X.append(X_c[idx])
        buffer_y.append(y_c[idx])
    R_replay.append(eval_all(model))
R_replay = np.array(R_replay)

print("Accuracy on each task after the final one (replay):")
print(np.round(R_replay[-1], 3))

Accuracy on each task after the final one (replay):
[0.917 0.889 0.927 0.935 0.962]

Upper bound: joint training

At each stage, retrain from scratch on everything seen so far. Nobody deploys this, that is the whole premise of the field, but I always run it because a continual method's score means nothing to me until I know the ceiling it should be measured against.

R_joint = []
for t in range(len(TASKS)):
    model = make_model()
    X_all = torch.cat([train_tasks[i][0] for i in range(t + 1)])
    y_all = torch.cat([train_tasks[i][1] for i in range(t + 1)])
    train(model, X_all, y_all)
    R_joint.append(eval_all(model))
R_joint = np.array(R_joint)

print("Accuracy on each task after the final stage (joint):")
print(np.round(R_joint[-1], 3))

Accuracy on each task after the final stage (joint):
[0.982 0.981 0.963 0.991 0.953]

Results

Two standard metrics, both cheap arithmetic on the accuracy matrix (Lopez-Paz and Ranzato, 2017; Chaudhry et al., 2018). Average accuracy is the mean over all tasks after training on the last one; forgetting is, for each task, the gap between the best accuracy it ever had and where it ended, averaged. I report both because they answer different questions, how good is the model now, and how much of what it once knew did it lose, and I have been misled before by looking at only the first.

def avg_accuracy(R):
    return R[-1].mean()

def forgetting(R):
    gaps = [R[:, j].max() - R[-1, j] for j in range(R.shape[1] - 1)]
    return float(np.mean(gaps))

strategies = [
    ("Naive fine-tuning", R_naive),
    ("EWC",               R_ewc),
    ("Replay",            R_replay),
    ("Joint (upper bound)", R_joint),
]

print(f"{'Strategy':<22}{'Avg accuracy':>14}{'Forgetting':>12}")
for name, R in strategies:
    print(f"{name:<22}{avg_accuracy(R):>13.1%}{forgetting(R):>11.1%}")

Strategy                Avg accuracy  Forgetting
Naive fine-tuning             20.0%      99.8%
EWC                           18.7%      99.8%
Replay                        92.6%       7.4%
Joint (upper bound)           97.4%       1.1%

import os
import matplotlib.pyplot as plt

os.makedirs("figs", exist_ok=True)

# Figure style: serif, muted colours, linestyle as second encoding, recessive grid
plt.rcParams.update({
    "figure.dpi": 150, "savefig.dpi": 200, "savefig.bbox": "tight",
    "font.family": "serif",
    "font.serif": ["STIXGeneral", "Times New Roman", "DejaVu Serif"],
    "mathtext.fontset": "stix",
    "font.size": 9.5, "axes.labelsize": 9.5, "axes.titlesize": 10,
    "xtick.labelsize": 8.5, "ytick.labelsize": 8.5,
    "text.color": "#1a1a1a", "axes.labelcolor": "#1a1a1a",
    "xtick.color": "#666666", "ytick.color": "#666666",
    "axes.edgecolor": "#666666", "axes.linewidth": 0.6,
    "axes.spines.top": False, "axes.spines.right": False,
    "axes.grid": True, "grid.color": "#dddddd", "grid.linewidth": 0.4,
    "axes.axisbelow": True, "legend.frameon": False,
})

STYLE = {
    "Naive fine-tuning":   dict(color="#0072B2", ls="-",  marker="o"),
    "EWC":                 dict(color="#D55E00", ls="--", marker="s"),
    "Replay":              dict(color="#009E73", ls="-",  marker="D"),
    "Joint (upper bound)": dict(color="#888888", ls=":",  marker=""),
}

# Accuracy on the first task as training moves through the sequence
stages = np.arange(1, 6)
fig, ax = plt.subplots(figsize=(6.6, 3.9))
s = STYLE["Joint (upper bound)"]
ax.plot(stages, R_joint[:, 0] * 100, ls=s["ls"], lw=1.2, color=s["color"], zorder=1)
for name, R in [("Naive fine-tuning", R_naive), ("EWC", R_ewc), ("Replay", R_replay)]:
    s = STYLE[name]
    ax.plot(stages, R[:, 0] * 100, ls=s["ls"], marker=s["marker"], ms=3.5,
            lw=1.4, color=s["color"], zorder=2)

# direct labels at the right edge, nudged apart if they collide
finals = sorted(
    [(R[-1, 0] * 100, n) for n, R in
     [("Naive fine-tuning", R_naive), ("EWC", R_ewc), ("Replay", R_replay),
      ("Joint (upper bound)", R_joint)]])
placed = []
for val, name in finals:
    ypos = val
    if placed and ypos - placed[-1] < 7:
        ypos = placed[-1] + 7
    placed.append(ypos)
    ax.annotate(name, (5, ypos), xytext=(7, 0), textcoords="offset points",
                va="center", fontsize=8.5, color=STYLE[name]["color"],
                annotation_clip=False)

ax.set_xticks(stages)
ax.set_xticklabels([f"after task {i}" for i in stages])
ax.set_ylim(-4, 104)
ax.set_ylabel("Accuracy on task 1 (digits 0 vs 1), %")
ax.set_title("Task 1 accuracy across the training sequence")
fig.subplots_adjust(right=0.76)
fig.savefig("figs/forgetting_curve.png")
plt.show()

# Full accuracy matrices: rows = after training task i, columns = tested on task j
fig, axes = plt.subplots(1, 4, figsize=(12, 3.2))
for ax, (name, R) in zip(axes, strategies):
    ax.imshow(R, cmap="Blues", vmin=0, vmax=1)
    for i in range(5):
        for j in range(5):
            ax.text(j, i, f"{R[i, j]:.0%}", ha="center", va="center", fontsize=7,
                    color="white" if R[i, j] > 0.6 else "#1a1a1a")
    # thin white separators between cells
    ax.set_xticks(np.arange(-0.5, 5), minor=True)
    ax.set_yticks(np.arange(-0.5, 5), minor=True)
    ax.grid(which="minor", color="white", linewidth=0.8)
    ax.tick_params(which="minor", length=0)
    ax.set_title(name)
    ax.set_xticks(range(5)); ax.set_xticklabels([f"T{j+1}" for j in range(5)])
    ax.set_yticks(range(5)); ax.set_yticklabels([f"after T{i+1}" for i in range(5)])
    ax.grid(False)
    for side in ("top", "right"):
        ax.spines[side].set_visible(True)
axes[0].set_ylabel("training stage")
fig.tight_layout()
fig.savefig("figs/accuracy_matrices.png")
plt.show()

# Final average accuracy, side by side
fig, ax = plt.subplots(figsize=(5.4, 3.3))
names = [n for n, _ in strategies]
vals = [avg_accuracy(R) * 100 for _, R in strategies]
ax.bar(range(len(names)), vals, width=0.5,
       color=[STYLE[n]["color"] for n in names], edgecolor="none")
for i, v in enumerate(vals):
    ax.text(i, v + 2, f"{v:.0f}%", ha="center", fontsize=8.5, color="#1a1a1a")
ax.set_xticks(range(len(names)))
ax.set_xticklabels(["Naive", "EWC", "Replay", "Joint"])
ax.set_ylim(0, 108)
ax.set_ylabel("Average accuracy over all five tasks, %")
ax.set_title("Final average accuracy by strategy")
ax.grid(axis="x")
fig.savefig("figs/final_accuracy.png")
plt.show()

What I take from this

Three things, and I hold them with different confidence. First, forgetting is not gradual decay; accuracy on old tasks does not drift down politely, it collapses within a task or two of the data disappearing, and no amount of reading prepared me for how abrupt it looks in your own run. Second, a small memory goes a long way: replay with 200 stored images recovers most of the gap to joint training, which matches the single most reliable finding in the empirical literature and is why replay, in some form, is the baseline I insist any clever alternative must beat. Third, and the one I most want you to carry out of here: method rankings are scenario-dependent. EWC looks respectable on task-incremental benchmarks and fell flat on its face in the class-incremental setting I used, so before importing any paper's conclusions, find which scenario its results were measured in.

Some experiments worth running next, in roughly the order I would run them myself: shrink the replay buffer and find where retention gives out; give each task its own output head and watch EWC start working, which is the scenario lesson happening in front of you; and reorder the tasks to see how much the sequence itself was carrying. For proper tooling beyond this toy, bigger benchmarks, more strategies, metric implementations you do not have to trust me for, use Avalanche, maintained by the ContinualAI community.

References

• McCloskey, M. and Cohen, N. J. (1989). Catastrophic interference in connectionist networks. Psychology of Learning and Motivation, 24.
• Kirkpatrick, J. et al. (2017). Overcoming catastrophic forgetting in neural networks. PNAS, 114(13). arXiv:1612.00796
• Lopez-Paz, D. and Ranzato, M. (2017). Gradient Episodic Memory for continual learning. NeurIPS. arXiv:1706.08840
• Chaudhry, A. et al. (2018). Riemannian walk for incremental learning. ECCV. arXiv:1801.10112
• van de Ven, G. M. and Tolias, A. S. (2019). Three scenarios for continual learning. arXiv:1904.07734