多头注意力机制

原创于 2025-08-14 09:02:25 发布 · 581 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能

一、自注意力机制

产生原因
- RNN 的局限性：在处理长序列数据（如长文本）时，RNN（如 LSTM、GRU）存在梯度消失或梯度爆炸问题，难以学习长距离依赖关系。例如，分析一篇长论文中开头假设与结尾结论的关联时，RNN 难以有效捕捉。
- 自注意力机制的优势：通过计算序列中每个位置与其他所有位置的关联程度，能够直接获取长距离依赖信息。处理文本时，可让一个单词直接 “关注” 到文本中其他任意位置的单词，不受距离限制，有效解决长序列依赖问题。
并行化解决
- CNN 的问题：CNN 的层级结构中，高层计算依赖底层输出，每一层计算需等待前一层结果，这种层级依赖限制了并行化程度。尽管同一层可进行一定并行计算，但随着层数增加，整体流程仍受束缚。
- 自注意力的改进：基于整个输入进行计算，且可并行处理，无需像 CNN 那样堆叠多层，大幅提升计算效率。
核心特点
- 长距离依赖捕获：通过计算序列中各元素间的相关性，捕捉全局依赖关系。
- 并行计算：不依赖序列顺序，可并行处理整个序列，显著提高计算效率。
- 动态权重分配：模型能动态关注序列中不同位置的重要信息，不再依赖固定的上下文向量。
- 灵活性：可处理不同长度的输入序列，不像卷积或 RNN 那样对输入结构有严格要求。
使用场景 语言含义极度依赖上下文，同一词或句子在不同语境中可能有完全不同的含义。例如：
- “货拉拉拉不拉拉布拉多要看拉布拉多在货拉拉上拉不拉 baba~” 中，多个 “拉” 的含义需结合上下文判断。
- 机器人第二法则 “机器人必须遵守人类给它的命令，除非该命令违背了第一法则” 中，“它” 指代机器人、“命令” 指代前文 “人类给它的命令”、“第一法则” 指代完整的机器人第一法则，这些均需结合上下文理解，此时需用到自注意力机制。
专业术语 自注意力机制通过查询向量（Query）、键向量（Key）、值向量（Value）实现序列元素间的信息交互和依赖建模：
- Q（Query）：表示当前查询者的位置，用于发出 “我想知道对我来说谁重要” 的问题。
- K（Key）：表示被查询者的身份，是所有位置给出的 “介绍信” 或 “标签”，用于说明自身信息。
- V（Value）：表示被查询者的实际信息，一旦被 “关注”，就会提供该信息。
- QKV 的意义：序列中每个 Token 都兼具 Q、K、V 三个角色。所有位置需通过 “查询 - 响应” 互动，单一角色表达能力有限；“我该关注谁” 是 “我” 与 “他们” 的交互，需分别建模（Q vs K）；最终融合的 V 信息可能与 Q・K 打分依据不同（如 K 强调结构特征，V 强调语义内容）。

二、实现过程

输入：输入为序列（如词向量序列），假设 $X=(x_1, x_2, \dots, x_n) \in \mathbb{R}^{n \times d}$ ，其中n为输入数量，d为输入维度。自注意力的目的是捕获n个实体间的关系。例如，“I Love Nature Language Processing” 经词向量转化（nn.Embedding）后作为输入，n为词的个数，d为每个词的维度。
词语关系：自注意力机制可识别词语间的关联，如 “我是一名学生，现在正在学自然语言处理” 中 “我” 与 “自然语言处理” 的关系。
线性变换：通过三个可学习的权重矩阵对输入X进行线性变换，得到 Q、K、V 矩阵：
- 查询向量（Q）：每个输入生成的查询向量，表示当前词的需求。通过权重矩阵映射到查询空间 $Q=XW_q$ ， $W_q$ 维度为d $\times$ $d_k$ ， $d_k$ 为查询向量维度，用于与 K 计算相似度（点积方式），确定当前词与其他词的相关性。
- 键向量（K）：每个输入生成的键向量，表示其能提供的信息。通过权重矩阵映射到键空间 $K=XW_k$ ， $W_k$ 维度为d $\times$ $d_k$ ，与 Q 计算点积生成注意力权重，点积越大则相关性越强。
- 值向量（V）：包含每个输入的实际信息内容，相关性决定信息被聚焦的程度。通过权重矩阵映射到值空间 $V=XW_v$ ， $W_v$ 维度为d $\times$ $d_v$ ， $d_v$ 为值向量维度），用于基于注意力得分进行加权求和，生成最终输出。
注意力得分：使用点积计算 Q 与 K 的相似度，除以缩放因子 $sqrt{d_k}$ 避免数值过大，得到注意力得分矩阵： $\text{Attention}(Q, K)$ = $\frac{QK^T}{\sqrt{d_k}}$ 。矩阵维度为n $\times$ n，元素(i, j)表示第i个元素与第j个元素的相似度。
归一化：按行对得分矩阵进行 softmax 操作，将注意力得分转换为概率分布（每行和为 1），得到注意力权重矩阵。公式为： $\hat{\alpha}_{1,i}$ = $\frac{\exp(\alpha_{1,i})}{\sum \exp(\alpha_{1,j})}$ ，其中 $\alpha_{1,i}$ 为第 1 个词语与第i个词语的原始注意力得分， $\hat{\alpha}_{1,i}$ 为归一化后的得分。
加权求和：将注意力权重矩阵与值矩阵V相乘，得到加权的值表示，即 $\text{Output} = \text{AttentionWeight} \times V$ 。

