softmax:建立hard-attention到soft-attention的桥梁

查询向量 $\boldsymbol{q}$ 与每个 $\boldsymbol{x}_{i}$ 计算相关性 $\alpha_{i}$ 后，为什么要使用softmax进行归一化？其他归一化方法不行？hard-attention和soft-attention在数学上有什么联系？这里提供一种数学解释。

由于引入softmax函数对注意力评分函数进行归一化并对编码序列进行加权平均容易让人误解为注意力机制就是加权平均，下面会对这个误解展开分析。

注意力机制简述

假设有一向量序列 $\boldsymbol{X} = [\boldsymbol{x}_{1}, \dots, \boldsymbol{x}_{n}] \in \mathbb{R}^{n \times d}$，有一个和任务相关的向量，称为查询向量 $\boldsymbol{q}$，那么注意力机制要做的事情分三步：

查询向量 $\boldsymbol{q}$ 与每个 $\boldsymbol{x}_{i}$ 计算相关性 $\alpha_{i}$，通过评分函数获得，即$\alpha_{i} = s(\boldsymbol{q}, \boldsymbol{x}_{i})$
使用softmax归一化相关性 $\alpha_{i}$，获得注意力分布
根据注意力分布计算向量序列的加权平均

这个过程用数学表示如下。首先注意力分布为，

$\begin{align} p(z=i|\boldsymbol{X},\boldsymbol{q}) &= \operatorname{softmax}(\alpha_{1},\dots,\alpha_{n}) \newline &= \frac{\exp(\alpha_{i})}{\displaystyle{\sum_{i=1}^{n}\exp(\alpha_{i})}} \newline &= \frac{\exp(s(\boldsymbol{q}, \boldsymbol{x}_{i}))}{\displaystyle{\sum_{i=1}^{n}\exp(s(\boldsymbol{q}, \boldsymbol{x}_{i}))}} \end{align}$

根据注意力分布计算向量序列的加权平均，获得注意力的输出，

$\operatorname{Attention}(\boldsymbol{X},\boldsymbol{q}) = \sum_{i=1}^{n}p(z=i|\boldsymbol{X},\boldsymbol{q})\boldsymbol{x}_{i}$

这就是注意力机制，严格来说说是soft的注意力机制。与之相反的hard注意力机制为，

那么soft的注意力机制为什么要使用softmax多相关性进行归一化？

soft Attention的导出

评分函数s计算查询向量 $\boldsymbol{q}$ 与每个 $\boldsymbol{x}_{i}$ 相关性$\alpha_{i} = s(\boldsymbol{q}, \boldsymbol{x}_{i})$，另$\boldsymbol{\alpha} = [\alpha_1, \dots, \alpha_n]$。$Q_{[i]}$表示向量$Q$的第$i$个分量。$\boldsymbol{x}_i
= \operatorname{\hat{A}ttention}(\boldsymbol{X},\boldsymbol{q})$表示在给定查询向量$\boldsymbol{q}$情况下，在向量序列 $\boldsymbol{X} = [\boldsymbol{x}_{1}, \dots, \boldsymbol{x}_{n}] \in \mathbb{R}^{n \times d}$中找到与其最相关的向量$\boldsymbol{x}_i$，也就是hard-attention，然后通过一系列推导soft化hard-attention，

$\begin{align} \boldsymbol{x}_i &= \operatorname{\hat{A}ttention}(\boldsymbol{X},\boldsymbol{q}) \newline &= \boldsymbol{x}_{\underset{i=1,\cdots,n}{\text{argmax}} \; s(\boldsymbol{q}, \boldsymbol{x}_{i})} \newline &= \sum_{i=1}^{n} \text{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; s(\boldsymbol{q}, \boldsymbol{x}_{i}) \Big)_{[i]} \times \boldsymbol{x}_i \newline &= \sum_{i=1}^{n} \text{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \alpha_i \Big)_{[i]} \times \boldsymbol{x}_i \newline &=\sum_{i=1}^{n} \text{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; \boldsymbol{\alpha} - \max(\boldsymbol{\alpha})\Big)_{[i]} \times \boldsymbol{x}_i \newline &\approx \sum_{i=1}^{n} \text{one-hot}\Big(\underset{i=1,\cdots,n}{\text{argmax}} \; e^{\boldsymbol{\alpha} - \displaystyle \log \sum_{i=1}^{n} e^{\alpha_i}} \Big)_{[i]} \times \boldsymbol{x}_i \newline &\approx \sum_{i=1}^{n} \frac{e^{\alpha_i}}{\displaystyle \sum_{i=1}^{n} e^{\alpha_i}} \boldsymbol{x}_i \newline &= \sum_{i=1}^{n} \operatorname{softmax}(\boldsymbol{\alpha})_{[i]} \times \boldsymbol{x}_i \newline &= \operatorname{Attention}(\boldsymbol{X},\boldsymbol{q}) \end{align}$

需要说明几点：

引入$\boldsymbol{\alpha} - \max(\boldsymbol{\alpha})$使得最大值为0，使得$e^0 = 1$，对应one-hot中的1
引入$e^x$是考虑到$e^0=1, 0 \lt e^{x|_{x \lt 0}} \lt 1$，更好适配one-hot特点
max不具有光滑性，被替换为其光滑近似logsumexp

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

注意力机制，本质上是想找与查询向量最相关的隐向量序列，即计算argmax，考虑到输出要对齐，实际上是想找，onehot(argmax)，比如的任务是中视野中找苹果，直接依据苹果的特征查找苹果。但是这个求解是无法计算梯度，所以才用onehot(argmax)的光滑版本softmax，以获得良好的梯度特性。也就是说，软性注意力机制中用到加权平均只不过是一种数学上的妥协，同时也带来很多有作用的side effect，但这些side effect不是注意力机制引入加权平均的本质原因。

因此，可以看到softmax是hard-attention到soft-attention的桥梁。

总结

本文从数学角度解释Attention为什么要使用softmax多相关性进行归一化，即加权平均只能说是一种数学技巧，而不是注意力机制的目标。同时也明白，softmax是建立hard-attention到soft-attention的桥梁。

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