Key Step Concept Unlearning (KSCU)

• KSCU is a selective concept erasure method for text-to-image diffusion models that fine-tunes only key denoising steps for effective unlearning.
• It divides the denoising process into task-dependent intervals (e.g., last 70% for NSFW) to balance the trade-off between concept removal and image quality.
• KSCU incorporates prompt augmentation, curriculum-based step selection, and targeted parameter updates to avoid over-unlearning and maintain related content.

Key Step Concept Unlearning (KSCU) is a concept-erasure method for text-to-image diffusion models that capitalizes on the stepwise sampling characteristic inherent in diffusion models during image generation. Rather than treating all denoising steps equally, KSCU focuses on pivotal steps with the most influence over the final outcome by dividing key steps for different concept unlearning tasks and fine-tuning the model only at those steps. Its stated objective is to improve the trade-off between unlearning effectiveness and generative retainability in settings such as NSFW content, unauthorized artistic styles, copyrighted artworks, and memorized public figures or objects (Zhang et al., 9 Jul 2025).

1. Definition and task-dependent scope

KSCU was introduced in the context of text-to-image diffusion safety, where concept unlearning is formulated as selective removal of a target concept from a pretrained model without retraining from scratch. The method is motivated by three limitations in prior work: previous methods generally ignore the stepwise structure of diffusion generation and optimize across all denoising steps; broad parameter updates often require larger cumulative parameter changes and can lead to over-unlearning and reduced quality on unrelated prompts; and robustness-oriented methods based on adversarial training are expensive and often further hurt generation quality (Zhang et al., 9 Jul 2025).

Its central claim is that the final content associated with a concept is disproportionately determined by only part of the denoising path. Earlier steps tend to recover coarse low-frequency structure, while later steps refine high-frequency details, and the later steps are argued to be especially important for concept-specific details. In this formulation, “key steps” are not necessarily a single timestep, but a subset or interval of denoising steps that has the strongest causal influence on whether the final image exhibits the concept to be erased (Zhang et al., 9 Jul 2025).

The key-step interval is task dependent rather than universal. The reported category-level settings are as follows.

Task	Key-step interval	Iterations
Class	last 70%	700
Style	last 50%	500
Instance	last 20%	200
NSFW	last 70%	750

This task dependence is one of the method’s defining features. For NSFW, training only the last 70% of denoising steps preserves about 95% of the unlearning accuracy obtained when training across all steps, which is presented as evidence that the first 30% of denoising steps contribute little to final nudity rendering (Zhang et al., 9 Jul 2025).

2. Diffusion-process formulation and key-step objective

KSCU is grounded in latent diffusion and classifier-free guidance. The latent diffusion model encodes an image $\mathbf{x}$ into latent space via an encoder $E$, yielding $\mathbf{z}=E(\mathbf{x})$, and diffusion then runs in latent space with forward process

$$q(\mathbf{z}_t | \mathbf{z}_{t-1}) = \mathcal{N}(\mathbf{z}_t; \sqrt{1-\beta_t} \mathbf{z}_{t-1}, \beta_t \mathbf{I}),$$

and reverse process

$$p_\theta(\mathbf{z}_{t-1} | \mathbf{z}_t) = \mathcal{N}(\mathbf{z}_{t-1}; \mu_\theta(\mathbf{z}_t, t), \Sigma_\theta(\mathbf{z}_t, t)).$$

Under classifier-free guidance, the guided prediction is

$$\hat{\epsilon}_\theta(\mathbf{x}_t, t, \mathbf{c}) = \epsilon_\theta(\mathbf{x}_t, t) + w \cdot (\epsilon_\theta(\mathbf{x}_t, t, \mathbf{c}) - \epsilon_\theta(\mathbf{x}_t, t)),$$

where $w$ is the guidance scale (Zhang et al., 9 Jul 2025).

The unlearning objective reverses the conditional contribution. The conditional-minus-unconditional term is interpreted as the direction associated with the target concept, and the method seeks to push the model in the opposite direction. The desired unlearned guided prediction is defined as

$$\hat{\epsilon}^*_\theta(\mathbf{x}_t, t, \mathbf{c}) = \epsilon_\theta(\mathbf{x}_t, t) - w \cdot (\epsilon_\theta(\mathbf{x}_t, t, \mathbf{c}) - \epsilon_\theta(\mathbf{x}_t, t)),$$

leading to the concept-erasure loss

$$\text{Loss}_{\text{unlearn}}=\left\| \epsilon_\theta^*(\mathbf{x}_t, t, \mathbf{c}) - \big( \epsilon_\theta(\mathbf{x}_t, t) - ( \epsilon_\theta(\mathbf{x}_t, t, \mathbf{c}) - \epsilon_\theta(\mathbf{x}_t, t) )\big)\right\|_2.$$

In plain language, the conditional prediction of the edited model is trained to match a reflected target: instead of moving toward the concept-conditioned direction, it should move away from it (Zhang et al., 9 Jul 2025).

