Diffusion Unlearning

Targeted concept unlearning in diffusion models via separable sparse-autoencoder latents.

This repository implements a concept-unlearning method for Stable Diffusion. It trains a sparse autoencoder (SAE) on UNet activations with a supervised contrastive (SupCon) objective over cross-attention patch masks, so that object-identity and artistic-style concepts land in separable sets of SAE latents. Unlearning then becomes latent selection: pick the latents that carry the target concept, scale them by a negative multiplier at inference, and leave everything else intact.

Three pipelines are provided:

Pipeline	Removes	Entry point
Object unlearning	an object class (e.g. Dogs)	bash/train_unlearning_object.sh
Style unlearning	an artistic style (e.g. Van_Gogh)	bash/train_unlearning_style.sh
Joint unlearning	a specific (object, style) pair	bash/train_unlearning_joint.sh

Built on top of SAeUron and evaluated on the UnlearnCanvas benchmark. See NOTICE for attribution.

---

1. Setup

1.1 Environment

Requires Python ≥ 3.10, a CUDA GPU, and accelerate for multi-GPU runs.

git clone <this-repo> Diffusion_Unlearning && cd Diffusion_Unlearning

pip install uv
uv venv --python 3.12
source .venv/bin/activate
uv pip install -r requirements.txt

# requirements.txt does not list torch/torchvision, so that you can pick the
# build matching your CUDA version. Both are required. For example:
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

Two further notes on dependencies:

• Triton accelerates the SAE decoder. It ships with the Linux torch wheel; if it is missing, SAE/utils.py falls back to an eager implementation and prints a notice. Force the fallback with SAE_DISABLE_TRITON=1.
• xformers is optional. Sampling enables memory-efficient attention only when is_xformers_available() returns true.

Configure accelerate once (or pass --num_processes through NUM_PROCESSES):

accelerate config

1.2 Checkpoints and benchmark data

The pipelines expect the UnlearnCanvas assets under checkpoints/:

checkpoints/
├── style50/                  # UnlearnCanvas-finetuned SD v1.4 (must contain model_index.json)
└── cls_model/
    ├── style50.pth           # style (theme) classifier — used as THEME_CKPT / STYLE_CKPT
    └── style50_cls.pth       # object (class) classifier — used as CLASS_CKPT

Two ways to get them:

Automatic (object/style pipelines only, uses gdown):

DOWNLOAD_CHECKPOINTS=true bash bash/train_unlearning_object.sh setup

Manual — download from the UnlearnCanvas project and place as above:

• Classifier checkpoints
• Diffusion model checkpoint

The joint pipeline deliberately refuses to download; run the object pipeline's setup once first, or point MODEL_NAME / CLASS_CKPT / THEME_CKPT at existing files.

Anchor prompts and the class/style vocabularies ship with this repo in UnlearnCanvas_resources/ — no benchmark images are needed for the three pipelines, because activations are collected from generated samples. The image benchmark is only required for the optional FID scripts.

1.3 Pretrained SAE checkpoints

Coming soon. Pretrained SAE checkpoints for the object, style, and joint pipelines are not released yet. Until then, train them yourself with the train mode of the three pipelines — see §2. Once released they will drop into sae-ckpts/<WANDB_PROJECT>/<RUN_NAME>_train_<dataset-basename>/<hookpoint>/, so the unlearn mode can be run directly with RUN_SWEEP=false.

---

2. Running the pipelines

All three take a mode as $1 and read defaults from bash/configs/<name>.env. Any variable in those files can be overridden inline from the environment.

bash bash/train_unlearning_object.sh [setup|train|unlearn|viz|all|resume]
bash bash/train_unlearning_style.sh  [setup|train|unlearn|viz|all|resume]
bash bash/train_unlearning_joint.sh  [setup|train|unlearn|all|resume]

