From Sampling to Optimization: A Unified Perspective on Differentiable and Non-Differentiable Optimization

从采样看优化：可导优化与不可导优化的统一视角 苏剑林 · 2020-06-23 · Mathematics optimization gradient optimizer sampling

Many readers are probably already aware that the inconsistency between the loss function and the evaluation metric is a classic phenomenon in machine learning. For instance, classification problems use cross-entropy as the loss function while the evaluation metric is accuracy or F1; text generation uses teacher-forcing cross-entropy as the loss function while the evaluation metric is BLEU, ROUGE, and so on. Ideally, we would directly optimize whatever metric we evaluate on, but evaluation metrics are typically non-differentiable, whereas most of us use gradient-based optimizers, which require the minimization target to be differentiable. This is the root of the inconsistency.

A few days ago I came across a paper on arXiv titled MLE-guided parameter search for task loss minimization in neural sequence modeling. As the name suggests, it studies how to directly optimize evaluation metrics for text generation. After reading it, I found the paper quite valuable — it in fact provides a new way of thinking about optimizing evaluation metrics, and its applicability is not limited to text generation. What's more, it even contains a unified perspective for understanding differentiable and non-differentiable optimization.

A Sampling Perspective

First, let's reconsider the optimization problem through the lens of sampling. Suppose the model's current parameters are $\theta$ and the optimization objective is $l(\theta)$. We want to determine the next update $\Delta\theta$. To do this, we first construct the distribution

\begin{equation}p(\Delta\theta|\theta)=\frac{e^{-[l(\theta + \Delta\theta) - l(\theta)]/\alpha}}{Z(\theta)},\quad Z(\theta) = \int e^{-[l(\theta + \Delta\theta) - l(\theta)]/\alpha} d(\Delta\theta)\end{equation}more

where $\alpha > 0$ is a hyperparameter. The meaning of this distribution is clear: we treat $\Delta\theta$ as a random variable, and the smaller $l(\theta + \Delta\theta)$ is, the greater the probability of the corresponding $\Delta\theta$. With this distribution in hand, we define the next update as its expectation

\begin{equation}\Delta\theta_* = \int p(\Delta\theta|\theta)\Delta\theta d(\Delta\theta) = \mathbb{E}_{\Delta\theta\sim p(\Delta\theta|\theta)}[\Delta\theta]\label{eq:delta}\end{equation}

In this perspective, we make no assumption about the differentiability of $l(\theta)$, so the above definition applies universally to both differentiable and non-differentiable optimization. Moreover, we can control the stability of the update by tuning $\alpha$: when $\alpha\to 0$, the definition of $p(\Delta\theta|\theta)$ implies that only the $\Delta\theta$ minimizing $l(\theta + \Delta\theta)$ has nonzero probability, meaning $\Delta\theta_*$ is the direction of steepest descent; when $\alpha\to +\infty$, $p(\Delta\theta|\theta)$ tends toward a uniform distribution, so $\Delta\theta_*$ tends toward zero — that is, maximally stable. By choosing an appropriate $\alpha$, in theory we can strike a good balance between "fast" and "stable" during optimization, which intuitively should yield better generalization.

Of course, up to this point the definition is still purely theoretical — we don't yet know the analytic form of $p(\Delta\theta|\theta)$, nor how to sample from it, let alone how to compute its expectation. Below we'll see how this theoretical form is gradually realized in both the differentiable and non-differentiable settings.

Differentiable Objectives

For a differentiable $l(\theta)$, although we cannot solve for $p(\Delta\theta|\theta)$ exactly, we can perform a Taylor expansion to obtain an approximate distribution, and from there estimate $\Delta\theta_*$. It turns out that expanding to first and second order respectively gives us gradient descent and Newton's method. In other words, gradient descent and Newton's method are, in some sense, both special cases of this perspective.

Gradient Descent

As a first attempt, let's assume $l(\theta)$ is first-order differentiable. Then the Taylor expansion gives

\begin{equation}l(\theta + \Delta\theta) - l(\theta)\approx \Delta\theta^{\top}\nabla_{\theta}l(\theta)\end{equation}

This is exactly $p(\Delta\theta|\theta)\sim e^{-\Delta\theta^{\top}\nabla_{\theta}l(\theta)/\alpha}$. If $\Delta\theta$ is unconstrained, normalization is impossible, so let's restrict $\Vert\Delta\theta\Vert\leq \epsilon$ and write $\nabla_{\theta}l(\theta)=g$. Then

