The Road to the Optimal Distribution: Minimization in Probability Space
When asked to find the minimum of a function, we typically start by taking the derivative and finding its zeros; in fortunate cases, one of these zeros happens to be the minimum point of the original function. For vector-valued functions, we replace the derivative with the gradient and find its zeros. When the gradient's zeros are hard to obtain analytically, we can use gradient descent to progressively approach the minimum.
All of the above are basic results from unconstrained optimization, with which many readers are probably already familiar. However, the subject of this post is optimization in probability space, i.e., cases where the input to the objective function is a probability distribution. This kind of optimization is considerably more complex, because the search space is no longer unconstrained: if we naively try to solve for the zero of the gradient or run gradient descent, the result is not guaranteed to be a valid probability distribution. So we need to find new methods of analysis and computation to ensure that the optimization result satisfies the properties of a probability distribution.
This has long been something of a headache for me, so recently I decided to "learn from the pain" and systematically study the problem of optimization over probability distributions. I've collected what I learned here for reference.more
Gradient Descent
Let's first revisit the relevant material on unconstrained optimization. Suppose our objective is
\begin{equation}\boldsymbol{x}_* = \mathop{\text{argmin}}_{\boldsymbol{x}\in\mathbb{R}^n} F(\boldsymbol{x})\end{equation}
Every high schooler knows that to find the extremum of a function, you take the derivative and set it to zero to find the critical points — this has become "common knowledge" for most people. But let me put a question to readers here: how many of you can actually prove this claim? In other words, why does the minimum of a function have anything to do with "the derivative being zero"?
The Search Perspective
We can explore this question from the perspective of search. Suppose the current best guess we know of is $\boldsymbol{x}$, denoted $\boldsymbol{x}_t$ — how do we determine whether $\boldsymbol{x}_t$ is the minimum point? We can flip the question around: if we can find $\boldsymbol{x}_{t+\eta}$ such that $F(\boldsymbol{x}_{t+\eta}) < F(\boldsymbol{x}_t)$, then $\boldsymbol{x}_t$ clearly cannot be the minimum. To this end, we can search over candidates of the form $\boldsymbol{x}_{t+\eta}$:
\begin{equation}\boldsymbol{x}_{t+\eta} = \boldsymbol{x}_t + \eta \boldsymbol{u}_t,\quad 0 < \eta \ll 1\end{equation}
When $F(\boldsymbol{x})$ is sufficiently smooth and $\eta$ is sufficiently small, we take the first-order approximation to be accurate enough, so we can use:
\begin{equation}F(\boldsymbol{x}_{t+\eta}) = F(\boldsymbol{x}_t + \eta \boldsymbol{u}_t) \approx F(\boldsymbol{x}_t) + \eta \boldsymbol{u}_t \cdot \nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t)\end{equation}
As long as $\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t)\neq 0$, we can choose $\boldsymbol{u}_t = -\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t)$ such that
\begin{equation}F(\boldsymbol{x}_{t+\eta}) \approx F(\boldsymbol{x}_t) - \eta \Vert\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t)\Vert^2 < F(\boldsymbol{x}_t)\end{equation}
This means that, for a sufficiently smooth function, the minimum can only be attained at a point satisfying $\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t) = 0$ or at infinity — which is exactly why the first step in finding an extremum is usually "set the derivative to zero." If $\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t) \neq 0$, we can always pick a sufficiently small $\eta$ and use
\begin{equation}\boldsymbol{x}_{t+\eta} = \boldsymbol{x}_t-\eta\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t)\label{eq:gd}\end{equation}
to obtain a point with a smaller value of $f$ — this is precisely gradient descent. Taking the limit $\eta\to 0$, we get the ODE:
\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = -\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t)\end{equation}
This is the "gradient flow" introduced in Gradient Flow: Exploring the Path Toward the Minimum, which can be viewed as the trajectory traced out by gradient descent in its search for the minimum.
Projected Descent
So far we've discussed unconstrained optimization; now let's briefly look at a simple extension of gradient descent to constrained optimization. Suppose the problem we face is:
\begin{equation}\boldsymbol{x}_* = \mathop{\text{argmin}}_{\boldsymbol{x}\in\mathbb{X}} F(\boldsymbol{x})\label{eq:c-loss}\end{equation}
where $\mathbb{X}$ is a subset of $\mathbb{R}^n$. For a rigorous theoretical treatment, one usually needs to require $\mathbb{X}$ to be a "bounded convex set," but for the purposes of a casual overview we can set these details aside for now.
If we still use gradient descent $\eqref{eq:gd}$ here, the biggest problem is that there's no guarantee $\boldsymbol{x}_{t+\eta}\in\mathbb{X}$ holds. But we can add one extra step — a projection operation:
\begin{equation}\Pi_{\mathbb{X}} (\boldsymbol{y}) = \mathop{\text{argmin}}_{\boldsymbol{x}\in\mathbb{X}}\Vert\boldsymbol{x}-\boldsymbol{y}\Vert\label{eq:project}\end{equation}
thereby forming "projected gradient descent":
\begin{equation}\boldsymbol{x}_{t+\eta} = \Pi_{\mathbb{X}}(\boldsymbol{x}_t-\eta\nabla_{\boldsymbol{x}_t}F(\boldsymbol{x}_t))\label{eq:pgd}\end{equation}
In plain terms, projected gradient descent first performs a gradient descent step, and then finds the point in $\mathbb{X}$ closest to the gradient-descent result and outputs that — which guarantees the output always lies within $\mathbb{X}$. In Making Model Training More Scientific (I): Convergence of Average Loss in SGD, we proved that, under certain assumptions, projected gradient descent can find the optimal solution to the constrained optimization problem $\eqref{eq:c-loss}$.
In the end, projected gradient descent turns the constrained optimization problem $\eqref{eq:c-loss}$ into a two-step procedure of "gradient descent + projection." The projection step $\eqref{eq:project}$ is itself a constrained optimization problem — although its objective is now fixed, it remains an unsolved sub-problem in general, requiring case-by-case $\mathbb{X}$-specific analysis, so it still calls for further investigation.
Discrete Distributions
This post focuses on optimization in probability space, i.e., search problems where the object being searched for must be a probability distribution. In this section we focus on discrete distributions, and denote the search space by $\Delta^{n-1}$, the set of all $n$-dimensional discrete probability distributions, i.e.,
\begin{equation}\Delta^{n-1} = \left\{\boldsymbol{p}=(p_1,p_2,\cdots,p_n)\left|\, p_1,p_2,\cdots,p_n\geq 0,\sum_{i=1}^n p_i = 1\right.\right\}\end{equation}
Our optimization objective is then
\begin{equation}\boldsymbol{p}_* = \mathop{\text{argmin}}_{\boldsymbol{p}\in\Delta^{n-1}} F(\boldsymbol{p})\label{eq:p-loss}\end{equation}
Lagrange Multipliers
For optimization problems with equality or inequality constraints, the standard approach is "Lagrange multipliers," which converts the constrained optimization problem $\eqref{eq:p-loss}$ into a weakly-constrained $\text{min-max}$ problem:
\begin{equation}\min_{\boldsymbol{p}\in\Delta^{n-1}} F(\boldsymbol{p}) = \min_{\boldsymbol{p}\in\mathbb{R}^n} \max_{\mu_i \geq 0,\lambda\in\mathbb{R}}F(\boldsymbol{p}) - \sum_{i=1}^n \mu_i p_i + \lambda\left(\sum_{i=1}^n p_i - 1\right)\label{eq:min-max}\end{equation}
Note that in this $\text{min-max}$ optimization, we've dropped the constraint $\boldsymbol{p}\in\Delta^{n-1}$ altogether, and the only remaining constraint is the fairly simple one on $\mu_i \geq 0$ in the $\max$ step. How do we prove that the right-hand optimization problem is equivalent to the left-hand one? It isn't actually that hard; we can understand it in three steps:
1. First, let's understand the meaning of $\text{min-max}$ on the right-hand side: $\min$ on the left, $\max$ on the right, meaning that we ultimately want to find as small a result as possible, but the objective function itself first requires taking a $\max$ over certain variables;
2. When $p_i < 0$, the $\max$ step must have $\mu_i\to\infty$, and in that case the resulting objective value is $\infty$; whereas if $p_i \geq 0$, then the $\max$ step must have $\mu_i p_i =0$, and in that case the objective value is finite — clearly the latter is smaller, so at the optimum of the right-hand side, $p_i\geq 0$ must hold; by the same argument we can show $\sum_{i=1}^n p_i = 1$ holds as well;
3. From step 2 we know that at the optimum of the right-hand side, $\boldsymbol{p}\in\Delta^{n-1}$ must hold, and the extra terms vanish, which makes it equivalent to the left-hand optimization problem.
Next we need the "Minimax Theorem":
If $\mathbb{X},\mathbb{Y}$ are two convex sets, $\boldsymbol{x}\in\mathbb{X},\boldsymbol{y}\in\mathbb{Y}$, and $f(\boldsymbol{x},\boldsymbol{y})$ is convex in $\boldsymbol{x}$ (for any fixed $\boldsymbol{y}$) and concave in $\boldsymbol{y}$ (for any fixed $\boldsymbol{x}$), then
\begin{equation}\min_{\boldsymbol{x}\in\mathbb{X}}\max_{\boldsymbol{y}\in\mathbb{Y}} f(\boldsymbol{x},\boldsymbol{y}) = \max_{\boldsymbol{y}\in\mathbb{Y}}\min_{\boldsymbol{x}\in\mathbb{X}} f(\boldsymbol{x},\boldsymbol{y})\end{equation}
The Minimax theorem gives a sufficient condition under which $\min,\max$ can be swapped. This introduces a new term, "convex set," which refers to a set such that the weighted average of any two points in the set stays within the set, i.e.,
\begin{equation}(1-\lambda)\boldsymbol{x}_1 + \lambda \boldsymbol{x}_2\in \mathbb{X},\qquad\forall \boldsymbol{x}_1,\boldsymbol{x}_2\in \mathbb{X},\quad\forall \lambda\in [0, 1]\end{equation}
As you can see, the condition of being a convex set isn't too restrictive — $\mathbb{R}^n,\Delta^{n-1}$ are all convex sets, and so is the set of all non-negative numbers, and so on.
For the objective function on the right-hand side of $\eqref{eq:min-max}$, it is linear in $\mu_i,\lambda$, so it satisfies the requirement of being concave in $\mu_i,\lambda$; and apart from the term $F(\boldsymbol{p})$, everything else is also linear in $\boldsymbol{p}$. So the convexity of the whole objective function in $\boldsymbol{p}$ is equivalent to the convexity of $F(\boldsymbol{p})$ in $\boldsymbol{p}$: that is, if $F(\boldsymbol{p})$ is convex in $\boldsymbol{p}$, then the $\min,\max$ in $\eqref{eq:min-max}$ can be swapped:
\begin{equation}\small\min_{\boldsymbol{p}\in\mathbb{R}^n} \max_{\mu_i \geq 0,\lambda\in\mathbb{R}}F(\boldsymbol{p}) - \sum_{i=1}^n \mu_i p_i + \lambda\left(\sum_{i=1}^n p_i - 1\right) = \max_{\mu_i \geq 0,\lambda\in\mathbb{R}} \min_{\boldsymbol{p}\in\mathbb{R}^n} F(\boldsymbol{p}) - \sum_{i=1}^n \mu_i p_i + \lambda\left(\sum_{i=1}^n p_i - 1\right)\end{equation}
This lets us first take the $\min$ over $\boldsymbol{p}$, which is an unconstrained minimization problem that can be solved by setting the gradient to zero and solving the resulting system of equations. The result will contain parameters $\lambda$ and $\mu_i$, which are finally determined via $p_i \geq 0$, $\mu_i p_i = 0$, and $\sum_{i=1}^n p_i = 1$.
Convex-Set Search
However, although Lagrange multipliers are regarded as the standard method for solving constrained optimization problems, they aren't very intuitive, and moreover they only give exact solutions via solving equations — they don't yield an iterative approximation algorithm analogous to gradient descent. So we shouldn't be satisfied with Lagrange multipliers alone.
From the search perspective, the key to solving optimization problems over probability space is to ensure that all trial points visited during the search remain within the set $\Delta^{n-1}$. In other words, suppose the current probability distribution is $\boldsymbol{p}_t\in \Delta^{n-1}$ — how do we construct the next trial point $\boldsymbol{p}_{t+\eta}$? It needs to satisfy two requirements: first, $\boldsymbol{p}_{t+\eta}\in \Delta^{n-1}$; second, we need to be able to control how close it is to $\boldsymbol{p}_t$ by controlling the magnitude of $\eta$. This is where the "convex set" property of $\Delta^{n-1}$ comes in handy — using this property, we can define $\boldsymbol{p}_{t+\eta}$ as
\begin{equation}\boldsymbol{p}_{t+\eta} = (1-\eta)\boldsymbol{p}_t + \eta \boldsymbol{q}_t,\quad \boldsymbol{q}_t\in \Delta^{n-1}\end{equation}
which gives us
\begin{equation}F(\boldsymbol{p}_{t+\eta}) = F((1-\eta)\boldsymbol{p}_t + \eta \boldsymbol{q}_t) \approx F(\boldsymbol{p}_t) + \eta(\boldsymbol{q}_t - \boldsymbol{p}_t)\cdot\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t)\end{equation}
Assuming the first-order approximation is accurate enough, finding the direction of steepest descent amounts to solving
\begin{equation}\mathop{\text{argmin}}_{\boldsymbol{q}_t\in\Delta^{n-1}}\,\boldsymbol{q}_t\cdot\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t) \end{equation}
This objective function turns out to be quite simple, and the answer is
\begin{equation}\boldsymbol{q}_t = \text{onehot}(\text{argmin}(\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t)))\end{equation}
Here $\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t)$ is a vector, and taking $\text{argmin}$ of a vector means finding the position of its smallest component. So the equation above says that $\boldsymbol{q}_t$ is a one-hot distribution, where the position of $1$ corresponds to the position of the smallest component of the gradient $\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t)$.
From this we see that the form of gradient descent in probability space is
\begin{equation}\boldsymbol{p}_{t+\eta} = (1 - \eta)\boldsymbol{p}_t + \eta\, \text{onehot}(\text{argmin}(\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t)))\end{equation}
and the condition for $\boldsymbol{p}_t$ to be a minimum of $F(\boldsymbol{p}_t)$ is:
\begin{equation}\boldsymbol{p}_t\cdot\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t) = (\nabla_{\boldsymbol{p}_t} F(\boldsymbol{p}_t))_{\min}\label{eq:p-min}\end{equation}
Here, $\min$ applied to a vector refers to returning its smallest component.
An Example
Let's take Sparsemax, introduced in The Road to Probability Distributions: A Survey of Softmax and Its Alternatives, as an example. Its original definition is
\begin{equation}Sparsemax(\boldsymbol{x}) = \mathop{\text{argmin}}\limits_{\boldsymbol{p}\in\Delta^{n-1}}\Vert \boldsymbol{p} - \boldsymbol{x}\Vert^2\end{equation}
where $\boldsymbol{x}\in\mathbb{R}^n$. It's not hard to see that, from the perspective of projected gradient descent discussed earlier, Sparsemax is exactly the "projection" operation from $\mathbb{R}^n$ onto $\Delta^{n-1}$.
Let's denote $F(\boldsymbol{p})=\Vert \boldsymbol{p} - \boldsymbol{x}\Vert^2$; its gradient with respect to $\boldsymbol{p}$ is $2(\boldsymbol{p} - \boldsymbol{x})$, so by equation $\eqref{eq:p-min}$, the equation satisfied by the minimum point is
\begin{equation}\boldsymbol{p}\cdot(\boldsymbol{p}-\boldsymbol{x}) = (\boldsymbol{p}-\boldsymbol{x})_{\min}\end{equation}
We adopt the convention that $x_i = x_j\Leftrightarrow p_i = p_j$ — here, an unbolded subscript such as $p_i$ denotes the $i$-th component of the vector $\boldsymbol{p}$ (i.e., a scalar), whereas the bolded subscript used earlier, such as $\boldsymbol{p}_t$, denotes the $t$-th iterate of $\boldsymbol{p}$ (i.e., still a vector) — please distinguish these carefully.
Under this convention, the equation above gives us
\begin{equation}p_i > 0 \quad \Leftrightarrow \quad p_i-x_i = (\boldsymbol{p}-\boldsymbol{x})_{\min}\end{equation}
Since $\boldsymbol{p}$ can be determined by $\boldsymbol{x}$, $(\boldsymbol{p}-\boldsymbol{x})_{\min}$ is a function of $\boldsymbol{x}$, which we denote by $-\lambda(\boldsymbol{x})$, giving us $p_i = x_i - \lambda(\boldsymbol{x})$ — but this only holds for $p_i > 0$; for $p_i=0$ we have $p_i-x_i > (\boldsymbol{p}-\boldsymbol{x})_{\min}$, i.e., $x_i - \lambda(\boldsymbol{x}) < 0$. Combining these two cases, we can write uniformly
\begin{equation}p_i = \text{relu}(x_i - \lambda(\boldsymbol{x}))\end{equation}
where $\lambda(\boldsymbol{x})$ is determined by requiring the components of $\boldsymbol{p}$ to sum to 1. For further details please refer to The Road to Probability Distributions: A Survey of Softmax and Its Alternatives.
Continuous Distributions
Having covered discrete distributions, let's now turn to continuous distributions. On the surface, continuous distributions might seem like merely the limiting version of discrete ones, so one might expect the results to be quite similar — but in fact their properties differ fundamentally, to the point that we need to build an entirely new methodology for continuous distributions.
The Objective Functional
Let's start with the objective function. As we know, continuous distributions are described by probability density functions, so the input to the objective here is a probability density function — and the objective function itself is no longer an ordinary function in this setting. We usually call it a "functional": a mapping from an entire function to a scalar. In other words, we need to find a probability density function that minimizes some target functional.
Although many people find "functional analysis chills the heart" (to borrow a common quip), in fact most of us have already encountered functionals, because mappings satisfying "input a function, output a scalar" are everywhere. For instance, the definite integral
\begin{equation}\mathcal{I}[f]\triangleq \int_a^b f(x) dx\end{equation}
is a mapping from a function to a scalar, so it too is a functional. In fact, essentially all the functionals we encounter in practical applications are built from definite integrals — for instance, the KL divergence between probability distributions:
\begin{equation}\mathcal{KL}[p\Vert q] = \int p(\boldsymbol{x})\log \frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}d\boldsymbol{x}\end{equation}
where, by default, the integral is taken over the whole space (all of $\mathbb{R}^n$). More general functionals may have integrands that also involve derivative terms, as in the principle of least action in theoretical physics:
\begin{equation}\mathcal{A}[x] = \int_{t_a}^{t_b} L(x(t),x'(t),t)dt\end{equation}
And the target functional we're about to minimize can, in general, be written as
\begin{equation}\mathcal{F}[p] = \int F(p(\boldsymbol{x}))d\boldsymbol{x}\end{equation}
For convenience, we can also define the functional derivative
\begin{equation}\frac{\delta\mathcal{F}[p]}{\delta p}(\boldsymbol{x}) = \frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\end{equation}
Compact Support
We also need a notation for the continuous probability space, whose basic definition is
\begin{equation}\mathbb{P} = \left\{p(\boldsymbol{x}) \,\Bigg|\, p(\boldsymbol{x})\geq 0(\forall\boldsymbol{x}\in\mathbb{R}^n),\int p(\boldsymbol{x})d\boldsymbol{x} = 1\right\}\end{equation}
It is not hard to prove that, if the limit $\lim_{\Vert\boldsymbol{x}\Vert\to\infty} p(\boldsymbol{x})$ of the probability density function $p(\boldsymbol{x})$ exists, then it must be the case that $\lim_{\Vert\boldsymbol{x}\Vert\to\infty} p(\boldsymbol{x}) = 0$ — a property we will use later in our derivations.
However, one can construct examples showing that not every probability density function has a well-defined limit at infinity. To sidestep such theoretical difficulties, when doing theoretical proofs we usually assume that the support of $p(\boldsymbol{x})$ is a compact set. This brings in two notions: support and compact set. The support is the set of all $\boldsymbol{x}$ for which $p(\boldsymbol{x}) > 0$ is nonzero, i.e.,
\begin{equation}\text{supp}(p) = \{\boldsymbol{x} | p(\boldsymbol{x}) > 0\}\end{equation}
The general definition of a compact set is fairly involved, but in $\mathbb{R}^n$, compactness is equivalent to being closed and bounded. So, in plain terms, assuming the support of $p(\boldsymbol{x})$ is compact directly gives $p(\boldsymbol{x})$ the property that "there exists a constant $C$ such that $p(\boldsymbol{x}) = 0$ holds whenever $\forall |\boldsymbol{x}| > C$," which simplifies the behavior of $p(\boldsymbol{x})$ at infinity and, fundamentally, sidesteps any discussion of $\lim_{\Vert\boldsymbol{x}\Vert\to\infty} p(\boldsymbol{x}) = 0$.
Theoretically, this is a rather strong assumption — it even rules out something as simple as the normal distribution (whose support is $\mathbb{R}^n$). But practically speaking, this assumption isn't too far-fetched, since we already noted that if the limit $\lim_{\Vert\boldsymbol{x}\Vert\to\infty} p(\boldsymbol{x})$ exists it must be zero, so beyond a certain range the density is effectively indistinguishable from zero anyway. There do exist examples where the limit fails to exist, but they generally have to be constructed rather artificially — for the kind of data we actually encounter in practice, the condition that the limit exists is basically always satisfied.
The Old Road Doesn't Work
Intuitively, one might expect optimization over continuous distributions to follow the same recipe as the discrete case — i.e., set $\boldsymbol{p}_{t+\eta}(\boldsymbol{x}) = (1 - \eta)\boldsymbol{p}_t(\boldsymbol{x}) + \eta \boldsymbol{q}_t(\boldsymbol{x})$ — since, just like the discrete case, the set of continuous probability density functions $\mathbb{P}$ is also a convex set. Let's substitute this into the target functional:
\begin{equation}\begin{aligned} \mathcal{F}[p_{t+\eta}] =&\, \int F(p_{t+\eta}(\boldsymbol{x}))d\boldsymbol{x} \\ =&\, \int F((1 - \eta)\boldsymbol{p}_t(\boldsymbol{x}) + \eta \boldsymbol{q}_t(\boldsymbol{x}))d\boldsymbol{x} \\ \approx&\,\int \left[F(p_t(\boldsymbol{x})) + \eta\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\Big(q_t(\boldsymbol{x}) - p_t(\boldsymbol{x})\Big)\right]d\boldsymbol{x} \end{aligned}\end{equation}
Assuming the first-order approximation suffices, the problem becomes
\begin{equation}\mathop{\text{argmin}}_{q_t\in \mathbb{P}}\int\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}q_t(\boldsymbol{x})d\boldsymbol{x}\end{equation}
This problem isn't hard to solve either, and the answer is analogous to the one-hot solution in the discrete case:
\begin{equation}q_t(\boldsymbol{x}) = \delta\left(\boldsymbol{x} - \mathop{\text{argmin}}_{\boldsymbol{x}'} \frac{\partial F(p_t(\boldsymbol{x}'))}{\partial p_t(\boldsymbol{x}')}\right)\end{equation}
Here $\delta(\cdot)$ is the Dirac delta function, representing the probability density of a point mass.
This all looks smooth so far, but in fact this path leads nowhere. First, the Dirac delta function is not a function in the ordinary sense — it's a generalized function (also a kind of functional); second, if we insist on viewing it as an ordinary function, the Dirac delta takes an infinite value at a single point, and given that infinite value, the "first-order approximation is good enough" assumption used in the derivation simply cannot hold.
Change of Variables
We could try to patch up the derivation from the previous section — for instance, by adding a constraint on $q_t(\boldsymbol{x}) \leq C$ to get a meaningful result — but such patchwork ultimately feels inelegant. But if we can't rely on the convex-set property, how should we construct the next trial distribution $\boldsymbol{p}_{t+\eta}(\boldsymbol{x})$?
This is where we need to make full use of the special properties of probability density functions — we can use a change of variables to transform one probability density function into another, which is a property unique to continuous distributions. Specifically, if $p(\boldsymbol{x})$ is a probability density function and $\boldsymbol{y}=\boldsymbol{T}(\boldsymbol{x})$ is an invertible transformation, then $p(\boldsymbol{T}(\boldsymbol{x}))\left|\frac{\partial \boldsymbol{T}(\boldsymbol{x})}{\partial\boldsymbol{x}}\right|$ is also a probability density function, where $|\cdot|$ denotes the absolute value of the matrix determinant.
Based on this property, we define the next trial distribution as
\begin{equation}\begin{aligned} p_{t+\eta}(\boldsymbol{x}) =&\, p_t(\boldsymbol{x} + \eta\boldsymbol{\mu}_t(\boldsymbol{x}))\left|\boldsymbol{I} + \eta\frac{\partial \boldsymbol{\mu}_t(\boldsymbol{x})}{\partial\boldsymbol{x}}\right| \\ \approx &\, \Big[p_t(\boldsymbol{x}) + \eta\boldsymbol{\mu}_t(\boldsymbol{x})\cdot\nabla_{\boldsymbol{x}} p_t(\boldsymbol{x})\Big]\left[1 + \eta\,\text{Tr}\frac{\partial \boldsymbol{\mu}_t(\boldsymbol{x})}{\partial\boldsymbol{x}}\right] \\[3pt] \approx &\, p_t(\boldsymbol{x}) + \eta\boldsymbol{\mu}_t(\boldsymbol{x})\cdot\nabla_{\boldsymbol{x}} p_t(\boldsymbol{x}) + \eta\, p_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\cdot\boldsymbol{\mu}_t(\boldsymbol{x}) \\[5pt] = &\, p_t(\boldsymbol{x}) + \eta\nabla_{\boldsymbol{x}}\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big] \\ \end{aligned}\end{equation}
We derived the same result in Rambling on Generative Diffusion Models (12): "Head-On" with the Diffusion ODE, where the approximate expansion of the determinant can be found in The Derivative of a Determinant.
Integral Transformation
Using this new $p_{t+\eta}(\boldsymbol{x})$, we obtain
\begin{equation}\begin{aligned} \mathcal{F}[p_{t+\eta}] =&\, \int F(p_{t+\eta}(\boldsymbol{x}))d\boldsymbol{x} \\ \approx&\, \int F\Big(p_t(\boldsymbol{x}) + \eta\nabla_{\boldsymbol{x}}\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big]\Big)d\boldsymbol{x} \\ \approx&\, \int \left[F(p_t(\boldsymbol{x})) + \eta\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\nabla_{\boldsymbol{x}}\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big]\right]d\boldsymbol{x} \\ =&\, \mathcal{F}[p_t] + \eta\int \frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\nabla_{\boldsymbol{x}}\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big] d\boldsymbol{x} \\ \end{aligned}\label{eq:px-approx}\end{equation}
Next, as in Deriving the Continuity Equation and the Fokker-Planck Equation via Test Functions, we need to derive an integral identity involving the probability density. First, we have
\begin{equation}\begin{aligned} &\,\int \nabla_{\boldsymbol{x}}\cdot\left[\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})} p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\right] d\boldsymbol{x} \\[5pt] =&\, \int \frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\nabla_{\boldsymbol{x}}\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big] d\boldsymbol{x} + \int \left(\nabla_{\boldsymbol{x}}\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\right)\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big] d\boldsymbol{x} \end{aligned}\end{equation}
By the divergence theorem, we have
\begin{equation}\int_{\Omega} \nabla_{\boldsymbol{x}}\cdot\left[\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})} p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\right] d\boldsymbol{x} = \int_{\partial\Omega} \left[\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})} p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\right]\cdot \hat{\boldsymbol{n}} dS\end{equation}
where $\Omega$ is the integration domain — here, the entirety of $\mathbb{R}^n$ — $\partial\Omega$ is its boundary (which for $\mathbb{R}^n$ is naturally at infinity), $\hat{\boldsymbol{n}}$ is the outward unit normal vector on the boundary, and $dS$ is the surface area element. Under the compact-support assumption, $p_t(\boldsymbol{x})=0$ vanishes at infinity, so the right-hand side above is really an integral of zero, giving a result of zero. Therefore, we have
\begin{equation}\int \frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\nabla_{\boldsymbol{x}}\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big] d\boldsymbol{x} = - \int \left(\nabla_{\boldsymbol{x}}\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\right)\cdot\big[p_t(\boldsymbol{x})\boldsymbol{\mu}_t(\boldsymbol{x})\big] d\boldsymbol{x}\end{equation}
Substituting into equation $\eqref{eq:px-approx}$ gives
\begin{equation}\mathcal{F}[p_{t+\eta}] \approx \mathcal{F}[p_t] - \eta\int \left(p_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\right)\cdot \boldsymbol{\mu}_t(\boldsymbol{x}) d\boldsymbol{x} \label{eq:px-approx-2}\end{equation}
Gradient Flow
Based on equation $\eqref{eq:px-approx-2}$, one simple choice that makes $\mathcal{F}[p_{t+\eta}] \leq \mathcal{F}[p_t]$ decrease is
\begin{equation}\boldsymbol{\mu}_t(\boldsymbol{x}) = \nabla_{\boldsymbol{x}}\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\end{equation}
The corresponding iterative scheme is
\begin{equation}p_{t+\eta}(\boldsymbol{x}) \approx p_t(\boldsymbol{x}) + \eta\nabla_{\boldsymbol{x}}\cdot\left[p_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\right] \end{equation}
We can also take the limit $\eta\to 0$ to get
\begin{equation}\frac{\partial}{\partial t}p_t(\boldsymbol{x}) = \nabla_{\boldsymbol{x}}\cdot\left[p_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\right] \end{equation}
or, more compactly,
\begin{equation}\frac{\partial p_t}{\partial t} = \nabla\cdot\left[p_t\nabla\frac{\delta \mathcal{F}[p_t]}{\delta p_t}\right] \end{equation}
This is exactly the Wasserstein gradient flow introduced in Gradient Flow: Exploring the Path Toward the Minimum — except here we arrived at the same result without ever introducing the notion of Wasserstein distance.
Since $p_{t+\eta}(\boldsymbol{x})$ is obtained from $p_t(\boldsymbol{x})$ via the transformation $\boldsymbol{x}\to \boldsymbol{x} + \eta \boldsymbol{\mu}_t(\boldsymbol{x})$, we can also write down the ODE governing the trajectory of motion of $\boldsymbol{x}$:
\begin{equation}\boldsymbol{x}_t = \boldsymbol{x}_{t+\eta} + \eta \boldsymbol{\mu}_t(\boldsymbol{x}_{t+\eta})\quad\Rightarrow\quad \frac{d\boldsymbol{x}_t}{dt} = -\boldsymbol{\mu}_t(\boldsymbol{x}_t) = -\nabla_{\boldsymbol{x}}\frac{\partial F(p_t(\boldsymbol{x}))}{\partial p_t(\boldsymbol{x})}\end{equation}
The meaning of this ODE is: starting from a sample $\boldsymbol{x}_0$ drawn from distribution $p_0(\boldsymbol{x})$, if we evolve it according to this ODE up to time $\boldsymbol{x}_t$, then the distribution followed by $\boldsymbol{x}_t$ at that time is exactly $p_t(\boldsymbol{x})$.
Summary
This post has systematically laid out methods for minimizing target functions over probability space, including the necessary conditions for attaining a minimum and iterative methods analogous to gradient descent. These results find frequent use in optimization and generative modeling (especially diffusion models), among other settings.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.