Gradient Flow: Exploring the Path Toward the Minimum

梯度流：探索通向最小值之路 苏剑林 · 2023-06-16 · Mathematics functional dynamics optimization gradient

In this article, we'll explore a concept called "Gradient Flow." Simply put, gradient flow connects the successive points visited while searching for a minimum via gradient descent into a trajectory that evolves over a (virtual) time variable; this trajectory is what we call the "gradient flow." In the latter half of the article, we'll focus on extending the concept of gradient flow to probability space, arriving at "Wasserstein gradient flow," which offers a new perspective for understanding the continuity equation, the Fokker-Planck equation, and related topics.

Gradient Descent

Suppose we want to search for the minimum of a smooth function $f(\boldsymbol{x})$. A common approach is gradient descent, which iterates as follows:

\begin{equation}\boldsymbol{x}_{t+1} = \boldsymbol{x}_t -\alpha \nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\label{eq:gd-d}\end{equation}

If $f(\boldsymbol{x})$ is convex in $\boldsymbol{x}$, then gradient descent will typically find the minimum point; otherwise, it will generally only converge to a "stationary point"—that is, a point where the gradient is zero—and in the more favorable case, converge to a local minimum. Here we won't draw a strict distinction between local minima and the global minimum, because in deep learning, even converging to a local minimum is already something to be grateful for. more

If we write $\alpha$ as $\Delta t$, and $\boldsymbol{x}_{t+1}$ as $\boldsymbol{x}_{t+\Delta t}$, and then consider the limit $\Delta t\to 0$, equation $\eqref{eq:gd-d}$ turns into an ODE:

\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = -\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\label{eq:gd-c}\end{equation}

The trajectory $\boldsymbol{x}_t$ obtained by solving this ODE is what we call the "Gradient Flow"—that is, the gradient flow is the trajectory traced out by gradient descent as it searches for a minimum. Given that $\eqref{eq:gd-c}$ holds, we also have:

\begin{equation}\frac{df(\boldsymbol{x}_t)}{dt} = \left\langle\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t),\frac{d\boldsymbol{x}_t}{dt}\right\rangle = -\Vert\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\Vert^2 \leq 0\end{equation}

This means that as long as $\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\neq\boldsymbol{0}$, when the learning rate is small enough, gradient descent will always move in the direction that decreases $f(\boldsymbol{x})$.

For more related discussion, see our earlier series on optimization algorithms, such as "A Dynamical Systems View of Optimization Algorithms (I): From SGD to Momentum" and "A Dynamical Systems View of Optimization Algorithms (III): A More Holistic Perspective".

The Steepest Direction

Why use gradient descent? The mainstream explanation is that "the negative gradient direction is the direction of locally fastest descent"—searching this phrase directly turns up plenty of material. This statement isn't exactly wrong, but it's a bit loose, because it fails to state its preconditions—being "fastest" necessarily involves a quantitative comparison, and only once we fix the metric of comparison can we determine what "fastest" actually means.

If we only care about the direction of fastest descent, the objective of gradient descent should be:

\begin{equation}\boldsymbol{x}_{t+1} = \mathop{\text{argmin}}_{\boldsymbol{x},\Vert\boldsymbol{x} - \boldsymbol{x}_t\Vert = \epsilon} f(\boldsymbol{x})\label{eq:gd-min-co}\end{equation}

Assuming a first-order approximation suffices, we have

\begin{equation}\begin{aligned} f(\boldsymbol{x})&\,=f(\boldsymbol{x}_t) + \langle \nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t),\boldsymbol{x} - \boldsymbol{x}_t\rangle\\ &\,\geq f(\boldsymbol{x}_t) - \Vert\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\Vert \Vert\boldsymbol{x} - \boldsymbol{x}_t\Vert\\ &\,= f(\boldsymbol{x}_t) - \Vert\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\Vert \epsilon\\ \end{aligned}\end{equation}

with equality holding when

\begin{equation}\boldsymbol{x} - \boldsymbol{x}_t = -\epsilon\frac{\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)}{\Vert\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\Vert}\quad\Rightarrow\quad\boldsymbol{x}_{t+1} = \boldsymbol{x}_t - \epsilon\frac{\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)}{\Vert\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\Vert}\label{eq:gd-d-norm} \end{equation}

