Reducing Memorisation in Generative Models via Riemannian Bayesian Inference

2026-01-30

AI Summary

This work proposes a novel Bayesian inference approach that mitigates the tendency of generative models to over-memorize training data at the expense of generalization. By introducing Riemannian geometry into the parameter spaces of flow matching and diffusion models for the first time, the method constructs a variational approximate posterior that adaptively captures the local geometric structure of the loss landscape and reflects the variability of the underlying data distribution. This geometrically informed posterior effectively suppresses memorization during sampling while preserving strong generalization capabilities. Theoretical analysis and empirical experiments demonstrate that the proposed Riemannian Bayesian inference framework significantly reduces memorization in generative models, offering a new paradigm for balancing memorization and generalization.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• memorisation
• generalisation
• generative models
• Bayesian inference
• Riemannian geometry

Innovation

• Riemannian Bayesian Inference
• Generative Models
• Memorisation Reduction
• Loss Geometry
• Approximate Posterior