Transformer学习与基础实验1——注意力机制

标签：Transformer self torch mask 实验 key attn 注意力

前置概念

自然语言处理（NLP）中，根据任务内容的不同，句子、段落等文本中需要更加关注的部分（不同的词语、句子等）也会不同。

在判断词在句子中的重要性时便使用了注意力机制，可以通过注意力分数来表达某个词在句子中的重要性，分数越高，说明该词对完成该任务的重要性越大。

计算注意力分数时，我们主要参考三个因素：query、key和value。

1. query：任务内容

2. key：索引/标签（帮助定位到答案）

3. value：答案

一个直观的例子：百度搜索中，输入的“搜索文本”即搜索的内容是query，搜索结果页面中各个结果标题（及其链接）为key，它与任务内容query相关，并能引导我们至具体的内容（value）。

一般在文本翻译中，我们希望翻译后的句子（目标序列）的所表达的意义和原始句子（源序列）的相似，所以进行注意力分数计算时，query一般和目标序列有关，key则与源序列有关。

（提出疑问Q1：那value和什么相关？）

计算注意力分数的方式

1. Additive Attention（加性注意力）

加性注意力机制通过一个带有非线性激活函数的神经网络来计算query和key之间的相似度。

①. 线性变换：首先，将query和每个key通过一个可学习的权重矩阵进行线性变换。

②. 拼接：然后，将变换后的query和key进行拼接。

③. 非线性激活：通过一个非线性激活函数（通常是tanh）处理拼接后的结果。

④. 加权求和：最后，通过另一个权重矩阵进行线性变换，并使用一个softmax函数来计算每个value的注意力权重。

数学公式表示为：

$e_i=v^T\tanh(W_qq+W_kk_i)$

代码实现中，使用批量计算时的计算流程图：

2. Scaled Dot-Product Attention（缩放点积注意力）

缩放点积注意力是由Ashish Vaswani等人在2017年提出的Transformer模型中使用的。这种方法直接计算query和key之间的点积，然后通过缩放和softmax函数来计算注意力权重。

①. 点积：计算query和每个key之间的点积QKT。在几何角度，点积（dot product）表示一个向量在另一个向量方向上的投影。换句话说，从几何角度上解读，点积代表了某个向量中的多少是和另一个向量相似的，两个向量点积的结果可以表示这两个向量相似的程度。

②. 缩放：将点积结果除以一个缩放因子（通常是key向量维度的平方根），以防止点积结果过大，导致梯度消失或爆炸。（提出疑问Q2：若点积结果过大，则分别如何导致梯度消失或爆炸的？缩放因子又是如何解决了这个问题的？）

③. softmax：通过softmax函数计算每个value的注意力权重。

数学公式表示：

${Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$

原论文的结构示意图：

点积注意力机制示意图（图片来源：Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017.）

计算流程详情图（注意暂时没有考虑“掩码”操作，上下同）：

pytorch代码实现缩放点积注意力

import torch
import torch.nn as nn
import torch.nn.functional as F

class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout_p=0.):
        super(ScaledDotProductAttention, self).__init__()
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, query, key, value, attn_mask=None):
        """
        缩放点积注意力机制
        参数:
            query: [batch_size, n_heads, seq_len_Q, d_k]
            key: [batch_size, n_heads, seq_len_K, d_k]
            value: [batch_size, n_heads, seq_len_K, d_v]
            attn_mask: [batch_size, seq_len_Q, seq_len_K]
        返回:
            output: [batch_size, n_heads, seq_len_Q, d_v]
            attn: [batch_size, n_heads, seq_len_Q, seq_len_K]
        """
        embed_size = query.size(-1)
        scaling_factor = embed_size ** -0.5  # 缩放因子

        # 计算缩放后的点积注意力分数
        attn = torch.matmul(query, key.transpose(-2, -1)) * scaling_factor

        if attn_mask is not None:
            # 将 attn_mask 中为 True 的位置的 attn 元素替换为 -1e9
            attn = attn.masked_fill(attn_mask, -1e9)

        # 通过 Softmax 函数将注意力分数转换为权重
        attn = F.softmax(attn, dim=-1)

        # 应用 Dropout
        attn = self.dropout(attn)

        # 使用注意力权重对 value 进行加权平均，生成最终的输出
        output = torch.matmul(attn, value)

        return output, attn

问题Q3：代码的计算中，两个矩阵Q和K（多个query和key向量放在一起）的向量维度一定要相同么？实际中Q和K的维度一定相同么？

问题Q4：注意力掩码attn_mask的意义是什么？在代码中是如何起作用的？

问题Q5：最终的输出output, attn分别有什么意义？后续的作用如何？

问题Q6：一般场景中，query, key, value的维度分别应是什么？各个维度的含义又是什么？

