# Wasserstein proximal operators describe score-based generative models and resolve memorization

arXiv: https://arxiv.org/abs/2402.06162 · HTML: https://mlsanity.com/paper/2402.06162/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 26 of its 42 references that are in the map.
- First posted on arXiv 9 Feb 2024.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 26 of 26)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Extracting Training Data from Diffusion Models](https://mlsanity.com/paper/2301.13188/) (2023)
- [Score-based Generative Modeling in Latent Space](https://mlsanity.com/paper/2106.05931/) (2021)
- [Fast Sampling of Diffusion Models with Exponential Integrator](https://mlsanity.com/paper/2204.13902/) (2022)
- [Diffusion Art or Digital Forgery? Investigating Data Replication in Diffusion Models](https://mlsanity.com/paper/2212.03860/) (2022)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Convergence of denoising diffusion models under the manifold hypothesis](https://mlsanity.com/paper/2208.05314/) (2022)
- [Score-Based Generative Modeling with Critically-Damped Langevin Diffusion](https://mlsanity.com/paper/2112.07068/) (2021)
- [Understanding and Mitigating Copying in Diffusion Models](https://mlsanity.com/paper/2305.20086/) (2023)
- [OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport](https://mlsanity.com/paper/2006.00104/) (2020)
- [On Memorization in Diffusion Models](https://mlsanity.com/paper/2310.02664/) (2023)
- [Score-Based Generative Models Detect Manifolds](https://mlsanity.com/paper/2206.01018/) (2022)
- [On the relation between optimal transport and Schr"odinger bridges: A stochastic control viewpoint](https://mlsanity.com/paper/1412.4430/) (2014)
- [Interacting particle solutions of Fokker-Planck equations through gradient-log-density estimation](https://mlsanity.com/paper/2006.00702/) (2020)
- [Score-based Generative Modeling Secretly Minimizes the Wasserstein Distance](https://mlsanity.com/paper/2212.06359/) (2022)
- [FP-Diffusion: Improving Score-based Diffusion Models by Enforcing the Underlying Score Fokker-Planck Equation](https://mlsanity.com/paper/2210.04296/) (2022)
- [A Good Score Does not Lead to A Good Generative Model](https://mlsanity.com/paper/2401.04856/) (2024)
- [How Much is Enough? A Study on Diffusion Times in Score-based Generative Models](https://mlsanity.com/paper/2206.05173/) (2022)
- [A mean-field games laboratory for generative modeling](https://mlsanity.com/paper/2304.13534/) (2023)
- [Wasserstein Proximal of GANs](https://mlsanity.com/paper/2102.06862/) (2021)
- [A Hamilton-Jacobi-based Proximal Operator](https://mlsanity.com/paper/2211.12997/) (2022)
- [A kernel formula for regularized Wasserstein proximal operators](https://mlsanity.com/paper/2301.10301/) (2023)
- [Lipschitz-regularized gradient flows and generative particle algorithms for high-dimensional scarce data](https://mlsanity.com/paper/2210.17230/) (2022)

## Built on by (most cited first, 14 of 14)

- [Memorization and Regularization in Generative Diffusion Models](https://mlsanity.com/paper/2501.15785/) (2025)
- [Beyond Log-Concavity and Score Regularity: Improved Convergence Bounds for Score-Based Generative Models in W2-distance](https://mlsanity.com/paper/2501.02298/) (2025)
- [Does Generation Require Memorization? Creative Diffusion Models using Ambient Diffusion](https://mlsanity.com/paper/2502.21278/) (2025)
- [Score-based generative models are provably robust: an uncertainty quantification perspective](https://mlsanity.com/paper/2405.15754/) (2024)
- [Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spaces](https://mlsanity.com/paper/2412.20556/) (2024)
- [Resolving Memorization in Empirical Diffusion Model for Manifold Data in High-Dimensional Spaces](https://mlsanity.com/paper/2505.02508/) (2025)
- [Combining Wasserstein-1 and Wasserstein-2 proximals: robust manifold learning via well-posed generative flows](https://mlsanity.com/paper/2407.11901/) (2024)
- [Nonlinear denoising score matching for enhanced learning of structured distributions](https://mlsanity.com/paper/2405.15625/) (2024)
- [Proximal optimal transport divergences](https://mlsanity.com/paper/2505.12097/) (2025)
- [HJ-sampler: A Bayesian sampler for inverse problems of a stochastic process by leveraging Hamilton-Jacobi PDEs and score-based generative models](https://mlsanity.com/paper/2409.09614/) (2024)
- [Manifold Generalization Provably Proceeds Memorization in Diffusion Models](https://mlsanity.com/paper/2603.23792/) (2026)
- [Shallow Diffuse: Robust and Invisible Watermarking through Low-Dimensional Subspaces in Diffusion Models](https://mlsanity.com/paper/2410.21088/) (2024)
- [Policy Transfer for Continuous-Time Reinforcement Learning: A (Rough) Differential Equation Approach](https://mlsanity.com/paper/2510.15165/) (2025)
- [Self-Consistent Generative Paths via Admissible Random Variational Transport](https://mlsanity.com/paper/2606.08953/) (2026)

## Abstract (arXiv)

We focus on the fundamental mathematical structure of score-based generative models (SGMs). We first formulate SGMs in terms of the Wasserstein proximal operator (WPO) and demonstrate that, via mean-field games (MFGs), the WPO formulation reveals mathematical structure that describes the inductive bias of diffusion and score-based models. In particular, MFGs yield optimality conditions in the form of a pair of coupled partial differential equations: a forward-controlled Fokker-Planck (FP) equation, and a backward Hamilton-Jacobi-Bellman (HJB) equation. Via a Cole-Hopf transformation and taking advantage of the fact that the cross-entropy can be related to a linear functional of the density, we show that the HJB equation is an uncontrolled FP equation. Second, with the mathematical structure at hand, we present an interpretable kernel-based model for the score function which dramatically improves the performance of SGMs in terms of training samples and training time. In addition, the WPO-informed kernel model is explicitly constructed to avoid the recently studied memorization effects of score-based generative models. The mathematical form of the new kernel-based models in combination with the use of the terminal condition of the MFG reveals new explanations for the manifold learning and generalization properties of SGMs, and provides a resolution to their memorization effects. Finally, our mathematically informed, interpretable kernel-based model suggests new scalable bespoke neural network architectures for high-dimensional applications.