\begin{equation}p(\Delta\theta|\theta) = \frac{e^{-\Delta\theta^{\top}g/\alpha}}{Z(g)},\quad Z(g)=\int_{\Vert\Delta\theta\Vert\leq\epsilon}e^{-\Delta\theta^{\top}g/\alpha}d(\Delta\theta)\end{equation}

Clearly $\Delta\theta_* = -\alpha\nabla_g \ln Z(g)$, so the key is to find $Z(g)$. Let the angle between $\Delta\theta$ and $g$ be $\eta$. Then

\begin{equation}Z(g)=\int_{\Vert\Delta\theta\Vert\leq\epsilon}e^{-\Vert\Delta\theta\Vert\times\Vert g\Vert \times (\cos\eta) / \alpha}d(\Delta\theta)\end{equation}

This is an integral over a high-dimensional sphere, and due to isotropy, once the norm $\Vert g\Vert$ is fixed, the whole integral is determined — that is, $Z(g)$ depends only on the norm of $g$, not on its direction, so we can also write it as $Z(\Vert g\Vert)$. We don't need to know the explicit form of $Z(\Vert g\Vert)$; it's enough to know that it's a function only of the norm $\Vert g\Vert$. In that case

\begin{equation}\Delta\theta_* = -\alpha\nabla_g \ln Z(g)= - \frac{Z'(\Vert g\Vert)}{Z(\Vert g\Vert)}\alpha\nabla_g\Vert g\Vert = - \frac{Z'(\Vert g\Vert)}{Z(\Vert g\Vert)}\frac{\alpha g}{\Vert g\Vert}\end{equation}

So the direction of $\Delta\theta_*$ is the direction of $-g$, i.e., the negative gradient direction — and this is exactly how we derive gradient descent. We can say it is the first-order approximation of equation $\eqref{eq:delta}$. Incidentally, it is possible to compute $Z(g)$ explicitly, though it's not an elementary function; see the discussion on Stack Exchange, Integral of exp over the unit ball, for the derivation.

Newton's Method

If $l(\theta)$ is second-order differentiable, we can expand to second order:

\begin{equation}l(\theta + \Delta\theta) - l(\theta)\approx \Delta\theta^{\top}\nabla_{\theta}l(\theta) + \frac{1}{2}\Delta\theta^{\top}\nabla_{\theta}^2 l(\theta) \Delta\theta\end{equation}

Writing $g=\nabla_{\theta}l(\theta),\mathcal{H}=\nabla_{\theta}^2 l(\theta)$, we get

\begin{equation}\begin{aligned} \log p(\Delta\theta|\theta)\sim&\, -\Delta\theta^{\top}g - \frac{1}{2}\Delta\theta^{\top} \mathcal{H} \Delta\theta\\ =&\, - \frac{1}{2}\left(\Delta\theta+\mathcal{H}^{-1}g\right)^{\top}\mathcal{H}\left(\Delta\theta+\mathcal{H}^{-1}g\right)+ \frac{1}{2}g^{\top} \mathcal{H}^{-1} g \end{aligned}\end{equation}

Clearly, since the exponent of $p(\Delta\theta|\theta)$ is quadratic in $\theta$, $p(\Delta\theta|\theta)$ is a Gaussian distribution, and the above expression tells us that this Gaussian has mean $-\mathcal{H}^{-1}g$ and covariance matrix $\mathcal{H}^{-1}$, so $\Delta\theta_*=-\mathcal{H}^{-1}g$. This result corresponds exactly to Newton's method, so we can say that Newton's method is the second-order approximation of equation $\eqref{eq:delta}$.

Non-Differentiable Objectives

For a non-differentiable $l(\theta)$, the Taylor-expansion approximation above is no longer possible, and in theory we can only estimate $\Delta\theta_*$ via direct sampling. The original paper proposes that we can use importance sampling to improve sampling efficiency and estimation accuracy — this is the core idea and main contribution of the paper.

Importance Sampling

Let me first briefly introduce the general concept of importance sampling. Suppose we have a probability distribution $p(x)$ and a function $f(x)$, and we want to estimate

\begin{equation}\int p(x)f(x)dx = \mathbb{E}_{x\sim p(x)}[f(x)]\end{equation}

This requires us to draw several samples $x_1,x_2,\dots,x_n$ from $p(x)$ and then compute $\frac{1}{n}\sum\limits_{i=1}^n f(x_i)$. However, there may be two difficulties here:

