结合代码理解各种注意力机制（一）：自注意力机制

transformer中最重要的就是注意力机制，从经典论文Attention is all you need出发，到后来的各种注意力机制的改进。本系列将手撕各种注意力机制，包括但不限于：

• self-attention(SA) 自注意力机制
• multi-head attention(MHA) 多头注意力机制
• multi-query attention(MQA) 分组注意力机制

在此系列的第一篇中，我们聚焦于经典的自注意力机制，mask的注意力机制，以及多头注意力机制的改进。

前置词典和词向量的构建

1. 词典dictionary构建

假设我们有一个句子：

The quick brown fox jumps over the lazy dog

然后对句子进行分词，构建词典。

text = "The quick brown fox jumps over the lazy dog"

dict = {token: i for i, token in enumerate(sorted(set(text.split())))}

dict

打印结果如下：

{‘The’: 0,
‘brown’: 1,
‘dog’: 2,
‘fox’: 3,
‘jumps’: 4,
‘lazy’: 5,
‘over’: 6,
‘quick’: 7,
‘the’: 8}

将原来的英文句子转换成对应词典中的索引

text_id = torch.tensor([dict[token] for token in text.split()])

得到的text_id结果如下：
tensor([0, 7, 1, 3, 4, 6, 8, 5, 2])

2. 词嵌入构建

利用torch.nn.Embedding构建词嵌入，在这里设置的词向量维度为5，方便演示。而实际中，大模型的词向量往往很高，比如在Gpt系列中词向量维度就是12800。

# 设置随机种子
torch.manual_seed(42)

# 构建词嵌入层
len_dict = len(dict) # 词典长度
dim_embedding = 5 # 词嵌入维度，可以自定义

embed = torch.nn.Embedding(len_dict, dim_embedding)
embedding_sentence = embed(text_id)

embedding_sentence

得到结果:

tensor([[ 1.9269,  1.4873,  0.9007, -2.1055,  0.6784],
        [ 0.5258, -0.4880, -0.4345, -1.3864, -1.2862],
        [-1.2345, -0.0431, -1.6047, -0.7521,  1.6487],
        [ 0.7624,  1.6423, -0.1596, -0.4974,  0.4396],
        [-0.7581,  1.0783,  0.8008,  1.6806,  1.2791],
        [-1.0892, -0.3553, -0.9138, -0.6581,  0.0780],
        [-1.4032,  0.0360, -0.0635,  0.6756, -0.0978],
        [ 1.2964,  0.6105,  1.3347, -0.2316,  0.6872],
        [-0.3925, -1.4036, -0.7279, -0.5594, -0.7688]],
       grad_fn=<EmbeddingBackward0>)

绘制了一张表格来表示各个词对应的词向量：

自注意力机制实现

流程说明

我们需要将每个词的embedding投影到三个空间中，分别表示query, key, value。

那么投影后的query, key, value向量维度dq, dk, dv是多少呢？需要注意的是，query和key的维度需要相同，而dv可以和dq, dk不同

注意力机制计算的公式如下：

可以根据这个计算流程图来计算：

代码实现

代码是经过结构化的，封装了一个self-attention的类，代码如下：

import torch
import torch.nn as nn
class SelfAttention(nn.Module):
    def __init__(self, d_emb, d_k, d_q, d_v):
        torch.manual_seed(0)
        super().__init__()
        self.d_k = d_k # 方便后面除以维度进行缩放
        self.get_query = nn.Linear(d_emb, d_q)
        self.get_key = nn.Linear(d_emb, d_k)
        self.get_value = nn.Linear(d_emb, d_v)

    def forward(self, embed):
        query = self.get_query(embed)
        key = self.get_key(embed)
        value = self.get_value(embed)

        attn_score = query @ key.T
        attn_score = torch.softmax(attn_score / self.d_k ** 0.5, dim = -1) # dim = -1表示最后一个维度，表示按列进行softmax
        attn_score = attn_score @ value
        
        return attn_score

使用这个模块的步骤如下：

d_emb = 5
d_q= 10
d_k = 10
d_v = 16

sa = SelfAttention(d_emb = d_emb, d_q= d_q, d_k=d_k, d_v=d_v)

attn_score = sa(embedding_sentence)
print(attn_score)
print(attn_score.shape)

得到的结果如下，具体的数值就略过，我们查看形状即可：

torch.Size([9, 16])

额外添加：因果自注意力机制

在transformer的decoder中，当前token只能关注到之前已经生成的token，而不能看到未来的token，因此，在计算注意力分数时，需要mask掉未来的信息。即注意力分数矩阵长这样：

