LogSumExp和Softmax的泰勒展开

最近看到论文《The Key to Going Linear: Analysis-Driven Transformer Linearization》里边直接对Softmax做近似展开来线性化Attention，有点收获，在此记录一下。

一些记号[#](https://spaces.ac.cn/archives/11814#%E4%B8%80%E4%BA%9B%E8%AE%B0%E5%8F%B7)

首先引入如下记号：

$$ (\text{1}) \mathbf{\mathit{x}} = \left(\right. x_{1} , x_{2} , \hdots , x_{n} \left.\right) \in \mathbb{R}^{n} , \bar{\mathbf{\mathit{x}}} = \frac{1}{n} \sum_{i = 1}^{n} x_{i} \\ (\text{2}) \text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) = log ⁡ \sum_{i = 1}^{n} e^{x_{i}} = log ⁡ n + log ⁡ \bar{e^{\mathbf{\mathit{x}}}} \\ (\text{3}) \text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) = \frac{e^{\mathbf{\mathit{x}}}}{\sum_{i = 1}^{n} e^{x_{i}}} = e^{\mathbf{\mathit{x}} - \text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)} $$

并约定向量的函数运算都是Hadamard意义下的，例如

$$ (\text{4}) e^{\mathbf{\mathit{x}}} = \left(\right. e^{x_{1}} , e^{x_{2}} , \hdots , e^{x_{n}} \left.\right) , \mathbf{\mathit{x}}^{2} = \left(\right. x_{1}^{2} , x_{2}^{2} , \hdots , x_{n}^{2} \left.\right) $$

特别地，要仔细留意$\bar{e^{\mathbf{\mathit{x}}}}$与$e^{\bar{\mathbf{\mathit{x}}}}$、$\bar{\mathbf{\mathit{x}}^{2}}$与$\left(\bar{\mathbf{\mathit{x}}}\right)^{2}$的区别：前者是先对每个分量做函数再平均，后者是先平均再做函数，二者一般不相等。 除了上述定义外，$\text{logsumexp}$和$\text{softmax}$还可以通过梯度联系起来：

$$ (\text{5}) \text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) = \nabla_{\mathbf{\mathit{x}}} \text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) $$

基本展开[#](https://spaces.ac.cn/archives/11814#%E5%9F%BA%E6%9C%AC%E5%B1%95%E5%BC%80)

接着我们对$\text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)$做展开。引入$\mathbf{\mathit{y}} = \mathbf{\mathit{x}} - \bar{\mathbf{\mathit{x}}}$，那么

$$ (\text{6}) \begin{matrix}\text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) = & log ⁡ n + \bar{\mathbf{\mathit{x}}} + log ⁡ \bar{e^{\mathbf{\mathit{y}}}} \\ = & log ⁡ n + \bar{\mathbf{\mathit{x}}} + log ⁡ \left(\right. \bar{1 + \mathbf{\mathit{y}} + \frac{\mathbf{\mathit{y}}^{2}}{2} + \frac{\mathbf{\mathit{y}}^{3}}{6} + \frac{\mathbf{\mathit{y}}^{4}}{24} + \hdots} \left.\right) \\ = & log ⁡ n + \bar{\mathbf{\mathit{x}}} + log ⁡ \left(\right. 1 + \frac{\bar{\mathbf{\mathit{y}}^{2}}}{2} + \frac{\bar{\mathbf{\mathit{y}}^{3}}}{6} + \frac{\bar{\mathbf{\mathit{y}}^{4}}}{24} + \hdots \left.\right) \\ = & log ⁡ n + \bar{\mathbf{\mathit{x}}} + \frac{\bar{\mathbf{\mathit{y}}^{2}}}{2} + \frac{\bar{\mathbf{\mathit{y}}^{3}}}{6} + \left(\right. \frac{\bar{\mathbf{\mathit{y}}^{4}}}{24} - \frac{\left(\right. \bar{\mathbf{\mathit{y}}^{2}} \left.\right)^{2}}{8} \left.\right) + \hdots\end{matrix} $$

最后一步用到了$log ⁡ \left(\right. 1 + t \left.\right) = t - t^{2} / 2 + t^{3} / 3 - \hdots$，如果有需要，还可以继续往下展开。如果读者不谙此道，也可以交给Kimi完成。 这里引入$\bar{\mathbf{\mathit{x}}}$的偏置是一个简化形式的小技巧，如果不引入，也能得到等价的结果，只不过会将$\mathbf{\mathit{y}}$的各个项显式展开：

$$ (\text{7}) \text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) = log ⁡ n + \bar{\mathbf{\mathit{x}}} + \underset{\bar{\mathbf{\mathit{y}}^{2}} / 2}{\underbrace{\left(\right. \frac{\bar{\mathbf{\mathit{x}}^{2}}}{2} - \frac{\left(\bar{\mathbf{\mathit{x}}}\right)^{ 2}}{2} \left.\right)}} + \underset{\bar{\mathbf{\mathit{y}}^{3}} / 6}{\underbrace{\left(\right. \frac{\bar{\mathbf{\mathit{x}}^{3}}}{6} - \frac{\bar{\mathbf{\mathit{x}}} \bar{\mathbf{\mathit{x}}^{2}}}{2} + \frac{\left(\bar{\mathbf{\mathit{x}}}\right)^{ 3}}{3} \left.\right)}} + \hdots $$

