Avoid Catastrophic Forgetting with Rank-1 Fisher from Diffusion Models

Official code for:

Avoid Catastrophic Forgetting with Rank-1 Fisher from Diffusion Models
Zekun Wang, Anant Gupta, Zihan Dong, Christopher J. MacLellan
ICLR 2026

OpenReview

This repository implements continual learning for class-conditional diffusion models with:

• EWC with diagonal Fisher (diag)
• EWC with rank-1 optimal Fisher approximation (rank1_opt) (main contribution)
• Generative replay
• Replay distillation (teacher-student noise prediction matching)

Training and evaluation are task-incremental: classes are split into tasks using group_size, and FID is evaluated on all seen tasks after each task.

• train-model.py: main entrypoint (config loading, continual loop, FID logging)
• example_config.json: runnable example config
• src/ddim.py: class-conditional DDIM model and sampler
• src/utils.py: dataset loaders, task split logic, training loop, FID evaluator
• src/parameter_scoring.py: gradient-based Fisher statistics and rank-1 coefficient
• src/ewc.py: EWC penalty over all past tasks
• src/gr.py: generative replay buffer built from frozen teacher generations
• src/fisher_analysis.py: Fisher analysis utilities

Python 3.8+ is recommended.

pip install torch torchvision diffusers torchmetrics tqdm datasets wandb

Default run (uses example_config.json):

python train-model.py

Run with a custom config:

python train-model.py --config path/to/your_config.json

All experiment options are read from JSON. Key fields:

• dataset: mnist | fmnist | cifar10 | imagenet32
• num_classes: number of classes to include
• group_size: number of classes per task (num_tasks = num_classes / group_size)
• use_ewc: enable EWC
• ewc_lambda: EWC regularization weight
• ewc_fisher_type: diag or rank1_opt
• use_generative_replay: enable generative replay
• gr_alpha, gr_pool_size_per_class, gr_num_inference_steps, gr_eta: replay controls
• use_distillation, gr_kl: replay distillation controls
• epochs, batch_size, lr, seed
• normalize, greyscale
• use_wandb, wandb_project, wandb_run_name
• output_dir: root folder for checkpoints, samples, and logs

The provided example_config.json runs MNIST with:

• num_classes=10, group_size=2 (5 tasks)
• rank-1 EWC (rank1_opt)
• replay + distillation enabled

• MNIST/FashionMNIST/CIFAR10 are downloaded via torchvision.
• ImageNet32 uses Hugging Face dataset benjamin-paine/imagenet-1k-32x32.
• Optional environment variables:
  • DATA_ROOT (default: ./data)
  • HF_DATASETS_CACHE (default: ./cache)

Under output_dir, each run writes to:

{dataset}-{ewc_fisher_type}{-gr}{-distill}-{seed}/

including:

• model-task{t}.pt: checkpoint after task t
• fisher-task{t}.pt: saved Fisher statistics (when EWC is enabled and t is not the final task)
• task_{t}/epoch_*/epoch_*_grid.png: sampled image grids logged during training

Set:

• use_ewc = true
• ewc_fisher_type = "rank1_opt"

Optional (as in demo config):

• use_generative_replay = true
• use_distillation = true

Then run:

python train-model.py --config example_config.json

@inproceedings{
wang2026avoid,
title={Avoid Catastrophic Forgetting with Rank-1 Fisher from Diffusion Models},
author={Zekun Wang and Anant Gupta and Zihan Dong and Christopher J. MacLellan},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=zCZcbRsc4g}
}