Selective-Finetuning-For-Targeted-and-Robust-Concept-Unlearning

TRUST: Targeted and Robust Concept Unlearning via Concept Neuron Localisation

This repository contains the official implementation for Selective Fine-Tuning for Targeted and Robust Concept Unlearning, a method for removing specific concepts (e.g. objects, styles, NSFW content) from text-to-image diffusion models — without retraining the whole network and without degrading the model’s general generation quality.

Instead of fine-tuning the entire U-Net, we first localise the small set of cross-attention “concept neurons” responsible for generating a target concept, and then selectively fine-tune only those parameters using one of two novel objectives:

CIP — Concept Influence Penalty, which directly suppresses the activations of the identified concept neurons (a strong, “hard” unlearning signal).
CSR — Concept Sensitivity Reduction, a softer regularizer that reduces the gradient sensitivity of the model with respect to the concept, while better preserving unrelated concepts.

Both objectives are paired with a preservation loss computed on a held-out set of unrelated concepts/images (e.g. MS-COCO) to limit collateral damage to the rest of the model.

How it works

Concept Neuron Discovery — For a target concept c_u (e.g. “Dog”), we generate sample images of the concept with the base model and compute a per-attention-head gradient-based saliency/importance score with respect to a concept-relevant loss (CLIP-similarity or noise-prediction error). Heads whose importance score exceeds μ + ξ·σ (z-score threshold, default ξ = 2.0) are marked as concept neurons.
Dynamic Mask-Guided Fine-tuning — Only the cross-attention Q/K/V projection weights belonging to the identified concept neurons are updated. The concept-neuron mask is recomputed at every training step, so the set of fine-tuned parameters adapts as the model changes during unlearning.
Unlearning Objective — One of CIP (concept_neurons_loss) or CSR (concept_weight_grad_loss) is combined with a preservation loss (preservation_loss) over a set of concepts/images that should remain unaffected (e.g. MS-COCO).

Pipeline overview

Setup

python3 -m venv cn_env
source cn_env/bin/activate
pip install -r requirements.txt
pip install -e .

Installing the package in editable mode (pip install -e .) makes concept_unlearning importable from anywhere, which is what the scripts under scripts/ rely on.

scripts/env_setup.sh configures the Hugging Face cache directories and other environment variables — adjust the paths inside it for your machine before running.

Usage

Run scripts from the repository root so that the relative default paths below (output directory, concept CSVs, etc.) resolve correctly.

The main entry point is scripts/start_finetuning.sh, which runs concept neuron discovery and selective fine-tuning end-to-end:

chmod +x scripts/start_finetuning.sh
./scripts/start_finetuning.sh -b 2 -s 2 -n 2 -q 2 -l noise_based_loss

Key options

Flag	Description	Default
`-m, --model_id`	Base diffusion model	`runwayml/stable-diffusion-v1-5`
`-c, --concept_to_remove`	Concept to unlearn	`Dog`
`-o, --output_dir`	Output directory for checkpoints/results	`outputs_finetuning`
`-l, --loss_function`	Unlearning objective: `concept_neurons_loss` (CIP), `concept_weight_grad_loss` (CSR), `noise_based_loss`, `clip_loss`, `grad_loss`, `preservation_loss`, `salun_loss`	`concept_neurons_loss`
`-b, --batch_size`	Training batch size	`5`
`-s, --steps_per_epoch`	Number of fine-tuning steps per epoch	`500`
`-n, --num_of_epochs`	Number of epochs	`3`
`-i, --num_of_inference_steps`	Diffusion sampling steps used during training/eval	`40`
`-r, --learning_rate`	Learning rate	`1e-4`
`-t, --threshold_concept_neuron`	Z-score threshold (`ξ`) for concept neuron selection	`2.0`
`-q, --num_of_sample_images_to_generate`	Number of images sampled per step for evaluation	`2`
`-k, --concept_guided`	Whether mask discovery is guided by the target concept	`concept_guided`
`-u, --unlearn_concepts_dataset_path`	Path to annotations CSV for the concept(s) to unlearn	`unlearning_concepts/annotations.csv`
`-p, --preserved_concepts_dataset_path`	Path to annotations CSV for concepts to preserve	`preserved_concepts/annotations.csv`
`-w, --compute_concept_neurons_based_on`	Saliency signal for neuron discovery: `clip_loss` or `noise_loss`	`noise_loss`
`-e, --run_evals_while_finetuning`	Run FID/CLIP/TIFA evaluations during fine-tuning	`True`
`-g, --gpu_id`	GPU ID	`$CUDA_VISIBLE_DEVICES`

Run ./scripts/start_finetuning.sh --help for the full list.

To compute the average CLIP score of a saved checkpoint against a set of COCO prompts:

chmod +x scripts/clip_score_computation.sh
./scripts/clip_score_computation.sh /path/to/saved/pipe

Examples

Conditional and concept-combination unlearning

Unlearning a concept (e.g. “Boy holding a Gun”, “Pregnant woman drinking alcohol”) while preserving the individual concepts it’s composed of:

Conditional unlearning examples

Nudity unlearning robustness

Comparison against SD 1.5, ESD, SalUn, and CoGFD on the I2P nudity benchmark:

Robustness on I2P nudity benchmark

Concept neuron dynamics during fine-tuning

Number of identified concept neurons over training steps (CIP), and the resulting Δ FID / Δ CLIP scores on preserved concepts compared to baselines:

Concept neurons vs steps FID and CLIP score comparison

Citation

If you use this code, please cite the associated paper:

@article{TRUST,
  title   = {Selective Fine-Tuning for Targeted and Robust Concept Unlearning},
  author  = {Mansi and Kori, Avinash and Toni, Francesca and Demetriou, Soteris},
  journal = {arXiv preprint arXiv:2602.07919},
  year    = {2026},
  url     = {https://arxiv.org/pdf/2602.07919}
}