On Function Smoothing: Differentiable Approximations to Non-Differentiable Functions

函数光滑化杂谈：不可导函数的可导逼近 苏剑林 · 2019-05-20 · Mathematics derivative function approximation smoothness

Generally speaking, what neural networks process is continuous floating-point numbers, and their standard outputs are also continuous values. But in real-world problems, we often need a discrete result — for example, in classification we want to output the correct category, and "category" is discrete while "category probability" is continuous. Similarly, many task evaluation metrics are actually discrete, such as accuracy and F1 for classification, or BLEU for machine translation, and so on.

Taking classification as an example again, a common evaluation metric is accuracy, while a common loss function is cross-entropy. There is indeed some correlation between decreasing cross-entropy and increasing accuracy, but it's not a strictly monotonic relationship. In other words, a drop in cross-entropy doesn't guarantee an increase in accuracy. Clearly, it would be ideal if we could directly use the negative of accuracy as the loss function, but accuracy is non-differentiable (it involves operations like $\text{argmax}$), so it can't be used directly.

There are generally two solutions to this problem: one is to bring in reinforcement learning, treating accuracy as the reward function — this is a bit like "using a sledgehammer to crack a nut"; the other is to try to find a smooth, differentiable approximation formula for accuracy. This post explores smooth approximations to common non-differentiable functions — sometimes we call this "smoothing," and sometimes "softening."

max

Most of what follows is built on the smooth approximation of the $\max$ operation. We have:

\begin{equation}\max(x_1,x_2,\dots,x_n) = \lim_{K\to +\infty}\frac{1}{K}\log\left(\sum_{i=1}^n e^{K x_i}\right)\end{equation}

Hence, by choosing a constant $K$, we get the approximation:

\begin{equation}\max(x_1,x_2,\dots,x_n) \approx \frac{1}{K}\log\left(\sum_{i=1}^n e^{K x_i}\right)\end{equation}

In models, we can often set $K=1$, which is equivalent to folding $K$ into the model itself, giving us the simplest form:

\begin{equation}\begin{aligned}\max(x_1,x_2,\dots,x_n) \approx&\, \log\left(\sum_{i=1}^n e^{x_i}\right) \\ \triangleq&\, \text{logsumexp}(x_1, x_2, \dots, x_n)\end{aligned}\label{eq:max-approx}\end{equation}

Here $\text{logsumexp}$ appears, which is a very common operator; in this context it is the smooth approximation of the $\max$ function. Indeed, the smooth approximation of $\max$ is actually $\text{logsumexp}$, not the literally similar-looking $\text{softmax}$. For related derivations, you can also refer to my earlier post In Search of a Smooth Maximum Function.

softmax

We just said that $\text{softmax}$ is not a smooth approximation of $\max$ — so what is it a smooth approximation of? It's actually the smooth approximation of $\text{onehot}(\text{argmax}(\boldsymbol{x}))$, i.e., first finding the position of the maximum value, then generating a vector of the same length with 1 at that position and 0 everywhere else, for example:

\begin{equation}[2, 1, 4, 5, 3]\quad \to \quad [0, 0, 0, 1, 0]\end{equation}

We can give a simple derivation from $\text{logsumexp}$ to $\text{softmax}$. Consider the vector $\boldsymbol{x}=[x_1, x_2, \dots, x_n]$, and then consider

\begin{equation}\boldsymbol{x}'=[x_1, x_2, \dots, x_n] - \max(x_1, x_2, \dots, x_n)\end{equation}

