Sample and Map from a Single Convex Potential: Generation using Conjugate Moment Measures
This work addresses generative modeling for machine learning practitioners by offering a novel method that is incremental but provides more intuitive solutions in specific cases.
The paper tackles the problem of generative modeling by proposing an alternative factorization using conjugate moment measures, which ties sampling and mapping more effectively than the canonical approach, and demonstrates improved results on simple examples like Gaussians and 1D distributions.
The canonical approach in generative modeling is to split model fitting into two blocks: define first how to sample noise (e.g. Gaussian) and choose next what to do with it (e.g. using a single map or flows). We explore in this work an alternative route that ties sampling and mapping. We find inspiration in moment measures, a result that states that for any measure $ρ$, there exists a unique convex potential $u$ such that $ρ=\nabla u \sharp e^{-u}$. While this does seem to tie effectively sampling (from log-concave distribution $e^{-u}$) and action (pushing particles through $\nabla u$), we observe on simple examples (e.g., Gaussians or 1D distributions) that this choice is ill-suited for practical tasks. We study an alternative factorization, where $ρ$ is factorized as $\nabla w^*\sharp e^{-w}$, where $w^*$ is the convex conjugate of a convex potential $w$. We call this approach conjugate moment measures, and show far more intuitive results on these examples. Because $\nabla w^*$ is the Monge map between the log-concave distribution $e^{-w}$ and $ρ$, we rely on optimal transport solvers to propose an algorithm to recover $w$ from samples of $ρ$, and parameterize $w$ as an input-convex neural network. We also address the common sampling scenario in which the density of $ρ$ is known only up to a normalizing constant, and propose an algorithm to learn $w$ in this setting.