1. We might not know how to sample from $p(x)$ at all.
2. Even if we do know how to sample from $p(x)$, in scenarios like VAEs, $p(x)$ is parameterized and requires gradients to be preserved, which direct sampling-based computation may not achieve.

In such cases, importance sampling may come to the rescue. It requires us to find a distribution $q(x)$ whose probability density expression we know and which is easy to sample from, and then rewrite the expression as follows:

\begin{equation}\int p(x)f(x)dx = \int q(x)\left[\frac{p(x)}{q(x)}f(x)\right]dx = \mathbb{E}_{x\sim q(x)}\left[\frac{p(x)}{q(x)}f(x)\right]\label{eq:is}\end{equation}

Now sampling is shifted to $q(x)$, and by our assumption, sampling from $q(x)$ is easy, and the analytic form of $q(x)$ is already known, so $\frac{p(x)}{q(x)}f(x)$ can also be computed. If $p(x)$ has parameters that require gradients, those gradients are also preserved. Clearly, the closer $q(x)$ is to $p(x)$, the more efficient the estimation; $q(x)$ represents a prior estimate of the "importance" of each sample from $p(x)$, which is why this approach is called importance sampling.

Given this, assuming $x_1,x_2,\dots,x_n\sim q(x)$, we have

\begin{equation}\mathbb{E}_{x\sim p(x)}[f(x)]\approx \frac{1}{n}\sum_{i=1}^n \frac{p(x_i)}{q(x_i)}f(x_i)\label{eq:is-2}\end{equation}

There is, however, a small remaining issue: both equation $\eqref{eq:is}$ and equation $\eqref{eq:is-2}$ require us to know the exact expression for $p(x)$, and sometimes we cannot even manage this — for example, in the case above, we only know that $p(\Delta\theta|\theta)$ is proportional to $e^{-[l(\theta + \Delta\theta) - l(\theta)]/\alpha}$, and its normalizing factor cannot be computed directly. In this case, we can rely on the relation

\begin{equation}1=\int p(x)dx=\int q(x)\left[\frac{p(x)}{q(x)}\right]dx=\mathbb{E}_{x\sim q(x)}\left[\frac{p(x)}{q(x)}\right]\approx\frac{1}{n}\sum_{i=1}^n \frac{p(x_i)}{q(x_i)}\end{equation}

That is, $\left[\frac{1}{n}\frac{p(x_1)}{q(x_1)},\frac{1}{n}\frac{p(x_2)}{q(x_2)},\dots,\frac{1}{n}\frac{p(x_n)}{q(x_n)}\right]$ should be approximately normalized. If we only know $p(x)\sim \rho(x)$ but not its normalizing factor, we can normalize it manually, and then equation $\eqref{eq:is-2}$ becomes

\begin{equation}\mathbb{E}_{x\sim p(x)}[f(x)]\approx \sum_{i=1}^n \frac{\rho(x_i)\big/q(x_i)}{\sum\limits_{i=1}^n \rho(x_i)\big/q(x_i)}f(x_i)\label{eq:is-3}\end{equation}

This avoids having to compute the normalizing factor at all.

Leveraging Differentiability

Now that we have all the mathematical tools ready, we can formally tackle our non-differentiable objective. Suppose $l(\theta)$ is the evaluation metric — say, average accuracy or average BLEU — which is the ultimate target we want to optimize, but it is non-differentiable. In most scenarios, however, we can find a differentiable (approximate) optimization objective $\tilde{l}(\theta)$, and we usually just apply gradient descent directly to optimize $\tilde{l}(\theta)$, which is what creates the inconsistency between the optimization objective and the evaluation metric.

But it has to be said that, in many cases, $\tilde{l}(\theta)$ really is a good approximation of $l(\theta)$ — in other words, $-\nabla_{\theta}\tilde{l}(\theta)$ really does point toward a reasonably reliable (though not optimal) update direction. This is where we can bring in importance sampling. Construct $q(\Delta\theta|\theta)$ as the normal distribution $\mathcal{N}(\Delta\theta; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2)$, and according to the importance-sampling equation $\eqref{eq:is-3}$, we get

