最近看到论文《The Key to Going Linear: Analysis-Driven Transformer Linearization》里边直接对Softmax做近似展开来线性化Attention,有点收获,在此记录一下。
一些记号
首先引入如下记号:
\begin{gather}\newcommand{logsumexp}{\mathop{\text{logsumexp}}}\newcommand{softmax}{\mathop{\text{softmax}}}
\boldsymbol{x} = (x_1, x_2, \cdots, x_n) \in \mathbb{R}^n,\qquad\overline{\boldsymbol{x}} = \frac{1}{n}\sum_{i=1}^n x_i \\
\logsumexp(\boldsymbol{x}) = \log\sum_{i=1}^n e^{x_i} = \log n + \log \overline{e^{\boldsymbol{x}}} \\
\softmax(\boldsymbol{x}) = \frac{e^{\boldsymbol{x}}}{\sum_{i=1}^n e^{x_i}} = e^{\boldsymbol{x} - \logsumexp(\boldsymbol{x})} \\
\end{gather}