Technical Summary: A Unified Framework for Diffusion Model Unlearning with f-Divergence

1. Problem Statement

Text-to-image (T2I) diffusion models (DMs), trained on large-scale internet-scraped datasets (e.g., LAION-5B), have demonstrated remarkable generative capabilities but also inherit undesirable properties, including the ability to generate Not-Safe-For-Work (NSFW) content, copyrighted artistic styles, and personal information. While retraining models on curated datasets is the theoretical solution for erasing specific concepts, it is computationally prohibitive and often impractical due to data deletion or resource constraints.

Existing unlearning algorithms typically rely on a paradigm where the model's generation for a "target" concept (to be erased) is aligned with an "anchor" concept (e.g., a null concept, superclass, or semantically related concept). This alignment is almost exclusively achieved by minimizing the Mean Squared Error (MSE) between the denoiser outputs of the target and anchor. Theoretically, this MSE minimization corresponds to minimizing the Kullback-Leibler (KL) divergence between two Gaussian distributions conditioned on the respective concepts. However, this approach limits the optimization landscape to a single divergence metric, potentially suboptimal for balancing unlearning efficacy with the preservation of generative fidelity for non-target concepts.

2. Methodology

The authors propose f-DMU, a unified framework that generalizes the unlearning objective from the standard KL divergence to any f-divergence. The framework is built upon the observation that the unlearning problem can be formulated as minimizing the divergence between the probability distributions of the original model $\Phi$ (conditioned on the anchor concept $c$) and the unlearned model $\hat{\Phi}$ (conditioned on the target concept $c^*$) at specific time steps of the diffusion process.

The core objective is defined as:
$$\min_{\hat{\Phi}} \mathbb{E}_{c^*, c, t} \left[ D_f \left( p_{\Phi}(x_{t-1}|x_t, c) \parallel p_{\hat{\Phi}}(x_{t-1}|x_t, c^*) \right) \right]$$

The framework offers two distinct approaches to solving this objective, depending on the properties of the chosen f-divergence:

A. Closed-Form Losses

For a subset of f-divergences where a closed-form expression exists for Gaussian distributions (specifically when variances are fixed, a standard practice in DMs), the authors derive specific loss functions. Key instances include:

• Squared Hellinger Distance ($H^2$): Yields a loss function involving an exponential decay of the MSE: $J_{H^2} \propto -\exp(-\text{MSE})$. This results in bounded gradients that diminish as the MSE increases, mitigating the impact of outliers and preventing abrupt model changes.
• Pearson $\chi^2$ Divergence: Yields a loss function involving an exponential growth of the MSE: $J_{\chi^2} \propto \exp(\text{MSE})$. This results in unbounded, rapidly growing gradients for large errors, leading to more aggressive updates.
• Jeffreys Divergence: Reduces to a standard MSE formulation under specific assumptions.

B. Variational Losses

For f-divergences lacking a closed-form expression (e.g., Total Variation, or when a more general formulation is desired), the authors utilize the variational representation of f-divergences. This transforms the problem into a min-max optimization game:
$$\min_{\hat{\Phi}} \max_{T} \mathbb{E} \left[ \mathbb{E}_{p_{\Phi}}[T(\Phi)] - \mathbb{E}_{p_{\hat{\Phi}}}[f^*(T(\hat{\Phi}))] \right]$$
Here, $T$ acts as a discriminator (parametrized by a neural network) estimating the divergence, while $\hat{\Phi}$ acts as the generator attempting to minimize it. This approach allows the framework to apply any f-divergence to the unlearning task.

3. Theoretical Analysis

The paper provides a rigorous theoretical analysis of the proposed framework:

• Gradient Dynamics: The authors analyze the gradient magnitudes of the closed-form losses. They demonstrate that $H^2$ and $\chi^2$ exhibit opposite behaviors compared to the standard KL/MSE. $H^2$ suppresses gradients for large errors (outliers), while $\chi^2$ amplifies them. This per-sample weighting is an inherent property of the f-divergence choice, distinct from standard learning rate schedulers.
• Convergence Properties: Using nonlinear dynamical systems theory, the authors prove the local exponential stability of the variational formulation around equilibrium points. They derive bounds on the convergence rate based on the eigenvalues of the Jacobian matrix. Crucially, they show that the choice of f-divergence impacts convergence speed; specifically, divergences with larger second derivatives of their Fenchel conjugates at equilibrium (e.g., $H^2$ and Jensen-Shannon) theoretically converge faster than others like KL or $\chi^2$.
• Mode-Seeking vs. Mode-Covering: The paper discusses how different divergences affect the distribution of generated images. $H^2$ and Jensen-Shannon exhibit "medium" mode-seeking properties, while KL and $\chi^2$ are "mode-covering," potentially assigning higher probability to areas between modes.

4. Experimental Results

The framework was evaluated on Stable Diffusion models (v1.4, v1.5, v2.1, XL, and 3) and FLUX, targeting the erasure of objects (e.g., R2D2), artistic styles (e.g., Van Gogh, Monet), and nudity.

• Gradient Stability: Empirical validation confirms that $H$-DMU (using $H^2$) maintains bounded gradients throughout training, avoiding the sudden degenerate image generation observed in MSE-based methods during intermediate fine-tuning stages.
• Unlearning Efficacy vs. Preservation:
  • Closed-Form: $H$-DMU consistently outperforms standard MSE in preserving non-target concepts (measured by Kernel Inception Distance, KID) while achieving comparable or better erasure of target concepts. Conversely, $\chi^2$-DMU achieves stronger erasure but often at the cost of generative quality and prior preservation.
  • Variational: Variational methods (e.g., $H_v$-DMU) tend to be more "aggressive," achieving rapid erasure in fewer iterations but often resulting in higher KID scores (lower image quality) due to coarse divergence estimates on small batches.
• Robustness: In the context of nudity erasure, $H$-DMU demonstrated superior robustness against adversarial prompts (e.g., Ring-A-Bell, MMA-Diffusion) compared to state-of-the-art methods like CAbl, DoCo, and SLD, generating significantly fewer NSFW images under attack.
• Multi-Concept Erasure: The framework successfully handled sequential erasure of 10 artistic styles, with $H$-DMU showing superior trade-offs between erasure and preservation compared to MSE.

5. Significance and Claims

The paper claims to present the first unified framework for diffusion model unlearning based on f-divergences. Its primary contributions are:

1. Generalization: Moving beyond the implicit KL/MSE assumption to a flexible family of divergences, allowing practitioners to select the optimal objective based on specific application goals (e.g., aggressive erasure vs. high-fidelity preservation).
2. Theoretical Grounding: Providing a theoretical analysis of gradient scaling and convergence properties, linking the choice of divergence to algorithmic stability and speed.
3. Practical Utility: Demonstrating that specific instances, particularly the closed-form Squared Hellinger distance ($H$-DMU), offer a superior default choice for most unlearning tasks by balancing effective concept removal with the preservation of the model's generative capabilities.

The authors emphasize that their framework is model-agnostic (applicable to flow-based models as well) and does not require retraining from scratch, offering a computationally efficient and theoretically sound alternative to existing concept erasure techniques.