文章目录
本文内容
本文基于李宏毅老师对 Self-Attention 的讲解,进行理解和补充,并结合Pytorch代码,最终目的是使得自己和各位读者更好的理解Self-Attention
李宏毅Self-Attention链接: https://www.youtube.com/watch?v=hYdO9CscNes
PPT链接见视频下方
通过本文的阅读,你可以获得以下知识:
- 什么是Self-Attention,为什么要用Self-Attention
- Self-Attention是如何做的
- Self-Attention这样如何设计的
- 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,i⋅ai=5⋅a1+2⋅a2+1⋅a3+2⋅a4
同理,对于 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 , 然后进行加权求和
如果按照上面这个式子做,还有两个问题:
- α \alpha α 之和不为1,这样会将输入向量放大或缩小
- 直接用输入向量 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 a⋅b=∣a∣∣b∣cosθ , 通过公式可以很容易得出结论:
- 两个向量夹角越小(越接近),其内积越大,相关性越高。反之,两个向量夹角越大,相关性越差,如果夹角为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=a1⋅a2 。 但如果直接这样,显然不好,例如,句子I saw a saw
的saw
和saw
相关性一定很高,这样不就错了嘛。
为了解决上面这个问题,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 Wq和Wk,我们就可以计算 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=q1⋅k2=(Wq⋅a1)⋅(Wk⋅a2)
上面这些内容可以汇总成如下图:
要计算 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(配角)的相关度,需要经历如下几步:
- 通过 W q W^q Wq ,计算 q 1 q^1 q1
- 通过 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
- 通过 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=Wv⋅a1 v2=Wv⋅a2 v3=Wv⋅a3 v4=Wv⋅a4
然后就可用 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,i′⋅vi=α1,1′⋅v1+α1,2′⋅v2+α1,3′⋅v3+α1,4′⋅v4
将求 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 的求解过程如下:
- 求出查询向量 q i q^i qi, 公式为 q i = W q ⋅ a i q^i = W^q \cdot a^i qi=Wq⋅ai
- 求出 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=Wk⋅aj
- 求出 α 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=qi⋅kj
- 将 α 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)
- 求出向量 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=Wv⋅aj
- 求出 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,j′⋅vj
其中, 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 Wq,Wk,Wv 矩阵的维度为 u × d u\times d u×d,它们三个的维度必须一致, u u u 是一个需要调的参数, u u u 也是 q 、 k 、 v 、 b q、k、v、b q、k、v、b 的维度(从后面公式可以看出),定义好 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×dq⋅Id×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)=Wk⋅I
同理,向量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)=Wv⋅I
得到了矩阵 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⎦⎥⎥⎥⎤=KT⋅Q=⎣⎢⎢⎢⎢⎡k1Tk2T⋮knT⎦⎥⎥⎥⎥⎤⋅(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×n⋅An×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)=V⋅softmax(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(dkQKT)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(dkQKT)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(nQn×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