Coffee: Controllable Diffusion Fine-tuning

TL;DR

Coffee is a novel method for controllable diffusion model fine-tuning that uses language-based regularization to prevent models from learning undesired concepts, enabling safer and more customizable image generation.

Contribution

It introduces a training-free approach that uses language prompts to control and prevent undesired concept learning during diffusion model fine-tuning.

Findings

Prevents models from learning specified undesired concepts.

Outperforms existing fine-tuning methods in experiments.

Requires no additional training for concept control.

Abstract

Text-to-image diffusion models can generate diverse content with flexible prompts, which makes them well-suited for customization through fine-tuning with a small amount of user-provided data. However, controllable fine-tuning that prevents models from learning undesired concepts present in the fine-tuning data, and from entangling those concepts with user prompts, remains an open challenge. It is crucial for downstream tasks like bias mitigation, preventing malicious adaptation, attribute disentanglement, and generalizable fine-tuning of diffusion policy. We propose Coffee that allows using language to specify undesired concepts to regularize the adaptation process. The crux of our method lies in keeping the embeddings of the user prompt from aligning with undesired concepts. Crucially, Coffee requires no additional training and enables flexible modification of undesired concepts by modifying textual descriptions. We evaluate Coffee by fine-tuning on images associated with user prompts paired with undesired concepts. Experimental results demonstrate that Coffee can prevent text-to-image models from learning specified undesired concepts during fine-tuning and outperforms existing methods. Code will be released upon acceptance.