Reducing Memorisation in Generative Models via Riemannian Bayesian Inference

TL;DR

The paper tackles memorisation in high-dimensional generative models by introducing a geometry-aware Bayesian approach to perturb model parameters. It leverages the Riemannian Laplace approximation to sample from a posterior that respects the loss geometry, producing an ensemble of generators whose outputs exhibit reduced memorisation while maintaining generalisation. The authors provide theoretical analysis linking perturbations to changes in the velocity field and substantiate these claims with toy and CIFAR-10 experiments, including an exact likelihood formulation for invertible generators and posterior predictive evaluation. The work demonstrates that respecting loss geometry during parameter perturbation yields a practical route to more robust, privacy-conscious generative modelling, scalable to flow-matching and diffusion architectures. Overall, the approach integrates differential geometry with Bayesian inference to navigate high-dimensional parameter spaces more effectively than Euclidean perturbations.

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.

Figures (14)

• Figure 1: We propose reducing memorisation of a trained generative model by perturbing the learnt model parameters ${\boldsymbol{\theta}}^\ast$ () by sampling from an approximate posterior distribution ${\boldsymbol{\theta}} \sim q\left({\boldsymbol{\theta}}\right)$, thus constructing new models. While samples from a Euclidean approximate posterior () reduce memorisation, accounting for the geometry () reduces memorisation without breaking the fit.
• Figure 2: The flat parameter space $\Theta$ and a two-dimensional surface $h(\Theta)=\mathcal{M}.$ The arrow ( ) represents the initial velocity vector ${\boldsymbol{v}}$ of ${\boldsymbol{\alpha}}(t)$ ( ) at the starting point ${\boldsymbol{\theta}}^*={\boldsymbol{\alpha}}(0)$, that is, $\dot{{\boldsymbol{\alpha}}}(0)={\boldsymbol{v}}$. By mapping ${\boldsymbol{\alpha}}(t)$ to $\mathcal{M}$ through $h$, we obtain the corresponding geodesic ${\boldsymbol{\gamma}}(t)=h({\boldsymbol{\alpha}}(t))$ ( ) on $\mathcal{M}.$ In this specific example we obtain a geodesic path lying in a low loss region.
• Figure 3: In the setting of Theorem \ref{['th:margin']}, in which all considered points lie on one straight line, the perturbation makes a generated point () wander straight towards the second closest training sample (). If the new point () falls into the green zone, it is no longer classified as memorised. The size of the green zone depends on $\lVert \widehat{{\boldsymbol{x}}} - {\boldsymbol{x}}^{\left(1\right)}(\widehat{{\boldsymbol{x}}})\rVert+\lVert \widehat{{\boldsymbol{x}}} - {\boldsymbol{x}}^{\left(1\right)}(\widehat{{\boldsymbol{x}}})\rVert,$ and is maximised if $\widehat{{\boldsymbol{x}}}$ is right between the two training samples.
• Figure 4: A 1D flow matching toy problem with a Gaussian base distribution $p_0$ and a GMM as the target distribution $p_\ast$. Left: the learnt conditional vector field $u_{{\boldsymbol{\theta}}}\left(x(t), t\right):\mathbb{R}\times[0,1]\rightarrow\mathbb{R}$ at the optimal parameters ${\boldsymbol{\theta}}^\ast$ for $x(t) \in [-3, 3]$ and $t\in[0,1]$. Each line ( ) is a trajectory $x(t)$ of a noise sample $x_0\sim p_0$ under $u_{{\boldsymbol{\theta}}^\ast}$, which yields $x(1)=\widehat{x}\sim p\left(\widehat{x}\right)$. This distribution overfits to the two fixed training points (), hence the generative model $g_{{\boldsymbol{\theta}}^\ast}$ has learnt to memorise rather than generalise to $p_\ast$. Right: we draw $S=1000$ models ${\boldsymbol{\theta}}^s \sim q\left({\boldsymbol{\theta}}|\mathcal{D}\right)$ from the Euclidean and Riemannian approximate posterior. Each model ${\boldsymbol{\theta}}^s$ gives a specific velocity field (as in left). We show the standard deviation per $\left(x(t),t\right)$-coordinate computed over the $S$ different velocity fields.
• Figure 5: Left: The memorisation ratio as a function of the distance threshold $c$ when generating data from the target distribution, the learnt distribution and the learnt distribution under the two Bayesian treatments. Right: The generalisation error computed as the $\operatorname{KL}$-divergence between the target distribution and the generated data distributions. For efficiency, we perform $50$ repetitions of computing $\operatorname{KL}$-divergences from a subset of $100$ generated data samples and plot the means and standard errors in the bar plot. The generated data distributions (most right) correspond to pushing noise samples from $p_0$ through the generator $g_{{\boldsymbol{\theta}}}$ using the learnt model ${\boldsymbol{\theta}}^\ast$ (top), using several models sampled from the Euclidean Laplace approximation (middle), and using several models sampled from the Riemannian Laplace approximation (bottom). We visualise the resulting distributions using kernel density estimation. See Appendix \ref{['app:experiments']} for details.

Theorems & Definitions (8)

• Definition 2.1
• Theorem 4.1
• Theorem 4.2
• Definition 1.1
• Definition 1.2
• proof
• proof
• Definition 3.1