Memorization Control in Diffusion Models from Denoising-centric Perspective

TL;DR

This work addresses memorization in diffusion models by adopting a denoising-centric perspective, revealing that uniform timestep sampling causes unequal learning contributions across denoising steps due to $SNR(t)$. The authors propose a learning-bias control via a confidence-interval based timestep sampling, with a Gaussian (and tail-mass redistribution) formulation that shifts learning toward later, lower-$SNR$ steps to reduce memorization while preserving generation quality. Theoretical analysis links gradient magnitude to $SNR(t)$ and demonstrates how the CI parameters $[c_l,c_h]$, along with $\,\mu,\sigma$, govern the memorization-generalization trade-off. Empirical results on image datasets and a 1D ECG dataset show that increasing the CI mean leads to distributions that align more closely with training data, validating the approach's generality and practical utility for domain-specific generation. This work provides a controllable, denoising-aware mechanism to mitigate memorization in diffusion models across modalities.

Abstract

Controlling memorization in diffusion models is critical for applications that require generated data to closely match the training distribution. Existing approaches mainly focus on data centric or model centric modifications, treating the diffusion model as an isolated predictor. In this paper, we study memorization in diffusion models from a denoising centric perspective. We show that uniform timestep sampling leads to unequal learning contributions across denoising steps due to differences in signal to noise ratio, which biases training toward memorization. To address this, we propose a timestep sampling strategy that explicitly controls where learning occurs along the denoising trajectory. By adjusting the width of the confidence interval, our method provides direct control over the memorization generalization trade off. Experiments on image and 1D signal generation tasks demonstrate that shifting learning emphasis toward later denoising steps consistently reduces memorization and improves distributional alignment with training data, validating the generality and effectiveness of our approach.

Figures (7)

• Figure 1: Illustration of the diffusion sampling trajectory, where early denoising steps play a critical role in determining whether the trajectory moves toward memorization or extrapolation. The trajectory is guided by the noise predicted by the denoising model at each timestep. Biases introduced in these early steps can accumulate over time, directing the sample toward the training data manifold or away from it. In this view, the diffusion model acts as a control knob for memorization by shaping the denoising dynamics along the sampling process.
• Figure 2: Illustration of the truncated normal timestep sampling. Probability mass outside the valid timestep range is treated as tail mass and redistributed to ensure full coverage of all timesteps.
• Figure 3: Distance of generated data and training data on (a) Pokémon and (b) Flower dataset with different $c_l$ and $c_h$ pairs (considering each diagonal) and the evaluation of model performance of different (c) mean location (1000 steps are divided into 3 categories - left, mid, and right) and (d) CI width on both datasets
• Figure 4: Difference of generated and training data on models with different mean location observed on 2 PCA components(a) and (b). Distribution of mean location at 100 and 1000 are selected to compare visually directly with the original (c, d, e, f) with fix CI=[700, 1000]
• Figure 5: Wasserstein and JS distance of generated and training data distribution on Pokemon (first row) and Flower (second row) data (C1 and C2 are 2 dimension of reduced dimension data)