Conjugate Moment Measures: A New Generative Modeling Paradigm That Unifies Sampling and Transport with a Single Convex Potential

One of the core challenges in generative modeling is how to organically connect "sampling from noise" with "mapping noise to data." Researcher Nina, at an academic seminar, presented a new work called Conjugate Moment Measures, which proposes using a single convex potential function W to simultaneously characterize both the sampling distribution and the transport map — offering a path to generative modeling that is fundamentally different from mainstream diffusion models.

Starting from the Moment Measure Decomposition

The theoretical starting point of this work is a theorem proved by Cordero-Erausquin and Klartag in 2015. The theorem states that any probability measure ρ with a first moment, zero barycenter, and not supported on a hyperplane can be decomposed into a special form — where u is a convex function, exp(-u) corresponds to what is known as a Gibbs distribution (or log-concave distribution), and the transport map in the decomposition is precisely the gradient of u.

The significance of this decomposition is: once the convex function u is known, one can first sample from the log-concave distribution associated with u, then use ∇u to transport the samples, thereby obtaining samples from ρ. Both the sampling process and the transport process are determined by the same convex potential u. A year later, Santambrogio proved the same conclusion using a variational problem over probability distributions, where the functional's first term is entropy and the second term comes from the dual of optimal transport. By proving the existence of a minimizer and deriving optimality conditions, the moment measure decomposition is exactly recovered.

Deep Connections with Optimal Transport

There is a tight bond between moment measures and optimal transport. Returning to the core of optimal transport — the Monge problem: given two distributions μ and ν, find a transport map T that pushes μ forward to ν while minimizing the transport cost. Brenier's theorem tells us that the gradient of a convex function is always an optimal transport map. Therefore, in the moment measure decomposition, ∇u is precisely the solution to the Monge problem between the Gibbs distribution and ρ. Interestingly, from the optimal transport perspective, u appears twice: once in the source distribution and once in the transport map.

Optimal Transport (OT) is a mathematical theory studying how to "move" one probability distribution to another at minimum cost, tracing back to Monge's earthmoving problem posed in 1781. The modern formulation was restated by Kantorovich in the mid-20th century as a linear program. Brenier's theorem (1991) is a landmark result in this theory: when the source distribution is absolutely continuous, the unique optimal map transporting the source to the target (under the squared Euclidean cost) must be the gradient of some convex function, called the Brenier potential. This reduces the optimal transport problem to solving for a convex function, offering great convenience for both numerical computation and theoretical analysis. In the context of generative modeling, OT provides a natural framework for "aligning" random noise to a data distribution — without stepwise denoising, but by directly constructing a globally optimal map. This is one of the core departures of this work from diffusion models.

Why Moment Measures Are Not Suited for Generative Modeling

Despite moment measures being an elegant answer, the authors argue that they are not well suited for generative modeling. By rewriting the functional studied by Cordero et al. using KL divergence, one can see that the u in the moment measure decomposition is the one that makes the Gibbs distribution associated with its Legendre transform (convex conjugate) u* closest to ρ.

The problem lies precisely here — what is close to ρ is the distribution associated with u*, while the prior Gibbs distribution associated with u is not guaranteed to be close to ρ, and in the Gaussian case may even be far away. When ρ is a Gaussian with covariance matrix Σ, the corresponding prior turns out to be a Gaussian with covariance Σ⁻¹. This means: if ρ is spread out (heavy-tailed), the prior is highly concentrated at zero; and vice versa. When the prior is extremely concentrated and the target is extremely spread out, the Lipschitz constant of the transport map explodes, causing disaster for numerical optimization. The same pathology arises for arbitrary distributions in marginal cases.

