神经网络中的Attention机制 pytorch 代码

注意力机制总述

引言
注意力分布
加权平均
注意力机制的变体
- 硬性注意力
- 键值对注意力
- 多头注意力
代码

引言

在计算能力有限的情况下，注意力机制作为一种资源分配方案，将有限的计算资源用来处理更为重要的信息，是解决信息超载问题的主要手段。

注意力分布

为了从N个输入向量[x1,…,xN]中选择出和某个特定任务相关的信息，我们需要引入一个和任务相关的表示，成为查询向量（Query Vector),并通过一个打分函数来计算每个输入向量和查询向量之间的相关性。
给定一个和任务相关的查询向量q,我们用注意力变量z∈[1,N]来表示被选择信息的索引位置，即z=n表示选择了第n个输入向量，为了方便计算，我们采用一种“软性”的信息选择机制，首先计算在给定q和X下，选择第i个输入向量的概率 α n \alpha_n αn,
α n = p ( z = n ∣ X , q ) = s o f t m a x ( s ( x n , q ) = e x p ( s ( x n , q ) ) ∑ j = 1 N e x p ( s ( x j , q ) ) \alpha_n=p(z=n|X,q) \\ = softmax(s(x_n,q) \\=\frac{exp(s(x_n,q))}{\sum_{j=1}^Nexp(s(x_j,q))} αn=p(z=n∣X,q)=softmax(s(xn,q)=∑j=1Nexp(s(xj,q))exp(s(xn,q))
其中 α n \alpha_n αn称为注意力分布（Attention Distribution），s(x,q)为注意力打分函数，可以使用以下几种方式来计算：
加点模型 s ( x , q ) = v T t a n h ( W x + U q ) s(x,q)=v^Ttanh(Wx + Uq) s(x,q)=vTtanh(Wx+Uq)
点积模型 s ( x , q ) = x T q s(x,q)=x^Tq s(x,q)=xTq
缩放点积模型 s ( x , q ) = x T q D s(x,q)=\frac{x^Tq}{\sqrt{D}} s(x,q)=D xTq
双线性模型 s ( x , q ) = x T W q s(x,q)=x^TWq s(x,q)=xTWq
其中W，U，v为可学习的参数，D为输入向量的维度。

加权平均

注意力分布 α n \alpha_n αn可以解释为在给定任务相关地查询q时，第n个输入向量受关注的程度。我们采用一种“软性”的信息选择机制对输入信息进行汇总，即
a t t ( X , q ) = ∑ n = 1 N α n x n = E z − p ( z ∣ X , q ) [ X z ] (1) att(X,q)=\sum_{n=1}^N \alpha_nx_n \\ =E_{z-p(z|X,q)}[X_z] \tag{1} att(X,q)=n=1∑Nαnxn=Ez−p(z∣X,q)[Xz](1)
上面公式1称为软性注意力机制（Soft Attention Mechanism）。下图是软性注意力机制的示例：

注意力机制可以单独使用，但更多地用作神经网络中的一个组件。

注意力机制的变体

硬性注意力

软性注意力选择的信息是所有输入向量在注意力分布下的期望，此外，还有一种注意力是只关注某一个输入向量，叫做硬性注意力（Hard Attention）
硬性注意力有两种实现方式：
（1）一种是选取最高概率的一个输入向量，即
a t t ( X , q ) = X n ~ att(X,q)=X_{\tilde n} att(X,q)=Xn~
其中 n ~ \tilde n n~为概率最大的输入向量的下标，即 n ~ = a r g m a x α n \tilde n=argmax\alpha_n n~=argmaxαn
（2）另一种硬性注意力可以通过在注意力分布式上随机采样的方式实现。
硬性注意力的一个缺点是基于最大采样或随机采样的方式来选择信息，使得最终的损失函数与注意力分布之间的函数关系不可导，因此无法使用反向传播算法进行训练，为了使用反向传播算法，一般使用软性注意力来代替硬性注意力。

键值对注意力

更一般地，我们可以用键值对（key-value pair）格式来表示输入信息，其中“键”用来计算注意力分布 α n \alpha_n αn,"值"用来计算聚合信息。
用（K，V）=[( k 1 k_1 k1, v 1 v_1 v1),…,( k n k_n kn, v n v_n vn)]表示N组输入信息，给定任务相关的查询向量q时，注意力函数为
a t t ( ( K , V ) , q ) = ∑ n = 1 N α n v n = ∑ n = 1 N e x p s ( k n , q ) ) ∑ j e x p ( s ( k j , q ) ) v n att((K,V),q)=\sum_{n=1}^N\alpha_nv_n \\ =\sum_{n=1}^N\frac{exps(k_n,q))}{\sum_j exp(s(k_j,q))}v_n att((K,V),q)=∑n=1Nαnvn=∑n=1N∑jexp(s(kj,q))exps(kn,q))vn
其中 s ( k n , q ) s(k_n,q) s(kn,q)为打分函数
下图给出键值对注意力机制的示例，当K=V时，键值对模式就等价于普通的注意力机制。

多头注意力

多头注意力（Muti-Head Attention）是利用多个查询Q=[ q 1 q_1 q1,…, q M q_M qM],来平行地计算从输入信息中选取多组信息，每个注意力都关注输入信息地不同部分。
a t t ( ( K , V ) , Q ) = a t t ( ( K , V ) , q 1 ) ⊕ . . . ⊕ a t t ( ( K , V ) , q M ) att((K,V),Q)=att((K,V),q_1)⊕...⊕att((K,V),q_M) att((K,V),Q)=att((K,V),q1)⊕...⊕att((K,V),qM)
其中⊕表示向量拼接

代码

光看这些理论上的东西估计你和我一样也是迷迷糊糊的，好像大致知道其意思，但是不知道具体实现的时候应该注意哪些细节，我推荐看一篇pytorch的官方Tutorial它讲的是在文本翻译的时候应用的注意力机制，它用的是软性注意力机制。
还有了解transformer的应该都知道transformer应用了多头注意力机制，不过pytorch用的是封装好的transformer函数，所以无法看到它的注意力到底是如何实现的，不过我看到有一篇将transformer的文章讲多头注意力机制讲的超详细，感觉都可以按他讲的复现了。
here