Selective Fine-Tuning for Targeted and Robust Concept Unlearning

AI summary

Overview

Research area: Machine unlearning and AI safety for text-to-image (T2I) diffusion models, specifically the targeted removal of harmful concepts from cross-attention layers.

Technical level: Advanced. The paper assumes familiarity with latent diffusion models, cross-attention mechanisms, classifier-free guidance, saliency maps, Hessian-based regularization, and standard generative evaluation metrics (FID, CLIP, TIFA).

Scope: The paper introduces TRuST (Targeted Robust Selective fine-Tuning), a framework that dynamically estimates concept-bearing neurons in a diffusion model's cross-attention layers and selectively fine-tunes them under two new regularization objectives to erase harmful concepts, concept combinations, and conditional concepts while preserving unrelated generation quality.

What This Paper Is About

Text-guided diffusion models such as Stable Diffusion v1.5 are widely used but can be exploited to generate harmful content, and prior work has identified 22 potential risks posed by T2I diffusion models (Bird et al., 2023). Existing concept unlearning methods either fine-tune the entire model, which is computationally expensive and untargeted, or rely on static neuron localization that becomes outdated as the model adapts. TRuST addresses both problems by dynamically re-identifying the concept neurons after every parameter update and applying selective fine-tuning with a Hessian-informed regularization, so that harmful concepts and harmful combinations of benign concepts can be erased without degrading the rest of the model.

Key Contributions

1. Dynamic concept neuron identification. A novel gradient-based method that finds the parameters encoding a target semantic concept inside the cross-attention key, query, and value projection matrices, using an alignment objective (the CLIPScore between a concept prompt and its generated image) rather than the predicted-noise correlation used by prior work. This localizes not just individual concepts but also concept combinations.

2. Two new unlearning objectives. The Concept Influence Penalty (CIP), a "hard" unlearning loss that penalizes the cardinality of high-influence concept neurons to induce sparsity, and Concept Sensitivity Reduction (CSR), a "soft" unlearning loss that minimizes the sensitivity of the predicted noise to those parameters and incorporates second-order curvature information, in each case paired with a preservation loss over benign concepts.

3. An end-to-end framework with dynamic mask-guided fine-tuning. The neuron masks are recomputed after every parameter update to account for representation shift during unlearning, and the framework is benchmarked against a broad set of state-of-the-art model editing and model steering baselines, plus ablation studies.

4. Extension to two harder unlearning tasks. Concept Combination Erasure (CCE) and the novel Conditional Concept Erasure (CoCE), where no prior baselines exist and TRuST is reported as a standalone result.

Main Findings

• Robustness to adversarial prompts. On the harmful concept "Nudity" evaluated with the I2P dataset, TRuST reduces Attack Success Rate to 0.11%, which the paper reports as a 45% ASR decrease compared to SalUn. The paper states degradations in ASR of 18.52%, 12.1%, 10.57%, 14.33%, and 1.89% over the UnlearnDiffAtk, MMA-Diffusion, Ring-a-Bell and P4D baselines (the paper lists five figures against four named techniques). In Table 2, TRuST with CIP reaches ASR of 0.0011 (I2P), 0.0027 (P4D), 0.0083 (Ring-a-Bell), 0.062 (MMA-Diffusion) and 0.0118 (UnlearnDiffAtk); TRuST with CSR reaches 0.0013, 0.0046, 0.0093, 0.077 and 0.0175 respectively.

• Utility preservation. TRuST keeps non-targeted generation quality nearly unchanged. The abstract reports ΔFID = 0.02; in Table 2, the ΔFID is 0.030 for CSR and 0.016 for CIP, and the text states the change is "as low as 0.016," described as better than any existing baseline and considerably better than SalUn. TIFA scores are 0.811 (CSR) and 0.808 (CIP) versus 0.813 for the unmodified SD-v1.5, and CLIP scores are 30.43 and 30.95 versus 31.3 for SD-v1.5. For comparison, SalUn reports ΔFID of 27.31 and a CLIP score of 26.2.

