多头注意力机制

目录

1、什么是多头注意力机制

2、多头注意力机制结构图

3、多头注意力机制的作用

4、代码实现多头注意力机制

---

1、什么是多头注意力机制

从多头注意力的结构图中，貌似这个所谓的多个头就是指多组线性变换，但是并不是，只使用了一组线性变换层，即三个变换张量对 Q、K、V 分别进行线性变换，这些变化不会改变原有张量的尺寸，因此每个变换矩阵都是方阵，得到输出结果后，多头的作用才开始显现，每一个头开始从词义层面分割输出的张量，也就是每一个头都先获得一组 Q、K、V进行注意力机制的计算，但是句子中的每个词的表示只获得一部分，也就是只分割了最后一维的词嵌入向量，这就是所谓的多头，将每个头获取的输入送到注意力机制中就形成了多头注意力机制。

2、多头注意力机制结构图

3、多头注意力机制的作用

这种结构的设计能让每个注意力机制去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能产生的偏差，让词义拥有来自多元的表达，实验表名可以从而提升模型效果

4、代码实现多头注意力机制

import torch
import torch.nn as nn
import numpy as np
import torch
import math
from torch.autograd import Variable
import torch
from torch.autograd import Variable
import math
import torch.nn as nn
# 构建Embedding类来实现文本嵌入层
class Embeddings(nn.Module):
    def __init__(self,vocab,d_model):
        """
        :param vocab: 词表的大小
        :param d_model: 词嵌入的维度
        """
        super(Embeddings,self).__init__()
        self.lut = nn.Embedding(vocab,d_model)
        self.d_model = d_model
    def forward(self,x):
        """
        :param x: 因为Embedding层是首层，所以代表输入给模型的文本通过词汇映射后的张量
        :return:
        """
        return self.lut(x) * math.sqrt(self.d_model)
class PositionalEncoding(nn.Module):
    def __init__(self,d_model,dropout,max_len=5000):
        """
        :param d_model: 词嵌入的维度
        :param dropout: 随机失活，置0比率
        :param max_len: 每个句子的最大长度，也就是每个句子中单词的最大个数
        """
        super(PositionalEncoding,self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len,d_model) # 初始化一个位置编码器矩阵，它是一个0矩阵，矩阵的大小是max_len * d_model
        position = torch.arange(0,max_len).unsqueeze(1) # 初始一个绝对位置矩阵 max_len * 1
        div_term = torch.exp(torch.arange(0,d_model,2)*-(math.log(1000.0)/d_model)) # 定义一个变换矩阵，跳跃式的初始化
        # 将前面定义的变换矩阵进行奇数、偶数的分别赋值
        pe[:,0::2] = torch.sin(position*div_term)
        pe[:,1::2] = torch.cos(position*div_term)
        pe = pe.unsqueeze(0)  # 将二维矩阵扩展为三维和embedding的输出（一个三维向量）相加
        self.register_buffer('pe',pe) # 把pe位置编码矩阵注册成模型的buffer，对模型是有帮助的，但是却不是模型结构中的超参数或者参数，不需要随着优化步骤进行更新的增益对象。注册之后我们就可以在模型保存后重加载时，将这个位置编码与模型参数一同加载进来
 
    def forward(self, x):
        """
        :param x: 表示文本序列的词嵌入表示
        :return: 最后使用self.dropout(x)对对象进行“丢弃”操作，并返回结果
        """
        x = x + Variable(self.pe[:, :x.size(1)],requires_grad = False)   # 不需要梯度求导，而且使用切片操作，因为我们默认的max_len为5000，但是很难一个句子有5000个词汇，所以要根据传递过来的实际单词的个数对创建的位置编码矩阵进行切片操作
        return self.dropout(x)
def subsequent_mask(size):
    """
    :param size: 生成向后遮掩的掩码张量，参数 size 是掩码张量的最后两个维度大小，它的最后两个维度形成一个方阵
    :return:
    """
    attn_shape = (1,size,size) # 定义掩码张量的形状
    subsequent_mask = np.triu(np.ones(attn_shape),k = 1).astype('uint8') # 定义一个上三角矩阵，元素为1，再使用其中的数据类型变为无符号8位整形
    return torch.from_numpy(1 - subsequent_mask) # 先将numpy 类型转化为 tensor，再做三角的翻转，将位置为 0 的地方变为 1，将位置为 1 的方变为 0