A second term regularizes the unconditional branch. The reported observation is that unconditional denoising itself can drift toward the target concept in later steps because it inherits the previous latent state as a prior, so KSCU adds

$$\text{Loss}_{\text{regular}} = \left\| \epsilon_\theta^*(\mathbf{x}_t, t) - \big( (1+\lambda_1) \epsilon_\theta(\mathbf{x}_t, t) - \lambda_1 \epsilon_\theta(\mathbf{x}_t, t, \mathbf{c}) \big)\right\|_2^2,$$

with

$E$0

for $E$1 decreasing from 1000 to 0, and final loss

$E$2

with $E$3 (Zhang et al., 9 Jul 2025).

3. Training pipeline, key-step table, and parameter selectivity

The timestep-selection mechanism is a predefined Key Step Table rather than a learned module. Given start step $E$4, end step $E$5, table length $E$6, and full loops before shift $E$7, the table is built by repeatedly appending intervals $E$8; after $E$9 repetitions, $\mathbf{z}=E(\mathbf{x})$0 is incremented by one up to $\mathbf{z}=E(\mathbf{x})$1. This creates a timestep list made of overlapping late-step intervals, so later timesteps are revisited more often. The reported effect is a curriculum biased toward later denoising steps while still covering a range instead of a single endpoint (Zhang et al., 9 Jul 2025).

Prompt handling is the third component. For a target concept $\mathbf{z}=E(\mathbf{x})$2, KSCU generates prompt variants using LLM-produced augmentation rules and then applies perturbations such as word shuffling, word removal, scrambled-character insertion, and Gaussian noise in text embedding space. This is presented as a cheaper alternative to adversarial training: instead of explicitly generating attack prompts, the method broadens the semantic support of the target concept during optimization (Zhang et al., 9 Jul 2025).

The parameter-update regime is task specific. For class, style, and instance unlearning, KSCU fine-tunes the model’s cross-attention layers. For NSFW unlearning, it updates all model components except the cross-attention module. The experimental setup uses Stable Diffusion 1.5 for Unlearn Canvas experiments and Stable Diffusion 1.4 for the rest unless otherwise stated, with DDIM, 50 denoising steps, batch size 1, Adam, and learning rate $\mathbf{z}=E(\mathbf{x})$3 (Zhang et al., 9 Jul 2025).

The evaluation covers four concept categories—Class, Style, NSFW, and Instance—using I2P for NSFW prompts and Unlearn Canvas for class and style unlearning. Main baselines are ESD, CA, FMN, EDiff, SalUn, and SHS. KSCU is therefore situated as a selective fine-tuning method rather than a closed-form edit or an inference-time steering method (Zhang et al., 9 Jul 2025).

4. Empirical performance and ablation structure

On Unlearn Canvas class unlearning, KSCU is not the absolute best in raw unlearning accuracy, but it achieves the strongest overall trade-off reported in the paper. The table gives ESD: UA 68.6, IRA 98.9, CRA 96.4, FID 26.1; EDiff: UA 78.0, IRA 94.1, CRA 90.1, FID 48.2; and KSCU: UA 70.0, IRA 99.0, CRA 97.1, FID 25.8. On style unlearning, KSCU matches the best unlearning accuracy while improving retention and image quality: ESD gives UA 100.0, IRA 85.7, CRA 99.0, FID 25.3, whereas KSCU gives UA 100.0, IRA 88.7, CRA 99.3, FID 23.1 (Zhang et al., 9 Jul 2025).

On NSFW/nudity unlearning with I2P and robustness evaluation, the main result is KSCU: UA 96.5, P4D 23.2, UDA 47.2, FID-30k 14.1. The principal comparison points are ESD: UA 88.2, P4D 39.3, UDA 73.7, FID-30k 15.4; EDiff: UA 93.9, P4D 29.6, UDA 52.5, FID-30k 19.8; and SalUn: UA 100.0, P4D 0.0, UDA 4.2, FID-30k 67.9. The stated interpretation is that KSCU combines very strong forgetting with the best FID, whereas SalUn is stronger on raw robustness but badly damages generation quality (Zhang et al., 9 Jul 2025).

The start-step ablation is the most direct validation of the key-step premise. For NSFW, the reported values are: KSCU $\mathbf{z}=E(\mathbf{x})$4: UA 79.7, P4D 59.2, UDA 85.3, FID 15.7; KSCU $\mathbf{z}=E(\mathbf{x})$5: UA 96.1, P4D 23.0, UDA 48.1, FID 14.3; KSCU $\mathbf{z}=E(\mathbf{x})$6: UA 96.0, P4D 26.7, UDA 54.5, FID 18.8; KSCU default $\mathbf{z}=E(\mathbf{x})$7: UA 96.5, P4D 23.2, UDA 47.2, FID 14.1. These results are used to argue that using too few late steps fails to erase the concept, while using all steps hurts image quality (Zhang et al., 9 Jul 2025).

