Dreaming: mitigating catastrophic forgetting via KL-regularized post-training on self-generated data

A study of catastrophic forgetting during online learning and a mitigation for it. A small transformer memorizes 10,000 key–value pairs ("old knowledge"), then is fine-tuned on 1,000 fresh pairs ("new knowledge") while a KL penalty against the frozen pretrained model computed on sequences sampled from the frozen model itself, hence "dreaming" pulls it back toward what it used to believe. We sweep the KL coefficient and measure retention of the old pairs against acquisition of the new ones. Full spec in dreaming.md.

Result

kl_coeff	retention (old)	acquisition (new)
(pretrained)	1.0000	0.0020
0	0.2911	1.0000
0.05	0.8830	1.0000
0.1	0.9569	1.0000
0.3	0.9828	0.9990
1	0.9760	0.9960
3	0.9697	0.9890
10	0.9665	0.9070

Three takeaways (default config, seed 0):

• Naive fine-tuning (kl_coeff = 0) is catastrophic: only 29% of the old pairs survive 50 epochs on the new data, even though the new data never contradicts the old (all keys are distinct).
• A little KL goes a long way: retention recovers steeply between 0 and 0.1 (0.29 → 0.88 → 0.96), and kl_coeff = 0.3 keeps >98% of old knowledge while still acquiring essentially all new pairs — nearly a free lunch at this scale.
• Too much KL trades the other way: by kl_coeff = 10 the constraint starts blocking acquisition (0.91) without buying more retention, which plateaus around 0.97 rather than returning to 1.0.

Numbers above are in results.json; the full training log is run.log.

Method

1. Data — each example is one sequence <key digits> KV <value digits> (e.g. 48711KV77), over a 12-token vocabulary: digits 0–9, a KV separator, and an input-only BOS. Keys are unique within and across the two sets, so old and new knowledge never conflict.
2. Pretrain — a ~1.1M-param decoder-only transformer (4 layers, d_model 128, 4 heads, RMSNorm pre-norm, RoPE, SwiGLU, untied head) trains with cross-entropy over all positions on the 10,000 old pairs. Loss on key positions matters: it lets the model generate keys itself, which sampling below requires. Pretraining reaches 100% exact-match on old values.
3. Post-train — freeze the pretrained model as a reference; train a copy on the 1,000 new pairs. Each step also draws 128 fresh full sequences from the frozen reference (ancestral sampling, temperature 1.0) and adds kl_coeff × KL(p_ref ‖ p_student) on those sequences — soft-label distillation of the model's own "dreams" back into it.
4. Evaluate — greedy-decode the value given BOS key KV; exact match over the full value string. Retention = accuracy on old pairs, acquisition = on new pairs.

Running it

Requires PyTorch and a GPU for reasonable speed (the full grid is ~3.5 min on CUDA; CPU works but is slow). Plotting needs matplotlib.

python dreaming.py                      # full default grid, writes results.json
python dreaming.py --kl-grid 0 0.3 1    # subset of the sweep
python dreaming.py --posttrain_epochs 100 --seed 1   # any Config field is a flag
python plot_results.py                  # results.json -> results.png

Sanity-check the setup quickly with a small run:

python dreaming.py --key_digits 4 --n_old_kvs 1000 --n_new_kvs 100 \
    --pretrain_epochs 5 --posttrain_epochs 11 --out smoke.json

Files

File	Purpose
dreaming.md	Experiment spec
dreaming.py	Model, training, sweep — the whole experiment
plot_results.py	Renders results.json → results.png
analyze_pretrain.py	Per-position loss breakdown + sample inspection after pretraining
results.json / run.log	Output of the reported run

Implementation decisions

Key implementation choices (documented in the spec and in dreaming.py's docstring):

• Optimizer: AdamW, lr 1e-3, weight decay 0, identical for both phases.
• key_digits = 5: 3-digit keys give only 1,000 unique keys, but the defaults need 10,000 + 1,000 = 11,000 of them.
• BOS token (id 11, input-only): prepended to every sequence so the model can generate keys unconditionally; the 11 task tokens are as specified.
• KL direction: forward KL(p_ref ‖ p_student) — penalizes the student for losing probability mass the reference has, averaged over all token positions of the sampled sequences.
• torch.compile (reduce-overhead on CUDA) is used for all training forwards; --compile false disables it. The same module instances are reused across the grid so compiled graphs stay warm.