\begin{equation} \Delta\theta_*=\mathbb{E}_{\Delta\theta\sim q(\Delta\theta|\theta)}\left[\frac{p(\Delta\theta|\theta)}{q(\Delta\theta|\theta)}\Delta\theta\right]\approx\sum_{i=1}^n \frac{e^{-[l(\theta + \Delta\theta_i) - l(\theta)]/\alpha}\big/\mathcal{N}(\Delta\theta_i; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2)}{\sum\limits_{i=1}^n e^{-[l(\theta + \Delta\theta_i) - l(\theta)]/\alpha}\big/\mathcal{N}(\Delta\theta_i; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2)}\Delta\theta_i \label{eq:sg}\end{equation}

where $\Delta\theta_1,\Delta\theta_2,\dots,\Delta\theta_n\sim\mathcal{N}(\Delta\theta; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2)$. Beyond this, $q(\Delta\theta|\theta)$ could also be a mixture model — the original paper uses:

\begin{equation}q(\Delta\theta|\theta)=\lambda \mathcal{N}(\Delta\theta; 0, \sigma^2) + (1-\lambda)\mathcal{N}(\Delta\theta; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2)\end{equation}

Readers may be curious about the number of samples needed. In the original paper's text generation task, choosing $n=4$ was already enough to yield a clear improvement, showing that once $q(\Delta\theta|\theta)$'s "guidance" is available, $n$ doesn't need to be very large.

Policy Gradient

In general, when one needs to directly optimize an evaluation metric, the common approach is the "policy gradient" method from reinforcement learning. So readers following along this far might wonder: how does the above method differ from policy gradient, and which is better?

Suppose the evaluation metric for a single sample is $l(y_t,y_p)$, where $y_t$ is the ground-truth label and $y_p$ is the prediction. Then the overall average metric is

\begin{equation}l(\theta)=\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[l\left(y_t, \mathop{\text{argmax}}_y p_{\theta}(y|x_t)\right)\right]\end{equation}

The non-differentiability stems from the $\mathop{\text{argmax}}$ operation. Policy gradient turns this into

\begin{equation}\tilde{l}(\theta)=\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[l\left(y_t,y\right)\right]\right]\end{equation}

and then uses (see A Chat About Reparameterization: From the Normal Distribution to Gumbel Softmax)

\begin{equation}\nabla_{\theta}\int p_{\theta}(x)f(x)dx = \int f(x)\nabla_{\theta}p_{\theta}(x)dx =\int p_{\theta}(x)f(x)\nabla_{\theta}\log p_{\theta}(x)dx\end{equation}

to obtain

\begin{equation}\nabla_{\theta}\tilde{l}(\theta)=\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[l\left(y_t,y\right)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]\right]\label{eq:pg}\end{equation}

This is the general form of policy gradient, also known as the REINFORCE estimator.

Equations $\eqref{eq:sg}$ and $\eqref{eq:pg}$ give two update directions derived from different angles — so where do they differ? The key difference lies in what is being sampled: in equation $\eqref{eq:sg}$, the sampling is over $\mathbb{E}_{\Delta\theta\sim q(\Delta\theta|\theta)}$, while in $\eqref{eq:pg}$, the sampling is over $\mathbb{E}_{y\sim p_{\theta}(y|x_t)}$. So the original paper's equation $\eqref{eq:sg}$ works by "sampling multiple sets of parameters, each producing one output sample" to compute the update, whereas policy gradient's equation $\eqref{eq:pg}$ works by "using a single set of parameters, but sampling multiple output samples" to compute the update.

In terms of computational cost, policy gradient should be somewhat cheaper, since sampling multiple output samples in step $\mathbb{E}_{y\sim p_{\theta}(y|x_t)}$ can be parallelized; in principle, sampling multiple sets of parameters to predict their respective samples in $\mathbb{E}_{\Delta\theta\sim q(\Delta\theta|\theta)}$ can also be parallelized, but this is not straightforward to implement. That said, policy gradient has its own share of problems — a classic one being that the gradient estimate has very high variance, which is why one typically needs to pretrain with an ordinary likelihood objective until reasonably converged before fine-tuning with policy gradient. The original paper's equation $\eqref{eq:sg}$, by contrast, aims to directly optimize the evaluation metric from start to finish, achieving importance sampling with the help of a differentiable objective — combining the strengths of both differentiable and non-differentiable optimization.

Summary

This post introduced a new perspective for understanding optimization algorithms, one under which the optimization of differentiable and non-differentiable objectives is unified: for differentiable objective functions, first- and second-order expansions under this perspective respectively yield gradient descent and Newton's method; for non-differentiable objective functions, we can use a differentiable approximation to perform importance sampling, thereby likewise accomplishing optimization of a non-differentiable objective.