求个梯度[#](https://spaces.ac.cn/archives/11814#%E6%B1%82%E4%B8%AA%E6%A2%AF%E5%BA%A6)

对于$\text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)$，我们直接利用恒等式$(\text{5})$，对式$(\text{6})$两边求导即得$\text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)$的展开式。为此，我们利用

$$ (\text{8}) \nabla_{\mathbf{\mathit{x}}} \bar{\mathbf{\mathit{x}}} = \frac{1}{n} , \nabla_{\mathbf{\mathit{x}}} \bar{\mathbf{\mathit{y}}^{m}} = \frac{m}{n} \left(\right. \mathbf{\mathit{y}}^{m - 1} - \bar{\mathbf{\mathit{y}}^{m - 1}} \left.\right) $$

得

$$ (\text{9}) \text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) = \frac{1}{n} \left(\right. 1 + \mathbf{\mathit{y}} + \left(\right. \frac{\mathbf{\mathit{y}}^{2}}{2} - \frac{\bar{\mathbf{\mathit{y}}^{2}}}{2} \left.\right) + \left(\right. \frac{\mathbf{\mathit{y}}^{3}}{6} - \frac{\bar{\mathbf{\mathit{y}}^{3}}}{6} - \frac{\bar{\mathbf{\mathit{y}}^{2}} \mathbf{\mathit{y}}}{2} \left.\right) + \hdots \left.\right) $$

注意如果截断有限项，那么右端依然能保持所有分量求和为1（从$\nabla_{\mathbf{\mathit{x}}} \bar{\mathbf{\mathit{y}}^{m}}$的梯度形式即可看出），但是无法保证每个分量的非负性，这一点在实际应用中需要留意（比如需要取对数求熵的时候）。 除了求梯度外，我们还可以用$\text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right) = e^{\mathbf{\mathit{x}} - \text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)} = e^{\mathbf{\mathit{y}} - \text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{y}} \left.\right)}$来展开，结合式$(\text{6})$得：

$$ (\text{10}) \begin{matrix}e^{\mathbf{\mathit{y}}} e^{- \text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{y}} \left.\right)} = & \left(\right. 1 + \mathbf{\mathit{y}} + \frac{\mathbf{\mathit{y}}^{2}}{2} + \frac{\mathbf{\mathit{y}}^{3}}{6} + \hdots \left.\right) exp ⁡ \left(\right. - log ⁡ n - \frac{\bar{\mathbf{\mathit{y}}^{2}}}{2} - \frac{\bar{\mathbf{\mathit{y}}^{3}}}{6} - \hdots \left.\right) \\ = & \frac{1}{n} \left(\right. 1 + \mathbf{\mathit{y}} + \left(\right. \frac{\mathbf{\mathit{y}}^{2}}{2} - \frac{\bar{\mathbf{\mathit{y}}^{2}}}{2} \left.\right) + \left(\right. \frac{\mathbf{\mathit{y}}^{3}}{6} - \frac{\bar{\mathbf{\mathit{y}}^{3}}}{6} - \frac{\bar{\mathbf{\mathit{y}}^{2}} \mathbf{\mathit{y}}}{2} \left.\right) + \hdots \left.\right)\end{matrix} $$

稀疏注意[#](https://spaces.ac.cn/archives/11814#%E7%A8%80%E7%96%8F%E6%B3%A8%E6%84%8F)

那这两个展开式有什么应用呢？首先是$\text{logsumexp} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)$的，它跟MoBA这类Block Sparse Attention的块打分有关。我们知道，Full Attention按$e^{\mathbf{\mathit{q}} \cdot \mathbf{\mathit{k}}}$给每个Token打分，那对于某个块$\mathcal{B}$，它的打分可以合理定义为每个Token的分数之和，等价于

$$ (\text{11}) s_{\mathcal{B}} = log ⁡ \underset{t \in \mathcal{B}}{\sum} e^{\mathbf{\mathit{q}} \cdot \mathbf{\mathit{k}}_{t}} $$

它的一阶近似是$log ⁡ \left|\right. \mathcal{B} \left|\right. + \mathbf{\mathit{q}} \cdot \bar{\mathbf{\mathit{k}}}$，其中$\bar{\mathbf{\mathit{k}}} = \frac{1}{\left|\right. \mathcal{B} \left|\right.} \underset{t \in \mathcal{B}}{\sum} \mathbf{\mathit{k}}_{t}$，这正好对应着MoBA用AvgPooling作为块内Landmark向量的经验做法。如果感觉一阶近似太粗糙，那么可以考虑高阶修正，根据式$(\text{6})$，二阶项是

