Preference-Calibrated Optimization with Score-Level Distribution Alignment for Text-to-Image Diffusion Model Unlearning
Abstract
While text-to-image diffusion models achieve remarkable generation quality, they inadvertently memorize sensitive content, necessitating machine unlearning to prevent undesired outputs. However, existing unlearning methods rely on suboptimal surrogate objectives rather than directly optimizing the unlearning goal, leading to fundamental objective mismatch. Moreover, these methods preserve model utility via surface-level constraints on model parameters or outputs, yet fail to capture the intrinsic generative dynamics of diffusion models, consequently triggering catastrophic forgetting. To address these challenges, we propose Preference-calibrated Optimization with Score-level Distribution Alignment (POSDA), a unified unlearning framework that harmonizes effective erasure with fine-grained structural preservation. Specifically, we reframe unlearning as a preference optimization problem by constructing a reward that explicitly quantifies the unlearning objective. Additionally, we introduce score-level distribution alignment to ensure the invariance of the underlying manifold topology of the unlearned model, thereby preventing distributional drift. Extensive experiments across object, style, and NSFW unlearning tasks demonstrate that POSDA achieves state-of-the-art erasure efficacy while maintaining superior model utility compared to existing methods.
Lay Summary
Text-to-image AI models can create high-quality images from written prompts, but they may also reproduce sensitive, unsafe, or otherwise unwanted content learned during training. A practical way to reduce this risk is machine unlearning, which aims to make a trained model stop generating selected concepts without retraining it from scratch. Existing methods often remove the target content only indirectly, and this can damage the model's ability to generate normal images. We propose POSDA, a new unlearning method that directly rewards the model for avoiding the unwanted concept while preserving how the model generates other content. The method also helps keep the model's internal generation process stable, so that removing one concept does not unnecessarily disrupt unrelated concepts. Experiments on object, artistic style, and unsafe-content removal show that POSDA more effectively suppresses unwanted outputs while better maintaining the model's general image-generation ability.