Memorization Control in Diffusion Models from Denoising-centric Perspective

2026-01-29

AI Summary

Diffusion models are prone to over-memorizing training samples, which compromises the consistency between the generated distribution and the training data. This work addresses this issue by analyzing the denoising process and reveals, for the first time, that biased sampling of timesteps is a key factor driving memorization. To mitigate this, the authors propose a signal-to-noise ratio–guided timestep reweighting strategy that dynamically adjusts the learning focus along the denoising trajectory based on confidence interval widths, thereby enabling an explicit trade-off between memorization and generalization. Experiments on both image and one-dimensional signal generation tasks demonstrate that shifting the learning emphasis toward later denoising steps significantly reduces memorization while improving sample quality and distribution alignment.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• memorization control
• diffusion models
• denoising
• distributional alignment
• timestep sampling

Innovation

• memorization control
• diffusion models
• denoising-centric perspective
• timestep sampling
• generalization-memorization trade-off