在代码中实现这个机制也很简单，需要用到2个函数：

1. torch.tril()函数，这个函数可以生成一个下三角矩阵，其中对角线以上的元素为1，对角线以下的元素为0。参数diagonal表示刚刚好位于对角线上的元素的值
2. torch.masked_fill()函数，这个函数可以对矩阵中的元素进行填充，参数mask.bool()表示将mask矩阵中的元素为True的元素填充为-inf，False的元素保持不变。

实现代码如下：

mask = torch.triu(torch.ones(attn_score.shape[0],attn_score.shape[0]),diagonal=1)
attn_score = attn_score.masked_fill(mask.bool(),float('-inf'))

mask矩阵为：

tensor([[0., 1., 1., 1., 1., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1., 1., 1., 1., 1.],
        [0., 0., 0., 1., 1., 1., 1., 1., 1.],
        [0., 0., 0., 0., 1., 1., 1., 1., 1.],
        [0., 0., 0., 0., 0., 1., 1., 1., 1.],
        [0., 0., 0., 0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.]])

attn_score为：

tensor([[-0.0394,    -inf,    -inf,    -inf,    -inf,    -inf,    -inf,    -inf,
            -inf],
        [-0.9510, -1.6893,    -inf,    -inf,    -inf,    -inf,    -inf,    -inf,
            -inf],
        [-2.9732, -3.2457, -0.2289,    -inf,    -inf,    -inf,    -inf,    -inf,
            -inf],
        [-0.5146, -1.5604,  0.4291,  0.2427,    -inf,    -inf,    -inf,    -inf,
            -inf],
        [ 1.9082, -0.0317,  3.1926,  0.7505,  2.3168,    -inf,    -inf,    -inf,
            -inf],
        [-1.7217, -2.1771, -0.6873, -0.7923,  1.2813, -1.1493,    -inf,    -inf,
            -inf],
        [ 0.2187, -0.9201,  0.7226,  0.0409,  1.5223,  0.2587,  0.7383,    -inf,
            -inf],
        [ 1.9516,  0.2456,  2.3482,  1.5573,  1.8757,  1.1425,  0.8713,  1.6937,
            -inf],
        [-0.8420, -1.8965, -0.7995, -0.2597,  1.3088, -1.2735, -0.3359,  0.2703,
         -1.4732]], grad_fn=<MaskedFillBackward0>)

完整的带mask的注意力机制代码如下：

import torch
import torch.nn as nn
class MaskSelfAttention(nn.Module):
    def __init__(self, d_emb, d_k, d_q, d_v):
        torch.manual_seed(0)
        super().__init__()
        self.d_k = d_k # 方便后面除以维度进行缩放
        self.get_query = nn.Linear(d_emb, d_q)
        self.get_key = nn.Linear(d_emb, d_k)
        self.get_value = nn.Linear(d_emb, d_v)

    def forward(self, embed):
        query = self.get_query(embed)
        key = self.get_key(embed)
        value = self.get_value(embed)

        attn_score = query @ key.T
        mask = torch.triu(torch.ones(attn_score.shape[0],attn_score.shape[0]),diagonal=1)
        attn_score = attn_score.masked_fill(mask.bool(),float('-inf'))
        attn_score = torch.softmax(attn_score / self.d_k ** 0.5, dim = -1) # 最后一个维度，表示按列进行softmax
        attn_score = attn_score @ value
        
        return attn_score

备注

但实际上，n = 9, d_emb = 5, embedding的形状为[9, 5]

而attn_score的形状为[9, 16]。
我们期望的是经过注意力机制之后，能生成一个和embedding形状相同的向量，称为△E,这样可以和原来的embedding相加，得到新的embedding： E = E + ΔE。 这也是add & normalize 中的add部分。

但目前attn_score的形状和embedding的形状不一致，怎么办呀？所以需要进行调整。
后面在实现的时候，通常还会有一个down_project层，维度为[d_k, d_emb]，将attn_score的形状调整到和embedding的形状一致。

理想状况下的self-attention的维度变化如下：

但实际上因为这个[d_emb, d_emb]的矩阵太大了，所以会拆成两个矩阵，一个up_project层,一个down_project层，中间的小维度就是d_v。而通常把第一个up_project层称为w_value。

d_v可以和dq, dk不同，因为它在作用时，是和n*n的矩阵进行相乘，和q,k的维度无关。

所以

参考文档
https://mp.weixin.qq.com/s/5TPYtEElfiSH8cHdu4uN7A
3b1b对self-attention的理解