The Legendre transform (convex conjugate) is a fundamental tool in convex analysis: given a convex function f, its convex conjugate is defined as f*(y) = sup_x {⟨x, y⟩ - f(x)}. The convex conjugate "flips" a function into the dual space, and for convex functions we have f** = f (double conjugate recovers the original). In the optimal transport context, the Kantorovich duality theorem transforms the original joint distribution optimization problem into an optimization over a pair of "dual potentials" (f, g), which are each other's convex conjugates. KL divergence (Kullback-Leibler divergence) is a standard tool for measuring the "gap" between two probability distributions, defined as KL(p‖q) = ∫ p log(p/q), equal to zero if and only if the two distributions are identical. By rewriting the original variational functional in a form containing KL divergence, the paper clearly reveals what the prior distribution in the moment measure decomposition is actually optimizing — a key step in understanding the core defect of "why the prior is far from the target distribution."

Conjugate Moment Measures: One Critical Modification

The newly proposed conjugate moment measure decomposition is formally very similar to the original decomposition, with a single difference: instead of using the gradient of the potential for transport, it uses the gradient of the Gaussian transform of the potential for transport, where W is still a convex function. This single change reverses the aforementioned pathology.

The proof strategy also differs from prior work. Rather than proving the equation directly, the authors apply convex duality by acting ∇W on both sides simultaneously, and instead prove an equivalent form. By Brenier's theorem, ∇W is the optimal transport map between ρ and its associated Gibbs distribution, i.e., W is the Brenier potential between the two. Furthermore, W is a fixed point of the mapping that "given potential V, returns the Brenier potential between ρ and its Gibbs distribution." The authors use the Schauder fixed-point theorem to prove this mapping has a fixed point, thereby guaranteeing the existence of a solution to the decomposition equation — this is also why the proof requires compactness assumptions (ρ absolutely continuous, supported on a compact convex set).

Revisiting the Gaussian case: when ρ has covariance Σ, the prior Gibbs distribution in the conjugate decomposition has covariance Σ^(1/3), rather than the original Σ⁻¹. That is, when ρ is spread out, the prior is also spread out; the two are better aligned, and the Lipschitz constant no longer explodes. This intuitively explains why the new decomposition is better suited for generative modeling.

The Schauder fixed-point theorem is a classical existence tool in functional analysis: if T is a continuous mapping on a compact convex set, then T must have a fixed point x* satisfying T(x*) = x*. It is the generalization of the finite-dimensional Brouwer fixed-point theorem to infinite-dimensional function spaces, widely used to prove the existence of solutions to differential equations and variational problems. In this paper's proof, the authors define the mapping "given potential V, return Brenier potential W" on a suitable compact convex set in function space, and by verifying the continuity and compactness conditions of this mapping, they can directly invoke the Schauder theorem to assert the existence of a fixed point, thereby guaranteeing that the conjugate moment measure decomposition equation has a solution. This proof strategy bypasses the difficulty of directly constructing an analytic solution to the decomposition equation and is a standard technique for handling such nonlinear operator equations.

How to Estimate the Potential Function W

The estimation algorithm comes directly from the fixed-point structure in the proof: initialize a potential function, then at each step update the potential with the Brenier potential between ρ and the current iterated Gibbs distribution, iterating until convergence.

The difficulty lies in estimating the Brenier potential (i.e., the optimal transport map) between two distributions at each step. In the one-dimensional case, the optimal transport map has a closed-form solution, expressible via the cumulative distribution function of the source distribution and the quantile function of the target distribution; this can be easily computed by approximating densities with histograms and calling Python libraries.

In one-dimensional experiments, the authors compare the two decompositions. The first example is four Gaussian mixtures highly concentrated at zero: the Gibbs distribution of the moment measure is extremely spread out (consistent with Gaussian theory), while the Gibbs distribution of the conjugate decomposition is better aligned with the data. The second example is three more spread-out Gaussian mixtures; the prior of the moment measure is still concentrated at zero, while the conjugate decomposition remains well aligned.

Higher Dimensions: Leveraging Convex Neural Networks and Amortization

In dimensions higher than one, the Brenier potential has no closed-form solution and requires neural network-based optimal transport solvers. The authors parameterize the potential function using Input Convex Neural Networks (ICNNs) — proposed by Brandon Amos, which guarantee convex output by applying non-negativity constraints to certain weights and using convex and non-decreasing activation functions — and train them with an optimal transport loss derived from the dual formulation.

