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

2024-07-22

European Conference on Computer Vision

AI Summary

Diffusion models achieve high-quality generation but pose privacy risks due to memorization of training data. Existing mitigation approaches are largely confined to text modalities or rely on data augmentation, lacking a general-purpose memory suppression framework for visual diffusion models. This paper proposes the first memory-mitigated training paradigm specifically designed for visual diffusion models. Our method integrates a loss-aware anti-gradient control mechanism with an iterative ensemble training strategy featuring sharded parameter aggregation—comprising mini-batch sample removal, intermittent parameter aggregation, and lightweight fine-tuning. Evaluated on four benchmark datasets, our approach significantly reduces memorization rates while marginally improving generation quality. It is computationally efficient and requires only minimal fine-tuning to adapt pre-trained models, enabling practical deployment without architectural modification.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Privacy
• Natural Language Processing: Generation

Application Category

• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Mitigate memorization in diffusion models
• Propose iterative ensemble training strategy
• Introduce anti-gradient control method

Innovation

• Iterative ensemble training strategy
• Anti-gradient control method
• Fine-tuning well-trained diffusion models