PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Gaurav Patel, Jun Fang, Greg Ver Steeg, Qiang Qiu, Sravan Sripada

arXiv:2610.10859v1cs.CVcs.LG

TL;DR

Few-step distilled text-to-image models can inherit or amplify unsafe content, while standard noise-prediction preference objectives are poorly matched to their generation dynamics. The paper introduces CePU, a consistency-based preference-unlearning framework, and reports effective targeted forgetting with utility retention across identity and nudity tasks, plus an object-level extension.

  • Problem

    Noise-prediction-based preference unlearning assumes full-step denoising trajectories and noise-based parameterization, which do not match few-step distilled models; unsafe content can also transfer or amplify during distillation.

    Source: §1 · p. 1

  • Method

    CePU replaces noise-prediction rewards with consistency-error rewards to align preference-driven unlearning with few-step generation dynamics.

    Source: §1 · p. 2

  • Results

    CePU achieves strong targeted forgetting while preserving overall generation utility, with Pareto-optimal trade-offs reported for identity removal and nudity unlearning under diverse red-teaming prompts.

    Source: §5 · p. 10

  • Takeaways & Limitations

    Direct unlearning on an FSD model achieves performance comparable to base-model unlearning followed by distillation and takes about 15 minutes on one A5000 GPU.

    Source: §1 · p. 2

    The consistency-based reward is a heuristic surrogate without a guaranteed sign or ordering relative to the likelihood reward; its practical validity is evaluated empirically.

    Source: §2.4 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Few-step distillation changes generation dynamics and can transfer or amplify unsafe content, while standard noise-prediction preference objectives are poorly matched to FSD models. CePU aligns preference-driven unlearning with few-step dynamics and is reported to forget targeted concepts while retaining utility.

1 Introduction

  • FSD models reduce sampling from 50–100 steps to 2–8 and commit early; undesirable residues can transfer from the base model and even be amplified.The paper argues that unlearning on the base model and then re-distilling has uncertain transfer and substantial computational overhead.Source: §1 · pp. 1–2
  • Noise-prediction preference objectives assume full-step trajectories and noise-based parameterization, making them poorly matched to FSD models and ineffective for unlearning.The paper identifies this objective dissonance as the key bottleneck.Source: §1 · p. 2
  • CePU achieves performance comparable to base-model unlearning followed by distillation, while directly unlearning on an FSD model takes ≈15 minutes on one A5000 GPU.The introduction reports evaluations on identity and nudity removal, with an extension to object-level unlearning.Source: §1 · p. 2
  • CePU replaces noise-prediction error with consistency error as a preference reward aligned with few-step generation dynamics.The consistency error compares final-timestep predictions from adjacent noisy timesteps.Source: §1 · p. 2
  • The paper reports effective unlearning across identity, nudity, and object-level settings, alongside analysis of utility–unlearning trade-offs across methods and model families.Nudity evaluation includes red-teaming prompts, and the introduction also notes an extension to object-level unlearning.Source: §1 · p. 2

2 Background, Observation, and Methodology

The section explains why noise-prediction-based preference objectives can mismatch few-step distilled models, then develops consistency-aware preference unlearning with a retention constraint. It also describes how SDEdit-generated counterfactual pairs provide preference supervision.

2.1 Unlearning as Preference Optimization

  • Standard diffusion DPO contrasts preferred and dispreferred samples through relative likelihoods, with diffusion formulations using noise-prediction error as a proxy.In unlearning, both likelihoods are conditioned on the undesirable-concept prompt so optimization favors the safe alternative in that context.Source: §2.1 · p. 3

2.2 Why Preference-based Unlearning Fails for Few-Step Distilled Models?

  • Noise-prediction error can cease to track reconstruction error under FSD parameterizations, weakening preference-based unlearning despite identical preference supervision.Using x0-error instead also risks ineffective unlearning because it does not preserve the few-step inductive bias.Source: §2.2 · p. 3

2.3 Rectifying Dissonance via Consistency

  • Consistency error provides a tractable surrogate that upper-bounds x0-error while preserving the few-step inductive bias.The section motivates consistency by decomposing sample-space error along the denoising trajectory.Source: §2.3 · p. 4

