# Redistribute Ensemble Training for Mitigating Memorization in Diffusion Models

arXiv: https://arxiv.org/abs/2502.09434 · HTML: https://mlsanity.com/paper/2502.09434/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 54 of its 67 references that are in the map.
- First posted on arXiv 13 Feb 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 54)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Deep Learning with Differential Privacy](https://mlsanity.com/paper/1607.00133/) (2016)
- [CARLA: An Open Urban Driving Simulator](https://mlsanity.com/paper/1711.03938/) (2017)
- [CLIPScore: A Reference-free Evaluation Metric for Image Captioning](https://mlsanity.com/paper/2104.08718/) (2021)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Extracting Training Data from Large Language Models](https://mlsanity.com/paper/2012.07805/) (2020)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [LAION-400M: Open Dataset of CLIP-Filtered 400 Million Image-Text Pairs](https://mlsanity.com/paper/2111.02114/) (2021)
- [Spatial Transformer Networks](https://mlsanity.com/paper/1506.02025/) (2015)
- [Random Erasing Data Augmentation](https://mlsanity.com/paper/1708.04896/) (2017)
- [StarGAN v2: Diverse Image Synthesis for Multiple Domains](https://mlsanity.com/paper/1912.01865/) (2019)
- [FlowNet: Learning Optical Flow with Convolutional Networks](https://mlsanity.com/paper/1504.06852/) (2015)
- [Diffusion Models in Vision: A Survey](https://mlsanity.com/paper/2209.04747/) (2022)
- [AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty](https://mlsanity.com/paper/1912.02781/) (2019)
- [Extracting Training Data from Diffusion Models](https://mlsanity.com/paper/2301.13188/) (2023)
- [Label-Efficient Semantic Segmentation with Diffusion Models](https://mlsanity.com/paper/2112.03126/) (2021)
- [Diffusion Art or Digital Forgery? Investigating Data Replication in Diffusion Models](https://mlsanity.com/paper/2212.03860/) (2022)
- [Unified Concept Editing in Diffusion Models](https://mlsanity.com/paper/2308.14761/) (2023)
- [Hard Prompts Made Easy: Gradient-Based Discrete Optimization for Prompt Tuning and Discovery](https://mlsanity.com/paper/2302.03668/) (2023)
- [Synthetic Data from Diffusion Models Improves ImageNet Classification](https://mlsanity.com/paper/2304.08466/) (2023)
- [Ablating Concepts in Text-to-Image Diffusion Models](https://mlsanity.com/paper/2303.13516/) (2023)
- [Forget-Me-Not: Learning to Forget in Text-to-Image Diffusion Models](https://mlsanity.com/paper/2303.17591/) (2023)
- [Effective Data Augmentation With Diffusion Models](https://mlsanity.com/paper/2302.07944/) (2023)
- [Red-Teaming the Stable Diffusion Safety Filter](https://mlsanity.com/paper/2210.04610/) (2022)

## Abstract (arXiv)

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/IETAGC.