Happy-LLM 注意力机制

2.1 注意力机制

2.1.1 什么是注意力机制

​ 随着自然语言处理（NLP）从传统的统计机器学习方法逐渐转向深度学习，文本表示方法也发生了显著变化。最初，文本表示主要依赖于统计学习模型，如向量空间模型和语言模型。后来，随着 Word2Vec 等单层神经网络的出现，文本表示开始进入深度学习时代。然而，传统的神经网络架构在处理序列数据时存在一些局限性。例如，全连接神经网络（FNN）虽然结构简单，但参数量大，计算效率低。卷积神经网络（CNN）虽然能够提取局部特征，但在处理长序列数据时效果不佳。循环神经网络（RNN）及其变体（如 LSTM）虽然能够处理序列数据，但存在两个主要问题：一是计算效率低，因为 RNN 的序列计算方式限制了并行计算的能力；二是难以捕捉长序列中的长距离依赖关系。

​ 为了解决这些问题，研究者们引入了注意力机制（Attention）。注意力机制最初是在计算机视觉（CV）领域提出的，其核心思想是将注意力集中在最重要的部分，而不是处理全部信息。在 NLP 中，注意力机制通过将重点放在文本中的关键部分（如某些单词或短语），能够更高效地处理文本数据。

​ 注意力机制的核心是三个变量：Query（查询向量）、Key（键向量） 和 Value（值向量）。通过计算 Query 和 Key 的相似度，可以得到一个权重，这个权重表示每个 Key 的重要性。然后，将这个权重与对应的 Value 相乘并求和，得到最终的输出。这个过程可以简单理解为：通过 Query 找到与之最相关的 Key，然后根据 Key 的重要性提取对应的 Value。

2.1.2 理解注意力机制

为了更好地理解注意力机制，我们可以用一个简单的例子来说明。假设我们有一个字典，字典的键（Key）和值（Value）如下：

{
    "apple": 10,
    "banana": 5,
    "chair": 2
}

如果我们要查找与“fruit”相关的值，我们可以将“fruit”作为 Query。通过计算 Query 和每个 Key 的相似度，我们可以得到每个 Key 的权重。例如，假设我们得到的权重如下：

{
    "apple": 0.6,
    "banana": 0.4,
    "chair": 0
}

那么，最终的输出值就是：

$$\text{output}=0.6\times 10+0.4\times 5+0\times 2=8$$

​ 这个过程可以通过数学公式来表示。假设 Query 和 Key 都是向量，我们可以通过计算 Query 和 Key 的点积来得到相似度，然后通过 Softmax 函数将相似度转换为权重。具体公式如下：

$$\text{scores}=\text{Query}\times {\text{Key}}^T$$，$$\text{weights}=\text{softmax}(\text{scores})$$，$$\text{output}=\text{weights}\times \text{Value}$$

为了防止计算结果过大或过小，我们通常会对 Query 和 Key 的点积结果进行缩放，具体公式为：

$$\text{attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V，d_k是Key的维度。$$

2.1.3 注意力机制的实现

在 PyTorch 中，我们可以用几行代码实现上述注意力机制：

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

def attention(query, key, value, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    p_attn = F.softmax(scores, dim=-1)
    return torch.matmul(p_attn, value), p_attn

2.1.4 自注意力机制

​ 自注意力机制（Self-Attention）是注意力机制的一种变体，它允许模型在处理序列数据时，同时考虑序列中的所有元素。在自注意力机制中，Query、Key 和 Value 都来自同一个输入序列。$$参数矩阵W_q、W_k和W_v定义，分别生成Query、Key和Value。$$自注意力机制的实现非常简单，只需要将 Query、Key 和 Value 的输入设置为同一个张量即可：

attention(x, x, x)

2.1.5 掩码自注意力机制

​ 掩码自注意力机制（Masked Self-Attention）是在自注意力机制的基础上引入了掩码（Mask），用于防止模型在计算时看到未来的信息。这在处理序列生成任务（如语言模型）时尤为重要，因为模型只能根据已知的历史信息来预测下一个元素。例如，假设我们有一个序列 [BOS, I, like, you, EOS]，在计算每个位置的注意力时，我们希望模型只能看到该位置之前的信息。为此，我们可以创建一个上三角矩阵作为掩码，将未来位置的值设置为负无穷大（-inf），这样在计算 Softmax 时，这些位置的权重就会趋近于零。掩码的生成代码如下：

mask = torch.full((1, seq_len, seq_len), float("-inf"))
mask = torch.triu(mask, diagonal=1)

在注意力计算中，我们将掩码与注意力分数相加，然后应用 Softmax 函数：

scores = scores + mask
scores = F.softmax(scores, dim=-1)

通过这种方式，模型在计算每个位置的注意力时，只能看到该位置之前的信息，从而实现了因果关系的建模。

2.1.6 多头注意力机制

​ 多头注意力机制（Multi-Head Attention）是 Transformer 架构中的一个关键组件，它允许模型从不同的角度同时学习输入数据的特征。具体来说，多头注意力机制将输入数据分成多个头（Head），每个头都独立地进行自注意力计算，然后将所有头的输出拼接起来，通过一个线性变换得到最终的输出。

多头注意力机制的公式如下：$$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W$$

其中，每个头的计算公式为：$${\text{head}}_i=\text{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$$

在 PyTorch 中，多头注意力机制的实现如下：

class MultiHeadAttention(nn.Module):
    def __init__(self, args, is_causal=False):
        super().__init__()
        assert args.dim % args.n_heads == 0
        self.n_local_heads = args.n_heads // args.model_parallel_size
        self.head_dim = args.dim // args.n_heads

        self.wq = nn.Linear(args.dim, self.n_local_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, self.n_local_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, self.n_local_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_local_heads * self.head_dim, args.dim, bias=False)
        self.attn_dropout = nn.Dropout(args.dropout)
        self.resid_dropout = nn.Dropout(args.dropout)

        if is_causal:
            mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
            mask = torch.triu(mask, diagonal=1)
            self.register_buffer("mask", mask)

    def forward(self, q, k, v):
        bsz, seqlen, _ = q.shape
        xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim).transpose(1, 2)
        xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim).transpose(1, 2)
        xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim).transpose(1, 2)

        scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
        if self.is_causal:
            scores = scores + self.mask[:, :, :seqlen, :seqlen]
        scores = F.softmax(scores, dim=-1)
        scores = self.attn_dropout(scores)
        output = torch.matmul(scores, xv)

        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        output = self.wo(output)
        output = self.resid_dropout(output)
        return output

参考资料：DataWhale Happy-LLM 2.1注意力机制