from torch import tensor, softmax
import math
def attention(query, key, value, mask = None, dropout = None):
    """
    :param query: 三个张量输入
    :param key: 三个张量输入
    :param value: 三个张量输入
    :param mask: 掩码张量
    :param dropout: 传入的 dropout 实例化对象
    :return:
    """
    d_model = query.size(-1)  # 得到词嵌入的维度，取 query 的最后一维大小
    scores = torch.matmul(query,key.transpose(-2,-1)) / math.sqrt(d_model)     # 按照注意力公式，将 query 和 key 的转置相乘，这里是将 key 的最后两个维度进行转置，再除以缩放系数，得到注意力得分张量 scores
    # query(2,8,4,64) key.transpose(-2,-1) (2,8,64,4) 进行矩阵乘法为 （2,8,4,4）
    if mask is not None:
        scores = torch.masked_fill(scores,mask == 0,-1e9)  # 使用 tensor 的 mask_fill 方法，将掩码张量和 scores 张量中每一个位置进行一一比较，如果掩码张量处为 0 ,则使用 -1e9 替换
        # scores = scores.masked_fill(mask == 0,-1e9)
    p_attn = softmax(scores, dim = -1) # 对 scores 的最后一维进行 softmax 操作，使用 F.softmax 方法，第一个参数是 softmax 对象，第二个参数是最后一个维度，得到注意力矩阵
    print('scores.shape ',scores.shape)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn,value),p_attn  # 返回注意力表示
 
import torch.nn as nn
import torch
import copy
class MultiHeadAttention(nn.Module):
    def __init__(self, head, embedding_dim , dropout=0.1):
        """
        :param head: 代表几个头的参数
        :param embedding_dim: 词向量维度
        :param dropout: 置零比率
        """
        super(MultiHeadAttention, self).__init__()
        assert embedding_dim % head == 0     # 确认一下多头的数量可以整除词嵌入的维度 embedding_dim
        self.d_k = embedding_dim // head  # 每个头获得词向量的维度
        self.head = head
        self.linears = nn.ModuleList([copy.deepcopy(nn.Linear(embedding_dim, embedding_dim)) for _ in range(4)])   # 深层拷贝4个线性层，每一个层都是独立的，保证内存地址是独立的，分别是 Q、K、V以及最终的输出线性层
        self.attn = None   # 初始化注意力张量
        self.dropout = nn.Dropout(p=dropout)
    def forward(self, query, key, value, mask=None):
        """
        :param query: 查询query [batch size, sentence length, d_model]
        :param key: 待查询key [batch size, sentence length, d_model]
        :param value: 待查询value [batch size, sentence length, d_model]
        :param mask: 计算相似度得分时的掩码（设置哪些输入不计算到score中）[batch size, 1, sentence length]
        :return:
        """
        if mask is not None:
            mask = mask.unsqueeze(1) # 将掩码张量进行维度扩充，代表多头中的第 n 个头
        batch_size = query.size(0)
        query, key, value = [l(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2) for l, x in zip(self.linears, (query, key, value))]  # 将1、2维度进行调换，目的是让句子长度维度和词向量维度靠近，这样注意力机制才能找到词义与句子之间的关系
        # 将每个头传递到注意力层
        x, self.attn = attention(query, key, value, mask=mask,
                                      dropout=self.dropout)
        # 得到每个头的计算结果是 4 维的张量，需要形状的转换
        # 前面已经将1,2两个维度进行转置了，所以这里要重新转置回来
        # 前面已经经历了transpose，所以要使用contiguous()方法，不然无法使用 view 方法
        x = x.transpose(1, 2).contiguous() \
            .view(batch_size, -1, self.head * self.d_k)
        return self.linears[-1](x)  # 在最后一个线性层中进行处理，得到最终的多头注意力结构输出

# 实例化参数
d_model = 512
dropout = 0.1
max_len = 60  # 句子最大长度
#-----------------------------------------词嵌入层
# 输入 x 是 Embedding层输出的张量，形状为 2 * 4 * 512
x = Variable(torch.LongTensor([[100,2,42,508],[491,998,1,221]]))
emb = Embeddings(1000,512)   # 嵌入层
embr = emb(x)
#-----------------------------------------位置编码
pe = PositionalEncoding(d_model, dropout,max_len)  # 位置编码
pe_result = pe(embr)
#-----------------------------------------多头注意力机制
head = 8
embedding_dim = 512
dropout = 0.2
query = key = value = pe_result
mask = torch.zeros(2,4,4)
mha = MultiHeadAttention(head,embedding_dim,dropout)
mha_result = mha(query,key,value,mask)
print(mha_result)
print(mha_result.shape)
 

scores.shape  torch.Size([2, 8, 4, 4])
tensor([[[ 7.7696, -2.7345,  2.4208,  ..., -2.6492,  1.5189,  4.8102],
         [ 3.3628, -3.0589,  4.1074,  ...,  0.7158, -1.8547,  5.7866],
         [ 3.2688, -2.4387,  3.8830,  ..., -1.7330, -7.4030,  8.4375],
         [ 2.4126, -4.9404, -0.9199,  ...,  0.9010,  0.4319,  6.2884]],

        [[ 2.6921,  2.3492, -7.3719,  ..., -0.8412, -5.3597, -5.0705],
         [ 7.8171,  2.8094, -5.2773,  ..., -1.8384, -4.4927, -4.3039],
         [ 1.3418,  5.0010, -6.6760,  ..., -0.6773, -8.4337, -2.2825],
         [ 2.4672,  3.3702, -3.9995,  ..., -4.1702, -2.3165, -3.0850]]],
       grad_fn=<AddBackward0>)
torch.Size([2, 4, 512])