2.4 Consistency-enforced Preference-driven Unlearning (CePU)

  • The consistency-aware preference objective samples a timestep and adjacent transition rather than evaluating the complete denoising trajectory.A Jensen-based per-step upper bound enables stochastic single-step training.Source: §2.4 · p. 5
  • CePU replaces the noise-prediction reward with relative step-to-step consistency against a reference model, aligning preference optimization with few-step generation dynamics.Its likelihood-reward surrogate is heuristic and has no guaranteed sign or ordering relative to the original likelihood; the method retains timestep-local terms because they can be estimated from adjacent transitions, unlike the O(m^2) cross-timestep interactions.Source: §2.4 · p. 5
  • CePU combines consistency-aware preference suppression with a retention regularizer that anchors non-target predictions to the reference model.The regularizer helps control the utility–unlearning trade-off while preserving few-step fidelity on non-target concepts.Source: §2.4 · pp. 5–6

2.5 Generating Paired Samples

  • SDEdit produces preferred counterfactuals from dispreferred images while preserving global scene structure and selectively changing concept-specific attributes.The paired samples supervise CePU and preference-based baselines using a standard, fully bootstrapped construction.Source: §2.3 · p. 4, §2.5 · p. 6

3 Experiments and Observations

Across identity and nudity tasks, CePU shows effective targeted forgetting while preserving utility; the experiments also extend to object-level unlearning and examine trade-offs and robustness.

3 Experiments and Observations

  • The benchmark compares preference-based, fine-tuning, and zero-shot baselines on three SD-based FSD models using four LCM sampling steps.Classifier-free guidance is disabled.Source: §3 · p. 6

3.1 Identity Unlearning

  • For identity unlearning, CePU attains superior target forgetting with the highest utility across the β Pareto curves.The qualitative comparison reports preserved non-target semantics, while competing methods can degrade realism, pose, hairstyle, and lighting.Source: Figure 3 · p. 7

3.2 Nudity Unlearning

  • CePU achieves superior DSR across DreamShaper-V8-LCM and SD-Turbo while maintaining low LPIPS and a more favorable DSR–utility frontier.Qualitative results also report improved suppression under implicit prompts, where baselines often leak unsafe content.Source: §3.2 · pp. 7–8, Figure 4 · p. 8, Figure 5 · p. 8

3.3 Object-level Unlearning

  • For object unlearning, CePU omits the positive preference term because there is no canonical preferred counterpart for a target object.This removes the need for coupled preference pairs; the consistency-based regularizer is retained to preserve adjacent concepts.Source: §3.3 · p. 8
  • On object-level unlearning of Imagenette [8] classes, CePU achieves Hmean 75.66% versus FADE [83] at 74.42%, with lower LPIPS (0.09 versus 0.12), indicating improved utility preservation.Source: §3.3 · p. 8, Table 3 · p. 9

3.4 Additional Analysis and Observations

  • Increasing λ improves non-target utility but weakens unlearning, while larger β weakens updates and improves utility retention.The analysis reports a strong identity balance at λ=100 and the best I2P nudity DSR–utility frontier at λ=10; CePU remains effective at one sampling step.Source: §3.4 · p. 9
  • Varying SDEdit edit strength τ only marginally affects I2P detections and LPIPS.SDEdit creates counterfactual pairs by perturbing an image at t=⌊τ × T⌋ and denoising back to t=0.Source: §3.4 · p. 9, Figure 8 · p. 9
  • CePU also unlearns effectively on SD-Turbo, which uses score-distillation and adversarial objectives rather than consistency-model distillation.The authors suggest this indicates CePU captures a more general property of few-step distilled models.Source: §3.4 · pp. 9–10

4 Related Works

Prior text-to-image unlearning work includes fine-tuning, pruning, multi-objective, meta-learning, adversarial, and zero-shot approaches. Preference-based methods extend DPO to diffusion and unlearning, while few-step models use varied distillation strategies.

