Replay and Retention Sets

Replay and retention sets are the two halves of handling catastrophic forgetting when a model is fine-tuned. A retention set measures what the adapted model has lost. Replay is a mitigation: the fine-tuning mixture includes examples of the earlier behavior so the model keeps being trained on it.

The problem

Fine-tuning computes gradient updates from the new data only. Nothing in that objective rewards preserving behavior that the data does not exercise, so performance on earlier tasks can fall sharply. This is catastrophic forgetting. For a language model it shows up in general knowledge, reasoning, instruction following, output formats, safety behavior, or other languages. Some change is intended: a behavior the fine-tuning data deliberately changes (a new house style, a new refusal policy) is the target of the adaptation and is not forgetting.

Retention set

A retention set is held-out data for the behaviors that should survive adaptation: a general benchmark or internal suite, format and schema checks, a safety or refusal suite, and a sample of earlier-task examples. It follows the same discipline as any evaluation set (no overlap with training, representative, big enough for the decision), plus two rules:

• It contains only behaviors where the old answer should survive. Deciding for each behavior whether the new model should differ from the base is a design decision made before the experiment.
• Score both the base model and the adapted model with the same harness. Report the absolute retention score and the drop from the base, per slice, since an average can hide a collapse in one behavior.

Retention should be tracked along the training curve and measured at the checkpoint that will ship, since forgetting speed depends on learning rate, method, and number of steps.

Replay

Replay (or rehearsal) adds earlier-task examples to the fine-tuning set, mixed with the new-task examples. It supplies the contrast the new data lacks: if the new task flips an answer on some inputs, replay examples show the model the cases where the old answer remains correct. Sources are a sample of earlier training data, general instruction data, or the model’s own outputs on general prompts, which is useful when the original data is unavailable. The main design choices are the replay fraction and which behaviors the replayed examples cover; behaviors without replay examples receive no direct rehearsal signal, although shared representations and other regularization can still help preserve them.

Worked example

This is a hypothetical toy, run with numpy: logistic regression on 20 features. Task A labels each input by the sign of \(w_A\cdot x\). Task B uses \(w_B\), which equals \(w_A\) with the signs of the first 10 weights flipped, so B conflicts with A on half the features. The model is first trained on A (98.6% accuracy on A, 59.8% on B). It is then fine-tuned on 300 B examples, with a varying number of A examples (the replay) added to the mix, and evaluated on 5,000 held-out inputs labeled for each task. The fixed seed, training steps, learning rate, and data generation are shown below. Here the 5,000 A examples are the retention set.

import numpy as np
rng = np.random.default_rng(7)
d = 20
wa = rng.normal(size=d); wb = wa.copy(); wb[:10] *= -1
train_a = rng.normal(size=(2000,d)); train_b = rng.normal(size=(300,d))
test = rng.normal(size=(5000,d))
labels = lambda x,w: (x @ w > 0).astype(float)
def fit(x,y,w,steps=300,lr=0.3):
    w = w.copy()
    for _ in range(steps):
        p = 1/(1+np.exp(-np.clip(x @ w,-30,30)))
        w -= lr * (x.T @ (p-y))/len(x)
    return w
acc = lambda w,truth: np.mean((test @ w > 0) == truth)
ya,yb = labels(test,wa),labels(test,wb)
base = fit(train_a,labels(train_a,wa),np.zeros(d),steps=600)
print(f"base A {acc(base,ya):.3f} B {acc(base,yb):.3f}")
replay = rng.normal(size=(900,d))
for n in (0,30,100,300,900):
    x = np.vstack([train_b,replay[:n]])
    y = np.concatenate([labels(train_b,wb),labels(replay[:n],wa)])
    w = fit(x,y,base)
    print(f"{n:3d} share {n/(300+n):.3f} B {acc(w,yb):.3f} A {acc(w,ya):.3f}")

replay examples	share of mixture	B accuracy	A accuracy (retention)
0	0%	.967	.615
30	9.1%	.927	.643
100	25%	.862	.706
300	50%	.761	.811
900	75%	.693	.901

With no replay the model learns B well and drops from 98.6% to 61.5% on A. As replay grows, retention rises and target accuracy falls. This is a trade-off set by the mixture: the model has limited capacity (20 weights) and the two tasks pull the same weights in opposite directions, so it cannot satisfy both. The toy shows the shape of the trade-off; the exchange rate for a real model and task has to be measured on its own retention set.

Reading the result

• Treat the target score and the retention score as a pair. Choosing a recipe by the target score alone selects the one that forgets most. A constraint set before the experiment, such as “retention may drop by at most 2 points”, turns the choice into a constrained one and exposes recipes that cannot meet it.
• If no replay setting meets the constraint, the adaptation is not worth shipping at that capacity and data; training settings alone do not supply missing contrast, although a lower learning rate, fewer steps, or regularization may reduce forgetting at the cost of adaptation. Re-test both scores after such changes.
• Replay data that does not resemble the behaviors you care about gives them no direct rehearsal, so protection is unlikely to be strong; shared representations may still help or hurt them indirectly, and only an independent retention set shows which. Match the replay to the retention set’s behaviors without copying its items; examples shared by both would make retention scores optimistic.
• Other mitigations (lower learning rates, fewer steps, lower-rank adapters, regularizing toward the base weights as in Kirkpatrick et al.) are also judged by the retention set.

References

• Kirkpatrick et al., “Overcoming catastrophic forgetting in neural networks” — https://arxiv.org/abs/1612.00796