三、多头注意力机制

基本概念：作为自注意力机制的扩展，核心思想是将 Q、K、V 分成多个头，每个头计算独立的注意力结果，拼接所有头的输出后通过线性变换得到最终输出。
实现过程（基于代码示例）：
- 数据准备与词嵌入：构建词汇表，将句子转换为索引张量后生成词嵌入（如句子 “I Love Nature Language Processing” 的嵌入形状为torch.size([5,512])。
- 参数定义：设头数为 8，词向量维度为 512，则每个头的维度为512/8=64。
- 生成 Q、K、V 矩阵：通过 3 个线性层分别映射输入，得到形状均为torch.size([5,512])的 Q、K、V 矩阵。
- 拆分多头：将 Q、K、V 的形状转换为[5, 8, 64]后转置为[8, 5, 64]，得到多头 Q、K、V形状均为torch.size([8,5,512])。
- 计算注意力得分：对 K 转置维度变为[8, 64, 5]，通过矩阵乘法计算 $(Q \cdot K^T) / \sqrt{\text{head\_dim}}$ ，得到形状为torch.size([8,5,5]的注意力得分。(可以实现词对词的映射关系，让这五个单词之间互相有联系)
- 计算注意力权重：对得分进行 softmax 操作（维度为 - 1），得到形状为torch.size([8,5,5]的注意力权重（每行和为 1）。

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

if __name__ == '__main__':
    torch.manual_seed(42)  # 设置随机种子，保证结果可复现

    # 1. 数据准备与词汇表构建
    document = "[UNK] I Love Nature Language Processing and you tell I Looking in my eyes"
    sentence = "I Love Nature Language Processing"  # 待处理句子
    vocab = {word: i for i, word in enumerate(set(document.split(" ")))}  # 构建词汇表
    vocab_len = len(vocab)
    dim = 512  # 词向量维度

    # 2. 生成词嵌入
    embedding = nn.Embedding(vocab_len, dim)
    # 将句子转换为索引张量，再生成嵌入矩阵
    sentence_embedding = embedding(
        torch.tensor([vocab[word] for word in sentence.split(" ")])
    )
    print("句子嵌入形状:", sentence_embedding.shape)  # 输出: torch.Size([5, 512])

    # 3. 定义多头注意力参数
    head_num = 8  # 注意力头数
    head_dim = dim // head_num  # 每个头的维度: 512/8=64

    # 4. 生成Q、K、V矩阵（使用3个线性层分别映射）
    # 注意：这里只需3个线性层（Q、K、V各一个），而非head_num个
    fc = nn.ModuleList([nn.Linear(dim, dim) for _ in range(3)])
    Q = fc[0](sentence_embedding)  # 查询矩阵
    K = fc[1](sentence_embedding)  # 键矩阵
    V = fc[2](sentence_embedding)  # 值矩阵
    print("Q矩阵形状:", Q.shape)  # 输出: torch.Size([5, 512])
    print("K矩阵形状:", K.shape)  # 输出: torch.Size([5, 512])
    print("V矩阵形状:", V.shape)  # 输出: torch.Size([5, 512])

    # 5. 拆分多头
    # 形状转换: [5, 512] → [5, 8, 64] → 转置为 [8, 5, 64]
    multi_head_Q = Q.reshape(-1, head_num, head_dim).transpose(0, 1)
    multi_head_K = K.reshape(-1, head_num, head_dim).transpose(0, 1)
    multi_head_V = V.reshape(-1, head_num, head_dim).transpose(0, 1)
    print("多头Q形状:", multi_head_Q.shape)  # 输出: torch.Size([8, 5, 64])
    print("多头K形状:", multi_head_K.shape)  # 输出: torch.Size([8, 5, 64])
    print("多头V形状:", multi_head_V.shape)  # 输出: torch.Size([8, 5, 64])

    # 6. 计算注意力得分（核心步骤）
    # 对K进行转置，使矩阵乘法维度匹配: [8, 5, 64] → [8, 64, 5]
    K_transposed = multi_head_K.transpose(1, 2)
    # 注意力得分 = (Q · K^T) / √(head_dim)
    attention_scores = torch.matmul(multi_head_Q, K_transposed) / math.sqrt(head_dim)
    print("注意力得分形状:", attention_scores.shape)  # 输出: torch.Size([8, 5, 5])
    # 打印第一个头的注意力得分（示例）
    print("\n第一个头的注意力得分矩阵:")
    print(attention_scores[0])  # 形状: [5, 5]，表示5个词之间的原始关注度

    # 7. 计算注意力权重（归一化）
    attention_weights = F.softmax(attention_scores, dim=-1)
    print("\n注意力权重形状:", attention_weights.shape)  # 输出: torch.Size([8, 5, 5])
    # 打印第一个头的注意力权重（每行和为1）
    print("\n第一个头的注意力权重矩阵:")
    print(attention_weights[0])

运行结果：