Efficient Lifelong Learning with A-GEM

In plain words

Train a neural network on task A, then on task B, and it usually forgets much of task A. This is called catastrophic forgetting, and it is the central problem of lifelong (or continual) learning, where a model sees a sequence of tasks one after another and each training example only once.

An earlier method, GEM, keeps a small notebook of examples from every past task. Before each update it checks, task by task, whether the update would make any past task worse, and if so it solves a small optimization problem to bend the update until no past task is hurt. That works, but it gets slow and memory-hungry as the number of tasks grows.

A-GEM simplifies the check. Instead of asking every past task separately, it grabs a random handful of notebook examples from all past tasks, computes one "don't make this worse" direction from them, and only intervenes if the new update points against it. The fix is a one-line projection: remove the part of the update that conflicts with that direction. The authors report the same or better average accuracy than GEM at a fraction of its time and memory, on par with simple regularization methods such as EWC.

The paper also argues for a stricter way to evaluate such methods (one pass over the data, hyperparameters tuned only on a few held-out tasks), adds a metric for how fast a model picks up a new task, and shows that describing each task by the attributes of its classes (beak shape, fur color and so on) helps every method learn new tasks faster.

TL;DR

A-GEM is a cheaper variant of Gradient Episodic Memory for single-pass continual learning. At each step it computes the gradient g on the current mini-batch and a reference gradient g_ref on a random batch drawn from the episodic memory of all past tasks; if gᵀg_ref < 0 it replaces g by g − (gᵀg_ref / g_refᵀg_ref) g_ref, otherwise it uses g unchanged §4 eq. 11. No quadratic program and no per-task gradient matrix are needed. On Permuted MNIST, Split CIFAR, Split CUB and Split AWA it matches or beats GEM in average accuracy and beats EWC, PI, MAS and RWalk, with training time close to EWC tab. 4 tab. 7. The paper also proposes a protocol (tune on 3 tasks, train and report on 17 in one pass), a Learning Curve Area (LCA) metric and an attribute-based joint-embedding model §2 §3 §5.

Problem

Most lifelong-learning work trains with several epochs per task and tunes hyperparameters by sweeping the whole task sequence repeatedly, which leaks the evaluation stream into model selection §2. Under a single pass, the regularization methods (EWC and relatives) barely beat plain fine-tuning §6.1 app. F. GEM works in the single-pass setting but needs, at every step, the gradients of all t−1 past tasks (a (t−1)×P matrix) and a QP solve; this becomes prohibitive as memory and task count grow §4 tab. 7.

Method

Setting §2. Tasks k = 1..T, each a dataset of triplets (x, t, y) with input x, task descriptor t and label y from a task-specific label set. The predictor is f_θ(x, t). Training examples of the evaluation stream are seen exactly once.

GEM recap §4. GEM stores an episodic memory M_k per task and solves, at task t:

\text{minimize}_\theta\ \ell(f_\theta, \mathcal{D}_t)\quad \text{s.t.}\quad \ell(f_\theta, \mathcal{M}_k) \le \ell(f_\theta^{t-1}, \mathcal{M}_k)\ \ \forall k<t eq. 6

In gradient form this becomes a projection with t−1 constraints:

\text{minimize}_{\tilde g}\ \tfrac12\lVert g-\tilde g\rVert_2^2\quad \text{s.t.}\quad \langle \tilde g, g_k\rangle \ge 0\ \ \forall k<t eq. 7

solved via its dual, a QP in t−1 variables:

\text{minimize}_v\ \tfrac12 v^\top G G^\top v + g^\top G^\top v \quad \text{s.t.}\quad v \ge 0 eq. 8

with G = −(g_1, …, g_{t−1}) ∈ ℝ^{(t−1)×P} and g̃ = Gᵀv* + g §4.

A-GEM §4. Constrain only the average memory loss over all past tasks:

\text{minimize}_\theta\ \ell(f_\theta, \mathcal{D}_t)\quad \text{s.t.}\quad \ell(f_\theta, \mathcal{M}) \le \ell(f_\theta^{t-1}, \mathcal{M}),\ \ \mathcal{M} = \cup_{k<t}\mathcal{M}_k eq. 9

\text{minimize}_{\tilde g}\ \tfrac12\lVert g-\tilde g\rVert_2^2\quad \text{s.t.}\quad \tilde g^\top g_{ref} \ge 0 eq. 10

where g_ref is the gradient on a batch (x_ref, y_ref) sampled at random from M. The closed-form solution when the constraint is violated:

\tilde g = g - \frac{g^\top g_{ref}}{g_{ref}^\top g_{ref}}\, g_{ref} eq. 11

The derivation (Lagrangian with multiplier α, z* = g + α g_ref, α* = −gᵀg_ref / g_refᵀg_ref) is in app. C.

