层层剖析,让你彻底搞懂Self-Attention的机制和原理

本文内容

本文基于李宏毅老师对 Self-Attention 的讲解,进行理解和补充,并结合Pytorch代码,最终目的是使得自己和各位读者更好的理解Self-Attention

李宏毅Self-Attention链接: https://www.youtube.com/watch?v=hYdO9CscNes
PPT链接见视频下方

通过本文的阅读,你可以获得以下知识:

  1. 什么是Self-Attention,为什么要用Self-Attention
  2. Self-Attention是如何做的
  3. Self-Attention这样如何设计的
  4. Self-Attention公式的细节

为什么要使用Self-Attention

假设现在一有个词性标注(POS Tags)的任务,例如:输入I saw a saw(我看到了一个锯子)这句话,目标是将每个单词的词性标注出来,最终输出为N, V, DET, N
在这里插入图片描述

这句话中,第一个saw为动词,第二个saw(锯子)为名词。如果想做到这一点,就需要保证机器在看到一个向量时,要同时考虑其上下文,并且,要能判断出上下文中每一个元素应该考虑多少。例如,对于第一个saw,要更多的关注I,而第二个saw,就应该多关注a

这个时候,就要Attention机制来提取这种关系:如果一个任务的输入是一个Sequence(一排向量),而且各向量之间有一定关系,那么就要利用Attention机制来提取这种关系

直观的感受下Self-Attention

在这里插入图片描述
该图描述了Self-Attention的使用。Self-Attention接受一个Sequence(一排向量,可以是输入,也可以是前面隐层的输出),然后Self-Attention输出一个长度相同的Sequence,该Sequence的每个向量都充分考虑了上下文

Self-Attenion是如何考虑上下文的

在这里插入图片描述
如图所示,每个输入都会和其他输入计算一个相关性分数,然后基于该分数,输出包含上下文信息的新向量

对于上图, a 1 a^1 a1需要与 a 1 , a 2 , a 3 , a 4 a^1,a^2,a^3,a^4 a1,a2,a3,a4 分别计算相关性分数 α 1 , 1 , α 1 , 2 , α 1 , 3 , α 1 , 4 \alpha_{1,1}, \alpha_{1,2}, \alpha_{1,3}, \alpha_{1,4} α1,1,α1,2,α1,3,α1,4需要和自己也计算一下), α \alpha α 的分数越高,表示两个向量的相关度越高

计算好 α 1 , ⋅ \alpha_{1,\cdot} α1, 后,就可以求出新的包含上下文信息的向量 b 1 b^1 b1,假设 α 1 , 1 = 5 , α 1 , 2 = 2 , α 1 , 3 = 1 , α 1 , 4 = 2 \alpha_{1,1}=5, \alpha_{1,2}=2, \alpha_{1,3}=1, \alpha_{1,4}=2 α1,1=5,α1,2=2,α1,3=1,α1,4=2,则:

b 1 = ∑ i α 1 , i ⋅ a i = 5 ⋅ a 1 + 2 ⋅ a 2 + 1 ⋅ a 3 + 2 ⋅ a 4 b_1 = \sum_{i}\alpha_{1,i} \cdot a^i = 5 \cdot a^1 + 2 \cdot a^2 + 1 \cdot a^3 + 2 \cdot a^4 b1=iα1,iai=5a1+2a2+1a3+2a4

同理,对于 b 2 b_2 b2,首先计算权重 α 2 , 1 , α 2 , 2 , α 2 , 3 , α 2 , 4 \alpha_{2,1}, \alpha_{2,2}, \alpha_{2,3}, \alpha_{2,4} α2,1,α2,2,α2,3,α2,4 , 然后进行加权求和

如果按照上面这个式子做,还有两个问题:

  1. α \alpha α 之和不为1,这样会将输入向量放大或缩小
  2. 直接用输入向量 a i a^i ai去乘的话,拟合能力不够好

对于问题1,通常的做法是将 α \alpha α 过一个Softmax(当然也可以选择其他的方式)

对于问题2,通常是将 a i a^i ai 乘个矩阵(做个线性变化),然后生成 v i v^i vi ,然后用 v i v^i vi 去乘 α \alpha α

如何计算相关性分数 α \alpha α

首先,复习下向量相乘。两个向量相乘(做内积),公式为: a ⋅ b = ∣ a ∣ ∣ b ∣ cos ⁡ θ a \cdot b = |a||b| \cos \theta ab=abcosθ , 通过公式可以很容易得出结论:

  • 两个向量夹角越小(越接近),其内积越大,相关性越高。反之,两个向量夹角越大,相关性越差,如果夹角为90°,两向量垂直,内积为0,无相关性

