Many readers are likely aware that the inconsistency between loss functions and evaluation metrics is a typical phenomenon in machine learning. For instance, in classification problems, the loss function is often cross-entropy, while the evaluation metrics are accuracy or F1-score. In text generation, the loss function is cross-entropy in a teacher-forcing format, while the evaluation metrics are BLEU, ROUGE, etc. Ideally, we would want to optimize the specific metric we are evaluating; however, evaluation metrics are usually non-differentiable. Since most of us use gradient-based optimizers, the minimization objective must be differentiable, which is the source of this 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 investigates how to directly optimize evaluation metrics for text generation. Upon reading it, I found this paper to be very valuable. In fact, it provides a new way of thinking about optimizing evaluation metrics, and its scope of application is not limited to text generation. Furthermore, it even contains a unified perspective for understanding differentiable and non-differentiable optimization.
Sampling Perspective
First, we can re-examine the optimization problem through the lens of sampling. Let the current model parameters be \theta and the optimization objective be l(\theta). We wish to determine the next update amount \Delta\theta. To this end, we first construct a distribution: 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) where \alpha > 0 is a hyperparameter. The meaning of this distribution is clear: we treat \Delta\theta as a random variable, where a \Delta\theta that results in a smaller l(\theta + \Delta\theta) has a higher probability of occurring. With this distribution, we define the next update amount as its expectation: \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}
In this perspective, we make no assumptions about the differentiability of l(\theta); therefore, the above definition is universal for both differentiable and non-differentiable optimization. Additionally, we can control the stability of the update by adjusting \alpha. When \alpha \to 0, it follows from the definition of p(\Delta\theta|\theta) that only the \Delta\theta that minimizes l(\theta + \Delta\theta) has a non-zero 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 0, which is the most stable. By choosing an appropriate \alpha, one can theoretically achieve a good balance between "speed" and "stability" in the optimization process, which intuitively leads to results with better generalization performance.
Of course, the definition so far is theoretical. We do not yet know the analytical form of p(\Delta\theta|\theta), nor do we know how to sample from it, let alone calculate its expectation. Below, we will see how this theoretical form is gradually applied to differentiable and non-differentiable scenarios.
Differentiable Objectives
For a differentiable l(\theta), although we cannot precisely solve for p(\Delta\theta|\theta), we can perform a Taylor expansion to obtain an approximate distribution and then estimate \Delta\theta_*. The results show that by expanding to the first and second order, we obtain gradient descent and Newton’s method, respectively. In other words, gradient descent and Newton’s method are, in a sense, just special cases within this perspective.
Gradient Descent
As a first attempt, let us assume l(\theta) is first-order differentiable. From the Taylor expansion, we get: l(\theta + \Delta\theta) - l(\theta)\approx \Delta\theta^{\top}\nabla_{\theta}l(\theta) This implies p(\Delta\theta|\theta)\sim e^{-\Delta\theta^{\top}\nabla_{\theta}l(\theta)/\alpha}. If \Delta\theta is unconstrained, normalization is impossible. Let us restrict \Vert\Delta\theta\Vert\leq \epsilon, and denote \nabla_{\theta}l(\theta)=g. Then: 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) It is clear that \Delta\theta_* = -\alpha\nabla_g \ln Z(g), so the key is to find Z(g). Let \eta be the angle between \Delta\theta and g, then: 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) This is an integral within a high-dimensional hypersphere. Due to isotropy, once the magnitude \Vert g\Vert is fixed, the entire integral result is fixed. That is, Z(g) only depends on the magnitude of g and is independent of its direction, so it can be written as Z(\Vert g\Vert). We do not need to know the specific form of Z(\Vert g\Vert); knowing it is merely a function of the magnitude \Vert g\Vert is sufficient. At this point: \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} Thus, the direction of \Delta\theta_* is the direction of -g, which is the opposite direction of the gradient. In this way, we have derived gradient descent. It can be said that it is a first-order approximation of Equation [eq:delta]. Additionally, it is possible to calculate Z(g) explicitly, though it is not an elementary function; the process can be found in the discussion on StackExchange: Integral of exp over the unit ball.
Newton’s Method
If l(\theta) is second-order differentiable, we can expand it to the second order: 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 Denoting g=\nabla_{\theta}l(\theta) and \mathcal{H}=\nabla_{\theta}^2 l(\theta), we have: \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} Clearly, because the exponent of p(\Delta\theta|\theta) is quadratic with respect to \theta, p(\Delta\theta|\theta) is a Gaussian distribution. The above equation implies that the mean of this Gaussian distribution is -\mathcal{H}^{-1}g and the covariance matrix is \mathcal{H}^{-1}. Therefore, \Delta\theta_*=-\mathcal{H}^{-1}g. This result corresponds to Newton’s method. Thus, Newton’s method can be considered a second-order approximation of Equation [eq:delta].
Non-Differentiable Objectives
For non-differentiable l(\theta), the Taylor expansion approximation mentioned above cannot be performed. Theoretically, we can only estimate \Delta\theta_* through 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
First, let’s 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: \int p(x)f(x)dx = \mathbb{E}_{x\sim p(x)}[f(x)] This usually requires us to draw several samples x_1, x_2, \dots, x_n from p(x) and then calculate \frac{1}{n}\sum_{i=1}^n f(x_i). However, two difficulties may arise:
1. We might not know how to sample from p(x) at all;
2. Even if we know how to sample from p(x), in scenarios like Variational Autoencoders, p(x) contains parameters and we need to preserve gradients, which direct sampling might not allow.
In such cases, importance sampling can help. It requires us to find a distribution q(x) for which we know the expression of the probability density and from which it is easy to sample. We then rewrite the integral as: \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} Now the sampling is shifted to q(x). According to our assumption, sampling from q(x) is easy, and the analytical form of q(x) is known, so \frac{p(x)}{q(x)}f(x) can be calculated. If p(x) has parameters requiring gradients, those gradients are preserved. Obviously, the closer q(x) is to p(x), the higher the estimation efficiency. q(x) represents a prior estimate of the "importance" of each sample in p(x), hence the name importance sampling.
Thus, assuming x_1, x_2, \dots, x_n \sim q(x), we have: \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} However, there is a small problem: Equation [eq:is] or [eq:is-2] requires us to know the exact expression of p(x). Sometimes we cannot even do that. For example, for the earlier p(\Delta\theta|\theta), we only know it is proportional to e^{-[l(\theta + \Delta\theta) - l(\theta)]/\alpha}, and its normalization factor cannot be calculated directly. In this case, we can use the relationship: 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)} This means that \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 the normalization factor, we can manually normalize it. Equation [eq:is-2] then becomes: \mathbb{E}_{x\sim p(x)}[f(x)]\approx \sum_{i=1}^n \frac{\rho(x_i)\big/q(x_i)}{\sum\limits_{j=1}^n \rho(x_j)\big/q(x_j)}f(x_i) \label{eq:is-3} This eliminates the need to calculate the normalization factor.
Leveraging Differentiability
At this point, all the mathematical tools we need are ready, and we can formally address our non-differentiable objective. Suppose l(\theta) is an evaluation metric, such as average accuracy or average BLEU. This is our ultimate optimization goal, but it is non-differentiable. However, in most scenarios, we can find a differentiable (approximate) optimization objective \tilde{l}(\theta). Usually, we directly optimize \tilde{l}(\theta) using gradient descent, which causes the inconsistency between the optimization objective and the evaluation metric.
Nevertheless, it must be said that in many cases \tilde{l}(\theta) is indeed a good approximation of l(\theta). In other words, -\nabla_{\theta}\tilde{l}(\theta) does point in a fairly reliable (though not optimal) update direction. This is where we can leverage importance sampling. We construct q(\Delta\theta|\theta) as a normal distribution \mathcal{N}(\Delta\theta; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2). According to Equation [eq:is-3] for importance sampling, we have: \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_{j=1}^n e^{-[l(\theta + \Delta\theta_j) - l(\theta)]/\alpha}\big/\mathcal{N}(\Delta\theta_j; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2)}\Delta\theta_i \label{eq:sg} where \Delta\theta_1, \Delta\theta_2, \dots, \Delta\theta_n \sim \mathcal{N}(\Delta\theta; -\nabla_{\theta}\tilde{l}(\theta), \sigma^2). Additionally, q(\Delta\theta|\theta) can use a mixture model. The original paper uses: 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) Readers might be concerned about the number of samples. In the text generation task, the original paper chose n=4, and the results showed significant improvement. This indicates that with the "guidance" of q(\Delta\theta|\theta), n does not need to be very large.
Policy Gradient
Generally, if one needs to directly optimize evaluation metrics, a common method is through "Policy Gradient" in reinforcement learning. Readers might wonder: what is the difference between the above method and policy gradient? Which is better?
Suppose the evaluation metric for a single sample is l(y_t, y_p), where y_t is the true label and y_p is the predicted result. Then the total average metric is: 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] The non-differentiability stems from the \mathop{\text{argmax}} operation. Policy gradient changes this to: \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] Then, using the identity (refer to "Talking about Reparameterization: From Normal Distribution to Gumbel Softmax"): \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 we obtain: \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} This is the general form of policy gradient, also known as the REINFORCE estimator.
Equations [eq:sg] and [eq:pg] represent two update directions derived from different perspectives. What is the difference between them? As we can see, the core difference lies in the sampling object: the sampling in Equation [eq:sg] is \mathbb{E}_{\Delta\theta\sim q(\Delta\theta|\theta)}, while the sampling in Equation [eq:pg] is \mathbb{E}_{y\sim p_{\theta}(y|x_t)}. Thus, the method provided in the original paper (Equation [eq:sg]) involves "sampling multiple sets of parameters, with each set outputting one sample" to calculate the update, whereas the policy gradient (Equation [eq:pg]) involves "only one set of parameters, but sampling multiple output samples" to calculate the update.
From a computational standpoint, policy gradient should be less intensive because the step \mathbb{E}_{y\sim p_{\theta}(y|x_t)} can be parallelized by sampling multiple samples. Theoretically, \mathbb{E}_{\Delta\theta\sim q(\Delta\theta|\theta)} could also be parallelized by sampling multiple sets of parameters to predict their respective samples, but this is harder to implement. However, policy gradient has its own issues, the most typical being the high variance of the gradient estimate. Consequently, it usually requires pre-training with a standard likelihood objective until it is "close enough" before using policy gradient for fine-tuning. In contrast, Equation [eq:sg] from the original paper attempts to directly optimize the evaluation metric from start to finish and leverages the differentiable objective to implement importance sampling, combining the advantages of both differentiable and non-differentiable optimization.
Summary
This article introduced a new perspective for understanding optimization algorithms, under which the optimization of differentiable and non-differentiable objectives is unified. For differentiable objective functions, by performing first-order and second-order expansions within this perspective, we can derive gradient descent and Newton’s method, respectively. For non-differentiable objective functions, by using differentiable approximations for importance sampling, we can similarly complete the optimization of non-differentiable objectives.
Please include the original address when reposting: https://kexue.fm/archives/7521
For more details on reposting, please refer to: "Scientific Space FAQ"