Forget-It-All: Multi-Concept Machine Unlearning via Concept-Aware Neuron Masking

Forget-It-All: Multi-Concept Machine Unlearning via Concept-Aware Neuron Masking

Kaiyuan Deng¹, Bo Hui³, Gen Li², Jie Ji², Minghai Qin⁵, Geng Yuan⁴, Xiaolong Ma¹

¹The University of Arizona, ²Clemson University, ³The University of Tulsa, ⁴University of Georgia, ⁵Western Digital Corporation

---

Introduction

Existing single-concept unlearning methods for text-to-image diffusion models break down in real-world settings that require erasing multiple, interrelated concepts at once: sequential application either re-acquires forgotten concepts or sacrifices generative quality, and recent multi-concept methods rely on fine-tuning, LoRA stacking, or LLM-derived concept graphs that are hyperparameter-sensitive and computationally heavy. We propose FIA (Forget-It-All), a training-free framework that unlearns arbitrary sets of concepts by pruning a small subset of neurons. FIA introduces a Contrastive Concept Saliency that quantifies each weight's contribution to a target concept, identifies Concept-Sensitive Neurons by combining temporal-sparsity and spatial-sparsity selection, and finally constructs per-concept masks that are fused under a concept-agnostic neuron retention rule so that broadly useful neurons are preserved while concept-specific ones are pruned. FIA needs only minimal hyperparameter tuning, achieves state-of-the-art unlearning at under 0.3% overall sparsity, and forgets each concept in roughly 11 seconds on a single A6000.

FIA simultaneously unlearns objects, NSFW content, and artistic styles while preserving overall generative quality (paper Figure 1).

Three-stage pipeline. Stage I computes Contrastive Concept Saliency from paired (concept, base) prompts, then aggregates over time and space to identify Concept-Sensitive Neurons. Stage II fuses per-concept masks while retaining Concept-Agnostic Neurons (paper Figure 2).

---

Installation

Prerequisites

• Anaconda
• NVIDIA GPU with CUDA 12.1 (or above) support — a single A6000 (48 GB) is sufficient; explicit and artist tasks fit on a 24 GB card

Setup

# Clone the repository
git clone https://github.com/<your-org>/FIA.git
cd FIA/official

# Create conda environment
conda create -n fia python=3.10 -y
conda activate fia

# Install PyTorch (CUDA 12.1)
pip install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu121

# Install remaining dependencies
pip install -r requirements.txt

Verify

python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')"

---

Usage

Concept Erasing

Multi-Object Unlearning (Imagenette)

python main.py \
    --config configs/object.yaml \
    --concepts "parachute, golf ball, garbage truck, cassette player, church, tench, \
                english springer, french horn, chain saw, gas pump" \
    --output_dir runs/object_10

Multi-Artistic-Style Unlearning

python main.py \
    --config configs/art.yaml \
    --concepts "Van Gogh, Monet, Pablo Picasso, Leonardo Da Vinci, Salvador Dali" \
    --output_dir runs/art_5

Explicit (NSFW) Content Unlearning

python main.py \
    --config configs/explicit.yaml \
    --concepts "naked" \
    --output_dir runs/explicit

Each pipeline run writes the unlearned UNet to runs/<task>/edited_unet.safetensors, the per-concept saliency tensors to runs/<task>/saliency/<concept>/, and a fusion_stats.json summarizing per-layer pruning ratios.

Evaluation

Object Erasure Evaluation

python evaluation/eval_object.py \
    --ckpt runs/object_10/edited_unet.safetensors \
    --csv  datasets/imagenette.csv \
    --out_dir runs/object_10/eval_imagenette

Nudity / Explicit Content Evaluation

python evaluation/eval_explicit.py \
    --ckpt runs/explicit/edited_unet.safetensors \
    --model_id CompVis/stable-diffusion-v1-4 \
    --eval_dataset i2p \
    --out_dir runs/explicit/eval_i2p \
    --max_prompts 300

Artistic Style Evaluation

python evaluation/eval_artist.py \
    --ckpt runs/art_5/edited_unet.safetensors \
    --artists "Van Gogh, Monet, Pablo Picasso, Leonardo Da Vinci, Salvador Dali" \
    --out_dir runs/art_5/eval_artist

COCO CLIP Score (Generative Quality)

python evaluation/eval_coco.py \
    --prompts_file datasets/coco_prompts.txt \
    --ckpt runs/object_10/edited_unet.safetensors \
    --model_id "runwayml/stable-diffusion-v1-5" \
    --max_prompts 200

Visual Sanity Check (paired base vs. edited)

python evaluation/inference.py \
    --ckpt runs/object_10/edited_unet.safetensors \
    --prompt "a photo of a french horn" \
    -n 4

One-Click Reproduction

The three shell scripts under scripts/ chain training and evaluation for each task:

bash scripts/run_object.sh
bash scripts/run_explicit.sh
bash scripts/run_art.sh

---

Citation

If you find this work useful, please cite our paper:

@article{deng2026fia,
  title   = {Forget-It-All: Multi-Concept Machine Unlearning via Concept-Aware Neuron Masking},
  author  = {Deng, Kaiyuan and Hui, Bo and Li, Gen and Ji, Jie and Qin, Minghai and Yuan, Geng and Ma, Xiaolong},
  journal = {arXiv preprint arXiv:2601.06163},
  year    = {2026}
}

---

Acknowledgements

We thank the authors of ESD, UCE, MACE, SPM, FMN, AC, SalUn, and ConceptPrune for releasing their code, which we use as baselines and evaluation harnesses. We also rely on NudeNet, the I2P benchmark, and Imagenette.