Prompt augmentation also matters. Without prompt augmentation, KSCU gives UA 92.5, P4D 30.3, UDA 51.8, FID 14.5; with prompt augmentation, the values are UA 96.5, P4D 23.2, UDA 47.2, FID 14.1. Efficiency is likewise part of the method’s practical profile: for style unlearning, KSCU halved the iteration count compared to ESD while still outperforming it, and for instance unlearning it reportedly achieves forgetting with only 20% of ESD’s iterations (Zhang et al., 9 Jul 2025).

5. Related formulations and broader generalizations

Although KSCU was introduced for diffusion models, later and adjacent work makes its underlying logic more explicit. In LLMs, Concept Unlearning was reformulated as removing a target node and its incident edges in a knowledge graph, with self-constructed triplets and explanatory sentences functioning as compact target representations. That work does not explicitly define “key steps,” but it states that the self-constructed triplets are high-impact knowledge units and that node-and-edge removal can be read as a KSCU-style abstraction of direct concept invocation and the factual supports that allow recovery of the concept (Yamashita et al., 19 Sep 2025).

A second line of work isolates what should survive unlearning. ReCARE defines CARE as co-occurring benign concepts that naturally co-occur with the erased target and therefore tend to be unintentionally suppressed during unlearning, but should be preserved. It introduces the CARE score and reports CARE preservation of 0.94 for nudity $\mathbf{z}=E(\mathbf{x})$8 person, 0.90 for Van Gogh $\mathbf{z}=E(\mathbf{x})$9 stars, and 0.97 for tench $$q(\mathbf{z}_t | \mathbf{z}_{t-1}) = \mathcal{N}(\mathbf{z}_t; \sqrt{1-\beta_t} \mathbf{z}_{t-1}, \beta_t \mathbf{I}),$$0 freshwater, making retained-concept discovery an explicit stage rather than an implicit side effect (Kim et al., 23 Jun 2026).

In diffusion models, subsequent methods shift the KSCU intuition from timestep selection toward representation selection. TRUST dynamically estimates target concept neurons in cross-attention projection matrices and selectively fine-tunes only those neurons, while PURE represents the target concept in cross-attention activation space captured along a short denoising trajectory and applies a single linear projector to cross-attention key and value weights. SAEParate, by contrast, organizes sparse latent representations into concept-specific clusters and suppresses selected latent features at inference time, with particularly strong gains in joint style-object unlearning (Mansi et al., 8 Feb 2026, Moon et al., 25 May 2026, Kim et al., 12 May 2026).

These developments suggest that KSCU can be understood at multiple granularities: denoising-step selection, neuron localization, cross-attention activation projection, and sparse latent feature suppression. The common thread is selective intervention on a compact internal support for the target concept rather than uniform editing across the full model.

6. Conceptual boundaries, evaluation criteria, and open problems

A central boundary for KSCU is the distinction between untraining and unlearning. One formulation defines untraining as removing the influence of a specific forget set, whereas unlearning generalizes beyond those examples to remove the broader concept or behavior they represent. Under that distinction, KSCU is concept unlearning only when it aims to suppress the broader underlying distribution rather than merely the listed prompts or examples (Triantafillou et al., 9 Apr 2026).

A second boundary concerns unlearning versus obfuscation. Work on probing-based evaluation for LLMs argues that many apparent unlearning methods actually add distracting associations rather than removing the privileged association itself. Its proposed criterion is distribution flattening over constrained answer spaces, together with refusal behavior and probing under Yes-No and MCQ settings, to test whether the model has become genuinely uncertain rather than merely redirected (Sun et al., 5 May 2025). This is directly relevant to KSCU because a diffusion model can likewise appear safe on literal prompts while still leaking the target concept under paraphrased or adversarial prompts.

Faithfulness of forgetting imposes a third requirement: interconnected knowledge that depends on the forgotten target should also be removed, while superficially similar but contextually irrelevant knowledge should be retained. FaithUn formalizes superficial unlearning as either failing to erase overlapping knowledge or unintentionally erasing non-overlapping knowledge, and KLUE addresses this by updating only knowledge-related neurons identified through attribution and regularized against answer-only bias (Yang et al., 26 Feb 2025). This gives KSCU a sharper evaluation target than direct target suppression alone.

Open problems also remain in difficulty estimation and adaptive intervention. Circuit-guided Unlearning Difficulty reports that easy-to-unlearn samples are associated with shorter, shallower interactions concentrated in earlier-to-intermediate parts of the original model, whereas hard samples rely on longer and deeper pathways closer to late-stage computation. A separate line of work identifies entanglement and memorization as two key factors affecting unlearning difficulty and proposes refinement of a heterogeneous forget set into homogenized subsets before unlearning (Cheng et al., 14 Jan 2026, Zhao et al., 2024). This suggests that a mature KSCU framework may need not only key-step selection, but also difficulty-aware scheduling, retained-concept discovery, and probing that can distinguish genuine erasure from surface-level suppression.