A probability distribution contains infinitely many degrees of freedom. Its mean is just a vector; its covariance is just a matrix. This compression is useful only if we understand what it loses. When the distribution evolves, can a retained statistic follow a closed equation without tracking everything else?
Usually a moment equation introduces a new moment, which introduces another, creating an endless hierarchy. But a particularly clean class of Wasserstein flows admits an exact answer. The criterion is geometric: the squared norm of the statistic's gradient must be affine in the statistic itself. It includes familiar first and second moments, but also nonpolynomial observables built from curved distance functions. The scalar result appears in Proposition 16.30 of Optimal Transport for Machine Learners [1].
Let the energy depend on one statistic
Let \(\Omega\subset\mathbb R^d\) be connected and open, and let \(\varphi\in C^2(\Omega)\). For a probability law \(\alpha\), define
\[ m(\alpha)=\int\varphi(x)\,d\alpha(x), \qquad \mathcal F_g(\alpha)=g(m(\alpha)). \]Here \(g\) is smooth on an interval containing the range of \(\varphi\). For example, \(g(m)=(m-m_\star)^2/2\) penalizes deviation from a desired moment. We begin with compactly supported laws, so all the integrals below are finite.
The formal Wasserstein gradient flow moves particles down the spatial gradient of the first variation of the energy. In this case, the first variation is \(g'(m)\varphi(x)\), so the particle velocity and continuity equation are
\[ \begin{aligned} v_t(x)&=-g'(m_t)\nabla\varphi(x),\\ \partial_t\alpha_t+\operatorname{div}(\alpha_tv_t)&=0. \end{aligned} \]This is the standard transport description of Wasserstein gradient flows [2]. We use it on time intervals where the characteristic flow exists and its support remains in \(\Omega\). Compact support and smooth velocities locally justify the following calculation; more general solutions require the corresponding integrability and weak-chain-rule hypotheses.
Testing the continuity equation with \(\varphi\) gives
\[ \dot m_t =-g'(m_t)\int|\nabla\varphi(x)|^2\,d\alpha_t(x). \]The difficulty is now visible. The energy depends on \(m_t\), but the derivative also depends on the average of \(|\nabla\varphi|^2\). When is that extra average already determined by \(m_t\)?
The exact closure criterion
By autonomous closure, we mean that for each energy \(g\), there is a function \(b_g\) such that \(\dot m=b_g(m)\) for every compactly supported initial law. Thus any two laws with the same moment must have the same instantaneous moment derivative. The right-hand side may depend on \(g\), but not on the rest of the distribution.
Theorem 1 (scalar closure). The instantaneous equation for \(m\) closes for every smooth \(g\) and every compactly supported probability measure on \(\Omega\) if and only if there are constants \(a,b\in\mathbb R\) such that
\[ |\nabla\varphi(x)|^2=a+b\varphi(x) \qquad(x\in\Omega). \]In that case, the closed equation is
\[ \dot m_t=-(a+bm_t)g'(m_t). \]The affine expression is automatically nonnegative on the range of admissible moments. In particular,
\[ \frac{d}{dt}g(m_t) =-(a+bm_t)|g'(m_t)|^2\le0. \]Proof: compare a two-point mixture with a single point
Sufficiency follows by integrating the affine identity. For necessity, it is enough to take \(g(s)=s\). Put \(\psi=|\nabla\varphi|^2\). Testing closure on Dirac masses shows that \(\psi(x)=q(\varphi(x))\) for some function \(q\): points with the same statistic must have the same derivative.
The image \(I=\varphi(\Omega)\) is an interval because \(\Omega\) is connected. Choose \(s_0,s_1\in I\), and points \(x_0,x_1\) realizing them. For \(t\in[0,1]\), compare the two-point law \((1-t)\delta_{x_0}+t\delta_{x_1}\) with a Dirac mass at any point whose statistic is \(s_t=(1-t)s_0+ts_1\). Their moments agree, so their derivatives agree:
\[ q((1-t)s_0+ts_1)=(1-t)q(s_0)+tq(s_1). \]This identity for every segment makes \(q\) affine on \(I\). If \(I\) is a singleton, \(\varphi\) is constant and the result is immediate. Thus \(\psi=a+b\varphi\). \(\square\)
Requiring closure for every \(g\) avoids accidental cancellations: a constant energy makes every statistic stationary, and would tell us nothing about \(\varphi\). In fact, the proof shows that testing the single energy \(g(s)=s\), but on all laws, is enough. Special families can have extra moment identities, but a closure valid only within such a family describes the evolution only if the flow also stays in that family.
Why a distance coordinate appears
On a connected region where \(\nabla\varphi\ne0\), choose \(s_0\) in the range of \(\varphi\) and set
\[ r(x)=\int_{s_0}^{\varphi(x)} \frac{ds}{\sqrt{a+bs}}. \]The chain rule turns the closure condition into the eikonal equation \(|\nabla r|=1\). Why is this a distance coordinate? Differentiating gives \(D^2r\,\nabla r=0\), so the integral curves of \(\nabla r\) are straight unit-speed lines, normal to the level sets of \(r\). Near a regular level set, \(r\), up to an additive constant, is its signed distance. Integrating \(d\varphi/dr=\sqrt{a+b\varphi}\) yields
\[ \varphi=s_0+\sqrt{a+bs_0}\,r+\frac b4r^2. \]Conversely, if \(|\nabla r|=1\) and
\[ \varphi=c+\kappa r+\frac{\lambda}{2}r^2, \]then direct differentiation gives
\[ |\nabla\varphi|^2 =\kappa^2-2\lambda c+2\lambda\varphi. \]So the statistic is quadratic in a distance-like coordinate, even if that coordinate is highly nonlinear in \(x\). This is a local description, valid before normal lines meet or a critical point is reached. In the first formula, the branch is the one where \(\sqrt{a+bs_0}+(b/2)r>0\). There is no claim that every smooth hypersurface has a globally smooth distance function.
Examples, and a genuine failure
For \(\varphi(x)=u\cdot x\), closure holds with \(a=|u|^2,b=0\). Particles translate along \(u\). For \(\varphi(x)=|x|^2\), one has \(a=0,b=4\); particles undergo a common radial scaling. Neither assertion requires a Gaussian initial law.
Now take an oriented smooth hypersurface \(\mathcal S\), or a local oriented patch. Its signed distance \(r_\mathcal S\) satisfies \(|\nabla r_\mathcal S|=1\) in a smooth tubular neighborhood. Thus \(r_\mathcal S+\lambda r_\mathcal S^2/2\) is another exactly closed statistic, for as long as the particles stay in that neighborhood. The normal direction varies from particle to particle, but each trajectory follows a straight normal line, and a single scalar governs the retained moment.
A counterexample is just as instructive. On the line, \(\varphi(x)=x^4\) has \(|\varphi'(x)|^2=16x^6\), which is not affine in \(x^4\). The laws
\[ \alpha=\delta_1,\qquad \beta=\tfrac12\delta_0+\tfrac12\delta_{2^{1/4}} \]both have fourth moment one. Their sixth moments are \(1\) and \(\sqrt2\), respectively. For \(g(s)=s\), their initial fourth-moment derivatives are therefore \(-16\) and \(-16\sqrt2\). No common scalar equation based on the fourth moment alone can describe both.
Numerical experimentation: different clouds, the same statistic
On an annulus about the unit circle, use \(\varphi(x)=|x|-1\) and \(g(m)=m^2/2\). Here \(a=1,b=0\), so \(m_t=m_0e^{-t}\). A particle with initial polar coordinates \((R_i,\theta_i)\) follows the exact trajectory
\[ R_i(t)=R_i-m_0(1-e^{-t}),\qquad \theta_i(t)=\theta_i. \]The experiment starts with two distinct clouds whose mean radius offset is \(m_0=0.2\). Their shapes stay different, yet their displayed statistics agree. All radii remain positive, so the signed-distance field is smooth throughout the displayed evolution.
Several statistics: a finite algebra of gradients
For \(m_i=\int\varphi_i\,d\alpha\) and energy \(g(m_1,\ldots,m_k)\), the same calculation gives
\[ \dot m_i =-\sum_j\partial_jg(m) \int\nabla\varphi_i\cdot\nabla\varphi_j\,d\alpha. \]Thus an exact finite system follows if every gradient product belongs to the affine span of \(1,\varphi_1,\ldots,\varphi_k\). Linear and quadratic coordinate functions have this property: their gradients are at most linear, and products of those gradients are at most quadratic. This explains mean–second-moment closure without an assumption that the evolving law is Gaussian.
Closed equations for observables and invariance of a family of distributions are related but distinct questions. The first can hold while the distribution retains complicated information which those observables never see.
References
- Gabriel Peyré. Optimal Transport for Machine Learners, Proposition 16.30 and §16.5. Book manuscript.
- Filippo Santambrogio. {Euclidean, Metric, and Wasserstein} Gradient Flows: an Overview. Bulletin of Mathematical Sciences 7, 87–154, 2017.
- Richard Jordan, David Kinderlehrer, and Felix Otto. The Variational Formulation of the Fokker–Planck Equation. SIAM Journal on Mathematical Analysis 29(1), 1–17, 1998.
The general gradient-flow viewpoint underlying the calculation has its classical variational formulation in [3]; the elementary closure proof above is specific to energies depending on finitely many observables.