A Unified Framework for Diffusion Model Unlearning with f-Divergence

Most existing methods for concept unlearning in text-to-image diffusion models minimize a mean squared error (MSE) loss between the denoiser outputs conditioned on a target and an anchor concept, which is implicitly the KL divergence between two Gaussians. We generalize this objective to any $f$-divergence, recovering MSE as the KL instance, and identify a family of $\alpha$-divergences whose Gaussian closed-form yields cheap, MSE-like training objectives. For the remaining $f$-divergences, we provide a min-max objective based on the variational formulation of the $f$-divergence. We theoretically analyze and numerically validate how different $f$-divergences impact the gradient magnitude and the convergence properties of the algorithm, affecting the quality of unlearning. For instance, we observe that the Hellinger closed-form instance consistently dominates MSE across multiple scenarios. More generally, the proposed unified framework offers a flexible paradigm for selecting the optimal divergence based on the application and user goal, allowing for finer control over the trade-off between unlearning efficacy and generative fidelity.

Modern AI models that generate images from text are trained on large datasets that often contain unsafe content and copyrighted material. This paper addresses the challenge of machine unlearning, which aims at removing from these models the ability to generate specific outputs without requiring a full retraining process. We propose $f$-DMU, a framework that replaces unwanted generations with safe, “anchor” concepts. Instead of offering a single tool like typical unlearning algorithms, $f$-DMU provides a theoretically grounded library of objectives. By using $f$-divergence, which can be intuitively viewed as a metric of dissimilarity between probability distributions, we allow developers to tailor the unlearning process to their specific needs. For instance, our framework provides guidelines on handling the trade-off between erasure of the target concepts and preservation of non-target concepts. Finally, this framework improves our understanding of some existing unlearning techniques by showing how they are specific instances within a single, unified system. Our work supports the safer and more reliable use of generative AI technologies.