Algorithm: A-GEM training  [alg. 2]
Input: model f_θ, task streams D_1..D_T (train/test), learning rate α
1. M ← {} ; A ← T×T zero matrix
2. for t = 1..T:
3.   for each mini-batch (x, y) in D_t^train:
4.     sample (x_ref, y_ref) from M
5.     g_ref ← ∇_θ ℓ(f_θ(x_ref, t), y_ref)
6.     g     ← ∇_θ ℓ(f_θ(x, t), y)
7.     if gᵀ g_ref ≥ 0: g̃ ← g
8.     else:            g̃ ← g − (gᵀ g_ref / g_refᵀ g_ref) · g_ref
9.     θ ← θ − α · g̃
10.  M ← UpdateEpsMem(M, D_t^train, T)   # add samples of task t to memory
11.  A[t, :] ← accuracy on every task's test set
Output: f_θ, A

Memory is populated per task with samples chosen uniformly at random §6. The code's actual procedure (see Paper vs code) differs in details.

Evaluation protocol §2 alg. 1. Split the task stream into D^CV (first T^CV tasks) and D^EV (remaining tasks); T^CV = 3, T = 20. Grid search hyperparameters with as many passes over D^CV as needed, pick h* by average accuracy (eq. 1) on D^CV's test sets, reset the model, then train once, in a single pass, over D^EV (17 tasks) and report metrics on D^EV.

Metrics §3. Let a_{k,i,j} be test accuracy on task j after training on the i-th mini-batch of task k, B_k the number of mini-batches of task k.

A_k = \frac1k \sum_{j=1}^{k} a_{k,B_k,j} eq. 1

F_k = \frac{1}{k-1}\sum_{j=1}^{k-1} f_j^k,\qquad f_j^k = \max_{l\in\{1,\dots,k-1\}} a_{l,B_l,j} - a_{k,B_k,j} eq. 2 eq. 3

Z_b = \frac1T \sum_{k=1}^{T} a_{k,b,k},\qquad \text{LCA}_\beta = \frac{1}{\beta+1}\sum_{b=0}^{\beta} Z_b eq. 4 eq. 5

LCA_0 is average 0-shot accuracy on a new task; LCA_β is the area under the early learning curve, high when 0-shot accuracy is good and learning is fast §3. Results use β = 10 fig. 1.

Joint embedding with task descriptors (the "-JE" models) §5. The task descriptor t^k is a C_k × A matrix of class attribute vectors. A feature extractor φ_θ(x) ∈ ℝ^D and an attribute embedding ψ_ω, an A × D parameter matrix, so ψ_ω(t^k) ∈ ℝ^{C_k×D} (each class embedding is a linear combination of attribute vectors) §5 fig. 9. Train with cross-entropy:

\ell_k(\theta,\omega) = \frac1N \sum_{i=1}^{N} -\log p(y_i^k \mid x_i^k, t^k;\theta,\omega) eq. 12

p(c \mid x_i^k, t^k;\theta,\omega) = \frac{\exp([\phi_\theta(x_i^k)\,\psi_\omega(t^k)^\top]_c)}{\sum_j \exp([\phi_\theta(x_i^k)\,\psi_\omega(t^k)^\top]_j)} eq. 13

This is orthogonal to A-GEM and is applied to every baseline §5. Integer task descriptors are used for Permuted MNIST and Split CIFAR; attribute descriptors for Split CUB and Split AWA §6.

Hyperparameters

Baseline settings chosen by the same grid search app. G:

Grids: EWC λ ∈ [1, 10, 100, 1000, 10000]; PI λ ∈ [0.001, 0.01, 0.1, 1, 10]; MAS λ ∈ [0.01, 0.1, 1, 10]; RWALK λ ∈ [0.1, 1, 10, 100, 1000] app. G. The -JE variants have their own optima, listed in app. G.

Dataset statistics tab. 1:

Results

Author-reported. Single pass over 17 evaluation tasks, mini-batch 10, means (± confidence interval) over 5 runs for MNIST/CIFAR and 10 runs for CUB/AWA. A_T in %, F_T and LCA_10 in [0, 1] §6.1.

Permuted MNIST and Split CIFAR tab. 4:

∞ means the model failed to train with the hyperparameters cross-validated on the first 3 tasks tab. 4.

Split CUB tab. 5 and Split AWA tab. 6, standard / joint-embedding (-JE):

Training time and memory tab. 7 (P parameters, B mini-batch size, H hidden state size, M episodic memory per task, T tasks; time is GPU training time):

GEM vs its variants tab. 3 (S-GEM projects onto one randomly chosen past task constraint per step app. D.3):

Average accuracy vs worst-case forgetting F_wst on the memory and on the test set tab. 2:

Other findings:

• A-GEM has the best A_T on all datasets except Permuted MNIST, where PROG-NN wins thanks to the large number of examples per task; PROG-NN runs out of memory on CUB and AWA §6.1.
• GEM violates its constraints at almost every step as tasks accumulate, while A-GEM's violation count plateaus much lower app. D.1 fig. 6.
• EWC only pulls ahead of VAN with a bigger network and multiple epochs (up to 10 on MNIST, 30 on CIFAR) app. F fig. 7 fig. 8.
• Attribute task descriptors raise A_T and LCA for every method; zero-shot accuracy of the -JE models improves as more tasks are learned, most for A-GEM-JE §6.1 fig. 4 fig. 5.

