Reducing Memorisation in Generative Models via Riemannian Bayesian Inference

TL;DR

This paper introduces a Riemannian Bayesian inference method for generative models that reduces memorisation by capturing the loss landscape's geometry, leading to better generalisation and less overfitting.

Contribution

It presents a novel Bayesian approach using Riemannian geometry to construct adaptive posteriors, improving memorisation control in complex generative models.

Findings

Reduces memorisation in generative models

Preserves model generalisation

Theoretically explains the benefits of geometric loss analysis

Abstract

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.