A tricky term in the loss is the convex conjugate (Legendre transform) of f: since f is convex, this term is a maximization of a concave function, solvable with gradient ascent, but recomputing and backpropagating at each step is costly. The authors adopt the amortization technique proposed by Amos: an additional MLP is introduced to predict the solution to the convex conjugate problem, serving as a warm start for the solver and thereby greatly accelerating computation. In practice, it was also found that even when only one gradient update is performed on the target per step — yielding an extremely "noisy" estimate of the optimal transport map — the algorithm still converges, making the final algorithm quite streamlined. During sampling, Langevin Monte Carlo (conceptually understood as noisy gradient descent on the potential W) is used to draw samples from the log-concave distribution, with step sizes adaptively adjusted according to the largest eigenvalue of W's Hessian.

Input Convex Neural Networks (ICNNs) were proposed by Brandon Amos et al. in 2017, specifically designed to parameterize convex functions. Their core constraints are twofold: first, all intermediate layer weights other than "skip connections" must be non-negative; second, activation functions must simultaneously satisfy convexity and monotone non-decreasingness (e.g., softplus). These two constraints together guarantee the convexity of the network output with respect to the input, allowing the gradient ∇W to be interpreted as a valid optimal transport map. Langevin Monte Carlo is a gradient-based MCMC sampling algorithm with the iteration scheme x_{t+1} = x_t - η∇U(x_t) + √(2η) ε, where ε is standard Gaussian noise and U is a potential energy function. For log-concave distributions (convex potential), Langevin dynamics theoretically guarantees convergence to the target distribution, and adaptive step size adjustment (here based on the largest eigenvalue of W's Hessian) is critical for convergence speed in practice.

Experimental Results and Unexpected Findings

In two-dimensional experiments, the authors train on five datasets and generate new samples; the generated distribution closely resembles the original, and the level sets of the potential W align well with the target distribution, further validating the theory.

On a higher-dimensional cartoon face dataset, the model also generates reasonable images. A particularly intriguing phenomenon is that the log-concave "noise" associated with W carries a significant amount of information — human face contours directly emerge in the noise, whereas typical generative models use uninformative Gaussian noise. Furthermore, specific characteristics of the final generated image (e.g., "red hair with sunglasses") can be predicted from a particular noise sample, because sampling and transport are tied together by the same W. The authors candidly acknowledge that the results do not reach SOTA for generative models, but are the best in the category of "generation with convex neural networks"; the experiments also reveal the bottleneck — ICNNs have limited expressiveness and are difficult to train.

Additionally, the authors derive the Monge-Ampère equation corresponding to the conjugate decomposition at the density level, enabling the model to access the density of the generated distribution like normalizing flows, thereby supporting tasks such as image reconstruction (gradient descent on corrupted pixels); they also propose an algorithm for estimating W when ρ is known only up to a normalizing constant (a common setup in sampling scenarios).

Open Questions

During the Q&A session, uniqueness became a focal point. The original moment measure has "uniqueness up to translation" under given conditions, but uniqueness is lost — and infinitely many solutions appear — once assumptions are relaxed. For the conjugate decomposition, the authors have not yet proved uniqueness, but in all practical examples the same potential function is consistently recovered, and in the Gaussian case the solution is even unique within the Gaussian family. She also mentioned currently attempting to move away from reliance on ICNNs — seeking new objective functions that do not require convex neural networks, or even allowing sampling and transport to use different potentials, so that the method can scale to more complex datasets.

Regarding questions about data with discontinuous clusters, the authors responded that assumptions such as absolute continuity, compactness, and convexity are required only for mathematical proof; in practice (e.g., image data often does not satisfy absolute continuity), the algorithm still runs and works effectively. This work is more like a bridge built between rigorous mathematics and practical generative modeling — its value lies not in surpassing benchmarks, but in providing a fundamentally new perspective in which sampling and mapping are naturally unified.