The Convergent Paths of Policy Gradient and Zeroth-Order Optimization
One huge reason for deep learning's success is that gradient-based optimization algorithms (SGD, Adam, etc.) can effectively solve most neural network models. However, being gradient-based requires the model to be differentiable. As research has progressed, we often encounter the need to optimize non-differentiable models — a typical example being direct optimization of evaluation metrics like accuracy, F1, or BLEU, or incorporating non-differentiable modules into neural networks (such as "skim reading" operations).
Gradient
This post will briefly introduce two effective methods for optimizing non-differentiable models: policy gradient, one of the key methods in reinforcement learning, and zeroth order optimization, which requires no gradients at all. On the surface, these are two entirely different optimization approaches, but this post will further demonstrate that for a broad class of optimization problems, they are essentially equivalent.more
Formal Description
Let's start by formally defining the problem we want to solve. Taking supervised learning as an example, with training data $(x_t,y_t)\sim\mathcal{D}$ and model $p_{\theta}(y|x)$, where $\theta$ is the parameter to be optimized with dimension $d$. Assuming the model itself is differentiable, its general form is $softmax(f_{\theta}(y|x)/\tau)$, where $\tau$ is called the temperature parameter, defaulting to $\tau=1$ unless otherwise noted. Suppose the ground truth label is $y_t$ and the predicted label is $y_p$; then the score for a single sample is denoted $r(y_t, y_p)$, and the training objective wants the total score to be as large as possible, i.e.,
\begin{equation}\theta = \mathop{\text{argmax}}_{\theta}\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[r\left(y_t, \mathop{\text{argmax}}_y p_{\theta}(y|x_t)\right)\right]\label{eq:base}\end{equation}
This looks quite complicated, but its meaning is actually intuitive and clear: we want to find the parameter $\theta$ such that the score $r(y_t,y_p)$ over the entire dataset is as large as possible, where $y_p=\mathop{\text{argmax}}_y p_{\theta}(y|x_t)$ means the model outputs whichever class has the highest probability at prediction time. Put simply, we want "the class $y$ with the highest predicted probability to be the class $y$ with the highest evaluation score."
This formulation corresponds to quite a lot of machine learning tasks — in NLP this includes text classification, sequence labeling, text generation, and even regression problems can be mapped onto it, making it fairly representative. The difficulty lies in the step $\mathop{\text{argmax}}_y$, which cannot provide a useful gradient, making it hard to directly optimize with gradient-based algorithms.
Policy Gradient
The idea behind policy gradient is straightforward: since the original objective $\eqref{eq:base}$ cannot be differentiated, we swap in a strongly-correlated, differentiable objective instead, such as
\begin{equation}\theta = \mathop{\text{argmax}}_{\theta}\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\sum_y p_{\theta}(y|x_t) r\left(y_t, y\right)\right]\label{eq:policy}\end{equation}
The Rearrangement Inequality
Clearly, the objective defined above contains no operators like $\mathop{\text{argmax}}_y$, so it is differentiable. The first thing we want to know, then, is what the relationship is between the expression above and the original objective $\eqref{eq:base}$ — and where the difference lies. To answer this we need the "rearrangement inequality" from mathematics:
Rearrangement Inequality For $a_1 \geq a_2 \geq \dots \geq a_n$ and $b_1 \geq b_2 \geq \dots \geq b_n$, and assuming $(c_1, c_2, \dots, c_n)$ is any permutation of $(b_1, b_2, \dots, b_n)$, we have
\begin{equation}\sum_{i=1}^n a_i b_i \geq \sum_{i=1}^n a_i c_i \geq \sum_{i=1}^n a_i b_{n+1-i}\end{equation}
That is, "sum of same-order products ≥ sum of shuffled-order products ≥ sum of reverse-order products."
The rearrangement inequality is a classic inequality, and its proof (typically via mathematical induction) is easy to find online, so we won't reproduce it here. From the rearrangement inequality we know that if the objective $\eqref{eq:policy}$ reaches its maximum, then $p_{\theta}(y|x_t)$ and $r\left(y_t, y\right)$ are in the same order — meaning we indeed achieve the goal that "the class $y$ with the highest predicted probability is the class $y$ with the highest evaluation score." But at the same time we also achieve "the class $y$ with the second highest predicted probability is the class $y$ with the second highest evaluation score," "the class $y$ with the third highest predicted probability is the class $y$ with the third highest evaluation score," and so on — which is more than what the original objective actually requires. So the objective $\eqref{eq:policy}$ is strongly correlated with the original objective, but demands more.
Note that the rearrangement inequality does not require $a_i,b_i$ to all be non-negative, so the actual scoring function $r(y_t, y)$ is allowed to be negative.
Estimating the Gradient via Sampling
Having established that the objective $\eqref{eq:policy}$ is workable, we can take its gradient:
\begin{equation}\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\sum_y \nabla_{\theta} p_{\theta}(y|x_t) r\left(y_t, y\right)\right]\label{eq:policy-grad-base}\end{equation}
Generally speaking, computing the gradient $\nabla_{\theta} p_{\theta}(y|x_t)$ is not difficult; the main difficulty lies in $\sum_y$, since it requires summing over all candidate classes, and in practice the number of candidate classes can be prohibitively large — a point also discussed previously in Musings on Reparameterization: From the Normal Distribution to the Gumbel Softmax. Therefore, a better approach is to convert it into a sampling-based estimate, i.e.,
\begin{equation}\begin{aligned} &\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\sum_y p_{\theta}(y|x_t)\frac{\nabla_{\theta} p_{\theta}(y|x_t)}{p_{\theta}(y|x_t)} r\left(y_t, y\right)\right]\\ =& \mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\sum_y p_{\theta}(y|x_t)r\left(y_t, y\right)\nabla_{\theta} \log p_{\theta}(y|x_t)\right]\\ =& \mathbb{E}_{(x_t,y_t)\sim\mathcal{D}, y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\nabla_{\theta}\log p_{\theta}(y|x_t)\right] \end{aligned}\end{equation}
In principle, this means we only need to sample an appropriate number of $y$ to estimate the expression above, and the resulting quantity is what's called the "policy gradient." With a gradient in hand, we can plug it into any off-the-shelf optimizer to perform the optimization.
Reducing Variance
We just said that adding a constant to $r(y_t, y)$ doesn't change the final result. However, it can change the efficiency of the sampling-based estimate — in statistical terms, it can change the variance of the sampling. Here's a simple example: $[4, 5, 6]$ and $[-10, 10, 15]$ both have a mean of 5 (representing the target we want to estimate), but their variances are 0.67 and 116.67 respectively — the latter's variance is far larger than the former's. If we only draw a single sample, the former's maximum deviation from the target is at most 1, while the latter's maximum deviation can reach 15, with a minimum deviation of 5. So even though in theory both give the same mean in the end, the former is far more efficient to estimate (fewer samples needed for the same estimation precision).
This simple example tells us that to improve estimation efficiency, we need an estimator with smaller variance. This is where we subtract a constant $b$ (called a baseline; "constant" here means it doesn't depend on $y$, though it can depend on $x$) from $r(y_t, y)$:
\begin{equation}\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[(r\left(y_t, y\right)-b)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]\label{eq:var-reduce}\end{equation}
The final result (the mean) doesn't change:
\begin{equation}\begin{aligned} &\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[(r\left(y_t, y\right)-b)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]\\ =&\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]-b\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[\nabla_{\theta}\log p_{\theta}(y|x_t)\right]\\ =&\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]-b\sum_y \nabla_{\theta} p_{\theta}(y|x_t)\\ =&\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]-b \nabla_{\theta} \sum_y p_{\theta}(y|x_t)\\ =&\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]-b \nabla_{\theta} 1\\ =&\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\nabla_{\theta}\log p_{\theta}(y|x_t)\right]\\ \end{aligned}\end{equation}
But its variance can change. We want to minimize the variance, and by $\mathbb{Var}[x]=\mathbb{E}[x^2]-\mathbb{E}[x]^2$ we know that minimizing the variance is equivalent to minimizing the second moment
\begin{equation}\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[(r\left(y_t, y\right)-b)^2\Vert\nabla_{\theta}\log p_{\theta}(y|x_t)\Vert^2\right]\end{equation}
This is just a matter of minimizing a quadratic function, and the optimal $b$ turns out to be:
\begin{equation}b = \frac{\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\Vert\nabla_{\theta}\log p_{\theta}(y|x_t)\Vert^2\right]}{\mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[\Vert\nabla_{\theta}\log p_{\theta}(y|x_t)\Vert^2\right]}\end{equation}
That is, take the weighted expectation of $r\left(y_t, y\right)$ with weights $\Vert\nabla_{\theta}\log p_{\theta}(y|x_t)\Vert^2$. But since computing the gradient for every candidate class would be expensive, we generally don't bother with this weighting and instead use a simplified version:
\begin{equation}b = \mathbb{E}_{y\sim p_{\theta}(y|x_t)}\left[r\left(y_t, y\right)\right]\end{equation}
Combining this with equation $\eqref{eq:var-reduce}$, we can see the idea is actually quite intuitive: it amounts to sampling several $y$ from $p_{\theta}(y|x_t)$, computing the mean $b$ of $r\left(y_t, y\right)$, and then doing gradient ascent (reinforcing the behavior) for samples above the mean and gradient descent (suppressing the behavior) for samples below it.
In a Nutshell
In short, policy gradient works by replacing a non-differentiable objective function — one involving an $\mathop{\text{argmax}}$ operation, say — with a differentiable objective $\eqref{eq:policy}$. In the language of reinforcement learning, $y$ is called the "policy," $p_{\theta}(y|x_t)$ is the "policy model" (decision model), and $r(y_t,y_p)$ is the "reward." Combined with sampling-based estimation and variance-reduction tricks, this yields an effective gradient estimate for the original model, enabling optimization.
Zeroth-Order Optimization
Zeroth-order optimization broadly refers to any optimization method that requires no gradient information; in most contexts, it specifically refers to optimization algorithms that estimate the parameter update direction based on sampling in parameter space combined with finite differences. Formally, it samples directly in parameter space and doesn't rely on any form of gradient, so in principle it applies to a very wide range of objectives. But precisely because it samples directly in parameter space — essentially a smarter form of grid search — it becomes quite inefficient when facing high-dimensional parameter spaces (as in deep learning), which severely limits its applicability.
Even so, that shouldn't stop us from learning the ideas behind zeroth-order optimization — an extra skill never hurts. Moreover, although deep learning models often have huge numbers of parameters, in most cases the majority of the modules we design are differentiable, with only a small portion being non-differentiable. It might therefore be possible to optimize the differentiable parts directly with gradient-based optimizers, while using zeroth-order optimization only for the non-differentiable parts — this is one of its practical use cases, an idea that appears in many NAS papers.
Zeroth-Order Gradient
Zeroth-order optimization doesn't require the gradient in the usual sense; instead, it defines a sampling-and-difference-based "substitute," which we'll call the "zeroth-order gradient":
For a scalar function $f(x)$, define its zeroth-order gradient at $x$ as
\begin{equation}\tilde{\nabla}_{x}f(x)=\mathbb{E}_{u\sim p(u)}\left[\frac{f(x + \varepsilon u) - f(x)}{\varepsilon}u\right]\label{eq:zero-grad}\end{equation}
where $\varepsilon$ is a small positive number specified in advance, and $p(u)$ is a pre-specified distribution with mean 0 and identity covariance matrix — typically we use the standard normal distribution.
As we can see, we only need to sample a number of points from $p(u)$ to estimate the zeroth-order gradient, and once we have it, we can treat it just like an ordinary gradient and plug it into a gradient-based optimizer — this is the basic idea behind zeroth-order optimization. In particular, if $f(x)$ is itself differentiable, then $f(x + \varepsilon u)=f(x)+\varepsilon u^{\top}\nabla_x f(x) + \mathcal{O}(\varepsilon^2)$, so when $\varepsilon\to 0$:
\begin{equation}\tilde{\nabla}_{x}f(x)=\int p(u) u \left(u^{\top}\nabla_x f(x)\right)du=\int p(u) \left(u u^{\top}\right)\nabla_x f(x)du=\nabla_x f(x)\end{equation}
That is, $\tilde{\nabla}_{x}f(x)$ equals the ordinary gradient, which confirms that $\tilde{\nabla}_{x}f(x)$ is indeed a reasonable generalization of the ordinary gradient.
There's a Baseline Here Too
Sharp-eyed readers may notice that in the definition of $\eqref{eq:zero-grad}$, since $p(u)$ has mean 0, the term $-f(x)$ doesn't actually affect the final result, i.e.,
\begin{equation}\tilde{\nabla}_{x}f(x)=\mathbb{E}_{u\sim p(u)}\left[\frac{f(x + \varepsilon u)}{\varepsilon}u\right] - \mathbb{E}_{u\sim p(u)}\left[\frac{f(x)}{\varepsilon}u\right]=\mathbb{E}_{u\sim p(u)}\left[\frac{f(x + \varepsilon u)}{\varepsilon}u\right]\label{eq:zero-grad-equal}\end{equation}
So what is the purpose of $-f(x)$? Just as with policy gradient, it's there to reduce variance. We can likewise introduce $b$ and minimize the second moment (equivalent to minimizing the variance)
\begin{equation}\mathbb{E}_{u\sim p(u)}\left[\left(\frac{f(x + \varepsilon u)-b}{\varepsilon}\right)^2\Vert u\Vert^2\right]\end{equation}
Solving for the optimal $b$ gives
\begin{equation}b=\frac{\mathbb{E}_{u\sim p(u)}\left[f(x + \varepsilon u)\Vert u\Vert^2\right]}{\mathbb{E}_{u\sim p(u)}\left[\Vert u\Vert^2\right]}\end{equation}
In practice, this can be estimated directly using a finite number of samples. In fact, if $f(x)$ is differentiable, the integral can be approximated using a Taylor expansion, giving the result $f(x)+\mathcal{O}(\varepsilon^2)$. From this angle, taking $b=f(x)$ directly is also a reasonable choice, which shows the necessity and rationale for introducing the $-f(x)$ term.
In a Nutshell
Zeroth-order optimization methods essentially define a reasonable generalization of the gradient based on finite differences. Since computing a finite difference doesn't require differentiability, this naturally applies to optimizing both differentiable and non-differentiable objectives. However, since the dimension of $u$ equals the dimension of the full parameter set $\theta$, for models with as many parameters as deep learning models have, the variance during the update process can be very large, making convergence difficult. So typically zeroth-order optimization is only used to optimize a small subset of a model's parameters, or as an auxiliary optimization technique (e.g., alternating updates like "differentiable objective + ordinary gradient + large learning rate" with "non-differentiable objective + zeroth-order gradient + small learning rate"). There has been some research on directly applying zeroth-order optimization in high-dimensional spaces (e.g., Gradientless Descent: High-Dimensional Zeroth-Order Optimization), but it hasn't been very successful so far.
Additionally, the unified perspective introduced earlier in Viewing Optimization Through Sampling: A Unified Perspective on Differentiable and Non-Differentiable Optimization can also be viewed as a form of zeroth-order optimization — it's a more unified generalization of common optimization ideas (through which gradient descent, Newton's method, zeroth-order gradients, etc., can all be derived). But in principle, it suffers from the same "chronic ailment" of high variance as zeroth-order gradients, a difficulty that zeroth-order optimization methods can hardly avoid.
Different in Appearance, United in Essence
On the surface, policy gradient and zeroth-order optimization do share many similarities: both require random sampling to estimate the gradient, both need ways to reduce variance, and so on. Of course, plenty of differences can also be listed — for instance, policy gradient samples over the policy space, while zeroth-order optimization samples over the parameter space; or that policy gradient is still fundamentally gradient-based, while zeroth-order optimization in principle needs no gradient at all.
So what exactly is the relationship between the two? In what follows, we will show that for the optimization problem $\eqref{eq:base}$ posed at the start of this post, the two are essentially equivalent. The proof strategy is to compute the zeroth-order gradient of the objective $\eqref{eq:base}$, and show that, after a series of simplifications, it is essentially the policy gradient.
Partitioning the Full Space
Let's denote
\begin{equation}\mathcal{R}_{\theta}=\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[r\left(y_t, \mathop{\text{argmax}}_y p_{\theta}(y|x_t)\right)\right]\end{equation}
Then, according to equation $\eqref{eq:zero-grad-equal}$, its zeroth-order gradient is
\begin{equation}\begin{aligned} \tilde{\nabla}_{\theta}\mathcal{R}_{\theta}=&\frac{1}{\varepsilon}\int \mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[r\left(y_t, \mathop{\text{argmax}}_y p_{\theta + \varepsilon u}(y|x_t)\right)\right] p(u) u du\\ =&\mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\frac{1}{\varepsilon}\int r\left(y_t, \mathop{\text{argmax}}_y p_{\theta + \varepsilon u}(y|x_t)\right)p(u) u du\right] \end{aligned}\end{equation}
which can be rewritten as
\begin{equation}\begin{aligned} \tilde{\nabla}_{\theta}\mathcal{R}_{\theta}=& \mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\frac{1}{\varepsilon}\sum_y\int_{\Omega_{y|x_t}} r\left(y_t, y\right)p(u) u du\right]\\ \Omega_{y|x_t} =& \left\{u\left|u\in\mathbb{R}^d, y=\mathop{\text{argmax}}_{\hat{y}} p_{\theta + \varepsilon u}(\hat{y}|x_t)\right.\right\} \end{aligned}\end{equation}
This looks complicated, but the idea is simple: different $u$ yield different prediction results $\mathop{\text{argmax}}_{y} p_{\theta + \varepsilon u}(y|x_t)$. We group together all the $u$ whose prediction result is the same $y$, denoting this group $\Omega_{y|x_t}$; the full space $\mathbb{R}^d$ is thereby partitioned into disjoint subsets $\Omega_{y|x_t}$. All integrals mentioned earlier without an explicitly stated integration domain are implicitly taken over the full space $\mathbb{R}^d$; once we partition the space, this becomes a sum of integrals over each of the partitioned subsets $\Omega_{y|x_t}$.
Indicator Function
Next, we use an "indicator function" trick, defining
\begin{equation}\chi(y|x_t, u)=\left\{\begin{aligned}1, u\in\Omega_{y|x_t}\\ 0, u\not\in\Omega_{y|x_t}\end{aligned}\right.\end{equation}
so that
\begin{equation}\begin{aligned} \frac{1}{\varepsilon}\sum_y\int_{\Omega_{y|x_t}} r\left(y_t, y\right)p(u) u du=\frac{1}{\varepsilon}\sum_y\int \chi(y|x_t, u) r\left(y_t, y\right)p(u) u du \end{aligned}\end{equation}
Recalling the meaning of $\Omega_{y|x_t}$: for any $u\in\Omega_{y|x_t}$, the model $p_{\theta + \varepsilon u}(\cdot|x_t)$'s prediction is $y$, and the indicator function outputs 1 in that case — so we find that in fact
\begin{equation}\chi(y|x_t, u)=\lim_{\tau\to 0} p_{\theta + \varepsilon u}(y|x_t)\end{equation}
where $\tau$ is the temperature parameter of the softmax (as noted at the beginning of the post). Based on this result, we can approximate $\chi(y|x_t, u)$ by $p_{\theta + \varepsilon u}(y|x_t)$, then substitute this into the expressions above to get
\begin{equation} \tilde{\nabla}_{\theta}\mathcal{R}_{\theta} \approx \mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\frac{1}{\varepsilon}\sum_y\int p_{\theta + \varepsilon u}(y|x_t) r\left(y_t, y\right)p(u) u du\right] \end{equation}
Approximating the Integral
Finally, since $p_{\theta + \varepsilon u}(y|x_t)$ is differentiable, we can use the expansion $p_{\theta + \varepsilon u}(y|x_t)\approx p_{\theta}(y|x_t) + \varepsilon u^{\top}\nabla_{\theta}p_{\theta}(y|x_t)$, substitute it into the equation above, and carry out the integral over $u$ to obtain
\begin{equation} \tilde{\nabla}_{\theta}\mathcal{R}_{\theta} \approx \mathbb{E}_{(x_t,y_t)\sim\mathcal{D}}\left[\sum_y r\left(y_t, y\right)\nabla_{\theta}p_{\theta}(y|x_t)\right] \end{equation}
Comparing this with equation $\eqref{eq:policy-grad-base}$, we can see that the right-hand side above is precisely the policy gradient.
So, the zeroth-order gradient — which on the surface looked quite different — ultimately points in essentially the same direction as the policy gradient. Truly a case of different in appearance but united in essence, of separate paths converging — one gets the sense that "there is only one correct answer." This can't help but remind the author of Leo Tolstoy's famous line:
Happy families are all alike; every unhappy family is unhappy in its own way.
The same seems to hold for parameter update directions (correct update directions all resemble one another)~
Summary
This post has introduced two approaches for handling non-differentiable optimization objectives: policy gradient and zeroth-order optimization. Each defines a new kind of "gradient" from a different angle to serve as the parameter update direction. On the surface, the two take different paths, but the discussion here shows that, when applied to the same kind of optimization problem as policy gradient, the update direction given by zeroth-order optimization is essentially equivalent to the policy gradient — different roads leading to the same destination. This post may also serve as introductory material on reinforcement learning for beginners.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.