Mode	Does
setup	Optional venv creation + checkpoint download/validation only.
train	Collect activations → train SAE → gather per-concept SAE latents.
unlearn	Hyper-parameter sweep → pick best params → sample → evaluate accuracy.
viz	UMAP / hypersphere cluster plots and selected-feature reports. (object, style)
all	train + unlearn + viz (joint: train + unlearn).
resume	Inspect the run directory and re-enter at the first incomplete stage.

Everything for one run lands in a self-contained namespace:

runs/<object|style|joint>_unlearning/<RUN_NAME>/
├── config/hparams.json     # resolved hyper-parameters for this run
├── logs/pipeline_*.log     # full tee'd stdout/stderr
├── activations/            # cached UNet activations (train_dataset/)
├── sae_activations/        # per-concept SAE latents (*.pkl)
├── sweep_results/          # sweep samples + per-config accuracy
├── eval_results/           # final samples and accuracy summaries
└── visualizations/         # UMAP, hypersphere, feature-overlap plots

SAE checkpoints are written outside the run dir, to sae-ckpts/<WANDB_PROJECT>/<RUN_NAME>_train_<dataset-basename>/<hookpoint>/. Pretrained ones are not published yet (see §1.3).

2.1 Object unlearning

RUN_NAME=obj_demo \
UNLEARN_TARGET=object \
TARGET_CLASS=Dogs \
NUM_PROCESSES=4 \
bash bash/train_unlearning_object.sh all

Stages: collect_activations → train_sae (SupCon on class labels) → gather_class_latents → sweep over PERCENTILES × MULTIPLIERS → find_best_params_cls_sweep → sample → run_acc_all_cls → check_params → viz.

Key knobs (bash/configs/train_unlearning_object.env):

Variable	Default	Meaning
UNLEARN_TARGET	object	object | style | both — which axis to unlearn.
TARGET_CLASS	(empty)	Restrict sampling to one class; empty = all 20.
HOOKPOINT	unet.up_blocks.1.attentions.1	UNet block the SAE is attached to.
SUPCON_WEIGHT	0.1	Weight of the SupCon term (0 = plain SAeUron SAE).
SUPCON_CONCEPT_TYPE	class	Which axis SupCon separates.
PERCENTILES	[99.999,99.995,99.99]	Latent-selection thresholds to sweep.
MULTIPLIERS	[-5.0,-1.0,-0.5,-0.1]	Negative scaling applied to selected latents.
RUN_SWEEP	true	Set false to reuse an existing class_params.pth.

2.2 Style unlearning

RUN_NAME=style_demo \
NUM_PROCESSES=6 \
bash bash/train_unlearning_style.sh all

Same shape as the object pipeline but with SUPCON_CONCEPT_TYPE=theme, HOOKPOINT=unet.up_blocks.1.attentions.2, and a single STYLE_PERCENTILE / STYLE_MULTIPLIER pair instead of a sweep (defaults 99.999 / -0.1). Evaluation runs run_acc_all_style.py.

2.3 Joint unlearning

RUN_NAME=joint_demo \
NUM_PROCESSES=4 \
bash bash/train_unlearning_joint.sh all

This is the pipeline the (object, style) disentanglement is built for. It trains with --joint_supcon True, which uses both label axes at once:

• strong positives — same class and same style, weight JOINT_SUPCON_STRONG_POS_WEIGHT
• hard negatives — matching on exactly one axis, weight JOINT_SUPCON_HARD_NEG_WEIGHT

so latents cannot collapse object and style into a single direction. It then gathers combination_latents_dict_<hookpoint>.pkl keyed by (class, theme) pairs, sweeps sweep_joint_unlearning.py, picks joint_best_params.pth, samples, and reports final_metrics/joint_metrics.pth.

Note the joint config defaults differ: K=64, EFFECTIVE_BATCH_SIZE=131072, USE_PRE_ENCODER_GELU=false, and a wider multiplier grid down to -11.0.

2.4 Resuming

resume probes the run directory for each stage's expected output and re-enters at the first one that is missing, so an interrupted run continues where it stopped:

RUN_NAME=obj_demo bash bash/train_unlearning_object.sh resume

With no RUN_NAME, resume picks the most recently modified run directory.

