Few-step diffusion models learn to forget unwanted image content safely
Enabling Preference-driven Unlearning in Few-step Distilled Text-to-Image Diffusion Models
Computer Vision and Pattern RecognitionMachine Learning
Summary
Text-to-image AI models that create pictures quickly can sometimes make harmful or unwanted images. The authors found that existing ways to make these models forget specific things don’t work well when the models are made to generate images faster with fewer steps. They created a new method that helps these quick models forget unwanted content, like nudity or certain identities, while still working fast and keeping good abilities for other images. Their tests show this method effectively removes targeted content without slowing down the model.
What this means in practice
- •For ai platform engineers: Enable fast text-to-image models to forget specific unwanted contents like nudity or identities directly without full retraining.
- •For content moderation teams: Remove harmful or inappropriate image generation in deployed AI models rapidly while preserving overall model performance.
Authors
Gaurav Patel, Jun Fang, Greg Ver Steeg, Qiang Qiu, Sravan Sripada
Abstract
Text-to-image diffusion models are increasingly distilled into few-step variants and being deployed to enable fast inference. However, their ability to generate harmful or undesired content poses significant safety risks. Data-driven unlearning methods suppress targeted generations by fine-tuning model weights using specialized unlearning objectives. Crucially, these objectives implicitly rely on multi-step denoising dynamics, an assumption that breaks down for few-step distilled (FSD) models, resulting in ineffective forgetting. Furthermore, performing unlearning on the non-distilled base model and subsequently re-distilling it to obtain an unlearned FSD model incurs substantial computational and time overhead, making it impractical in many settings. Hence, we address this limitation with a preference-driven unlearning framework that revisits Direct Preference Optimization (DPO) for diffusion models. We show that standard DPO and its unlearning derivatives, formulated around noise-prediction error, transfer poorly to FSD models due to their altered generation dynamics. To overcome this, we introduce a modified preference optimization formulation explicitly aligned with the few-step generation properties, enabling direct concept removal in FSD models while preserving few-step efficiency and maintaining strong retention of desirable (non-targeted) capabilities. We evaluate our framework primarily on identity and NSFW (nudity) removal tasks and also extend our method to object-level unlearning. Extensive experiments demonstrate consistent and effective forgetting, and strong retention performance, establishing our method as a practical and principled solution for unlearning in FSD models.