$$ (\text{12}) \frac{1}{2 \left|\right. \mathcal{B} \left|\right.} \underset{t \in \mathcal{B}}{\sum} \left(\right. \mathbf{\mathit{q}} \cdot \left(\right. \mathbf{\mathit{k}}_{t} - \bar{\mathbf{\mathit{k}}} \left.\right) \left.\right)^{2} = \frac{1}{2} \mathbf{\mathit{q}}^{\top} \underset{\mathbf{\Sigma}}{\underbrace{\left(\right. \frac{1}{\left|\right. \mathcal{B} \left|\right.} \underset{t \in \mathcal{B}}{\sum} \left(\right. \mathbf{\mathit{k}}_{t} - \bar{\mathbf{\mathit{k}}} \left.\right) \left(\right. \mathbf{\mathit{k}}_{t} - \bar{\mathbf{\mathit{k}}} \left.\right)^{\top} \left.\right)}} \mathbf{\mathit{q}} $$

即二阶近似是$log ⁡ \left|\right. \mathcal{B} \left|\right. + \mathbf{\mathit{q}} \cdot \bar{\mathbf{\mathit{k}}} + \mathbf{\mathit{q}}^{\top} \mathbf{\Sigma} \mathbf{\mathit{q}} / 2$，其中$\mathbf{\Sigma}$正好是块内$\mathbf{\mathit{k}}_{t}$的协方差矩阵。进一步地，我们可以考虑对角近似来降低计算成本。这一高阶修正的思路，跟SPLA本质上是相同的。后续相关工作还有HiLS，它另外学了一个非等权平均的中心向量来取代$\bar{\mathbf{\mathit{k}}}$。

线性注意[#](https://spaces.ac.cn/archives/11814#%E7%BA%BF%E6%80%A7%E6%B3%A8%E6%84%8F)

至于$\text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)$展开式的应用，自然是线性化Attention了。在《Transformer升级之路：5、作为无限维的线性Attention》中，我们总结了三种线性化的思路，它们都是对$e^{\mathbf{\mathit{q}} \cdot \mathbf{\mathit{k}}}$做近似。然而，近似$e^{\mathbf{\mathit{q}} \cdot \mathbf{\mathit{k}}}$后还要归一化，倒不如直接对归一化后的$\text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)$做近似。

记$\mathbf{\mathit{x}} = \left(\right. \mathbf{\mathit{q}} \cdot \mathbf{\mathit{k}}_{1} , \hdots , \mathbf{\mathit{q}} \cdot \mathbf{\mathit{k}}_{n} \left.\right)$，那么有一阶近似：

$$ (\text{13}) \text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)_{i} \approx \frac{1}{n} + \mathbf{\mathit{q}} \cdot \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right) $$

这正是开头论文所讨论的方案，它本身也很直观：平均向量$\bar{\mathbf{\mathit{k}}}$为注意力的相对大小提供了一个基准，相似度比$\bar{\mathbf{\mathit{k}}}$大的Token，要加大注意力，否则减少。为了提高精度，Based将$e^{\mathbf{\mathit{q}} \cdot \mathbf{\mathit{k}}}$展开到了二阶（此时刚好保证非负，参考这里），但从我们的视角看，考虑$\text{softmax} ⁡ \left(\right. \mathbf{\mathit{x}} \left.\right)$的二阶近似更科学，二阶项是

$$ (\text{14}) \frac{1}{2 n} \left[\right. \left(\right. \mathbf{\mathit{q}} \cdot \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right) \left.\right)^{2} - \frac{1}{n} \sum_{j = 1}^{n} \left(\right. \mathbf{\mathit{q}} \cdot \left(\right. \mathbf{\mathit{k}}_{j} - \bar{\mathbf{\mathit{k}}} \left.\right) \left.\right)^{2} \left]\right. $$

其中$\left(\right. \mathbf{\mathit{q}} \cdot \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right) \left.\right)^{2}$可以通过引入外积来变成内积：

$$ (\text{15}) \left(\right. \mathbf{\mathit{q}} \cdot \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right) \left.\right)^{2} = \mathbf{\mathit{q}}^{\top} \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right) \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right)^{\top} \mathbf{\mathit{q}} = \left(\langle \mathbf{\mathit{q}} \mathbf{\mathit{q}}^{\top} , \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right) \left(\right. \mathbf{\mathit{k}}_{i} - \bar{\mathbf{\mathit{k}}} \left.\right)^{\top} \rangle\right)_{F} $$

所以跟Based类似，将Softmax Attention的Softmax截断到二阶近似，结果也是线性注意力。

文章小结[#](https://spaces.ac.cn/archives/11814#%E6%96%87%E7%AB%A0%E5%B0%8F%E7%BB%93)

本文分别推导了LogSumExp和Softmax的泰勒展开式，并讨论了它们的两个潜在应用。

_转载到请包括本文地址：[https://spaces.ac.cn/archives/11814](https://spaces.ac.cn/archives/11814 "LogSumExp和Softmax的泰勒展开")_

_更详细的转载事宜请参考：_[《科学空间FAQ》](https://spaces.ac.cn/archives/6508#%E6%96%87%E7%AB%A0%E5%A6%82%E4%BD%95%E8%BD%AC%E8%BD%BD/%E5%BC%95%E7%94%A8 "《科学空间FAQ》")