We can see that the resulting update direction is exactly the negative gradient direction, which is why it's called the direction of locally fastest descent. However, don't forget this result was obtained under the constraint $\Vert\boldsymbol{x} - \boldsymbol{x}_t\Vert = \epsilon$, where $\Vert\cdot\Vert$ is the norm of Euclidean space; if we change the definition of the norm, or simply change the constraint altogether, the result will differ. So, to be precise, we should say: "In Euclidean space, the negative gradient direction is the direction of locally fastest descent."

An Optimization Perspective

Equation $\eqref{eq:gd-min-co}$ is a constrained optimization problem, which is somewhat cumbersome both to generalize and to solve. Moreover, the solution to $\eqref{eq:gd-min-co}$ is $\eqref{eq:gd-d-norm}$, not the original gradient descent update $\eqref{eq:gd-d}$. In fact, it can be shown that the optimization objective corresponding to $\eqref{eq:gd-d}$ is

\begin{equation}\boldsymbol{x}_{t+1} = \mathop{\text{argmin}}_{\boldsymbol{x}} \frac{\Vert\boldsymbol{x} - \boldsymbol{x}_t\Vert^2}{2\alpha} + f(\boldsymbol{x})\label{eq:gd-min}\end{equation}

That is, we fold the constraint into the objective as a penalty term, so that we no longer need to solve a constrained problem, and generalization becomes easier too. Moreover, even with this extra $\frac{\Vert\boldsymbol{x} - \boldsymbol{x}_t\Vert^2}{2\alpha}$ term added, we can still guarantee that optimizing the above objective won't push us toward a worse point, since substituting $\boldsymbol{x} = \boldsymbol{x}_t$ makes the objective function above exactly equal to $f(\boldsymbol{x}_t)$, so the optimum $\min_{\boldsymbol{x}}$ can be no larger than $f(\boldsymbol{x}_t)$.

When $\alpha$ is small enough, the first term dominates, so for the first term to become small enough $\Vert\boldsymbol{x} - \boldsymbol{x}_t\Vert$ must be small, meaning the optimal point must be very close to $\boldsymbol{x}_t$. We can therefore expand $f(\boldsymbol{x})$ around $\boldsymbol{x}_t$, obtaining

\begin{equation}\boldsymbol{x}_{t+1} = \mathop{\text{argmin}}_{\boldsymbol{x}} \frac{\Vert\boldsymbol{x} - \boldsymbol{x}_t\Vert^2}{2\alpha} + f(\boldsymbol{x}_t)+\langle\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t),\boldsymbol{x}-\boldsymbol{x}_t\rangle\end{equation}

This is now simply the minimization of a quadratic function, and the solution is exactly equation $\eqref{eq:gd-d}$.

Clearly, besides the squared norm, we could consider other regularization terms, leading to different gradient descent schemes. For example, Natural Gradient Descent uses the KL divergence as its regularization term:

\begin{equation}\boldsymbol{x}_{t+1} = \mathop{\text{argmin}}_{\boldsymbol{x}} \frac{KL(p(\boldsymbol{y}|\boldsymbol{x})\Vert p(\boldsymbol{y}|\boldsymbol{x}_t))}{\alpha} + f(\boldsymbol{x})\end{equation}

where $p(\boldsymbol{y}|\boldsymbol{x})$ is some probability distribution associated with $f(\boldsymbol{x})$. To solve the above, we again expand around $f(\boldsymbol{x})$, expanding $f(\boldsymbol{x})$ to first order as before; but the KL divergence is special—its first-order expansion vanishes (see here), so we must expand it at least to second order, giving

\begin{equation}\boldsymbol{x}_{t+1} = \mathop{\text{argmin}}_{\boldsymbol{x}} \frac{(\boldsymbol{x}-\boldsymbol{x}_t)^{\top}\boldsymbol{F}(\boldsymbol{x}-\boldsymbol{x}_t)}{2\alpha} + f(\boldsymbol{x}_t)+\langle\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t),\boldsymbol{x}-\boldsymbol{x}_t\rangle\end{equation}

Here $\boldsymbol{F}$ is the Fisher information matrix; we won't go into the computational details, which can also be found here. Now the equation above is again essentially a quadratic minimization problem, with the result

\begin{equation}\boldsymbol{x}_{t+1} = \boldsymbol{x}_t -\alpha \boldsymbol{F}^{-1}\nabla_{\boldsymbol{x}_t}f(\boldsymbol{x}_t)\end{equation}