4 Related Works

  • Diffusion unlearning spans noise-prediction fine-tuning, counterfactual objectives, prior-preserving objectives, pruning, multi-objective optimization, meta-learning, adversarial erasure, and zero-shot updates.The evaluated baselines are selected from methods implemented in HuggingFace Diffusers because some CompVis implementations prevent use of HuggingFace FSD checkpoints.Source: §4 · p. 10
  • Diffusion-DPO, DUO, and SafetyDPO apply preference comparisons to diffusion and unlearning, but had not been evaluated on FSD models in the cited work.PSO was adapted as an FSD baseline although it was not designed for unlearning.Source: §4 · p. 10
  • FSD models compress sampling to 2–8 steps using consistency, score-distillation with adversarial losses, distribution matching, score matching, or mean-flow approaches.The passage describes these as distinct strategies for few-step generation.Source: §4 · p. 10

5 Conclusion

The paper identifies a mismatch between noise-prediction preference objectives and FSD generation dynamics, then introduces CePU’s consistency-based rewards. It reports effective forgetting with preserved utility across identity and nudity tasks, and an extension to object-level unlearning.

5 Conclusion

  • Noise-prediction error-based preference objectives mismatch the few-step generation dynamics of FSD models.The conclusion frames this as a core issue in post-distillation unlearning.Source: §5 · p. 10
  • CePU replaces noise-prediction rewards with consistency-error rewards aligned with FSD parameterization.The derivation links preference optimization to per-timestep consistency as a stable, tractable surrogate for unlearning.Source: §5 · p. 10
  • CePU achieves strong targeted forgetting while preserving overall generation utility, with Pareto-optimal trade-offs reported for identity and nudity unlearning.The framework is also extended to object-level unlearning.Source: §5 · p. 10

B Probabilistic Sampling Framework for Diffusion Models

The section places diffusion samplers in a shared marginal-preserving framework, then derives the likelihood-ratio connection between clean-sample and noise-prediction errors.

B Probabilistic Sampling Framework for Diffusion Models

  • DDPM, DDIM, and LCM samplers arise as instances of one reverse process induced by non-Markovian forward distributions that preserve the same marginals.The framework preserves forward marginals without requiring a particular joint factorization.Source: B Probabilistic Sampling Framework for… · p. 17
  • The trainable reverse process predicts the clean signal from a noisy input and uses that estimate in reverse transitions, adding Gaussian noise at the final step.This construction mirrors the non-Markovian inference family while remaining tractable.Source: B Probabilistic Sampling Framework for… · pp. 17–18
  • Accelerated sampling evaluates a subsequence of diffusion times, preserving the marginals while reversing only the selected trajectory.The subsequence has S ≪ T steps.Source: B Probabilistic Sampling Framework for… · p. 18
  • The Gaussian clean-sample conditional is a separate modeling assumption and a coarse approximation of a generally non-Gaussian, multimodal posterior.The likelihood-ratio results in Lemma 4 and Proposition 2 are conditional on this assumption; Diffusion-DPO uses the same choice [85].Source: B Probabilistic Sampling Framework for… · p. 18

C.1 Lemmas

  • Strict convexity of softplus gives a per-step Jensen inequality, with equality only when all inputs are equal.The lemma establishes that softplus is increasing as well as strictly convex.Source: §C.1 · p. 19
  • Under the Gaussian assumption, the implicit likelihood-ratio reward equals a timestep-weighted difference in noise-prediction errors.The positive timestep weight rescales the reward without changing its sign; setting it constant changes per-timestep scale [85].Source: §C.1 · pp. 20–21

C.2 Corollaries

The corollaries bound a multi-step consistency objective and show how choices of variance and parameterization recover classical and few-step samplers.

C.2 Corollaries

  • The per-timestep Jensen bound upper-bounds the consistency preference loss by the average of its per-step losses, with equality when all per-step gaps are equal.The equality condition is stated as uniform consistency gaps across steps.Source: §C.2 · p. 22
  • Appropriate variance and parameterization choices recover DDPM, deterministic DDIM, and the latent consistency or LCM scheduler from the unified reverse update.The corollary identifies these samplers as special cases of the same update.Source: §C.2 · p. 22