掩码的生成和使用

在处理数据时，尤其是文本数据，为了统一长度，会使用占位符补齐了一些稍短的文本。

"Hello world!" --> ['<bos>' , 'hello', 'world', '!', '<eos>', '<pad>', '<pad>' ]

这些占位符<pad>没有特别的意义，在实际计算中不应该参与到注意力分数的计算中。为此可以在注意力机制中加入 padding 掩码，即识别输入序列中的占位符，保证计算时这些位置对应的注意力分数为0。实际实现中，占位符<pad>的数字表示一般是固定的，根据该数值找到占位符的位置（索引），使用一个掩码矩阵，通常将占位符对应位置的元素标记为1，具体的实现结合以下代码和演示理解

import torch
import torch.nn.functional as F

def get_attn_pad_mask(seq_q, seq_k, pad_idx):
    _, len_q = seq_q.size()
    _, len_k = seq_k.size()

    pad_attn_mask = seq_k.eq(pad_idx)
    pad_attn_mask = pad_attn_mask.unsqueeze(1).expand(-1, len_q, len_k)
    return pad_attn_mask

if __name__ == '__main__':
    q = k = torch.tensor([[1, 3, 4, 5, 2, 0, 0]], dtype=torch.float32)
    pad_idx = 0
    mask = get_attn_pad_mask(q, k, pad_idx)
    print(mask)
    print(q.shape, mask.shape)

输出结果：

生成的掩码具体是如何在注意力分数的计算中发挥其作用？下述演示可以帮助理解：

if __name__ == '__main__':
    q = k = torch.tensor([[1, 3, 4, 5, 2, 0, 0]], dtype=torch.float32)
    pad_idx = 0
    mask = get_attn_pad_mask(q, k, pad_idx)
    print(mask)
    print(q.shape, mask.shape)

    attn = torch.ones(mask.shape)
    attn = attn.masked_fill(mask, -1e9)
    attn_weights = F.softmax(attn, dim=-1)
    print(attn_weights)

输出结果：

自注意力机制与多头注意力机制

自注意力机制

自注意力机制中，我们只关注当前输入的文本本身，查看每个“单词”（token）对于周边单词（token）的重要性。这样可以很好地理清句子中的逻辑关系，如代词指代。

具体地，在'While the dog was running, it hurt itself.' 这句话中，'it' 指代句中的 'the dog'，所以自注意力会赋予 'the'、'dog' 更高的注意力分值。

自注意力分数的计算还是可以使用上述的缩放点积注意力公式，只不过这里的query, key和value都变成了句子本身点乘各自权重。（其实也可以不乘以一个权重，直接将输入的序列作为Q、K、V，思考两种实现有什么区别？实际中可能会有什么不同？）

给定序列X∈ $R^{n\cdot d\_model}$ ，序列长度为n，维度为d_model。在计算自注意力时， $Q = W^{^{Q}}X, K = W^{^{K}}X, V = W^{^{V}}X$

其中，序列中位置为i的词与位置为j的词之间的自注意力分数为：

自注意力机制的代码演示：

class SelfAttention(nn.Module):
    def __init__(self, d_model, dropout_p=0.):
        super(SelfAttention, self).__init__()
        self.attn = ScaledDotProductAttention(dropout_p)
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, x, attn_mask=None):
        """
        自注意力机制
        参数:
            x: [batch_size, n, d_model]
            attn_mask: [batch_size, n, n]
        返回:
            output: [batch_size, n, d_model]
            attn: [batch_size, n, n]
        """
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        output, attn = self.attn(Q, K, V, attn_mask)

        return output, attn

if __name__ == '__main__':
    d_model = 64
    self_attn = SelfAttention(d_model)

    batch_size = 1
    n = 32
    x = torch.ones((batch_size, n, d_model))
    attn_mask = torch.ones((batch_size, n, n), dtype=torch.bool)
    output, attn = self_attn(x, attn_mask)
    print(output.shape, attn.shape)

输出：torch.Size([1, 32, 64]) torch.Size([1, 32, 32])

问题Q7：如何理解自注意力机制通过计算注意力分数让输入序列中的每个元素（单词）能够在生成上下文表示时，关注序列中的其它元素（单词）？具体的意义是什么？

多头注意力机制

多头注意力可以看作是上述自注意力机制（输入序列矩阵乘以权重矩阵之后再作为Q、K、V的做法）的扩展，它可以使模型通过不同的方式关注输入序列的不同部分，从而提升模型的特征表示能力和最终训练效果。

图片来源：Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017.

代码实现时（和下述pytorch实现相对应），多头注意力机制流程图（注意和上图的对应关系）：

（理解形式）：

