A major reason for the immense success of deep learning is that gradient-based optimization algorithms (such as SGD, Adam, etc.) can effectively solve most neural network models. However, since these are based on gradients, they require the model to be differentiable. As research deepens, we often encounter the need to solve non-differentiable models. Typical examples include directly optimizing evaluation metrics like Accuracy, F1, or BLEU, or incorporating non-differentiable modules (such as "skipping" operations) into neural networks.
This article briefly introduces two effective methods for solving non-differentiable models: Policy Gradient, one of the core methods in reinforcement learning, and Zeroth-Order Optimization, which requires no gradients at all. On the surface, these appear to be two entirely different optimization approaches. However, this article will further demonstrate that for a large class of optimization problems, the two are essentially equivalent.
Formal Description
First, let us formally define the problem we need to solve. Taking supervised learning as an example, let the training data be (x_t, y_t) \sim \mathcal{D} and the model be p_{\theta}(y|x), where \theta represents the parameters to be optimized with dimension d. Assuming the model itself is differentiable, its general form is \text{softmax}(f_{\theta}(y|x)/\tau), where \tau is the temperature parameter (defaulting to \tau=1 unless otherwise specified). If the true label is y_t and the predicted label is y_p, the score for a single sample is denoted as r(y_t, y_p). The training objective is to maximize the total score: \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} While this looks complex, its meaning is intuitive: we want to find parameters \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 the label with the highest probability during prediction. Simply put, we want "the y with the highest predicted probability to be the y with the highest evaluation score."
This formulation corresponds to many machine learning tasks. In NLP, it includes text classification, sequence labeling, text generation, and even regression problems. The difficulty lies in the fact that the \mathop{\text{argmax}}_y step does not provide effective gradients, making it difficult to optimize directly using gradient-based algorithms.
Policy Gradient
The idea behind Policy Gradient is straightforward: since the original objective [eq:base] is not differentiable, we replace it with a strongly related objective that is differentiable, such as: \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}
Rearrangement Inequality
Clearly, the objective defined above does not contain operators like \mathop{\text{argmax}}_y, so it is differentiable. What is the relationship between this and the original objective [eq:base]? The answer lies in 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), then: \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} In other words, "the sum of products of similarly ordered sequences \geq mixed-order sum \geq reverse-order sum."
The rearrangement inequality is a classic result, and its proof (usually by induction) is easily found online. According to this inequality, if the objective [eq:policy] reaches its maximum, then p_{\theta}(y|x_t) and r(y_t, y) must be similarly ordered. This means we indeed achieve the goal that "the y with the highest predicted probability is the y with the highest score." However, it also simultaneously aims for "the y with the second highest probability to be the y with the second highest score," and so on. These additional constraints are not strictly required by the original objective. Thus, objective [eq:policy] is strongly related to the original objective but is more demanding.
Note that the rearrangement inequality does not require a_i, b_i to be non-negative, so the scoring function r(y_t, y) can potentially be negative.
Sampling Gradient Estimation
Once we determine that objective [eq:policy] is feasible, we can compute its gradient: \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} Generally, computing \nabla_{\theta} p_{\theta}(y|x_t) is not difficult, but \sum_y is the main challenge because it requires summing over all candidate categories, which may be prohibitively large. (Related discussions appeared in my previous post "Talking about Reparameterization: From Normal Distribution to Gumbel Softmax"). Therefore, a better approach is to convert it into a sampling estimation form: \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} In principle, we only need to sample an appropriate number of y values to estimate the above expression. The result is known as the "Policy Gradient." With the gradient, we can use standard optimizers to complete the optimization.
Reducing Variance
As mentioned, adding a constant to r(y_t, y) does not change the final theoretical result. However, it can change the efficiency of the sampling estimation—in statistical terms, it changes the variance. For example, consider [4, 5, 6] and [-10, 10, 15]. Both have a mean of 5 (representing our target), but their variances are 0.67 and 116.67, respectively. If we sample only one point, the former has a maximum error of 1, while the latter can have an error of 15. Thus, while the theoretical mean is the same, the estimation efficiency of the former is much higher.
This simple example tells us that to improve estimation efficiency, we must find an estimator with lower variance. We can subtract a constant b (called a baseline, where "constant" means it does not depend on y, though it can depend on x): \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} The final result (mean) remains unchanged: \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} However, the variance may change. To minimize variance, we minimize the second moment (since \text{Var}[X] = \mathbb{E}[X^2] - \mathbb{E}[X]^2): \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] This is a simple quadratic minimization problem. The optimal b is: 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]} This is a weighted expectation of r(y_t, y) with weights \Vert \nabla_{\theta} \log p_{\theta}(y|x_t) \Vert^2. Since obtaining the gradient for every candidate category is costly, we often ignore the weights and use a simplified version: b = \mathbb{E}_{y\sim p_{\theta}(y|x_t)} \left[ r\left(y_t, y\right) \right] Combining this with equation [eq:var-reduce], the intuition is clear: sample several y from p_{\theta}(y|x_t), calculate the average score b, and perform gradient ascent for those y with scores above the average (reinforcement) and gradient descent for those below (weakening).
In a Nutshell
Simply put, when encountering non-differentiable objectives like \mathop{\text{argmax}}, Policy Gradient replaces them with a differentiable objective [eq:policy]. In reinforcement learning terms, y is the "policy," p_{\theta}(y|x_t) is the "decision model," and r(y_t, y_p) is the "reward." Combined with sampling estimation and variance reduction techniques, it provides an effective gradient estimate for the original model.
Zeroth-Order Optimization
Zeroth-order optimization refers to optimization methods that do not require gradient information. Generally, it refers to algorithms that estimate the parameter update direction based on parameter sampling and finite differences. Formally, it samples directly in the parameter space and does not rely on any form of gradient, making it theoretically capable of solving a wide range of objectives. However, because it samples in the parameter space (essentially a smarter grid search), its efficiency is quite low in high-dimensional spaces (like deep learning), limiting its use cases.
Nevertheless, the ideas behind zeroth-order optimization are worth learning. While deep learning models have many parameters, most modules are differentiable. Perhaps only a small part is non-differentiable. We can use gradient optimizers for the differentiable parts and zeroth-order optimization for the rest—an idea seen in many NAS (Neural Architecture Search) papers.
Zeroth-Order Gradient
Zeroth-order optimization does not use gradients in the usual sense, but it defines a "substitute" based on sampling and differences, which we call the "zeroth-order gradient":
For a scalar function f(x), define its zeroth-order gradient at x as: \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} where \varepsilon is a small positive constant and p(u) is a pre-specified distribution with mean 0 and identity covariance matrix (usually a standard normal distribution).
By sampling several points from p(u), we can estimate the zeroth-order gradient and use it in standard gradient-based optimizers. Notably, if f(x) is differentiable, then f(x + \varepsilon u) = f(x) + \varepsilon u^{\top} \nabla_x f(x) + \mathcal{O}(\varepsilon^2). As \varepsilon \to 0: \tilde{\nabla}_{x} f(x) = \int p(u) u \left( u^{\top} \nabla_x f(x) \right) du = \int p(u) (u u^{\top}) \nabla_x f(x) du = \nabla_x f(x) Thus, \tilde{\nabla}_{x} f(x) is indeed a reasonable generalization of the standard gradient.
Also has a Baseline
Astute readers may notice that in definition [eq:zero-grad], since the mean of p(u) is 0, the term -f(x) does not affect the final theoretical result: \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} So what is the role of -f(x)? Like in Policy Gradient, it is used to reduce variance. We can introduce b and minimize the second moment: \mathbb{E}_{u\sim p(u)} \left[ \left( \frac{f(x + \varepsilon u) - b}{\varepsilon} \right)^2 \Vert u \Vert^2 \right] The optimal b is: 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]} In practice, this can be estimated with finite samples. If f(x) is differentiable, a Taylor expansion shows b \approx f(x) + \mathcal{O}(\varepsilon^2), justifying the choice of b=f(x).
In a Nutshell
Zeroth-order optimization defines a reasonable generalization of the gradient based on finite differences. Since differences do not require differentiability, it applies to both differentiable and non-differentiable objectives. However, since the dimension of u is the dimension of all parameters \theta, the variance in deep learning models is huge, making convergence difficult. Thus, it is usually used for a small subset of parameters or as an auxiliary method. There is research on applying zeroth-order methods directly to high-dimensional spaces (e.g., "Gradientless Descent: High-Dimensional Zeroth-Order Optimization"), but it is not yet widely successful.
Additionally, the unified perspective introduced in "Optimization from Sampling: A Unified View of Differentiable and Non-Differentiable Optimization" can also be seen as a zeroth-order method. While it generalizes gradient descent, Newton’s method, etc., it suffers from the same high-variance issues inherent to zeroth-order methods.
Different Paths, Same Destination
Policy Gradient and Zeroth-Order Optimization share many similarities: both use random sampling to estimate gradients and both require variance reduction. Of course, there are differences: Policy Gradient samples the policy space, while Zeroth-Order samples the parameter space; Policy Gradient is fundamentally gradient-based, while Zeroth-Order theoretically is not.
What is the relationship between them? We will now prove that for the optimization problem [eq:base] presented at the beginning, the two are essentially equivalent. The proof involves calculating the zeroth-order gradient of objective [eq:base] and showing that, after simplification, it is basically the Policy Gradient.
Partitioning the Full Space
Let: \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] According to equation [eq:zero-grad-equal], its zeroth-order gradient is: \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} This can be transformed into: \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 \mid u \in \mathbb{R}^d, y = \mathop{\text{argmax}}_{\hat{y}} p_{\theta + \varepsilon u}(\hat{y}|x_t) \right\} \end{aligned} The idea is simple: different u lead to different predictions \mathop{\text{argmax}}_y p_{\theta + \varepsilon u}(y|x_t). We group all u that yield the same prediction y into a set \Omega_{y|x_t}. Thus, the full space \mathbb{R}^d is partitioned into disjoint subsets \Omega_{y|x_t}. The integral over the full space becomes the sum of integrals over these subsets.
Indicator Function
Next, we use an "indicator function" trick. Define: \chi(y|x_t, u) = \begin{cases} 1, & u \in \Omega_{y|x_t} \\ 0, & u \notin \Omega_{y|x_t} \end{cases} Then: \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 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) predicts y. We can observe that: \chi(y|x_t, u) = \lim_{\tau \to 0} p_{\theta + \varepsilon u}(y|x_t) where \tau is the temperature parameter in the softmax. Based on this, we can approximate \chi(y|x_t, u) with p_{\theta + \varepsilon u}(y|x_t) and substitute it back: \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]
Approximating the Integral
Finally, since p_{\theta + \varepsilon u}(y|x_t) is differentiable, we 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) and integrate over u: \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] Comparing this to equation [eq:policy-grad-base], we see that the right side is exactly the Policy Gradient.
Thus, despite their apparent differences, the zeroth-order gradient ultimately points in a direction similar to the Policy Gradient. They truly are different paths to the same destination. It reminds me of Leo Tolstoy’s famous quote:
All happy families are alike; each unhappy family is unhappy in its own way.
The same applies to parameter update directions (all correct update directions are alike)!
Summary
This article introduced two solutions for non-differentiable optimization: Policy Gradient and Zeroth-Order Optimization. They define new "gradients" from different perspectives. While they take different paths, our exploration shows that for the same optimization problem, the update direction provided by zeroth-order optimization is essentially equivalent to Policy Gradient. Furthermore, this article serves as an introductory reference for beginners in reinforcement learning.
Original address: https://kexue.fm/archives/7737