<jupyter_start><jupyter_text>Content1. [load package](1-load-package)2. [首先理解：$Softmax(XX^T)X$](2--softmaxxxtx)3. [$QKV$矩阵](3-qkv矩阵)4. [集成self-attention类](4-集成self-attention类)5. [多头数量的确定](5-headnums的确定)6. [17篇注意力机制PyTorch实现，包含MLP、Re-Parameter系列热门论文](https://www.jiqizhixin.com/articles/2021-08-03-11)    Reference: * [超详细图解Self-Attention的那些事儿](https://mp.weixin.qq.com/s?__biz=MzI5MDUyMDIxNA==&mid=2247579430&idx=1&sn=3c63a42410e2107f1f91f861dc25c3cb&chksm=ec1d5adfdb6ad3c953722280aae160d9bf8a73d4cb4374de196d2a2b5bb2b53ac6b7eecbc568&mpshare=1&scene=1&srcid=05019UXsKNB17Qv7uzylnUDI&sharer_sharetime=1651341683033&sharer_shareid=d50902e5393c9beaed62a664f9b58b17rd)* [Attention Is All You Need(注意力模型）](https://zhuanlan.zhihu.com/p/44731789)* [Attention 和self-attention](https://zhuanlan.zhihu.com/p/109496099)* [超细节的BERT/Transformer知识点](https://zhuanlan.zhihu.com/p/132554155)![](https://img2020.cnblogs.com/blog/1470684/202103/1470684-20210309221818623-1103910065.png)  1. load package<jupyter_code>import torch
from torch import nn
from math import sqrt
import numpy as np
import pandas as pd
import seaborn as sns
sns.set(style='darkgrid', font='SimHei')
if torch.cuda.is_available():
    print('Cuda is available')
    cuda0 = torch.device('cuda:0')
    X = torch.tensor([
        [1,2,1,2,1],
        [1,1,3,2,1],
        [3,1,2,1,1]
    ], dtype=torch.float64, device=cuda0)# 字向量定义
else:
    X = torch.tensor([
        [1,2,1,2,1],
        [1,1,3,2,1],
        [3,1,2,1,1]
    ], dtype=torch.float64)<jupyter_output>Cuda is available<jupyter_text>现在有组字向量* 早[1,2,1,2,1]* 上[1,1,3,2,1]* 好[3,1,2,1,1]* 它们构成了一个句子[早上好]* 那么这个句子的向量则为一个矩阵，命名为$X$:  ```[        [1,2,1,2,1],        [1,1,3,2,1],        [3,1,2,1,1]    ]  ```<jupyter_code>X<jupyter_output><empty_output><jupyter_text>2. $ Softmax(XX^T)X $  2.1 矩阵内积$ (XX^T) $的意义* 求矩阵内每个字向量之间的相互投影，即相关程度<jupyter_code>X_dot = X @ X.T
X_dot
df_corr = pd.DataFrame(X_dot.cpu().T, columns=['早', '上', '好'], index=['早', '上', '好'])
ax = sns.heatmap(df_corr, annot=True)<jupyter_output><empty_output><jupyter_text>2.2 $ Softmax() $的意义![](https://mmbiz.qpic.cn/sz_mmbiz_jpg/gYUsOT36vfq8debSx7RJV8NahibXzTMAYfoTEVeRf6or6JibQ9KuG1kH4DblLQlxrsgKecEoxpkSUgj1d6kjpDBQ/640?wx_fmt=jpeg&wxfrom=5&wx_lazy=1&wx_co=1)* 将进行归一化处理，转换成概率形式，使每一个字向量在其他字向量的投影的和为1<jupyter_code># torch.matmul(X, X.T)
# torch.mm(X, X.T)
softmax = nn.Softmax(dim=-1)
sft_X_dot = softmax(X @ X.T)
sft_X_dot
df_corr_weight = pd.DataFrame(sft_X_dot.cpu(), columns=['早', '上', '好'], index=['早', '上', '好'])
ax = sns.heatmap(df_corr_weight, annot=True)<jupyter_output><empty_output><jupyter_text>2.3 最后一步$ Softmax(XX^T)X $的意义![](https://mmbiz.qpic.cn/sz_mmbiz_jpg/gYUsOT36vfq8debSx7RJV8NahibXzTMAYNhw7lCUTCEpoUKNH2CEsGGqCPgNBQ1rkeYVI6cwzF4DYeX9ULXQMXA/640?wx_fmt=jpeg&wxfrom=5&wx_lazy=1&wx_co=1)* 将“早”在每个字上的加权相关（即注意力）与原矩阵$X$的一个列向量相乘，求出一个新的行向量* 在新的行向量中，每一个维度的数值都是由三个词向量在这一维度的数值加权求和得来的，这个新的行向量就是“早”字词向量经过注意力机制加权求和之后的表示* $Softmax(XX^T)$为attetnion score，但是已经很难表示原本的句子，因此还需要再重新乘于$X$，得到加权后的结果。* 单独一个$X$矩阵只能包含了word embedding内每个字的独立信息，无法包括字与字之间的信息，相互之间没什么关系。但是经过与attention score相乘后，$V$中每个token的向量（即一个单词的word embedding向量），在$N$(例如$N=5$在本例子)维的每个维度上（每一列）上，都会对其他token做出调整（关注度不同）。与$V$相乘这一步，相当于提纯，让每个单词关注该关注的部分。<jupyter_code>weight_sft_X_dot = sft_X_dot@X
weight_sft_X_dot<jupyter_output><empty_output><jupyter_text>图中右半部分的颜色深浅，其实就是我们上图中黄色向量中数值的大小，意义就是单词之间的相关度（回想之前的内容，相关度其本质是由向量的内积度量的）！ ![](https://mmbiz.qpic.cn/sz_mmbiz_jpg/gYUsOT36vfq8debSx7RJV8NahibXzTMAYawvlPOUnPibj0icbBKccm4A6iaC9QiaAbWtuXGhk3XVAz4AHttjhzUq4FA/640?wx_fmt=jpeg&wxfrom=5&wx_lazy=1&wx_co=1)  3. $QKV$矩阵* [Transformer中K 、Q、V的设置以及为什么不能使用同一个值](https://www.cnblogs.com/jins-note/p/14508523.html) ![](https://mmbiz.qpic.cn/sz_mmbiz_jpg/gYUsOT36vfq8debSx7RJV8NahibXzTMAYZrGVNq0AaDiamGVwhTsENJgxJCzgpAkgQhnQibDvFOHWVnFrAUNgqiccg/640?wx_fmt=jpeg&wxfrom=5&wx_lazy=1&wx_co=1) 其实，许多文章中所谓的$QKV$矩阵、查询向量之类的字眼，根本其来源都是是$X$与矩阵的乘积，**本质上都是$X$的线性变换。*** 为什么不直接使用$X$而要对其进行线性变换？* 当然是为了提升模型的拟合能力，矩阵$W$都是可以训练的，起到一个缓冲的效果。* 如果你真正读懂了前文的内容，读懂了$Softmax(XX^T)$这个矩阵的意义，相信你也理解了所谓查询向量一类字眼的含义。  3.1 $\sqrt d_k$的意义* $d_k$为$K$的dimension* 假设$Softmax(QK^T)$的均值为0，方差为1，那么$d$为$Softmax(QK^T)$的方差，  * $d$的值越大会导致$Softmax(QK^T)$的分布会趋于陡峭(分布的方差大，分布集中在绝对值大的区域)  * 控制$Softmax(QK^T)$的分布：    * $Softmax(QK^T)$的分布会和$d$有关。$Softmax(QK^T)$中每一个元素除以$d$后，方差又变为1。    * 利用$\sqrt d_k$对$Softmax(QK^T)$的分布“陡峭”程度与$d$解耦，从而使得训练过程中梯度值保持稳定* 利用$Q$输入与key-value memories做检索，找到和问题相似的memory的key计算相关性分数，然后对value embedding进行加权求和，得到一个输出向量。衍生出了self-attention里的$Q$，$K$，$V$表示，在self-attention里的把X映射到$QKV$  3.1 $QKV$权重矩阵$W_Q$、$W_K$和$W_V$的维度探索![](https://img2020.cnblogs.com/blog/1470684/202103/1470684-20210309221742405-1927745420.png)* 现有一$3*4$维矩阵$X$,即矩阵包含3个字，其word embedding的大小为5，想要重新编码为字向量长度为10的矩阵* 即$Softmax(QK^T/\sqrt d_k)V$的输出维度是$3*10$  * 思路：首先$Softmax(QK^T)$肯定是一个$3*3$的矩阵，因为它表示了每个字之间相互的注意力（相关程度）  * 若要求得输出结果为$3*10$的矩阵，那么$V$的维度就一定是$3*10$  * 而$V=XW_V$，已知$X$的维度是$3*5$，那么$W_V$的维度一定是$5*10$* 而对于$W_Q$和$W_K$的维度，因为要保证输出的$QK^T$的维度一定是$3*3$  * 则$Q$和$K$的维度保持一致并且一定是$3*N$  * 以$Q$为例:    * $Q=XW_Q$，已知$X$的维度是$3*5$，则$W_Q$的维度一定是$5*N$<jupyter_code>output_size = 10
embedding_size = 5
dim_q = dim_k = 8
X<jupyter_output><empty_output><jupyter_text>3.1.1 初始化各个权重矩阵<jupyter_code># 随机初始化WQ
wq = torch.randn((embedding_size, dim_q), device=cuda0, dtype=torch.float64)
# 随机初始化WK
wk = torch.randn((embedding_size, dim_k), device=cuda0, dtype=torch.float64)
# 随机初始化WV
wv = torch.randn((embedding_size, output_size), device=cuda0, dtype=torch.float64)<jupyter_output><empty_output><jupyter_text>3.1.2 求出$Q$矩阵<jupyter_code>Q = X @ wq
Q<jupyter_output><empty_output><jupyter_text>3.1.3 求出$K$矩阵<jupyter_code>K = X @ wk
K<jupyter_output><empty_output><jupyter_text>3.1.4 求出$V$矩阵<jupyter_code>V = X @ wv
V<jupyter_output><empty_output><jupyter_text>3.1.5 求出$QK^T$矩阵<jupyter_code>attention_score = Q @ K.T
attention_score<jupyter_output><empty_output><jupyter_text>3.1.6 $QK^T/\sqrt{d_k}$：利用$d_k$对$QK^T$进行解耦，防止$QK^T$的变化过大<jupyter_code>scaled_attention_score = attention_score/sqrt(dim_k)
scaled_attention_score<jupyter_output><empty_output><jupyter_text>3.1.7 利用Softmax()对$QK^T/\sqrt{d_k}$进行归一化处理<jupyter_code>soft_scaled_att_score = softmax(scaled_attention_score)
soft_scaled_att_score<jupyter_output><empty_output><jupyter_text>3.1.8 将注意力矩阵$Softmax(QK^T/\sqrt{d_k})$乘上$V$，对原句进行表达<jupyter_code>attention_output = soft_scaled_att_score@V
attention_output<jupyter_output><empty_output><jupyter_text>3.2 当batch_size>1时  3.2.1 定义一个例子<jupyter_code>X2 = torch.tensor([[[1,2,1,2,1], # 早
                    [1,1,3,2,1], # 上
                    [3,1,2,1,1]], # 好
                    
                    [[3,1,2,1,1], # 好
                    [1,1,3,2,1], # 早
                    [1,2,1,2,1]] # 上
                ], dtype=torch.float64, device=cuda0)
X2, X2.size()
batch_size = X2.size()[0]
batch_size<jupyter_output><empty_output><jupyter_text>3.2.2 先实现$Softmax(XX^T)X$<jupyter_code>X2_dot = torch.bmm(X2, X2.permute(0, 2, 1))
X2_dot
soft_X2_dot = softmax(X2_dot)
soft_X2_dot
torch.bmm(soft_X2_dot, X2)<jupyter_output><empty_output><jupyter_text>3.2.3 实现$Softmax(QK^T/\sqrt{d_k})V$<jupyter_code>embedding_size = 5
dim_q = dim_k = 8
output_size = 10<jupyter_output><empty_output><jupyter_text>3.2.3.1 定义qkv权重<jupyter_code># q = nn.Linear(embedding_size, dim_q)
# k = nn.Linear(embedding_size, dim_k)
# v = nn.Linear(embedding_size, output_size)
wq = torch.randn(batch_size, embedding_size, dim_q, device=cuda0, dtype=torch.float64)
wk = torch.randn(batch_size, embedding_size, dim_k, device=cuda0, dtype=torch.float64)
wv = torch.randn(batch_size, embedding_size, output_size, device=cuda0, dtype=torch.float64)

Q = torch.bmm(X2, wq)
Q
K = torch.bmm(X2, wk)
K
V = torch.bmm(X2, wv)
V
QK_T = torch.bmm(Q, K.permute(0, 2, 1))
QK_T
scaled_QK_T = QK_T/sqrt(dim_k)
scaled_QK_T
soft_scaled_QK_T = softmax(scaled_QK_T)
soft_scaled_QK_T
batch_attention_output = torch.bmm(soft_scaled_QK_T, V)
batch_attention_output<jupyter_output><empty_output><jupyter_text>4. 集成self-attention类  4.1 单头self-attention<jupyter_code>class Self_Attention(nn.Module):
    def __init__(self, embedding_size, dim_k, output_size):
        super(Self_Attention, self).__init__()
        # input (batch_size, seq_len, embedding_size)
        # q (batch_size, embedding_size, dim_k)
        # k (batch_size, embedding_size, dim_k)
        # v (batch_size, embedding_size, output_size)
        # self.batch_size = batch_size
        self.embedding_size = embedding_size
        self.dim_q = self.dim_k = dim_k
        self.dim_v = output_size
        self.output_size = output_size
        self.norm_factor = 1/sqrt(self.dim_k)

        self.q = nn.Linear(self.embedding_size, self.dim_q)
        self.k = nn.Linear(self.embedding_size, self.dim_k)
        self.v = nn.Linear(self.embedding_size, self.dim_v)

    def forward(self, x):
        Q = self.q(x) # dim: batch_size * seq_len * dim_q
        K = self.k(x) # dim: batch_size * seq_len * dim_k
        V = self.v(x) # dim: batch_size * seq_len * output_size(dim_v)

        attention_score = torch.bmm(Q, K.permute(0, 2, 1)) * self.norm_factor # dim: batch_size * seq_len * seq_len
        soft_attention_score = nn.Softmax(dim=-1)(attention_score) # dim: batch_size * seq_len * seq_len
        output = torch.bmm(soft_attention_score, V) # dim: batch_size * seq_len * output_size(dim_v)

        return output
batch_size = 2
seq_len = 3
embedding_size = 4
x = torch.randn(batch_size, seq_len, embedding_size) # dim: batch_size=2, seq_len=3, embedding_size=4
x, x.size()
self_attention = Self_Attention(embedding_size=embedding_size, dim_k=10, output_size=5)
res = self_attention(x)
res, res.size()<jupyter_output><empty_output><jupyter_text>4.2 多头self-attention![](https://pic2.zhimg.com/80/v2-ed923dba9f3c3f32decd2b8a6283da41_720w.jpg)<jupyter_code>head_nums = 2
X2, X2.shape
Q_mh = torch.bmm(X2, wq).reshape(-1, # batch_size
                                 X2.shape[0], # seq_len
                                #  X2.shape[1], # embedding_size
                                 head_nums, # head number
                                 dim_q//head_nums # segmented by head_nums
                                )
Q_mh, Q_mh.size()
# K_mh = torch.bmm(X2, wk).reshape(-1, X2.shape[0], X2.shape[1], dim_k//head_nums)
K_mh = torch.bmm(X2, wk).reshape(-1, X2.shape[0], head_nums, dim_k//head_nums)
K_mh, K_mh.size()
# V_mh = torch.bmm(X2, wv).reshape(-1, X2.shape[0], X2.shape[1], output_size//head_nums)
V_mh = torch.bmm(X2, wv).reshape(-1, X2.shape[0], head_nums, output_size//head_nums)
V_mh, V_mh.size()
# Q_mh · K_mh.T
atten_score_mh = nn.Softmax(dim=-1)(Q_mh @ K_mh.permute(0, 1, 3, 2)/sqrt(dim_k))
atten_score_mh, atten_score_mh.size()
atten_score_mhV_mh = (atten_score_mh @ V_mh)
atten_score_mhV_mh, atten_score_mhV_mh.size()
# concat multi-head
res_mh = atten_score_mhV_mh.reshape(X2.shape[0], X2.shape[1], -1)
res_mh, res_mh.size()
class Self_Attention_multi_head(nn.Module):
    # input: batch_size, seq_len, embedding_size
    # q: batch_size, embedding_size, dim_q
    # k: batch_size, embedding_size, dim_k
    # v: batch_size, embedding_size, dim_v
    def __init__(self, embedding_size, dim_k, output_size, head_nums):
        super(Self_Attention_multi_head, self).__init__()
        assert dim_k % head_nums == 0
        assert output_size % head_nums == 0
        self.embedding_size = embedding_size
        self.dim_q = self.dim_k = dim_k
        self.dim_v = output_size
        self.q = nn.Linear(self.embedding_size, self.dim_q)
        self.k = nn.Linear(self.embedding_size, self.dim_k)
        self.v = nn.Linear(self.embedding_size, self.dim_v)
        self.norm_factor = 1/sqrt(dim_k)
        self.head_nums = head_nums

    def forward(self, x):
        Q = self.q(x).reshape(-1, x.shape[0], self.head_nums, self.dim_q//self.head_nums) 
        # Q dim: batch_size, multi-segment(head_nums), seq_len, embedding_size/head_nums
        K = self.k(x).reshape(-1, x.shape[0], self.head_nums, self.dim_k//self.head_nums)
        # K dim: batch_size, multi-segment(head_nums), seq_len, embedding_size/head_nums
        V = self.v(x).reshape(-1, x.shape[0], self.head_nums, self.dim_v//self.head_nums)
        # V dim: batch_size, multi-segment(head_nums), seq_len, outputs_size/head_nums

        QK = Q @ K.permute(0, 1, 3, 2)
        # QK dim: batch_size, multi-segment(head_nums), seq_len, seq_len
        scaled_QK = QK * self.norm_factor
        atten_mh = nn.Softmax(dim=-1)(scaled_QK)
        # atten_mh dim: batch_size, multi-segment(head_nums), seq_len, seq_len
        atten_mhV = atten_mh @ V
        # atten_mhV dim: batch_size, multi-segment(head_nums), seq_len, output_size/head_nums

        output = atten_mhV.reshape(x.shape[0], x.shape[1], -1)
        # concat multi-segment
        # output dim: batch_size, seq_len, output_size
        return output
batch_size = 2
seq_len = 3
embedding_size = 4
x = torch.randn(batch_size, seq_len, embedding_size) # dim: batch_size=2, seq_len=3, embedding_size=4
x, x.size()
dim_k = 6
output_size = 10
head_nums = 2
self_attention_mh = Self_Attention_multi_head(embedding_size=embedding_size, 
                                              dim_k=dim_k, 
                                              output_size=output_size,
                                              head_nums=head_nums)
res = self_attention_mh(x)
res, res.size()<jupyter_output><empty_output><jupyter_text>5. head_nums的确定 * Reference: [Transformers Explained Visually (Part 3): Multi-head Attention, deep dive](https://towardsdatascience.com/transformers-explained-visually-part-3-multi-head-attention-deep-dive-1c1ff1024853): This means that separate sections of the Embedding can learn different aspects of the meanings of each word, as it relates to other words in the sequence. This allows the Transformer to capture richer interpretations of the sequence.![](https://miro.medium.com/max/1400/1*75EUBJLaqAMcDjgwWVh-_A.png)<jupyter_code>batch_size = 2
seq_len = 3
embedding_size = 100
x2 = torch.randn(batch_size, seq_len, embedding_size)
head_nums = 4
dim_k = 64
output_size = 100
self_attention_mh = Self_Attention_multi_head(embedding_size=embedding_size,
                                              dim_k=dim_k,
                                              output_size=output_size,
                                              head_nums=head_nums)
res = self_attention_mh(x2)
res.size()
X = torch.tensor([[
                      [1, 2, 3, 4, 5, 6],
                      [5, 6, 7, 8, 9, 10],
                      [9, 10, 11, 12, 12, 14],
                  ], 
                  [
                      [1, 2, 3, 4, 5, 6],
                      [1, 2, 3, 4, 5, 6],
                      [1, 2, 3, 4, 5, 6],
                  ]], dtype=torch.float64, device=cuda0)
seq_len = 3
embed_size = 6
batch_size = 2
head_num = 2
head_X = X.reshape(-1, X.shape[0], head_num, embed_size//head_num)
head_X
head_XXTX = (head_X @ head_X.permute(0, 1, 3, 2)) @ head_X
head_XXTX
concat_X = head_XXTX.reshape(batch_size, seq_len, embed_size)
concat_X<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14462175335514454, 'fraction_numerical': 0.05141216532015758, 'mean_word_length': 4.391648668106551, 'pattern_counts': {'":': 0, '<': 84, '<?xml version=': 0, '>': 85, 'https://': 15, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '21', 'dir': 'jupyter-structured-clean-dedup', 'id': '198882', 'license': 'MIT', 'n_tokens_mistral': 7688, 'n_tokens_neox': 7501, 'n_words': 814, 'path': 'Transformer/self_attention.ipynb', 'repo_name': 'AlanConstantine/GradientDescent', 'size': '14977'}