Reducing Memorisation in Generative Models via Riemannian Bayesian Inference
This work addresses a key challenge in generative AI for improving model reliability, though it appears incremental as it builds on existing Bayesian and geometric methods.
The paper tackles the problem of balancing memorization and generalization in generative models by introducing a Riemannian Bayesian inference approach that reduces memorization while preserving generalization, with empirical demonstrations showing reduced memorization.
Modern generative models can produce realistic samples, however, balancing memorisation and generalisation remains an open problem. We approach this challenge from a Bayesian perspective by focusing on the parameter space of flow matching and diffusion models and constructing a predictive posterior that better captures the variability of the data distribution. In particular, we capture the geometry of the loss using a Riemannian metric and leverage a flexible approximate posterior that adapts to the local structure of the loss landscape. This approach allows us to sample generative models that resemble the original model, but exhibit reduced memorisation. Empirically, we demonstrate that the proposed approach reduces memorisation while preserving generalisation. Further, we provide a theoretical analysis of our method, which explains our findings. Overall, our work illustrates how considering the geometry of the loss enables effective use of the parameter space, even for complex high-dimensional generative models.