通过上面的结论,很容易想到,要计算 a 1 a^1 a1 a 2 a^2 a2 的相关性,直接做内积即可,即 α 1 , 2 = a 1 ⋅ a 2 \alpha_{1,2} = a_1 \cdot a_2 α1,2=a1a2 。 但如果直接这样,显然不好,例如,句子I saw a sawsawsaw相关性一定很高,这样不就错了嘛。

为了解决上面这个问题,Self-Attention又额外“训练”了两个矩阵 W q W^q Wq W k W^k Wk

  • W q W^q Wq 负责对“主角”进行线性变化,将其变换为 q q q,称为query
  • W k W^k Wk 负责对“配角”进行线性变化,将其变换为 k k k,称为key

有了 W q 和 W k W^q和W^k WqWk,我们就可以计算 a 1 a^1 a1 a 2 a^2 a2 的相关分数 α 1 , 2 \alpha_{1,2} α1,2了,即:

α 1 , 2 = q 1 ⋅ k 2 = ( W q ⋅ a 1 ) ⋅ ( W k ⋅ a 2 ) \alpha_{1,2} = q^1 \cdot k^2 = (W^q \cdot a^1 )\cdot (W^k \cdot a^2) α1,2=q1k2=(Wqa1)(Wka2)

上面这些内容可以汇总成如下图:
在这里插入图片描述
要计算 a 1 a^1 a1(主角)与 a 1 , a 2 , a 3 , a 4 a^1, a^2, a^3, a^4 a1,a2,a3,a4(配角)的相关度,需要经历如下几步:

  1. 通过 W q W^q Wq ,计算 q 1 q^1 q1
  2. 通过 W k W^k Wk,计算 k 1 , k 2 , k 3 , k 4 k^1, k^2, k^3, k^4 k1,k2,k3,k4
  3. 通过 q q q k k k , 计算 α 1 , 1 , α 1 , 2 , α 1 , 3 , α 1 , 4 \alpha_{1,1}, \alpha_{1,2}, \alpha_{1,3}, \alpha_{1,4} α1,1,α1,2,α1,3,α1,4

上图并没有把 k 1 k^1 k1 画出来,但实际计算的时候,需要计算 k 1 k1 k1,即需要计算 a 1 a^1 a1和其自身的相关分数。

α \alpha α 归一化

还记得上面提到的, α \alpha α之和不为1,所以,在上面的到了 α 1 , ⋅ \alpha_{1, \cdot} α1, 后,还需要过一下Softmax,将 α 1 , ⋅ \alpha_{1, \cdot} α1,归一化。如下图:

在这里插入图片描述

最终,会将归一化后的 α 1 , ⋅ ′ \alpha'_{1, \cdot} α1, 作为 a 1 a^1 a1 与其它向量的相关分数。 同理, a 2 , a 3 , . . . a^2, a^3, ... a2,a3,... 向量与其他向量的相关分数也这么求。

不一定非要用Softmax,你开心想用什么都行,说不定效果还不错,也不一定非要归一化。 只是通常是这么做的

整合上述内容

求出了相关分数 α ′ \alpha ' α,就可以进行加权求和计算出包含上下文信息的向量 b b b 了。还记得上面提到过,如果直接用 a a a α ′ \alpha ' α 进行加权求和,泛化性不够好,所以需要对 a a a 进行线性变换,得到向量 v v v,所以Self-Attention还需要训练一个矩阵 W v W^v Wv 用于对 a a a 进行线性变化,即:

v 1 = W v ⋅ a 1          v 2 = W v ⋅ a 2           v 3 = W v ⋅ a 3             v 4 = W v ⋅ a 4 v^1 = W^v \cdot a^1 ~~~~~~~~v^2 = W^v \cdot a^2~~~~~~~~~v^3 = W^v \cdot a^3~~~~~~~~~~~v^4 = W^v \cdot a^4 v1=Wva1        v2=Wva2         v3=Wva3           v4=Wva4

然后就可用 v v v α ′ \alpha ' α 进行加权求和,得到 b b b 了。

