Redistribute Ensemble Training for Mitigating Memorization in Diffusion Models

TL;DR

This work addresses the privacy risk of memorization in diffusion models by introducing a visual-modality, shard-based memorandum mitigation framework. It combines Iterative Ensemble Training (IET) with Anti-Gradient Control (AGC), Threshold-Aware Augmentation (TAA), and Memory Samples Redistribute (MSR) to selectively skip memorization-prone samples while ensuring they are relearned across shards. The method achieves substantial memorization reductions across multiple datasets and also improves or preserves image quality, including notable gains when fine-tuning diffusion models like Stable Diffusion. By demonstrating compatibility with existing inference-time mitigation strategies, the approach offers a scalable, practical path to privacy-aware diffusion modeling with broad potential applications.

Abstract

Diffusion models, known for their tremendous ability to generate high-quality samples, have recently raised concerns due to their data memorization behavior, which poses privacy risks. Recent methods for memory mitigation have primarily addressed the issue within the context of the text modality in cross-modal generation tasks, restricting their applicability to specific conditions. In this paper, we propose a novel method for diffusion models from the perspective of visual modality, which is more generic and fundamental for mitigating memorization. Directly exposing visual data to the model increases memorization risk, so we design a framework where models learn through proxy model parameters instead. Specially, the training dataset is divided into multiple shards, with each shard training a proxy model, then aggregated to form the final model. Additionally, practical analysis of training losses illustrates that the losses for easily memorable images tend to be obviously lower. Thus, we skip the samples with abnormally low loss values from the current mini-batch to avoid memorizing. However, balancing the need to skip memorization-prone samples while maintaining sufficient training data for high-quality image generation presents a key challenge. Thus, we propose IET-AGC+, which redistributes highly memorizable samples between shards, to mitigate these samples from over-skipping. Furthermore, we dynamically augment samples based on their loss values to further reduce memorization. Extensive experiments and analysis on four datasets show that our method successfully reduces memory capacity while maintaining performance. Moreover, we fine-tune the pre-trained diffusion models, e.g., Stable Diffusion, and decrease the memorization score by 46.7\%, demonstrating the effectiveness of our method. Code is available in: https://github.com/liuxiao-guan/IET_AGC.

Figures (9)

• Figure 1: Prior methods focus solely on the captions associated with the memorized images, such as caption augmentation. In contrast, our approach takes a more generalizable framework by considering aspects from the visual modality.
• Figure 2: Threshold-Aware Augmentation (TAA) collaborated with Anti-Gradient Control. We apply three different treatments based on the comparison between the sample's loss ratio and the skipping threshold.
• Figure 3: Comparison of the training losses between memorized and non-memorized images.
• Figure 4: Framework overview of our method. During the training stage, we train multiple proxy models on several data shards. Besides, we selectively skip samples based on their training loss and track how often each sample is skipped in each shard. During the interaction stage, there are two main parts: first, the proxy models are aggregated into a new model, and its weights are distributed as initial weights for the next training phase; second, each shard redistributes its top $P$ skipped samples to the next shard, assigning the last shard to the first. In the next training stage, each shard resumes training with the updated data and model.
• Figure 5: The proposed model update procedure (AGC with TAA). During training, we dynamically update and maintain a memory bank of losses at each timestep. For each sample's loss ratio $\frac{Loss}{l_t}$, we compare it with $\lambda$ and $R\lambda$ to update the loss, considering three cases: for losses less than $\lambda$, we skip the sample and update its skip times; for losses between $\lambda$ and $R\lambda$, we augment the sample and retrain to obtain a new loss; for losses greater than $R\lambda$, we keep the loss unchanged.