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

TL;DR

This paper introduces a novel training framework for diffusion models that reduces memorization and privacy risks by using iterative ensemble training and anti-gradient control, improving privacy and efficiency.

Contribution

The paper proposes a new visual modality-based training method with ensemble and anti-gradient techniques to mitigate memorization in diffusion models, applicable to fine-tuning.

Findings

Reduces memory capacity while maintaining or improving performance.

Effective in limiting memorization on four datasets.

Applicable to fine-tuning with limited epochs.

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.