3. LCM Sampler

  • For the LCM sampler, the model directly predicts x0 with a consistency parametrization, and the chosen variance makes the reverse-step mean simplify to that prediction scaled by √α_(t−1).This yields the latent consistency scheduler used for few-step distilled models.Source: §3 · pp. 24–25

C.3 Propositions

The propositions connect prediction error to consistency increments and decompose the exact likelihood-ratio reward into local and cross-timestep contributions, motivating a local surrogate with bounded guarantees.

C.3 Propositions

  • The prediction error is bounded by √m times the root-sum-square of the consistency increments across m steps.The bound is stated under the endpoint and modeling assumptions of Proposition 1.Source: §C.3 · p. 25
  • Under the stated assumptions, the exact implicit reward decomposes into timestep-local consistency terms and cross-timestep interactions.The local term is a difference of squared consistency-increment norms; the cross term couples increments from distinct timesteps.Source: §C.3 · p. 25
  • CePU retains only the local reward because cross-timestep interactions are not directly available from a single sampled adjacent transition.The local term can therefore be evaluated with the stochastic per-timestep sampling procedure used during DPO training.Source: §C.3 · p. 27
  • The local surrogate is not guaranteed to match the exact reward or its sign; sign agreement is sufficient when 2|X| < |L|, and the one-step decomposition is exact.The authors state that Proposition 2 motivates the reward rather than providing a convergence or optimality guarantee.Source: §C.3 · p. 27

D Directly Unlearning on FSD Model

Direct CePU unlearning on a few-step distilled model is reported as comparable to unlearning the base model and distilling it, while avoiding the distillation cost. The section also specifies the training and evaluation setup.

D Directly Unlearning on FSD Model

  • ∼15 minutes on one A5000 GPU for direct CePU unlearning compares with ∼24 hours on 4×A5000 GPUs for distillation, while achieving comparable unlearning effectiveness and utility preservation.The comparison is against unlearning the base SD1.5 model and then distilling it; longer distillation may further improve generation quality.Source: D Directly Unlearning on FSD Model · pp. 27–28

E Training and Evaluation Details

  • Preference-driven methods use rank-32 LoRA and AdamW, with β = 100, learning rate 1 × 10−4, batch size 4, and 500 training iterations for nudity and identity removal.LoRA is not applied to cross-attention layers, and the consistency-error timestep offset is k = 20.Source: E Training and Evaluation Details · p. 28
  • Evaluation primarily generates images in 4 steps using the LCM scheduler with classifier-free guidance disabled.Source: E Training and Evaluation Details · p. 28

E.1 Training Dataset Creation

The section describes counterfactual training data, consistency-error computation, and evaluation protocols for identity, nudity, and object-level unlearning. CePU achieves strong object removal while retaining non-target and semantically adjacent concepts.

E.1 Training Dataset Creation

  • SDEdit creates preferred counterparts at τ = 0.6, preserving global scene semantics while changing concept-specific attributes; object-level examples instead use sampled class prompts.The same paired-data protocol is used across preference-based methods.Source: §E.1 · p. 28

E.2 Computing the Consistency Error

  • Adjacent noisy states reuse the same Gaussian noise; gradients update only the noisier-step prediction, with the lower-noise prediction as a stop-gradient target and the reference model frozen.The construction is used for both the consistency-based preference loss and the retention loss.Source: §E.2 · pp. 28–29

E.3 Identity Unlearning

  • Identity pairs edit target-identity images into neutral counterparts; disjoint 64-prompt tests assess target accuracy, identity-agnostic LPIPS, and recognition of four other celebrities.Other-identity retention is evaluated but not explicitly enforced during unlearning.Source: §E.3 · p. 29

E.4 Nudity Unlearning

  • Nudity training uses 800 prompt pairs with nudity terms removed from preferred prompts; held-out red-team sets assess NudeNet DSR, while benign LAION prompts assess CLIP Score and LPIPS.The red-team evaluations include adversarial and sanitized variants; their prompts are not used for training.Source: §E.4 · p. 31

