Efficient Fine-Tuning and Concept Suppression for Pruned Diffusion Models

TL;DR

The paper tackles the challenge of deploying efficient diffusion models by integrating pruning, distillation, and concept unlearning into a single bilevel optimization framework. The lower level performs standard diffusion fine-tuning with distillation to restore generation quality, while the upper level directs the model away from unwanted concepts, enabling selective suppression without sacrificing fidelity. The proposed approach, solvable with first-order bilevel methods, outperforms two-stage baselines on artist-style erasure and NSFW content removal, while maintaining strong generation quality on unrelated concepts. This enables safer, more practical deployment of diffusion models in resource-constrained settings. The framework is plug-in compatible with various pruning and unlearning methods, making it broadly applicable for controlled diffusion in real-world applications.

Abstract

Recent advances in diffusion generative models have yielded remarkable progress. While the quality of generated content continues to improve, these models have grown considerably in size and complexity. This increasing computational burden poses significant challenges, particularly in resource-constrained deployment scenarios such as mobile devices. The combination of model pruning and knowledge distillation has emerged as a promising solution to reduce computational demands while preserving generation quality. However, this technique inadvertently propagates undesirable behaviors, including the generation of copyrighted content and unsafe concepts, even when such instances are absent from the fine-tuning dataset. In this paper, we propose a novel bilevel optimization framework for pruned diffusion models that consolidates the fine-tuning and unlearning processes into a unified phase. Our approach maintains the principal advantages of distillation-namely, efficient convergence and style transfer capabilities-while selectively suppressing the generation of unwanted content. This plug-in framework is compatible with various pruning and concept unlearning methods, facilitating efficient, safe deployment of diffusion models in controlled environments.

Figures (8)

• Figure 1: Comparison of images generated in the styles of Claude Monet (top row) and Mary Cassatt (bottom row), both impressionist artists, using Stable Diffusion, a pruned model finetuned using standard knowledge distillation, and our proposed controlled fine-tuning method. While both the Stable Diffusion and distilled models generate images in Monet's style, our method effectively unlearns Monet’s Impressionist style while preserving Cassatt’s distinct Impressionist features. This indicates our approach’s advantage in selectively suppressing specific styles while maintaining high fidelity to the other unremoved concepts and features.
• Figure 2: Comparison of generative quality and style adherence: Row 1: The original Stable Diffusion 2.1 model. Row 2: A pruned version fine-tuned with 20,000 iterations of combined DDPM and distillation loss. Row 3: A pruned version fine-tuned with 20,000 iterations of our proposed bilevel fine-tuning approach, removing styles of Van Gogh, Monet, and Picasso. Our bilevel method is successful in retaining generative quality and style diversity while suppressing undesirable concepts. See the \ref{['sec:app-exp_sub:prompts']} for prompts used.
• Figure 3: Why can a two-stage approach (fine-tuning followed by forgetting) be suboptimal? If fine-tuning yields $\hat{\theta}$, initializing the concept unlearning parameters with $\hat{\theta}$ and optimizing the concept unlearning loss (\ref{['eq:concept_ablation']}) results in $\theta'$, which is suboptimal for both fine-tuning and for concept unlearning. In contrast, our bilevel method (\ref{['eq:constrained']}) produces the optimal solution $\theta^{*}$, achieving better performance for both fine-tuning and unlearning.
• Figure 4: Effect of Distillation: Adding distillation to the fine-tuning process of a pruned model significantly accelerates convergence, especially in resource-constrained settings.
• Figure 5: Effect of Pruning: Pruning enables significantly faster convergence compared to random initialization, making it an excellent choice for training a small diffusion model.