• Both objectives beat baselines, with different profiles. CIP is more robust to adversarial prompts and better at erasing target concepts but slightly reduces output fidelity for benign concepts; CSR preserves realism and semantic coherence better, especially for concept combinations, but is less aggressive at unlearning.

• Concept combination erasure (CCE). In the CLIP-score-versus-fine-tuning-step analysis, TRuST retains individual concept alignment across fine-tuning, outperforming the state-of-the-art CoGFD, and achieves effective unlearning of concept combinations in nearly half the number of steps required by CoGFD.

• Conditional concept erasure (CoCE). Using the conditional prompt "Cat on the Table," TRuST unlearns that combination while preserving the semantically related prompt "Cat under the Table." High TIFA scores for individual concepts such as "Cat" and "Table" indicate the unlearning does not spill over to adjacent concepts; the paper notes that TIFA scores of 1.0 for individual prompts are expected because they measure basic presence detection.

• Scalability. Averaged Unlearning Accuracy and Retaining Accuracy are reported as outperforming all baselines as the number of erased concepts grows. For four concept combinations, the paper reports only a 0.02 drop in TIFA score and less than a 1-point increase in FID (Appendix E, Table 7).

• Efficiency. Table 3 reports 60 fine-tuning steps, 350 images and 0.25 hours for TRuST (CSR), and 100 steps, 350 images and 0.72 hours for TRuST (CIP), against 1300 steps / 800 images / 2 hours for SalUn, 150 steps / 0.5 hours for CoGFD, 1000 steps / 540 images / 2.5 hours for ESD, and 80,000 steps / 10^7 images / 150,000 hours for retraining. The text describes the CSR strategy as cutting fine-tuning to 60 steps — less than half of CoGFD's and a reduction from 1300 to 60 steps for standard concept unlearning — and also states that TRuST with CIP loss reaches better performance than other state-of-the-art baselines in "just 15 minutes of finetuning," which is two times faster than the then-current state of the art (note that Table 3 lists 0.25 hours for CSR and 0.72 hours for CIP).

• Noise-based grounding is inefficient. The paper reports that localization grounded in predicted noise rather than concept alignment requires more than 5x the fine-tuning steps and 8x the wall-clock time to unlearn targeted concepts, even when improved with dynamic localization.