b 1 = ∑ i α 1 , i ′ ⋅ v i = α 1 , 1 ′ ⋅ v 1 + α 1 , 2 ′ ⋅ v 2 + α 1 , 3 ′ ⋅ v 3 + α 1 , 4 ′ ⋅ v 4 b^1 = \sum_i \alpha'_{1,i} \cdot v^i = \alpha'_{1,1} \cdot v^1 + \alpha'_{1,2} \cdot v^2 + \alpha'_{1,3} \cdot v^3 + \alpha'_{1,4} \cdot v^4 b1=iα1,ivi=α1,1v1+α1,2v2+α1,3v3+α1,4v4

将求 b 1 b^1 b1 的整个过程可以归纳为下图:

在这里插入图片描述
用更正式的话描述一下整个过程:

有一组输入序列 I = ( a 1 , a 2 , ⋯   , a n ) I = (a^1, a^2, \cdots, a^n) I=(a1,a2,,an),其中 a i a^i ai 为向量, 将序列 I I I 通过Self-Attention,可以将其转化为另外一个序列 O = ( b 1 , b 2 , ⋯   , b n ) O = (b^1, b^2, \cdots, b^n) O=(b1,b2,,bn),其中向量 b i b^i bi 是由向量 a i a^i ai 结合其上下文得出的 b i b^i bi 的求解过程如下:

  1. 求出查询向量 q i q^i qi, 公式为 q i = W q ⋅ a i q^i = W^q \cdot a^i qi=Wqai
  2. 求出 k 1 , k 2 , ⋯   , k n k^1,k^2, \cdots, k^n k1,k2,,kn,公式为 k j = W k ⋅ a j k^j = W^k \cdot a^j kj=Wkaj
  3. 求出 α i , 1 , α i , 2 , ⋯   , α i , n \alpha_{i,1}, \alpha_{i,2}, \cdots, \alpha_{i,n} αi,1,αi,2,,αi,n , 公式为 α i , j = q i ⋅ k j \alpha_{i,j}=q^i\cdot k^j αi,j=qikj
  4. α i , 1 , α i , 2 , ⋯   , α i , n \alpha_{i,1}, \alpha_{i,2}, \cdots, \alpha_{i,n} αi,1,αi,2,,αi,n 进行归一化得到 α i , 1 ′ , α i , 2 ′ , ⋯   , α i , n ′ \alpha'_{i,1}, \alpha'_{i,2}, \cdots, \alpha'_{i,n} αi,1,αi,2,,αi,n,公式为 α i , j ′ = S o f t m a x ( α i , j ; α i , ⋅ ) = exp ⁡ ( α i , j ) / ∑ t exp ⁡ ( α i , t ) \alpha'_{i,j} = Softmax(\alpha_{i,j};\alpha_{i,\cdot}) = \exp(\alpha_{i,j})/\sum_t \exp(\alpha_{i,t}) αi,j=Softmax(αi,j;αi,)=exp(αi,j)/texp(αi,t)
  5. 求出向量 v 1 , v 2 , ⋯   , v n v^1, v^2, \cdots, v^n v1,v2,,vn, 公式为: v j = W v ⋅ a j v^j=W^v \cdot a^j vj=Wvaj
  6. 求出 b i b^i bi, 公式为 b i = ∑ j α i , j ′ ⋅ v j b^i = \sum_j \alpha'_{i,j} \cdot v^j bi=jαi,jvj

其中, W q , W k , W v W^q, W^k, W^v Wq,Wk,Wv 都是训练出来的


到这里Self-Attention的面纱已经揭开,但还没有结束,因为上面的步骤如果写成代码,需要大量的for循环,显然效率太低,所以需要进行向量化,能合并成向量的合成向量,能合并成矩阵的合成矩阵

向量化

向量 a a a 的矩阵化,假设列向量 a i a^i ai 维度为 d d d,显然可以将输入转化为矩阵 I I I,公式为:

I d × n = ( a 1 , a 2 , ⋯   , a n ) I_{d\times n} = (a^1, a^2, \cdots, a^n) Id×n=(a1,a2,,an)

接下来定义 W q , W k , W v W^q,W^k, W^v WqWk,Wv 矩阵的维度为 u × d u\times d u×d,它们三个的维度必须一致, u u u 是一个需要调的参数, u u u 也是 q 、 k 、 v 、 b q、k、v、b qkvb 的维度(从后面公式可以看出),定义好 W q W^q Wq 的维度后,就可以将 q q q 矩阵化了,

向量 q q q 的矩阵化,公式为:

Q u × n = ( q 1 , q 2 , ⋯   , q n ) = W u × d q ⋅ I d × n Q_{u\times n} = (q^1, q^2, \cdots, q^n) = W^q_{u\times d} \cdot I_{d\times n} Qu×n=(q1,q2,,qn)=Wu×dqId×n

