Enabling Preference-driven Unlearning in Few-step Distilled Text-to-Image Diffusion Models

AI summary

Overview

• Research area: Computer vision / generative modeling — specifically safety-oriented machine unlearning for few-step distilled (FSD) text-to-image diffusion models.
• Technical level: Advanced. The paper builds on Direct Preference Optimization, consistency models, and parameterization theory from diffusion literature.
• Scope: A single paper proposing a preference-based unlearning framework (CePU) that removes targeted concepts directly from few-step distilled text-to-image models without re-distillation, evaluated on identity removal, nudity suppression, and (as claimed) object-level unlearning.

What This Paper Is About

Text-to-image diffusion models are increasingly compressed into "few-step" versions that generate images in 2–8 sampling steps instead of 50–100, for speed. These distilled models can still produce harmful content such as celebrity likenesses or nudity, and existing unlearning methods — which fine-tune the model to "forget" a concept — were designed around the long, multi-step denoising behavior of the original base model. The paper shows that this mismatch ("objective dissonance") makes noise-prediction-based DPO-style unlearning ineffective on few-step models, and proposes a consistency-based preference objective that works directly on the distilled model itself.

Key Contributions

1. Identification of an objective dissonance. The authors show that Diffusion-DPO and its unlearning derivatives score preferences using noise-prediction (epsilon) error, which is only equivalent to the sample-space reconstruction error under a specific parameterization. Few-step distilled models such as Latent Consistency Models use timestep-dependent skip and output coefficients, breaking that equivalence and weakening unlearning.

2. CePU, a consistency-enforced preference-driven unlearning framework. Instead of noise-prediction error, the reward signal is the consistency error — the discrepancy between the model's final-timestep predictions at two adjacent noisy timesteps — which is the inductive bias that few-step models are actually trained to satisfy. A derived objective, C-DPO, plugs this signal into the DPO preference formulation, and a retention regularizer keeps non-targeted concepts anchored to the reference model.

3. Direct post-distillation unlearning. Rather than unlearning the base model and re-distilling, CePU operates on the deployed FSD artifact. The paper reports that this achieves performance comparable to unlearn-then-redistill (Appendix D) while taking roughly 15 minutes on a single A5000 GPU — against cited re-distillation costs of 32 A100 GPU hours (Luo et al.) and 36 hours on 72 A100 GPUs (Yin et al.).

4. Broad evaluation with trade-off analysis. Experiments span identity removal (Angelina Jolie, Brad Pitt), nudity suppression across red-teaming prompt sets, and an extension to object-level unlearning, with Pareto-frontier analysis over the DPO temperature parameter beta for utility versus forgetting.

Main Findings

• Standard preference unlearning underperforms on FSD models. The paper attributes weaker forgetting by DPO, DUO, and SafetyDPO on few-step models to the reward signal no longer tracking sample-space reconstruction behavior, illustrated in the LPIPS-versus-timesteps comparison where the FSD variant shows steep early decay and commits within a few steps.

• Identity unlearning (DreamShaper-V8-LCM, Angelina Jolie). CePU reaches a target identity accuracy of 0, non-target accuracy of 55.08, a harmonic mean of 71.03, and LPIPS of 0.0287, versus the unmodified FSD baseline's 60.94 target accuracy, 67.97 non-target accuracy, 49.61 harmonic mean, and LPIPS of 0.

• Identity unlearning (DreamShaper-V8-LCM, Brad Pitt). CePU reaches target accuracy 0, non-target accuracy 58.59, harmonic mean 73.89, and LPIPS 0.0290.

• Identity unlearning (DreamShaper-V7-LCM). For Angelina Jolie, CePU reaches target accuracy 18.75, non-target accuracy 49.61, harmonic mean 61.60, LPIPS 0.0176; for Brad Pitt, 7.81, 57.03, 70.47, and 0.0158 respectively. These are described as best or second-best among preference-based methods, with consistently lower LPIPS.

• Nudity unlearning (DreamShaper-V8-LCM). NudeNet detections on red-teaming prompts drop from a total of 3120 for the unmodified FSD baseline to 130 for CePU, a Defense Success Rate of 95.83, with LPIPS 0.0834 and CLIP score 0.3040. Among preference-based methods, CePU has the lowest detection total (158 for PSO, 217 for SafetyDPO, 260 for DPO, 301 for DUO) and the highest DSR on this model.

• Nudity unlearning (SD-Turbo). CePU produces 6 total detections with a DSR of 99.32, LPIPS 0.0794, and CLIP score 0.3122, versus 876 detections for the unmodified baseline. SafetyDPO is close on this model at 16 detections and a DSR of 98.17.

• Perceptual utility is preserved better than by many baselines. The paper reports that competing methods often degrade realism or attributes such as pose, hairstyle, and lighting, while CePU retains these.

• Existence of a utility–unlearning trade-off. Pareto curves across beta values of 100, 250, 500, and 1000 (and additionally 2000 for nudity) show that increasing beta favors utility retention through stronger alignment to the reference model, and that CePU attains a more favorable frontier than other preference-based methods.

