Training Data Protection with Compositional Diffusion Models

TL;DR

Compartmentalized Diffusion Models enable perfect selective forgetting and continual learning for large-scale diffusion models, allow serving customized models based on the user's access rights, and reduce memorization.

Abstract

We introduce Compartmentalized Diffusion Models (CDM), a method to train different diffusion models (or prompts) on distinct data sources and arbitrarily compose them at inference time. The individual models can be trained in isolation, at different times, and on different distributions and domains and can be later composed to achieve performance comparable to a paragon model trained on all data simultaneously. Furthermore, each model only contains information about the subset of the data it was exposed to during training, enabling several forms of training data protection. In particular, CDMs enable perfect selective forgetting and continual learning for large-scale diffusion models, allow serving customized models based on the user's access rights. Empirically the quality (FID) of the class-conditional CDMs (8-splits) is within 10% (on fine-grained vision datasets) of a monolithic model (no splits), and allows (8x) faster forgetting compared monolithic model with a maximum FID increase of 1%. When applied to text-to-image generation, CDMs improve alignment (TIFA) by 14.33% over a monolithic model trained on MSCOCO. CDMs also allow determining the importance of a subset of the data (attribution) in generating particular samples, and reduce memorization.

Figures (7)

• Figure 1: Compositional diffusion models. We train two diffusion models on two disjoint data distributions: Birds (CUB-200 wahcub) and Stanford Cars krause20133d. (Left) Image generated by naively composing the models by averaging their output. The sample images are distorted and contain elements of both distributions within the same image. (Center) Sample images generated by our method using the same models. The quality is substantially higher and the samples do not conflate the two distributions. (Right) The key idea is illustrated graphically at each step of the reverse diffusion process, where we show the estimated optimal weights $w_i$ to assign to each component. At the beginning, the two components are weighted equally, also shown in the plot below, but as the model approaches convergence, the weights increasingly favor only one of the two models, based on the likelihood that it was trained on a data distribution closest to the current sample.
• Figure 2: Forgetting/Continual Learning with Compartmentalized DMs: Relative change (wrt single shard) in the FID score as we continually add more shards (class-conditional generation with 8-splits). FID scores for our method continuously decreases as we add more shards compared to naive averaging which may result in incorrect mixture of vector fields. Classifier (our) based ensembling enables correct selection of models from the mixture of ensembles, and therefore results in decrease in FID with increase in data. This plot (when viewed from right to left) also shows the increase in FID, given a request to forget an entire shard.
• Figure 3: Credit attribution with CDMs Compartmentalized DMs enable us to provide credit attribution corresponding to samples belonging to different models. Plot shows average credit attribution when shards are from different domains (CUB200 wahcub and Cars krause20133d) and when shards are uniform split of the same domain. For different source domains, CDMs selects the appropriate domain model during backward diffusion, resulting in higher source weight for one model compared to another. For same source domains, CDMS assigns equal weight for majority of backward process, until the end when it selects one source model.
• Figure 4: Out-of-Distribution Coverage and Memorization: Top (A): CDMs enable out of distribution (OOD) generalization by combining shard belonging to different domains. Figure (A) shows that SD2.1 fine-tuned (FT) on AE6+ produces unrealistic images of the bird, and does not follow the prompt correctly, however adding a shard corresponding to birds (SD2.1 FT CUB), with appropriate mixture weights enables the CDMs to generalize OOD, and produce more realistic birds (eg $(w_0,w_1)=(0.5,0.5)$ or $(w_0,w_1)=(0.25,0.75)$. CDMs have the flexibility to add diffusion score from different data sources to improve OOD generalization. Bottom (B): CDMs can also help in reducing memorization. Figure (B) shows that SD2.1 fine-tuned (FT) on AE6+ often memorizes training data carlini2023extracting, eg Big Ben tower $(w_0,w_1)=(1,0)$, which can often violate the privacy rights of the user data. We can prevent memorization in diffusion models by using a mixture of models like in CDMs, which will merge diffusion flows from multiple sources preventing generation of memorized training samples at inference (eg $(w_0,w_1)=(0.5,0.5)$ or $(w_0,w_1)=(0.25,0.75)$). We show that CDMs also aid improving the diversity of the generated images, while preserving the inherent concept presented in the input prompt.
• Figure 5: Diffusion model composition with our method vs naive averaging

Theorems & Definitions (2)

• Proposition 3.1
• proof