Unlearning in Diffusion Models: A Unified Framework with KL Divergence and Likelihood Constraints

2026-05-29

AI Summary

This work addresses the fundamental tension in diffusion models between preserving model utility and effectively removing undesirable information during targeted unlearning. The authors propose a unified constraint optimization framework that minimizes deviation from the pretrained model while explicitly enforcing separation from the to-be-forgotten distribution through constraints based on reverse/forward KL divergences and likelihood. They establish, for the first time, three equivalent optimization formulations, prove their strong duality, and derive closed-form optimal solutions. Implemented via a primal-dual algorithm, the approach significantly outperforms weight fine-tuning baselines in both concept and data unlearning tasks: KL-constrained variants achieve superior forgetting efficacy, whereas likelihood-constrained variants better preserve legitimate concepts, thereby advancing both the theoretical understanding and practical implementation of machine unlearning.

Abstract

Unlearning in diffusion models aims to remove undesirable data or concepts while preserving the utility of pretrained models -- two fundamentally conflicting objectives. We propose a principled constrained optimization framework that formulates unlearning as minimizing the deviation from a pretrained model, subject to explicit separation constraints from the unlearning distributions. Specifically, we formulate three constrained optimization problems based on reverse and forward KL divergences, and likelihood constraints. The first two generalize existing approaches for concept and data unlearning, while the third offers a novel and natural formulation for unlearning. Despite the nonconvexity of the KL constraints, we establish strong duality for all three problems, enabling us to explicitly characterize their optimal solutions as unlearning targets and develop primal-dual algorithms for each formulation. Experimental results demonstrate that our KL-constrained approach achieves superior retention-unlearning tradeoffs compared to weight-based baselines for concept and data unlearning, and that our likelihood-based approach matches unlearning effectiveness while better preserving retained concepts compared to baselines.

Problem

• unlearning
• diffusion models
• KL divergence
• likelihood constraints
• model utility

Innovation

• diffusion models
• machine unlearning
• constrained optimization
• KL divergence
• strong duality