max函数光滑逼近：一种与softmax相关的形式

一个关于max函数光滑逼近，其特例居然和均值、tanh函数、Logistics函数相关！

这个逼近其实在过去的文章中也略有提及，不过今天那它来单独分析一下。

max & argmax

设有向量

$$\boldsymbol{x} = [x_{1}, \cdots,x_{n}]$$

$\max(\boldsymbol{x})$就是要找向量中元素的最大值。$\arg \max( \boldsymbol{x})$要找向量中元素的最大值的标量，就是所在的位置。假设第$x_j$最大，那么

$$j = \arg \max (\boldsymbol{x})$$

有时候我们希望这个“标量”能够对齐原来向量的长度，即表达成$\operatorname{one-hot}\left(\arg \max( \boldsymbol{x}) \right)$，

$$\operatorname{one-hot}\left(\arg \max( \boldsymbol{x}) \right) = [0, \dots, 1, \dots, 0]$$

等式右侧第$j$个分量为$1$，其他为$0$​。比如圆周率的头几个数，

$$\boldsymbol{x} = [3,1,4,1,5,9,2]$$

结果为

$$6 = \arg \max \left( [3,1,4,1,5,{\color{red} {9}},2] \right)$$

one-hot形式为，

$$[0,0,0,0,0,1,0] = \operatorname{one-hot}\big(\arg \max \left([3,1,4,1,5,{\color{red} {9}},2]\right) \big)$$

于是$\max(\boldsymbol{x})$​可以使用$\arg \max( \boldsymbol{x})$​表示，

$$\begin{align} \max(\boldsymbol{x}) &= \boldsymbol{x} \otimes \operatorname{one-hot}\left(\arg \max( \boldsymbol{x}) \right) \newline &= \sum_{i=1}^{n} x_i \times \operatorname{one-hot}[\arg \max(\boldsymbol{x})]_{[i]} \end{align}$$

这里$\operatorname{one-hot}[\arg \max(\boldsymbol{x})]_{[i]}$表示$\operatorname{one-hot}[\arg \max(\boldsymbol{x})]$的第$i$个分量。

argmax的光滑逼近

一下推导使用到$\max(\boldsymbol{x})$​的一光滑形式，但这不是我们要的最终形式，

$$\max(\boldsymbol{x}) \approx \frac{1}{\alpha} \log( \sum_{i=1}^{n}e^{\alpha x_{i}})$$

上式右边部分称为Logsumexp。这个从凸优化的角度，考虑到Logsumexp的凸性，因此有，

$$\begin{align} f(x_{1},x_{2},\dots,x_{n}) &= \frac{1}{\alpha}\log(\sum_{i=1}^{n}e^{\alpha x_{i}}) \newline &= \frac{1}{\alpha}\log(n \times \frac{1}{n}\sum_{i=1}^{n}e^{\alpha x_{i}}) \newline &\le \frac{\log(n)}{\alpha} + \frac{1}{\alpha}\log(\sum_{i=1}^{n}\frac{1}{n} \times e^{\alpha x_{i}}) \newline &\le \frac{\log(n)}{\alpha} + \frac{1}{n} \sum_{i=1}^{n}x_{i} \newline &\le \frac{\log(n)}{\alpha} + \max(x_{1},x_{2},\dots,x_{n}) \end{align}$$

另一方面，假设有数据$x_1, \dots, x_n$，其中$x_j$最大，有

$$\begin{align} f(x_{1},x_{2},\dots,x_{n}) &= \frac{1}{\alpha}\log(\sum_{i=1}^{n}e^{\alpha x_{i}}) \newline &\ge \frac{1}{\alpha}\log(e^{\alpha x_{j}}) \newline &= \max(x_{1},\dots,x_{n}) \end{align}$$

因此有，

$$\max(x_{1},\dots,x_{n}) \le \frac{1}{\alpha}\log(\sum_{i=1}^{n}e^{\alpha x_{i}}) \le \frac{\log(n)}{\alpha} + \max(x_{1},x_{2},\dots,x_{n})$$

通过夹逼定理有，

$$\lim_{\alpha \to + \infty} \frac{1}{\alpha}\log(\sum_{i=1}^{n}e^{\alpha x_{i}}) = \max(x_{1},\dots,x_{n})$$

同样我们假设$\boldsymbol{x} = [x_{1}, \cdots,x_{n}]$​中$x_j$​最大，那么one-hot形式第$j$​个分量为$1$​，其他为$0$​，有推导，