E.5 Object-level Unlearning

  • Object training targets one class while treating the others as retained; evaluation reports target and non-target accuracy, their harmonic mean, and LPIPS.The setup covers Imagenette and fine-grained object categories.Source: §E.5 · p. 32

F.1 Additional Object-Level Unlearning on Imagenette.

  • CePU completely removes all 10 Imagenette targets (Atarget = 0.00), with Aother = 60.85, LPIPS = 0.0900, and the highest Hmean = 75.66.The reported Hmean is the highest across methods.Source: §F.1 · pp. 32–33

F.2 Adjacent Concept Preservation

  • CePU removes German Shepherds (Atarget = 0) while outperforming FADE on adjacent-breed accuracy (77.76% vs. 75.28%) and LPIPS (0.0568 vs. 0.0790).Adjacent concepts include Malinois, Rottweiler, Norwegian Elkhound, Labrador Retriever, and Golden Retriever.Source: §F.2 · p. 33

F.3 Extended Identity Unlearning Results.

Extended identity results compare preference-based methods across β values on DS-V8-LCM and DS-V7-LCM. CePU traces a stronger removal-retention frontier across both backbones.

F.3 Extended Identity Unlearning Results.

  • Lower β favors stronger forgetting but can reduce retention of other identities; higher β improves retention while weakening target suppression.CePU’s results span both DS-V8-LCM and DS-V7-LCM.Source: §F.3 · p. 33
  • Across DS-V8-LCM and DS-V7-LCM, CePU achieves a stronger Atarget–Aother Pareto frontier than prior preference-based approaches.On DS-V8-LCM it fully removes the target at all reported β values, with the best Hmean at β = 1000; on the harder DS-V7-LCM task it remains robust and competitive.Source: §F.3 · pp. 33–34, Table I · p. 34

F.4 Extended Nudity Unlearning Results.

The extended nudity results compare preference-based methods across β values on three few-step backbones. CePU reports strong suppression across all three, with β settings that trade safety against utility.

F.4 Extended Nudity Unlearning Results.

  • On DS-V8-LCM, CePU at β = 100 reduces NudeNet detections from 3120 to 61 (98.04% DSR); on DS-V7-LCM, it reduces 2554 to 68 (97.34% DSR).These settings provide the strongest suppression on those two backbones.Source: §F.4 · p. 35, Table L · p. 36
  • On SD-Turbo, CePU at β = 250 achieves 5 detections and 99.43% DSR; at β = 500 it reports 6 detections and 99.32% DSR with substantially improved LPIPS and CLIP alignment.The β = 250 setting gives the strongest overall suppression; β = 500 offers a more favorable safety–utility balance.Source: §F.4 · pp. 35–36
  • Across the three backbones, lower β strongly suppresses unsafe generations, whereas higher β better preserves the original model distribution but weakens unlearning.CePU lies on a more favorable safety–utility frontier than prior preference-based baselines.Source: §F.4 · pp. 35–36, Table L · p. 36

G Limitations

The paper identifies limits in its tested scope, evaluation proxies, and hyperparameter coverage, alongside broader risks from selective concept removal and unlearning misuse.

G Limitations

  • Evaluation covers identity and nudity on three FSD models; violence, hate, self-harm, multilingual or codemixed prompts, and simultaneous multi-concept unlearning remain outside this study.The authors identify these as directions for follow-up work.Source: G Limitations · p. 37
  • Automatic safety and utility proxies may be imperfect or biased, and fixed hyperparameters were tuned on limited tasks rather than systematically tested across deployment regimes.The authors suggest human studies, richer safety benchmarks, and sensitivity analysis as extensions.Source: G Limitations · p. 37

H Broader Impacts

  • Selective removal may suppress legitimate expression or amplify representational bias, and unlearning may be misused for censorship or fail under adversarial prompting.The paper calls for careful evaluation, transparency, and human oversight.Source: H Broader Impacts · p. 37

Related papers

Loading 2610.10859v1…