Loss curvature predicts memorization in diffusion models over training
Loss Curvature as a Sampling-Free Proxy for Memorization in Diffusion Models
Machine LearningArtificial Intelligence
Summary
Diffusion models can sometimes memorize their training data, which is a challenge for generating new content. This paper explores how a mathematical property called loss curvature, related to how the model's error changes, can indicate when memorization starts in a model during training. The authors also define a new way to understand if a specific example is memorized, linking it to this curvature measure. Their findings provide a way to estimate memorization without needing extra data samples.
What this means in practice
- •For machine learning engineers: Monitor loss curvature during diffusion model training to detect early memorization and adjust training to improve generalization.
- •For ai ethics compliance teams: Use input loss curvature metrics to flag models likely to reproduce sensitive training data and manage data privacy risks.
Authors
Shashank Sule, Xinyue Yu, Boonam Shin, Yuxuan Liu, Abolfazl Hashemi, Hayden Schaeffer
Abstract
Understanding the mechanism behind why diffusion models memorize their training data remains an open question in generative modeling. In this paper, we examine this question from two perspectives: (a) the model-level perspective concerning the onset of memorization in a model during training, and (b) the sample-level perspective concerning the proclivity of an individual sample to be memorized by a given model. We show that the so-called input loss curvature, or the Frobenius norm of the Hessian of the diffusion model's training loss, is an effective metric for quantifying both perspectives on memorization. In particular, we empirically demonstrate that input curvatures averaged across training and validation samples can identify when a model is likely to start regurgitating its training data when tracked over training time. However, this metric is not sufficient for measuring an individual sample's likelihood of being memorized. To address this perspective, we introduce a counterfactual definition of memorization in diffusion models, extending the analogous notion from supervised learning proposed by [Feldman, 2019]. We prove that this counterfactual memorization is controlled by the input curvature of the leave-one-out model, thus establishing a direct theoretical link between a sample's local loss landscape and its generation likelihood. We validate our derived bounds through several empirical experiments and show that input-dependent curvature provides a proxy for memorization.