• Three challenges motivate the design. Catastrophic interference (removing an unsafe concept perturbs neighboring representations, harming benign-concept utility), saliency shift during optimization (the number of concept neurons changes after every fine-tuning step under SalUn's static assumption), and computational efficiency (slow convergence on entangled or combinatorial concepts).

Methodology in Plain English

The starting point is the observation that harmful outputs often come not from a single bad word but from combinations of individually harmless ones — the paper's example is "child" and "beer," which are benign alone but harmful in combination such as "child drinking beer." Because diffusion models compose concepts readily, unlearning has to be precise enough to sever a harmful association without erasing the benign ingredients.

TRuST works inside the cross-attention layers of a latent diffusion model, where text and image information are aligned. The authors define a "concept neuron" as a single parameter in the key, query, or value projection matrices that carries information about a concept. To find which parameters matter, they generate an image from the concept prompt, score how well the image matches the prompt using CLIPScore, and take the gradient of that score with respect to the parameters. Parameters whose gradient magnitude exceeds a threshold set as ξ times the standard deviation plus the mean of the gradient matrix (with ξ = 2.0 in all experiments) are marked as concept neurons, forming a saliency mask.

Unlearning then proceeds on two alternative paths. The Concept Influence Penalty directly penalizes the number of concept neurons that remain influential, pushing the model toward a sparse set — the paper calls this hard unlearning. Concept Sensitivity Reduction instead minimizes how sensitive the predicted noise is to those parameters, so their influence is weakened rather than removed — the paper calls this soft unlearning, and because gradients of gradients are involved, the Hessian of the noise-prediction loss governs the local curvature implicitly. Both objectives are combined with a preservation loss that keeps the model faithful on a set of benign prompts, guarding against unintended forgetting caused by classifier-free guidance sharing parameters.

A key practical decision is that the concept neuron set is recomputed after every parameter update rather than fixed before training. This tracks the representation drift that occurs during unlearning and prevents the method from overfitting to an outdated subset of neurons. Experiments use Stable Diffusion v1.5, with 1000 real-world images from the MS COCO 30k validation set for the preservation loss and a COCO subset with the target concept removed for evaluating retained concepts.

Why This Matters

Impact on research. The paper reframes concept localization as a dynamic rather than static problem and replaces noise-prediction-based saliency with a semantic alignment objective, offering a general recipe for neuron-level editing. It also introduces Conditional Concept Erasure as a task with no prior baselines, and provides a seven-dimension comparison table positioning prior methods across methodology-driven dimensions (weight modification, training-free operation, anchor-free operation, targeted tuning of concept neurons or layers) and use-case dimensions (multiple concepts, concept combinations, conditional concepts), with TRuST as the only method satisfying all of the latter three.

Real-world applications.

• Safety filtering for publicly released image generators, where harmful concepts and adversarial prompts must be neutralized in the deployed weights rather than by output filters that can be paraphrased around.
• Publisher and platform compliance workflows, so that models trained on web-scale data such as LAION can be adjusted after the fact without full retraining.
• Moderation of harmful combinations and contextual associations, such as prompts that are innocuous word-by-word but harmful together.
• Cost-constrained deployment, where erasure must fit within small compute and data budgets.

Industry relevance. Full retraining is listed at 80,000 steps, 10^7 images and 150,000 hours, and full fine-tuning baselines such as SalUn at 1300 steps and 800 images. TRuST's reported 60 steps, 350 images and 0.25 hours make post-hoc safety editing of large generative models substantially more practical for teams that cannot retrain foundation models. The authors also note TRuST is architecture-agnostic because it operates on gradients and saliency in projection weights rather than a specific diffusion backbone.

Future Directions

• Extending the approach beyond diffusion to other generative paradigms, which the authors name explicitly: flow matching models, stochastic interpolants, and large language models.
• Establishing baselines and benchmarks for Conditional Concept Erasure, the task the paper introduces without prior comparators.
• Further characterizing when to prefer CIP versus CSR, given their different trade-offs between adversarial robustness and fidelity on benign concepts (the paper points to Appendix A.5 for an extended discussion).
• Investigating the sensitivity of the gradient threshold hyperparameter ξ, where the paper fixes ξ = 2.0 and defers ablation details to Appendix A.11.3.

Target Audience

Researchers and practitioners working on machine unlearning, generative model safety, and diffusion model interpretability, particularly those who need to modify released T2I models without full fine-tuning. It is also relevant to safety engineers and policy-facing teams who must reason about adversarial robustness of erasure methods, and to readers interested in neuron-level model editing and localization techniques. Because the methods involve Hessian-based regularization and cross-attention internals, the paper is best suited to readers with an intermediate-to-advanced background in deep generative modeling.

Authors’ abstract

Text guided diffusion models are used by millions of users, but can be easily exploited to produce harmful content. Concept unlearning methods aim at reducing the models' likelihood of generating harmful content. Traditionally, this has been tackled at an individual concept level, with only a handful of recent works considering more realistic concept combinations. However, state of the art methods depend on full finetuning, which is computationally expensive. Concept localisation methods can facilitate selective finetuning, but existing techniques are static, resulting in suboptimal utility. In order to tackle these challenges, we propose TRUST (Targeted Robust Selective fine Tuning), a novel approach for dynamically estimating target concept neurons and unlearning them through selective finetuning, empowered by a Hessian based regularization. We show experimentally, against a number of SOTA baselines, that TRUST is robust against adversarial prompts, preserves generation quality to a significant degree, and is also significantly faster than the SOTA. Our method achieves unlearning of not only individual concepts but also combinations of concepts and conditional concepts, without any specific regularization.