Iterative Ensemble Training with Anti-Gradient Control for Mitigating Memorization in Diffusion Models

TL;DR

This work tackles privacy risks in diffusion models caused by memorization of training data. It introduces Iterative Ensemble Training with Anti-Gradient Control (IET-AGC), which trains multiple diffusion models on separate data shards, aggregates their parameters to form a global model, and uses a memory-bank driven anti-gradient strategy to skip low-loss memorized samples. Across CIFAR-10/100, AFHQ-DOG, and LAION-10k, IET-AGC substantially reduces memorization (MQ reductions up to ~87%), while maintaining or improving generation quality (FID) and enabling efficient fine-tuning. The approach offers a generic, scalable pathway to privacy-preserving diffusion models for both unconditional and text-conditional generation, with practical guidelines on hyperparameters and analysis of skipping behavior.

Abstract

Diffusion models, known for their tremendous ability to generate novel and high-quality samples, have recently raised concerns due to their data memorization behavior, which poses privacy risks. Recent approaches for memory mitigation either only focused on the text modality problem in cross-modal generation tasks or utilized data augmentation strategies. In this paper, we propose a novel training framework for diffusion models from the perspective of visual modality, which is more generic and fundamental for mitigating memorization. To facilitate forgetting of stored information in diffusion model parameters, we propose an iterative ensemble training strategy by splitting the data into multiple shards for training multiple models and intermittently aggregating these model parameters. Moreover, practical analysis of losses illustrates that the training loss for easily memorable images tends to be obviously lower. Thus, we propose an anti-gradient control method to exclude the sample with a lower loss value from the current mini-batch to avoid memorizing. Extensive experiments and analysis on four datasets are conducted to illustrate the effectiveness of our method, and results show that our method successfully reduces memory capacity while even improving the performance slightly. Moreover, to save the computing cost, we successfully apply our method to fine-tune the well-trained diffusion models by limited epochs, demonstrating the applicability of our method. Code is available in https://github.com/liuxiao-guan/IET_AGC.

Figures (7)

• Figure 1: Overview of our IET-AGC method. (a) Iterative Ensemble Training (IET): we divide the dataset $D$ into $K$ different data shards. Each shard $D_i$ trains a separate diffusion model $\theta_i$. After a period of training, these models are merged by averaging and repeating this training strategy. (b) Anti-Gradient Control (AGC): during the training process, we dynamically update and maintain a memory bank of losses at each time step. For loss values smaller than $\lambda$ times the corresponding memory bank, we exclude these losses to prevent the model from memorizing such images.
• Figure 2: Comparison of the losses between memorized and non-memorized images. The solid line represents the averaged losses of memorized and non-memorized images, while the two dashed lines represent the losses of the 15th and 85th percentile data respectively.
• Figure 3: The similar grid image of the default training and our method. Odd-numbered columns represent images from the training set, while even-numbered columns represent images from the generation set that has the smallest $\ell_{2}$ distance to the corresponding training set image. The images in the grid are arranged in ascending order of $\ell_{2}$ distance, and the selected images from both methods correspond to the same position in terms of $\ell_{2}$ distance.
• Figure 4: Most skipped images vs least skipped images.
• Figure 5: Distribution of skipped image counts.