$$\begin{align} [0, \dots, 1, \dots, 0] &= \operatorname{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; x_i \Big) \newline &= \operatorname{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \big[x_i - \max(\boldsymbol{x}) \big] \Big) \newline &= \operatorname{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \exp\big[x_i - \max(\boldsymbol{x}) \big] \Big) \newline &\approx \operatorname{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \exp\Big[x_i - \frac{1}{\alpha}\log(\sum_{i=1}^{n} e^{\alpha x_i}) \Big] \Big) \newline &= \operatorname{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \exp \frac{1}{\alpha} \Big[\alpha x_i - \log(\sum_{i=1}^{n} e^{\alpha x_i}) \Big] \Big) \newline &= \operatorname{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \exp \Big[\alpha x_i - \log(\sum_{i=1}^{n} e^{\alpha x_i}) \Big] \Big) \newline &= \operatorname{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \frac{e^{\alpha x_i}}{\displaystyle \sum_{i=1}^{n} e^{\alpha x_i}} \Big) \newline &\approx [\frac{e^{\alpha x_1}}{\displaystyle \sum_{i=1}^{n} e^{\alpha x_i}}, \dots, \frac{e^{\alpha x_n}}{\displaystyle \sum_{i=1}^{n} e^{\alpha x_i}}] \newline &= \operatorname{softmax}(\alpha \boldsymbol{x}) \newline \end{align}$$

需要说明几点：

• 引入$\big[x_i - \max(\boldsymbol{x}) \big]$使得最大值为0，使得$e^0 = 1$，对应one-hot中的1
• 引入$e^x$是考虑到$e^0=1, 0 \lt e^{x|_{x \lt 0}} \lt 1$，并拉大$[x_1, x_2, \dots, x_n]$间的距离，更好适配one-hot特点
• max不具有光滑性，被替换为其光滑近似logsumexp，$\operatorname{logsumexp}(x_{1},x_{2},\dots,x_{n}) = \log(\sum_{i=1}^{n}e^{x_{i}})$

理解好这三点就明白上述推导过程。

max的光滑逼近

于是，接着第一部分的推导，

$$\begin{align} \max(\boldsymbol{x}) &= \boldsymbol{x} \otimes \operatorname{one-hot}\left(\arg \max( \boldsymbol{x}) \right) \newline &= \sum_{i=1}^{n} x_i \times \operatorname{one-hot}[\arg \max(\boldsymbol{x})]_{[i]} \newline &\approx \sum_{i=1}^{n} \frac{e^{\alpha x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}} x_{i} \newline &= \frac{\displaystyle \sum_{i=1}^{n}x_{i} e^{\alpha x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}} \end{align}$$

于是，这就是我们要找的$\max(\boldsymbol{x})$光滑形式，

$$\max(\boldsymbol{x})\approx \frac{\displaystyle \sum_{i=1}^{n}x_{i} e^{\alpha x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}}$$

有三个关键点：

• 如果$\alpha = 0$，则上式右侧为均值
• 如果$\alpha = +\infty$，则上式右侧为$\max(\boldsymbol{x})$
• 如果$\alpha = -\infty$，则上式右侧为$\min(\boldsymbol{x})$

第三点使用$\min(\boldsymbol{x}) = - \max(-\boldsymbol{x})$性质容易理解。

特例

取$n=2, x_1 = 0, x_2 = x$​有特例，

$$\begin{align} \max(\boldsymbol{x}) &\approx \max\{0, x\} \newline &= \frac{\displaystyle \sum_{i=1}^{n}x_{i} e^{\alpha x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}} \newline &= \frac{x e^{\alpha x}}{1 + e^{ \alpha x}} = x \sigma(\alpha x) \end{align}$$

这也是$\operatorname{relu}(x) = \max{0, x}$的光滑逼近，也就是激活函数Swish。

取$n=2, x_1 = -x, x_2 = x$​​，有特例，

$$\begin{align} \max(\boldsymbol{x}) &=\max\{-x, x\} \newline &\approx \frac{\displaystyle \sum_{i=1}^{n}x_{i} e^{\alpha x_{i}}}{\displaystyle \sum_{i=1}^{n}e^{\alpha x_{i}}} \newline &= \frac{x(e^{\alpha x} - e^{-\alpha x})}{e^{\alpha x} + e^{-\alpha x}} \newline &= x \tanh(\alpha x) \end{align}$$

这是$f(x) = |x| = =\max{-x, x}$的光滑逼近。

总结

其实这个$\max(\boldsymbol{x})$函数的光滑形式比较优雅，这里单独拿出来说一说。

转载请包括本文地址：https://allenwind.github.io/blog/11886
更多文章请参考：https://allenwind.github.io/blog/archives/