Attention的矩阵表示及理解

Attention两篇文章链接：其中一个是Luong,提的dot product attention， 另一个是Vaswali的scaled dot product attention ， 也就是大名鼎鼎的attention is all you need。

说到attention不再过多赘述，论文中的公式推导感觉比较简单，结合自己的理解写一下矩阵层面的表示。数学好的可以跳过。

在attention is all you need这篇文章中，他是这么写的：
$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$，而Luong那篇文章中，公式比较多且分散。
但无论如何，大致总结是（看下图），先算attention score（Q，K相乘），再用softmax算distribution，再把distribution和hidden state相乘获得attention output（最上面那个MatMul），再把output和另一个hidden相加（concat）。

这里以seq2seq模型中的attention为例。

首先我们有encoder hidden state的一个序列: $H=[h_1,h_2,...h_N]$
然后有 $t$ 时刻的decoder state $s^t$, 所有时刻的decoder state就是$S=[s^1,s^2,...,s^t]$。
每次用所有的encoder hidden state去和当前时刻的decoder state相乘（dot product）

对于$t$时刻而言的attention score就是用$e^t=[h_1^T{s}^t,h_2^T{s}^t,...,h_N^T{s}^t]$，
但实际在计算中，我们是把整个decoder hidden state和encoder hidden state乘起来，而不是像循环一样对每个时刻都依次计算
$$E=\left[\begin{array}{cccc}{h}_1^T{s}^1& h_2^T{s}^1,& ...,& h_N^T{s}^1\\ ⋮& ⋮& \ddots & ⋮\\ h_1^T{s}^{t-1}& h_2^T{s}^{t-1},& ...,& h_N^T{s}^{t-1}\\ h_1^T{s}^t& h_2^T{s}^t,& ...,& h_N^T{s}^t\end{array}\right]=\left[\begin{array}{c}{s}^1\\ s^2\\ ⋮\\ s^t\end{array}\right]\cdot \left[\begin{array}{ccccc}{h}_1^T& h_2^T& \cdots & & h_N^T\end{array}\right](1)$$

这里看不懂没关系，先看下一步如何计算attention distribution的计算，这里其实就是用softmax把$e^t$中的每个元素映射到[0,1]之间并且总和为1:
$${\alpha }^t=softmax(e^t)$$

如果用矩阵的方式表示其实就是
$$softmax(E)=[{\alpha }^1,{\alpha }^2,...,{\alpha }^t{]}^T$$ 其中e的右上标是时刻$t$的意思, 而不是指数。其中每一个$\alpha$都是一个1堆概率，长度为N，，也就是encoder hidden state的长度，并且每一个$\alpha$总和为1。

下一步是将attention distribution和encoder hidden state相乘获得attention output, 即把每一个$\alpha$中的每一个概率和对应位置的encoder hidden state相乘，再求和：
$$a_t={\Sigma }_{i=1}^N{\alpha }_i^t{h}_i=[{\alpha }_1^t,{\alpha }_2^t,{\alpha }_3^t,...{\alpha }_i^t]\cdot \left[\begin{array}{c}{h}_1\\ h_2\\ h_3\\ ⋮\\ h_i\end{array}\right]$$
这个操作对应于下图的encoder recurrent layer指向attention distribution的那条红线，以及Attention distribution指向Attention output的黑线以及最上面的三角形。

把每个时刻的$a_t$都放到一起的话：
$$\begin{gathered} A=[a_1,a_2,a_3,...,a_t{]}^T=\left[\begin{array}{c}{a}_1\\ a_2\\ ⋮\\ a_t\end{array}\right]=\left[\begin{array}{cccc}{\alpha }_1^1& {\alpha }_2^1& \cdots & {\alpha }_i^1\\ {\alpha }_1^2& {\alpha }_2^2& \cdots & {\alpha }_i^2\\ ⋮& ⋮& \ddots & ⋮\\ {\alpha }_1^t& {\alpha }_2^t& \cdots & {\alpha }_i^t\end{array}\right]\cdot \left[\begin{array}{c}{h}_1\\ h_2\\ ⋮\\ h_i\end{array}\right] \\ =Softmax(E)\cdot H^T \\ =Softmax(S\cdot H^T)\cdot H^T \end{gathered}$$

A中的每一个$a_t$是当前$t$时刻的attention output，需要和对应时刻的decoder state的s^t执行concat操作。如下图

写成公式就是$concat[a_t,s_t]$，然后对每一个时刻的decoder state都执行这样的操作，其实就是concat(A, S)，其中A就是$[a_1,a_2,...,a_t]=\left[\begin{array}{cccc}{\Sigma }_{i=1}^N{\alpha }_i^1{h}_i& {\Sigma }_{i=1}^N{\alpha }_i^2{h}_i& ...& {\Sigma }_{i=1}^N{\alpha }_i^t{h}_i\end{array}\right]$, 而S就是$[s_1,s_2,...,s_t]$

所以最后的attention就等于$concat(Softmax(S\cdot H^T)\cdot H^T,S)$

翻译成代码，我们只需要知道S，H，并且写一个矩阵乘法，softmax和concat即可，在pytorch中分别可以通过torch.bmm(), F.softmax(), torch.cat()实现。

遇到不同的attention score计算方式，只需要修改Softmax中的计算即可，比如我要用scaled dot product， softmax层里面就是$S\cdot H^T/\sqrt{n}$其中n是input的长度, 其他地方不需要修改。