Selective Fine-Tuning for Targeted and Robust Concept Unlearning

TL;DR

TRUST is a dynamic, Hessian-regularized selective fine-tuning method that effectively unlearns harmful concepts in diffusion models, improving robustness and efficiency over existing static approaches.

Contribution

The paper introduces TRUST, a novel dynamic approach for targeted concept unlearning that outperforms state-of-the-art methods in robustness, speed, and flexibility without additional regularization.

Findings

TRUST effectively unlearns individual and combined concepts.

It is more robust against adversarial prompts.

TRUST is significantly faster than existing methods.

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.