• Unsafe generations occur even without explicit prompts for FSD models. The paper reports elevated pre-unlearning detection rates in FSD models relative to their base models, which it uses as justification for a stringent evaluation setting.

Methodology in Plain English

The researchers start from the observation that few-step distilled models generate differently from their base models: they commit to an output early rather than refining it over many steps. Preference-based unlearning methods such as Diffusion-DPO work by comparing how well a model predicts a "preferred" (safe) image versus a "dispreferred" (unsafe) one, using the error in noise prediction as the score. The authors prove that this score is only equivalent to actual sample-space reconstruction error under a particular parameterization, and that few-step models — for example LCM-style models that use timestep-dependent coefficients — do not satisfy that condition. Using the sample-space error directly does not help either, because it ignores the step-to-step consistency that few-step models rely on.

Their fix replaces the reward with a consistency error: compare what the model predicts as the final clean output at one noisy timestep with what it predicts at a slightly less noisy adjacent timestep, and penalize disagreement. They derive a decomposition showing the relative likelihood splits into a timestep-local consistency term and a cross-timestep interaction term; they keep the local term because it can be estimated from a single sampled adjacent pair, and use a Jensen-inequality upper bound so that a single random timestep can be sampled per training step — preserving the stochastic training style of DPO. The resulting C-DPO loss suppresses the dispreferred concept, while an added retention regularizer compares the updated model's few-step prediction on non-target prompts against the frozen reference model, with a weight lambda that controls the forgetting-versus-utility trade-off. In practice, the two adjacent noisy states come from forward-noising the same clean image with a shared noise realization, and the lower-noise prediction is treated as a stop-gradient target.

Training pairs are built without human annotation: an unsafe image is generated for a dispreferred prompt, then SDEdit is applied at a mid-range noise strength with a preferred counterfactual prompt to produce a matching safe image that keeps the layout but drops the targeted concept. Evaluation uses the LCM scheduler with 4 sampling steps and classifier-free guidance disabled, on DreamShaper-V8-LCM, DreamShaper-V7-LCM, and SD-Turbo, against preference-based baselines (DPO, DUO, PSO, SafetyDPO/AlignGuard), non-preference fine-tuning baselines (ESD variants, CA, FADE), and training-free zero-shot methods (UCE, SAFREE).

Why This Matters

Research impact. The paper reframes a practical failure — preference-based unlearning not working on distilled models — as a formal objective mismatch grounded in parameterization theory (Lemma 1, Propositions 1 and 2). This gives the field a principled reason to revisit the reward signal used in diffusion DPO, not just its hyperparameters, and argues that post-distillation unlearning is a legitimate target rather than an afterthought to base-model unlearning.

Real-world applications.

• Removing specific celebrity likenesses from open-weight, fast-generation image models deployed in consumer apps.
• Suppressing nudity and other NSFW content in few-step models used behind prompt filters that can be bypassed through obfuscation or multilingual prompts.
• Compliance-driven model editing where a model owner must demonstrate a specific concept has been removed after deployment.
• Extending the same mechanism to object-level or brand-level removal as an artist- or trademark-rights tool.

Industry relevance. The reported cost of roughly 15 minutes on a single A5000 GPU, versus the cited 32 A100 GPU hours and 36 hours on 72 A100 GPUs for re-distillation pipelines, matters for teams that ship distilled models and cannot afford to re-run distillation for every concept-removal request. The retention regularizer's adjustable weight also gives a controllable dial for the safety-versus-quality trade-off that product teams must manage.

Future Directions

• Strengthening the surrogate's theoretical status. The paper states that the consistency-based reward is a heuristic surrogate rather than an approximation with guaranteed sign or ordering, and leaves its practical validity to empirical evaluation. A tighter analysis of when the omitted cross-timestep term is small enough to ignore is an open question.
• More systematic preference-pair construction. The authors note that paired dataset generation is orthogonal to their contribution and is handled by SDEdit, pointing to complementary work on better preference-pair construction for concept removal.
• Broader concept coverage. Object-level unlearning is presented as an extension; the truncated paper text does not report its quantitative results, and the framework's behavior on other concept types (styles, copyright characters, violence) is not covered.
• Generalization across distillation families. Evaluation covers DreamShaper-V8-LCM, DreamShaper-V7-LCM, and SD-Turbo. Whether the consistency reward transfers to other few-step distillation methods and step counts beyond the 4-step LCM scheduler setting is not established.

Target Audience

Researchers and engineers working on diffusion model safety, machine unlearning, and preference optimization, particularly those deploying few-step or consistency-distilled text-to-image models. It is most useful to readers already comfortable with DPO for diffusion models and with consistency model training; readers new to these topics will need the background sections and appendices to follow the derivations. Practitioners building content-moderation pipelines for fast generative image systems will find the cost and retention trade-off analysis the most directly applicable part.

Authors’ 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.