Paper vs code

Checked against facebookresearch/agem at 4542149.

• g_ref batch size on CIFAR. Paper: 1300 §6. Code default is 256 conv_split_cifar.py and the replication script does not override --eps-mem-batch replicate_results_iclr19.sh. MNIST (256), CUB (128) and AWA (128) match fc_permute_mnist.py conv_split_cub.py conv_split_awa.py.
• Memory population. Paper: samples chosen uniformly at random per task §6; Alg. 2 adds |M|/T samples after each task alg. 2. Code: while training on a task, each incoming example is written into a fixed-size ring buffer slot for its class (mem_size slots per class), so the memory holds the last mem_size examples per class of the shuffled stream; the slots of task t become visible to g_ref sampling only after task t ends conv_split_cifar.py conv_split_cifar.py fc_permute_mnist.py. Because data is shuffled first conv_split_cifar.py, this is a random subset in effect.
• Small memory. When the filled memory has at most eps_mem_batch examples, the code uses all of it for g_ref; otherwise it samples without replacement conv_split_cifar.py. The paper only says the stated batch sizes apply "when the episodic memory is sufficiently filled" §6.
• First task. Alg. 2 samples from an empty memory on task 1 alg. 2; the code uses an ordinary SGD step on the first task conv_split_cifar.py model.py.
• Multi-head loss on memory. In split (multi-head) setups, each memory example's cross-entropy is computed only over the logits of its own task's classes; per-task losses are summed and divided by the batch size model.py model.py conv_split_cifar.py. The paper writes f_θ(x_ref, t) with the current task id t alg. 2, which would be wrong for multi-head outputs.
• Joint-embedding bias. The code adds a learned per-class bias to the logits φ(x)ψ(t)ᵀ model.py; eq. 13 has no bias eq. 13.
• Projection. Matches eq. 11; the projection triggers only when the dot product is strictly negative, with no epsilon in the denominator model.py.
• Learning rates for A-GEM match app. G: MNIST 0.1 and CIFAR 0.03 come from the replication script, CUB 0.03 and AWA 0.01 are hard-coded replicate_results_iclr19.sh conv_split_cub.py conv_split_awa.py.

Pitfalls

• Use the memory gradient with the right heads. For multi-head tasks, g_ref must be computed with each memory example routed to its own task's output classes, not to the current task's head (see Paper vs code).
• Speed-up claim vs table. The text says A-GEM is "about 100 times faster" than GEM §6.1, but tab. 7 gives 3442 s vs 477 s on MNIST and 5238 s vs 449 s on CIFAR, i.e. roughly 7× and 12×. Expect the table's numbers, not the text's.
• Not truly free. Each step costs two backward passes (one on the current batch, one on a memory batch of 128 to 1300 examples), versus one for plain SGD; A-GEM takes about 2.6× (MNIST) and 4.3× (CIFAR) the training time of VAN tab. 7. Training memory holds an extra P-sized gradient vector (2·P vs P) tab. 7.
• Projection only removes conflict on average. A-GEM bounds the average memory loss, not each task's loss. GEM has lower worst-case forgetting on the memory itself (e.g. CIFAR memory accuracy 97.1% vs 72.1%) tab. 2 §4. If per-task guarantees matter, A-GEM is weaker.
• g_refᵀg_ref near zero. Eq. 11 divides by ‖g_ref‖²; neither paper nor code guards against a near-zero reference gradient model.py.
• Alg. 2's UpdateEpsMem is underspecified. It sets s = |M|/T and adds s samples alg. 2, which is zero for an empty memory; follow the per-task budgets in §6 (or the per-class slots in the code) instead.
• Protocol matters. Hyperparameters are tuned on the first 3 tasks only, and results are on the other 17 §2. PI fails to train on MNIST under this protocol tab. 4. Results tuned on the full stream are not comparable.
• Split AWA reuses classes. Tasks sample classes with replacement, each class's training data is split into disjoint parts per occurrence, and the class weights are re-initialized in each head §6. Split AWA's training images per task are not given tab. 1.
• Permuted MNIST step count. app. D.1 states 5500 gradient steps per task at batch 10 (55000 examples), while tab. 1 lists 60000 training images per task; how the remainder is used is not stated.
• LCA indexing. Z_b starts at b = 0, the accuracy before the first mini-batch of the task (0-shot) eq. 4 eq. 5; the code evaluates after every mini-batch for the first 20 iterations, then every 50 conv_split_cifar.py.
• Weight decay, momentum and LR schedule are not stated; the replication script uses plain SGD at a constant rate.

Lineage

• Builds on: GEM (Lopez-Paz & Ranzato, 2017, arXiv:1706.08840): same episodic memory and gradient-projection idea, but t−1 per-task constraints and a QP are replaced by one averaged constraint with a closed-form projection §4.
• Compared against: EWC (arXiv:1612.00796), PI, MAS, RWalk, iCaRL, Progressive Networks, GEM, and a multi-task upper bound §6.
• Metrics A and F come from RWalk (Chaudhry et al., 2018); LCA is new here §3.