2.5 Weights & Biases

Logging is on by default (WANDB_MODE=online, project sae_stable-diffusion-v1-4). Disable with WANDB_MODE=offline, or set WANDB_ENTITY / WANDB_PROJECT to redirect. Pipeline-stage events are logged through scripts/wandb_pipeline_logger.py; if wandb is not installed the pipelines warn and continue.

---

3. Repository layout

Diffusion_Unlearning/
├── bash/
│   ├── train_unlearning_object.sh   # object pipeline
│   ├── train_unlearning_style.sh    # style pipeline
│   ├── train_unlearning_joint.sh    # joint (object × style) pipeline
│   ├── configs/*.env                # per-pipeline defaults, all env-overridable
│   └── extra/                       # unmaintained helper scripts (see its README)
├── SAE/
│   ├── sae.py                       # SAE module (BatchTopK / TopK, aux-k loss)
│   ├── trainer.py                   # training loop + SupCon / joint-SupCon losses
│   ├── config.py                    # SaeConfig, TrainConfig, CacheActivationsRunnerConfig
│   ├── hooked_sd_noised_pipeline.py # SD pipeline with activation hooks
│   ├── hooked_scheduler.py          # scheduler exposing intermediate latents
│   ├── cache_activations_runner_unlearn_canvas.py
│   ├── unlearning_utils.py          # latent selection / intervention helpers
│   ├── kernels.py, utils.py
│   └── __init__.py
├── scripts/                         # ~65 entry points (see §4)
├── utils/hooks.py                   # forward-hook utilities used at sampling time
├── UnlearnCanvas_resources/
│   ├── const.py                     # 20 object classes, 50 styles (+ `Seed_Images` = unstyled)
│   └── anchor_prompts/finetune_prompts/sd_prompt_<Class>.txt
├── configs/semantic_prompt_variants.py
├── requirements.txt
├── LICENSE                          # MIT
└── NOTICE                           # upstream attribution (Apache-2.0 portions)

---

4. Scripts

Core pipeline (invoked by the three bash/ entry points)

Stage	Script
Collect UNet activations	collect_activations_unlearn_canvas.py
Train SAE	train.py
Gather latents (object / style / joint)	gather_sae_acts_ca_prompts_cls.py, gather_sae_acts_ca_prompts.py, gather_sae_acts_ca_prompts_joint.py
Sweep	sweep_cls_distr.py, sweep_joint_unlearning.py
Pick best params	find_best_params_cls_sweep.py, find_best_params_joint_sweep.py
Sample with intervention	sample_unlearning_cls_distr.py, sample_unlearning_distr.py, sample_unlearning_joint_distr.py
Evaluate accuracy	run_acc_all_cls.py, run_acc_all_style.py, run_acc_all_cls_sweep.py, run_joint_sweep_eval.py, eval_joint_final.py, accuracy_unlearncanvas_*.py, avg_accuracy_*.py
Visualize	plot_sae_latent_clusters.py, extract_selected_features_from_{cls,style}_latents_pkl.py
Sanity checks / logging	check_params.py, wandb_pipeline_logger.py

---

5. Citation

The paper describing this method is not published yet; citation information will be added here once it is available.

In the meantime, please cite the works this builds on:

@article{cywinski2025saeuron,
  title   = {SAeUron: Interpretable Concept Unlearning in Diffusion Models with
             Sparse Autoencoders},
  author  = {Cywi\'nski, Bartosz and Deja, Kamil},
  journal = {arXiv preprint arXiv:2501.18052},
  year    = {2025}
}

@article{zhang2024unlearncanvas,
  title   = {UnlearnCanvas: A Stylized Image Dataset to Benchmark Machine
             Unlearning for Diffusion Models},
  author  = {Zhang, Yihua and others},
  journal = {arXiv preprint arXiv:2402.11846},
  year    = {2024}
}

6. License

Released under the MIT License. Portions derived from SAeUron remain under Apache-2.0 — see NOTICE, which must be retained when redistributing.