同理,向量k的矩阵化,公式为:

K u × n = ( k 1 , k 2 , ⋯   , k n ) = W k ⋅ I K_{u\times n} = (k^1, k^2, \cdots, k^n) = W^k \cdot I Ku×n=(k1,k2,,kn)=WkI

同理,向量v的矩阵化,公式为:

V u × n = ( v 1 , v 2 , ⋯   , v n ) = W v ⋅ I V_{u\times n} = (v^1, v^2, \cdots, v^n) = W^v \cdot I Vu×n=(v1,v2,,vn)=WvI

得到了矩阵 Q Q Q K K K,那么就很容易得出相关分数 α \alpha α 的矩阵了,

相关分数 α \alpha α 的矩阵为

A n × n = [ α 1 , 1 α 2 , 1 ⋯ α n , 1 α 1 , 2 α 2 , 2 ⋯ α n , 2 ⋮ ⋮ ⋮ α 1 , n α 2 , n ⋯ α n , n ] = K T ⋅ Q = [ k 1 T k 2 T ⋮ k n T ] ⋅ ( q 1 , q 2 , ⋯   , q n ) A_{n\times n} = \begin{bmatrix} \alpha_{1,1} & \alpha_{2,1} & \cdots &\alpha_{n,1} \\ \alpha_{1,2} & \alpha_{2,2} & \cdots &\alpha_{n,2} \\ \vdots & \vdots & &\vdots \\ \alpha_{1,n} & \alpha_{2,n} & \cdots &\alpha_{n,n} \\ \end{bmatrix} = K^T \cdot Q =\begin{bmatrix} {k^1}^T \\ {k^2}^T \\ \vdots \\ {k^n}^T \end{bmatrix} \cdot (q^1, q^2, \cdots, q^n) An×n=α1,1α1,2α1,nα2,1α2,2α2,nαn,1αn,2αn,n=KTQ=k1Tk2TknT(q1,q2,,qn)

进一步, α ′ \alpha ' α 的矩阵为

A n × n ′ = softmax ( A ) = [ α 1 , 1 ′ α 2 , 1 ′ ⋯ α n , 1 ′ α 1 , 2 ′ α 2 , 2 ′ ⋯ α n , 2 ′ ⋮ ⋮ ⋮ α 1 , n ′ α 2 , n ′ ⋯ α n , n ′ ] A'_{n\times n} = \textbf{softmax}(A) = \begin{bmatrix} \alpha'_{1,1} & \alpha'_{2,1} & \cdots &\alpha'_{n,1} \\ \alpha'_{1,2} & \alpha'_{2,2} & \cdots &\alpha'_{n,2} \\ \vdots & \vdots & &\vdots \\ \alpha'_{1,n} & \alpha'_{2,n} & \cdots &\alpha'_{n,n} \\ \end{bmatrix} An×n=softmax(A)=α1,1α1,2α1,nα2,1α2,2α2,nαn,1αn,2αn,n

A ′ A' A 有了, V V V 有了,那就可以对输出向量 b b b 进行矩阵化了,

输出向量b的矩阵化,公式为:

O u × n = ( b 1 , b 2 , ⋯   , b n ) = V u × n ⋅ A n × n ′ = ( v 1 , v 2 , ⋯   , v n ) ⋅ [ α 1 , 1 ′ α 2 , 1 ′ ⋯ α n , 1 ′ α 1 , 2 ′ α 2 , 2 ′ ⋯ α n , 2 ′ ⋮ ⋮ ⋮ α 1 , n ′ α 2 , n ′ ⋯ α n , n ′ ] O_{u\times n} = (b^1, b^2, \cdots, b^n) = V_{u\times n} \cdot A'_{n\times n} = (v^1, v^2, \cdots, v^n) \cdot \begin{bmatrix} \alpha'_{1,1} & \alpha'_{2,1} & \cdots &\alpha'_{n,1} \\ \alpha'_{1,2} & \alpha'_{2,2} & \cdots &\alpha'_{n,2} \\ \vdots & \vdots & &\vdots \\ \alpha'_{1,n} & \alpha'_{2,n} & \cdots &\alpha'_{n,n} \\ \end{bmatrix} Ou×n=(b1,b2,,bn)=Vu×nAn×n=(v1,v2,,vn)α1,1α1,2α1,nα2,1α2,2α2,nαn,1αn,2αn,n

将上面全部整合起来,就可以的到,整合后的公式