This is what's known as "Natural Gradient Descent."

A Primer on Functionals

Equation $\eqref{eq:gd-min}$ allows us not only to generalize the regularization term, but also to generalize the optimization target itself—for instance, to functionals.

The word "functional" (泛函) might sound intimidating, but in fact long-time readers of this blog will have encountered it many times already. Simply put: an ordinary multivariate function takes a vector as input and outputs a scalar, whereas a functional takes a function as input and outputs a scalar—for example, the definite integral:

\begin{equation}\mathcal{I}[f] = \int_a^b f(x)dx\end{equation}

For any function $f$, the result of computing $\mathcal{I}[f]$ is a scalar, so $\mathcal{I}[f]$ is a functional. Another example is the KL divergence mentioned earlier, defined as

\begin{equation}KL(p\Vert q) = \int p(\boldsymbol{x})\log \frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}d\boldsymbol{x}\end{equation}

Here the integral is understood to be over the whole space by default. If we fix $p(\boldsymbol{x})$, this becomes a functional of $q(\boldsymbol{x})$, because $q(\boldsymbol{x})$ is a function—feed in any function satisfying the required conditions, and $KL(p\Vert q)$ will output a scalar. More generally, the $f$-divergence introduced in "An Introduction to f-GAN: A Production Line for GAN Models" is likewise a kind of functional. These are all relatively simple functionals; more complex functionals may involve derivatives of the input function, such as the principle of least action in theoretical physics.

Below, we'll mainly be concerned with functionals whose domain is the set of all probability density functions—that is, functionals that take a probability density as input and output a scalar.

The Flow of Probability

Suppose we have a functional $\mathcal{F}[q]$ whose minimum we want to compute. Following the logic of gradient descent, as long as we can find some notion of its "gradient," we can iterate in the negative direction of that gradient.

To pin down the iteration scheme, let's continue our earlier line of thought and generalize equation $\eqref{eq:gd-min}$, where $f(\boldsymbol{x})$ is naturally replaced by $\mathcal{F}[q]$. What should the first-term regularizer be replaced with? In equation $\eqref{eq:gd-min}$ it was the squared Euclidean distance, so it's natural to guess that here too it should be replaced by some squared distance. For probability distributions, a distance with particularly good properties is the Wasserstein distance (more precisely, the "2-Wasserstein distance"):

\begin{equation}\mathcal{W}_2[p,q]=\sqrt{\inf_{\gamma\in \Pi[p,q]} \iint \gamma(\boldsymbol{x},\boldsymbol{y}) \Vert\boldsymbol{x}-\boldsymbol{y}\Vert^2 d\boldsymbol{x}d\boldsymbol{y}}\end{equation}

We won't go into detail about it here; interested readers should consult "From Wasserstein Distance and Duality Theory to WGAN". If we further replace the Euclidean distance in equation $\eqref{eq:gd-min}$ with the Wasserstein distance, the final objective becomes

\begin{equation}q_{t+1} = \mathop{\text{argmin}}_{q} \frac{\mathcal{W}_2^2[q,q_t]}{2\alpha} + \mathcal{F}[q]\end{equation}

Unfortunately, the author is not able to give a concise derivation of the solution to the above objective—indeed, even the author himself hasn't fully grasped the derivation—so I'll simply state the solution, following Introduction to Gradient Flows in the 2-Wasserstein Space and {Euclidean, Metric, and Wasserstein} Gradient Flows: an overview:

\begin{equation}q_{t+1}(\boldsymbol{x}) = q_t(\boldsymbol{x}) + \alpha \nabla_{\boldsymbol{x}}\cdot\left(q_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\frac{\delta \mathcal{F}[q_t(\boldsymbol{x})]}{\delta q_t(\boldsymbol{x})}\right)\end{equation}

or, taking the limit,

\begin{equation}\frac{\partial q_t(\boldsymbol{x})}{\partial t} = \nabla_{\boldsymbol{x}}\cdot\left(q_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\frac{\delta \mathcal{F}[q_t(\boldsymbol{x})]}{\delta q_t(\boldsymbol{x})}\right)\end{equation}

This is the "Wasserstein Gradient Flow," where $\frac{\delta \mathcal{F}[q]}{\delta q}$ is the variational derivative of $\mathcal{F}[q]$; for a functional given by a definite integral, the variational derivative is just the derivative of the integrand:

\begin{equation}\mathcal{F}[q] = \int F(q(\boldsymbol{x}))d\boldsymbol{x} \quad\Rightarrow\quad \frac{\delta \mathcal{F}[q(\boldsymbol{x})]}{\delta q(\boldsymbol{x})} = \frac{\partial F(q(\boldsymbol{x}))}{\partial q(\boldsymbol{x})}\end{equation}

Some Examples

According to "An Introduction to f-GAN: A Production Line for GAN Models", the $f$-divergence is defined as

\begin{equation}\mathcal{D}_f(p\Vert q) = \int q(\boldsymbol{x}) f\left(\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}\right)d\boldsymbol{x}\end{equation}

Fixing $p$ and letting $\mathcal{F}[q]=\mathcal{D}_f(p\Vert q)$, we obtain

\begin{equation}\frac{\partial q_t(\boldsymbol{x})}{\partial t} = \nabla_{\boldsymbol{x}}\cdot\Big(q_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\big(f(r_t(\boldsymbol{x})) - r_t(\boldsymbol{x}) f'(r_t(\boldsymbol{x}))\big)\Big)\label{eq:wgd}\end{equation}

where $r_t(\boldsymbol{x}) = \frac{p(\boldsymbol{x})}{q_t(\boldsymbol{x})}$. Following "Deriving the Continuity Equation and Fokker-Planck Equation via Test Functions", the equation above takes the form of a continuity equation, so the ODE

\begin{equation}\frac{d\boldsymbol{x}}{dt} = -\nabla_{\boldsymbol{x}}\big(f(r_t(\boldsymbol{x})) - r_t(\boldsymbol{x}) f'(r_t(\boldsymbol{x}))\big)\end{equation}

lets us sample from the distribution $q_t$. Based on our earlier discussion, equation $\eqref{eq:wgd}$ is the Wasserstein gradient flow minimizing the $f$-divergence between $p,q$ and itself; the $f$-divergence is zero precisely when $t\to\infty$, i.e., $q_t=p$, so at $t\to\infty$ the above ODE realizes sampling from the distribution $p$. However, this result is currently only of formal interest and has no practical use, because it would require us to know the closed form of the distribution $p$, and also to solve for $q_t$ from equation $\eqref{eq:wgd}$, before we could even compute the right-hand side of the ODE and thereby carry out the sampling—this is an extremely difficult computation that generally cannot be carried out in practice.

A relatively simpler example is the (reverse) KL divergence, in which case $f=-\log$; substituting into equation $\eqref{eq:wgd}$ gives

\begin{equation}\begin{aligned}\frac{\partial q_t(\boldsymbol{x})}{\partial t} =&\, - \nabla_{\boldsymbol{x}}\cdot\left(q_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\log \frac{p(\boldsymbol{x})}{q_t(\boldsymbol{x})}\right)\\ =&\, - \nabla_{\boldsymbol{x}}\cdot\Big(q_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\big(\log p(\boldsymbol{x}) - \log q_t(\boldsymbol{x})\big)\Big)\\ =&\, - \nabla_{\boldsymbol{x}}\cdot\big(q_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\log p(\boldsymbol{x})\big) + \nabla_{\boldsymbol{x}}\cdot\nabla_{\boldsymbol{x}} q_t(\boldsymbol{x}) \end{aligned}\end{equation}

Comparing again with the results of "Deriving the Continuity Equation and Fokker-Planck Equation via Test Functions", this is exactly a Fokker-Planck equation, corresponding to the SDE:

\begin{equation}d\boldsymbol{x} = \nabla_{\boldsymbol{x}}\log p(\boldsymbol{x}) dt + \sqrt{2}dw\end{equation}

That is, if we know $\log p(\boldsymbol{x})$, then we can use the equation above to sample from $p(\boldsymbol{x})$. Compared with the earlier ODE approach, this dispenses with the need to solve for $q_t(\boldsymbol{x})$, making it a relatively more usable scheme.

Summary

This article introduced the concept of "gradient flow" that arises when searching for a minimum via gradient descent, including its extension from vector-space gradient flow to Wasserstein gradient flow in probability space, and its connections to the continuity equation, the Fokker-Planck equation, and ODE/SDE-based sampling.