LoyalDiffusion: A Diffusion Model Guarding Against Data Replication

TL;DR

This work addresses privacy risks in diffusion models arising from memorization of training data. It introduces LoyalDiffusion, which combines a replication-aware U-Net (RAU-Net) that replaces direct skip-connection transfers with information transfer blocks and a two-stage training strategy that applies RAU-Net only at timesteps where image fidelity is less sensitive, thereby reducing replication without sacrificing quality. The framework also integrates with data-centric strategies (GC&DF) to further mitigate memorization, and its effectiveness is demonstrated on SD v2.1 fine-tuned on LAION-2B, achieving substantial replication reductions while maintaining competitive FID and CLIP scores. A Bing-based RepliBing evaluation suggests replication scores approach a real-world lower bound, underscoring the practical significance of the approach for privacy-preserving diffusion models. Overall, LoyalDiffusion provides a novel model-centric avenue for replication mitigation that complements existing data-centric methods and highlights the role of timesteps in memorization dynamics.

Abstract

Diffusion models have demonstrated significant potential in image generation. However, their ability to replicate training data presents a privacy risk, particularly when the training data includes confidential information. Existing mitigation strategies primarily focus on augmenting the training dataset, leaving the impact of diffusion model architecture under explored. In this paper, we address this gap by examining and mitigating the impact of the model structure, specifically the skip connections in the diffusion model's U-Net model. We first present our observation on a trade-off in the skip connections. While they enhance image generation quality, they also reinforce the memorization of training data, increasing the risk of replication. To address this, we propose a replication-aware U-Net (RAU-Net) architecture that incorporates information transfer blocks into skip connections that are less essential for image quality. Recognizing the potential impact of RAU-Net on generation quality, we further investigate and identify specific timesteps during which the impact on memorization is most pronounced. By applying RAU-Net selectively at these critical timesteps, we couple our novel diffusion model with a targeted training and inference strategy, forming a framework we refer to as LoyalDiffusion. Extensive experiments demonstrate that LoyalDiffusion outperforms the state-of-the-art replication mitigation method achieving a 48.63% reduction in replication while maintaining comparable image quality.

Figures (6)

• Figure 1: Comparison of replication scores between our proposed LoyalDiffusion with and without dataset optimization (DO) and prior methods including MC somepalli2023understanding, GN somepalli2023understanding, RC somepalli2023understanding, CWR somepalli2023understanding, GC&DF li2024mitigate.
• Figure 2: Overview of the proposed LoyalDiffusion framework. (a) Standard U-Net architecture. (b) Replication-Aware U-Net (RAU-Net) where the skip connection (SC) 3 & 4 being modified to one single Conv 3×3 layer. (c) Training strategy for LoyalDiffusion where the training data optimization block includes generalized captioning and dual fusion li2024mitigate. (d) Image generation process using LoyalDiffusion.
• Figure 3: Similarity between a training image and the generated image that replicates it and CLIP score calculating using the caption during the diffusion generating process. It shows that similarity and CLIP score do not change too much during the early inference interval with large timesteps.
• Figure 4: Images in the first row are from training dataset that are replicated. Images in the second row are generated using fine-tuned baseline model, which show some replication with training images. Images in the third row are generated by our best performance model, with less content replicated with training images.
• Figure 5: Examples of generated images with different FIDs. The images come from the following models: column (1) MP 3&4 from Table \ref{['tab:various-SC-result']}; column (2) GC&DF from Table \ref{['tab:compare-prior']}; column (3) Baseline from Table \ref{['tab:compare-prior']}; column (4) $w_{lat} = 0.1, w_{emb}=0.5,$ and $\tau = 300$ from Table \ref{['tab:result-with-GCDF']}; column (5) Remove skip connection 4 from Table \ref{['tab:raunet-result-2']}; column (6) Two-stage result with $\tau = 100$ from Table \ref{['tab:two-stage-result_1']}; and (7) Multi-Conv 1&4 from Table \ref{['tab:various-SC-result']}.