that is, subtracting the overall maximum from every entry. This new vector has the maximum at the same position as the original vector, i.e., $\text{onehot}(\text{argmax}(\boldsymbol{x}))=\text{onehot}(\text{argmax}(\boldsymbol{x}'))$. Without loss of generality, suppose the entries of $x_1,x_2,\dots,x_n$ are pairwise distinct; then the maximum of the new vector is clearly 0, and all other entries are negative. This lets us consider

\begin{equation}e^{\boldsymbol{x}'}=[e^{x_1 - \max(x_1, x_2, \dots, x_n)}, e^{x_2 - \max(x_1, x_2, \dots, x_n)}, \dots, e^{x_n - \max(x_1, x_2, \dots, x_n)}]\end{equation}

as an approximation to $\text{onehot}(\text{argmax}(\boldsymbol{x}'))$, since the maximum is 0, so the corresponding position is $e^0=1$, while the rest are negative and, after exponentiation, will be close to 0.

Finally, substituting the approximation $\eqref{eq:max-approx}$ into the expression above and simplifying gives us

\begin{equation}\begin{aligned}\text{onehot}(\text{argmax}(\boldsymbol{x}))=&\,\text{onehot}(\text{argmax}(\boldsymbol{x}'))\\ \approx&\, \left(\frac{e^{x_1}}{\sum\limits_{i=1}^n e^{x_i}}, \frac{e^{x_2}}{\sum\limits_{i=1}^n e^{x_i}}, \dots, \frac{e^{x_n}}{\sum\limits_{i=1}^n e^{x_i}}\right)\\ \triangleq&\,\text{softmax}(x_1, x_2, \dots, x_n) \end{aligned}\end{equation}

argmax

$\text{argmax}$ refers to directly giving the index (an integer) of the position of the maximum value in a vector, for example

\begin{equation}[2, 1, 4, 5, 3]\quad \to \quad 4\end{equation}

Here we follow the usual convention of indexing from 1, so the returned result is 4; but in programming languages, indexing generally starts from 0, so the returned result would usually be 3.

If we want a smooth approximation to $\text{argmax}$, we naturally hope it outputs a floating-point number close to 4. To construct such an approximation, first notice that $\text{argmax}$ is actually equal to

\begin{equation}\text{sum}\Big(\underbrace{[1, 2, 3, 4, 5]}_{\text{order vector [1, 2, ..., n]}}\,\, \otimes\,\,\underbrace{[0, 0, 0, 1, 0]}_{\text{onehot}(\text{argmax}(\boldsymbol{x}))}\Big)\end{equation}

that is, the inner product of the array $[1, 2, \dots, n]$ and $\text{onehot}(\text{argmax}(\boldsymbol{x}))$. Constructing a softened version of $\text{argmax}$ is then simple: just replace $\text{onehot}(\text{argmax}(\boldsymbol{x}))$ with $\text{softmax}(\boldsymbol{x})$, giving

\begin{equation}\text{argmax} (\boldsymbol{x}) \approx \sum_{i=1}^n i\times \text{softmax}(\boldsymbol{x})_i\end{equation}

Accuracy

Most of the approximations discussed above are derived by first expressing the correct form in terms of one-hot vectors, and then using softmax to approximate the one-hot operation, thereby obtaining a smooth approximation. Using this same idea, we can also derive smooth approximations for many other operators, such as accuracy.

For simplicity, let's introduce the notation $\boldsymbol{1}_k$, denoting the one-hot vector with a 1 at position $k$. Suppose that in a classification problem the target class is $i$ and the predicted class is $j$; then we can consider the one-hot vectors $\boldsymbol{1}_i$ and $\boldsymbol{1}_j$, and take their inner product

\begin{equation}\langle \boldsymbol{1}_i, \boldsymbol{1}_j\rangle = \left\{\begin{aligned}&1,\,\,(i=j)\\ &0,\,\,(i\neq j)\end{aligned}\right.\end{equation}

That is, when the two classes are the same, the inner product is exactly 1, and when they differ, the inner product is exactly 0. So the inner product of the one-hot vectors of the target class and the predicted class exactly defines a "prediction correct" counting function. With this counting function in hand, we can compute accuracy:

\begin{equation}\text{accuracy}=\frac{1}{|\mathcal{B}|}\sum_{\boldsymbol{x}\in\mathcal{B}}\langle \boldsymbol{1}_i(\boldsymbol{x}), \boldsymbol{1}_j(\boldsymbol{x})\rangle\end{equation}

where $\mathcal{B}$ denotes the current batch, and the above expression computes the accuracy over one batch. However, in neural networks, to ensure differentiability, the final output can only be a probability distribution (the result after softmax), so the smooth approximation of accuracy replaces the one-hot vector of the predicted class with the probability distribution:

\begin{equation}\text{accuracy}\approx \frac{1}{|\mathcal{B}|}\sum_{\boldsymbol{x}\in\mathcal{B}}\langle \boldsymbol{1}_i(\boldsymbol{x}), p(\boldsymbol{x})\rangle\end{equation}

Similarly, we can derive smooth approximations to metrics like recall and F1. Taking binary classification as an example, suppose $p(\boldsymbol{x})$ is the predicted probability of the positive class, and $t(\boldsymbol{x})$ is the label (0 or 1) of sample $\boldsymbol{x}$; then the smooth approximation of the F1 score for the positive class is:

\begin{equation}\text{positive class F1}\approx\frac{2 \sum\limits_{\boldsymbol{x}\in\mathcal{B}}t(\boldsymbol{x}) p(\boldsymbol{x})}{\sum\limits_{\boldsymbol{x}\in\mathcal{B}}\big[t(\boldsymbol{x}) + p(\boldsymbol{x})\big]}\end{equation}

The accuracy approximation formula derived this way is differentiable, so its negative can be used directly as a loss. However, in practice, during sampling-based estimation, it's a biased estimator of F1 (since the denominator also involves summation over the batch), and this can sometimes disturb the optimization trajectory or even cause divergence. So, generally speaking, it's better not to use it directly from the start — instead, train first with ordinary cross-entropy until reasonably converged, and then fine-tune using the negative F1 as the loss.

softkmax

$\text{softmax}$ is a smooth approximation of "set the position of the maximum to 1 and everything else to 0" — so what about a smooth approximation of "set the positions of the top $k$ values to 1 and everything else to 0"? We might call this $\text{soft-}k\text{-max}$.

I haven't managed to construct a simple closed form for $\text{soft-}k\text{-max}$, but it can be built recursively:

Given input $\boldsymbol{x}$, initialize $\boldsymbol{p}^{(0)}$ as the all-zero vector;
Compute $\boldsymbol{x} = \boldsymbol{x} - \min(\boldsymbol{x})$ (ensuring all elements are non-negative);
For $i=1,2,\dots,k$, do:
$\boldsymbol{y} = (1 - \boldsymbol{p}^{(i-1)})\otimes\boldsymbol{x}$;
$\boldsymbol{p}^{(i)} = \boldsymbol{p}^{(i-1)} + \text{softmax}(\boldsymbol{y})$
Return $\boldsymbol{p}^{(k)}$.

As for why this works, it becomes clear once you replace $\text{softmax}(\boldsymbol{y})$ with $\text{onehot}(\text{argmax}(\boldsymbol{y}))$ and recurse: essentially, we first compute $\max$, then subtract off the position corresponding to $\max$, so that the second-largest value becomes the new maximum, and then reapply $\text{softmax}$, recursing accordingly.

Summary

Function smoothing is a fairly interesting piece of mathematics that shows up often in machine learning. On the one hand, it's a technique for making certain operations differentiable, allowing models to be optimized directly via backpropagation without having to "bring in" reinforcement learning. On the other hand, in some cases it can also enhance a model's interpretability, since the corresponding non-differentiable function is often quite interpretable — after training with the smoothed version, one can sometimes revert to the non-differentiable version to interpret the model's outputs.

Of course, appreciated purely as an instance of mathematical beauty, it's also quite a pleasure to behold~