O = Attention ( Q , K , V ) = V ⋅ softmax ( K T Q ) O = \textbf{Attention}(Q, K, V) = V\cdot \textbf{softmax}(K^T Q) O=Attention(Q,K,V)=Vsoftmax(KTQ)

如果你看过其他文章,你应该会看到真正的最终公式如下:

 Attention  ( Q , K , V ) = softmax ⁡ ( Q K T d k ) V \text { Attention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V  Attention (Q,K,V)=softmax(dk QKT)V

其实我们的公式和这个公式只差了一个转置和 d k \sqrt{d_k} dk 。转置不比多说,就是表示方式不同。

原公式的 Q , K , V Q,K,V Q,K,V以及输出 O O O,对应我们公式的 Q T , K T , V T Q^T,K^T,V^T QT,KT,VT O T O^T OT

d k d_k dk是什么,为什么要除以 d k \sqrt{d_k} dk

首先, d k d_k dk是输入向量的数量,也就是上面的 n n n 。而矩阵相乘会放大原有矩阵的标准差,放大的倍数约为 d k \sqrt{d_k} dk ,为了将标准差缩放回原来的大小,所以要除以 d k \sqrt{d_k} dk

例如,假设 Q u × n Q_{u \times n} Qu×n K u × n T K^T_{u\times n} Ku×nT 的均值为0,标准差为1。则矩阵 Q K T QK^T QKT 的均值为0,标准差为 n \sqrt{n} n ,矩阵相乘使得其标准差放大了 n \sqrt{n} n 倍,这个 n n n Q u × n Q_{u\times n} Qu×n中的 n n n

矩阵的均值就是把所有的元素加起来除以元素数量,方差同理。

可以通过以下代码验证这个结论(数学不好,只能通过实验验证结论了,哭):

Q = np.random.normal(size=(123, 456)) # 生成均值为0,标准差为1的 Q和K
K = np.random.normal(size=(123, 456))
print("Q.std=%s, K.std=%s, \nQ·K^T.std=%s, Q·K^T/√d.std=%s" 
      % (Q.std(), K.std(), 
         Q.dot(K.T).std(), Q.dot(K.T).std() / np.sqrt(456)))
Q.std=0.9977961671085275, K.std=1.0000574599289282,
Q·K^T.std=21.240017020263437, Q·K^T/√d.std=0.9946549289466212

通过输出可以看到,Q和K的标准差都为1,但是两矩阵相乘后,标准差却变为了 21.24, 通过除以 d k \sqrt{d_k} dk ,标准差又重新变为了 1

再看另一个例子,该例子Q和K的标准差是随机的,更符合真实的情况:

Q = np.random.normal(loc=1.56, scale=0.36, size=(123, 456)) # 生成均值为随机,标准差为随机的 Q和K
K = np.random.normal(loc=-0.34, scale=1.2, size=(123, 456))
print("Q.std=%s, K.std=%s, \nQ·K^T.std=%s, Q·K^T/√d.std=%s" 
      % (Q.std(), K.std(), 
         Q.dot(K.T).std(), Q.dot(K.T).std() / np.sqrt(456)))
Q.std=0.357460640868945, K.std=1.204536717914841, 
Q·K^T.std=37.78368871510589, Q·K^T/√d.std=1.769383337989377

可以看到,最开始Q的标准差为 0.35 0.35 0.35, K的标准差为 1.20 1.20 1.20,结果矩阵相乘后标准差达到了 37.78 37.78 37.78, 经过缩放后,标准差又回到了 1.76 1.76 1.76

代码实战:Pytorch定义SelfAttention模型

接下来使用Pytorch来定义SelfAttention模型,这里使用原论文中的公式:

 Attention  ( Q , K , V ) = softmax ⁡ ( Q K T d k ) V \text { Attention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V  Attention (Q,K,V)=softmax(dk QKT)V

这里为了使代码定义逻辑更清晰,下面我将各个部分的维度标记出来:

O n × u =  Attention  ( Q n × u , K n × u , V n × u ) = softmax ⁡ ( Q n × u K u × n T n ) V n × u          ( 1 ) O_{n\times u} = \text { Attention }(Q_{n\times u}, K_{n\times u}, V_{n\times u})=\operatorname{softmax}\left(\frac{Q_{n\times u} K^{T}_{u\times n}}{\sqrt{n}}\right) V_{n\times u} ~~~~~~~~(1) On×u= Attention (Qn×u,Kn×u,Vn×u)=softmax(n Qn×uKu×nT)Vn×u        (1)

其中,各个变量定义为:

  • n n n:input_num,输入向量的数量,例如,你一句话包含20个单词,则该值为20
  • u u u:output_vector_dim,输出向量的维度,对应公式中的 u u u,即输出向量 b b b的维度,例如,经过Attention后的输出向量 b b b,如果你想让他的维度为 15 15 15,则该值为 15 15 15

上述公式中, Q , K , V Q,K,V Q,K,V是通过矩阵 W q , W k , W v W^q,W^k,W^v Wq,Wk,Wv和输入向量 I I I 计算出来的,而一般对于要训练的矩阵,代码中一般使用线性层来表示,详情可参考:Pytorch nn.Linear的基本用法,所以最终 Q Q Q 矩阵的计算公式为:

Q n × u = I n × d W d × u q          ( 2 ) Q_{n \times u} = I_{n\times d} W^q_{d\times u} ~~~~~~~~(2) Qn×u=In×dWd×uq        (2)

K , V K,V K,V 矩阵同理。其中

  • d d d:input_vector_dim: 输入向量的维度,例如你将单词编码为了10维的向量,则该值为10

有了公式(1)和(2),就可以定义SelfAttention模型了,代码如下:

class SelfAttention(nn.Module):
    def __init__(self, input_vector_dim: int, input_num:int, output_vector_dim=None):
        """
        初始化SelfAttention,包含如下关键参数:
        input_vector_dim: 输入向量的维度,对应上述公式中的d,例如你将单词编码为了10维的向量,则该值为10
        input_num: 输入向量的数量,对应公式中的n,例如,你一句话包含20个单词,则该值为20
        output_vector_dim: 输出向量的维度,对应公式中的u,即b的维度,例如,经过Attention后的输出向量b,如果你想让他的维度为15,则该值为15,若不填,则取input_vector_dim
        """
        super(SelfAttention,self).__init__()
        
        self.input_vector_dim = input_vector_dim
        self.input_num = input_num
        self.output_vector_dim = output_vector_dim
        if not self.output_vector_dim:
            self.output_vector_dim = input_vector_dim
        
        """
        实际写代码时,常用线性层来表示需要训练的矩阵,方便反向传播和参数更新
        W的维度为 input_vector_dim x output_vector_dim
        """
        self.W_q = nn.Linear(input_vector_dim, output_vector_dim, bias=False)
        self.W_k = nn.Linear(input_vector_dim, output_vector_dim, bias=False)
        self.W_v = nn.Linear(input_vector_dim, output_vector_dim, bias=False)
        
        # 这个是根号下d_k
        self._norm_fact = 1 / np.sqrt(input_num)
        
    def forward(self, x):
        """
        进行前向传播:
        x: 输入向量,size为(batch_size, input_num, input_vector_dim)
        """
        
        batch_size = x.size(0)
            
        # 通过W_q, W_k, W_v矩阵计算出,Q,K,V
        # Q,K,V矩阵的size为 (batch_size, input_num, output_vector_dim)
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        
        # permute用于变换矩阵的size中对应元素的位置,
        # 即,将K的size由(batch_size, input_num, output_vector_dim),变为(batch_size, output_vector_dim,input_num)
        # 0,1,2 代表各个元素的下标,即变换前,batch_size所在的位置是0,input_num所在的位置是1
        K_T = K.permute(0,2,1)
        
        # bmm是batch matrix-matrix product,即对一批矩阵进行矩阵相乘
        # bmm详情参见:https://pytorch.org/docs/stable/generated/torch.bmm.html
        atten = nn.Softmax(dim=-1)(torch.bmm(Q,K_T)) * self._norm_fact
        
        # 最后再乘以 V
        output = torch.bmm(atten, V)        
        
        return output

接下来使用一下,定义50个一批,输入向量维度为3, 一次输入5个向量,欲经过Attention层后,编码成5个4维的向量

model = SelfAttention(3, 5, 4)
model(torch.Tensor(50,5,3)).size()
torch.Size([50, 5, 4])

Attention模型一般作为整体模型的一部分,是套在其他模型中使用的,最经典的莫过于Transformer





参考资料

李宏毅Self-Attention: https://www.youtube.com/watch?v=hYdO9CscNes

超详细图解Self-Attention: https://zhuanlan.zhihu.com/p/410776234

Pytorch nn.Linear的基本用法:https://blog.csdn.net/zhaohongfei_358/article/details/122797190

猜你喜欢

转载自blog.csdn.net/zhaohongfei_358/article/details/122861751