（注：所有的这些模型结构图提供清晰的、可二次编辑的PPT绘制文件或后缀为.drawio文件，见文末）

①多头注意力通过对输入的embedding乘以不同的权重参数 $W^Q$ 、 $W^K$ 和 $W^V$ ，将其映射到多个小维度空间中，称之为“头”（head）；

（问题Q8：这里的多个小维度空间中的“多个”具体是指多少个？）

②每个头部并行计算自己的自注意力分数，得到多组（n_heads）自注意力分数；

③在获得多组自注意力分数后，我们将结果拼接到一起；

④ $W^O$ 为可学习的权重参数，用于将拼接后的多头注意力输出映射回原来的维度（有可能在映射前的维度已经和原维度一致），得到多头注意力的最终输出。

简单来说，在多头注意力中，每个头部可以“解读”输入内容的不同方面，比如：捕捉全局依赖关系、关注特定语境下的词元词性、识别词和词之间的语法关系等。

现给出多头注意力机制的特例实现（自注意力机制的多头计算，即输入的Q, K, V矩阵是相同的，也即直接使用输入序列），实际上更一般化的多头注意力机制的实现需要考虑许多参数，具体可以看pytorch的官方实现的源码（transformer的实现源码）。

（问题Q9：代码实现时是如何实现映射到多个“头部”并使每个“头部”进行所谓的并行计算？可以有多种不同的实现么？）

import torch
import torch.nn as nn
import torch.nn.functional as F
from scaled_dot_product_attention import ScaledDotProductAttention


class MultiHeadAttention(nn.Module):
    def __init__(self, d_key, n_heads, dropout_p=0.):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.d_key = d_key
        self.head_dim = d_key // n_heads

        if self.head_dim * n_heads != d_key:
            raise ValueError("d_key must be divisible by n_heads")

        self.W_Q = nn.Linear(d_key, d_key, bias=False)
        self.W_K = nn.Linear(d_key, d_key, bias=False)
        self.W_V = nn.Linear(d_key, d_key, bias=False)
        self.W_O = nn.Linear(d_key, d_key, bias=False)
        self.attention = ScaledDotProductAttention(dropout_p=dropout_p)

    def forward(self, query, key, value, attn_mask):
        batch_size = query.size(0)

        Q = self.W_Q(query).view(batch_size, -1, self.n_heads, self.head_dim)
        K = self.W_K(key).view(batch_size, -1, self.n_heads, self.head_dim)
        V = self.W_V(value).view(batch_size, -1, self.n_heads, self.head_dim)

        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        if attn_mask is not None:
            attn_mask = attn_mask.unsqueeze(1).expand(-1, self.n_heads, -1, -1)

        context, attn = self.attention(Q, K, V, attn_mask)

        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.head_dim)
        output = self.W_O(context)

        return output, attn

if __name__ == '__main__':
    batch_size = 1
    seq_len = 40
    d_key = 32
    n_heads = 8
    dropout_p = 0.1

    query = torch.ones((batch_size, seq_len, d_key))
    key = torch.ones((batch_size, seq_len, d_key))
    value = torch.ones((batch_size, seq_len, d_key))
    attn_mask = torch.ones((batch_size, seq_len, seq_len), dtype=torch.bool)

    multi_head_attn = MultiHeadAttention(d_key, n_heads, dropout_p)
    output, attn = multi_head_attn(query, key, value, attn_mask)
    print(output.shape, attn.shape)

    attn_builtin = nn.MultiheadAttention(d_key, n_heads, dropout_p, batch_first=True)  # 官方库中的接口
    output_builtin, attn_builtin = attn_builtin(query, key, value, torch.ones((1, seq_len), dtype=torch.bool))
    print(output_builtin.shape, attn_builtin.shape)

输出：

torch.Size([1, 40, 32]) torch.Size([1, 8, 40, 40])
torch.Size([1, 40, 32]) torch.Size([1, 40, 40])

关注微信公众号——分享之心，后台回复：Transformer，获取该系列实验的所有源码（包含mindspore和pytorch两个框架的版本)、文档、模型结构图（部分帮助理解的流程图）文件（.drawio文件、PPT文件）。

注：对于源码，部分源文件中的“通过sys模块添加系统路径使得可以正确加载自定义的模块”部分需要根据实际运行机器的路径进行修改，本地安装好第三方依赖包（如pytorch、jieba等）后可以直接运行。

预告：

Transformer学习与基础实验2——Transformer结构

Transformer学习与基础实验3——transformer应用举例演示——英汉翻译（1. 数据准备与处理）

Transformer学习与基础实验4——英汉翻译（2. 模型构建、训练、推理）

标签：Transformer,self,torch,mask,实验,key,attn,注意力
From： https://blog.csdn.net/qq_61784003/article/details/144853958