Transformer中的注意力机制(Attention)

丁希希哇2026-10-17 17:06

相关文章

李沐《动手学深度学习》注意力机制

【论文精读】Transformer：Attention Is All You Need

Transformer 的架构原理

文章目录

一、注意力机制(Attention)

注意力机制是理解现代大模型（如GPT、BERT、LLaMA）的绝对核心。没有注意力机制，就没有今天的大模型。

注意力机制解决了什么问题？

  1. 传统 RNN/LSTM 的瓶颈
    在注意力机制出现前，处理序列（如翻译句子）主要靠RNN。但是RNN有两个致命问题：
    • 长距离依赖丢失：句子一长，前面的信息传到后面就忘了
    • 无法并行：必须按顺序计算，GPU 再强也跑不快。
  2. 注意力机制的核心思想：动态聚焦
    注意力 = 加权平均。权重越大，说明当前词越需要参考那个词的信息。

注意力机制的数学公式：

注意力机制的核心，是让每个位置根据当前需求（Q），动态地衡量其他所有位置（K）的重要性，并按重要性加权融合它们的信息（V），从而获得一个融入了上下文的、动态的表示。

Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q,K,V)=\text{softmax}(\frac{QK^T}{\sqrt{d_k}})V Attention(Q,K,V)=softmax(dk QKT)V

为了更直观地理解 Q、K、V 三个角色的分工，可以借助「图书馆检索」的类比：

角色 类比 作用
Query (Q) 你的搜索词 当前词在问："我需要什么信息？"
Key (K) 每本书的标签 每个词在说："我是什么信息？"
Value (V) 每本书的内容 每个词实际携带的信息

Q、K、V 都是输入向量 X 乘以不同的可学习权重矩阵得到的：

Q = X W Q , K = X W K , V = X W V Q=XW^Q, K=XW^K, V=XW^V Q=XWQ,K=XWK,V=XWV

  • 为什么Q、K 要使用独立的权重矩阵：
    为了让模型有能力从不同角度表达输入序列中的依赖关系。如果 Q 和 K 完全对称，匹配就会失去方向性。
  • 为什么需要 Q、K、V（查询、键、值）三个矩阵：
    如果只有 Q 和K，我们只能知道每个词与其他词的相关性，但无法获得与这些相关性对应的具体内容。V 是在相关性确定后，被权重加权的部分，生成的结果便是模型对输入序列的理解。
  • 为什么要除以 d k \sqrt{d_k} dk ：
    d k d_k dk 是 Key 的维度，当维度很大时，点积结果会很大，导致softmax进入饱和区（梯度极小）。除以 d k \sqrt{d_k} dk 是为了缩放，让梯度稳定。

注意力机制的多种变体：

  • Self-Attentiohn：Q、K、V 都来自同一个序列 → 自己看自己；
  • Masked Self-Attention：加个掩码，不让看未来 → 自回归生成；
  • Cross-Attention：Q 来自一个序列，K、V 来自另一个 → 跨序列对齐；
  • Multi-Head：做多次注意力，每次关注不同模式 → 多视角；
  • FlashAttention：数学不变，只是算得更快更省显存。

二、自注意力机制

（一）Self-Attention

自注意力（Self-Attention）就是：序列中的每个位置，都用同一套 Q、K、V 规则，去关注同一个序列里的所有位置（包括自己），从而得到一个融合了全局上下文的新表示。

自注意力的三个核心特性：

  1. 全局感受野： 每个词直接和所有词交互，不管隔多远。一个 1000 词的句子，第 1 个词和第 1000 个词也能直接建立联系。这是它碾压 RNN 的关键。
  2. 动态权重： 权重不是固定的，而是根据当前输入实时算出来的。同一句话换个语境，注意力分布就变了。
  3. 并行计算： 所有位置的 Q、K、V 可以同时算，没有先后依赖。这是 Transformer 能在大规模 GPU 上训练的根本原因。

自注意力本身对词序不敏感。因为它是加权求和，打乱词序，结果可能一样。所以必须额外加上位置编码（Positional Encoding），把"第几个词"的信息注入进去。否则"猫追老鼠"和"老鼠追猫"在模型眼里没区别。

自注意力 vs 全连接：

全连接层也让每个位置看到所有位置，自注意力与全连接层的主要区别在于权重是否动态：

  • 全连接的权重是训练好后固定的，与输入无关
  • 自注意力的权重是根据输入实时计算出来的，随内容变化

（二）Masked Self-Attention

Masked Self-Attention（掩码自注意力）是 GPT 系列大模型的心脏。

掩码自注意力的核心思想： 在计算注意力时，强制每个位置只能看到它自己和它左边的词，右边的词全部屏蔽掉。

Attention ( Q , K , V ) = softmax ( Q K T d k + M ) V \text{Attention}(Q,K,V)=\text{softmax}(\frac{QK^T}{\sqrt{d_k}}+M)V Attention(Q,K,V)=softmax(dk QKT+M)V

其中 M M M 是掩码矩阵，它是一个上三角矩阵，将对角线以上的位置设为 − ∞ -\infty −∞，对角线及以下的位置设为 0 0 0：

M i j = { 0 , i ≥ j − ∞ , i < j M_{ij} = \begin{cases} 0, & i \ge j \\ -\infty, & i < j \end{cases} Mij={0,−∞,i≥ji<j

这样在 softmax 之前，未来位置的分数会被压成 0，从而被完全屏蔽。

训练时的并行性：

虽然逻辑上每个位置只能看左边，但是训练时所有位置可以同时计算。因为掩码矩阵是预先知道的，一次矩阵乘法就能算出所有位置的注意力，然后统一加掩码、softmax。不需要像 RNN 那样一步步来。

推理时的 KV Cache：

生成时，模型一个词一个词地输出，每一步都要重新计算前面所有词的K、V，非常浪费。KV Cache就是把之前算过的K、V缓存起来，每步只算新词的K、V，然后拼接。

  • 第 1 步：缓存 K₁, V₁
  • 第 2 步：缓存 K₁K₂, V₁V₂
  • 第 3 步：缓存 K₁K₂K₃, V₁V₂V₃

推理时不需要掩码了，因为新词本来就没有未来词，只需和缓存的 K、V 做注意力即可。

三、交叉注意力机制

交叉注意力就是 Q 来自一个序列，K 和 V 来自另一个序列的注意力机制。

类型 Q 来源 K、V 来源 建模对象
自注意力 序列 A 序列 A 序列内部关系
交叉注意力 序列 A 序列 B 两个序列之间的关系

为什么需要交叉注意力？

交叉注意力的作用：让一个序列在生成时，动态地从另一个序列中提取相关信息。它让模型能够：

  • 翻译时对齐源语言和目标语言；
  • 图文任务中对齐图像区域和文本词；
  • 语音识别中对齐音频帧和文字。

交叉注意力的数学机制：

交叉注意力的公式和自注意力完全一样：

Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q,K,V)=\text{softmax}(\frac{QK^T}{\sqrt{d_k}})V Attention(Q,K,V)=softmax(dk QKT)V

唯一的区别在于 Q、K、V 的来源：

Q = X decoder W Q , K = X encoder W K , V = X encoder W V Q=X_{\text{decoder}}W^Q, K=X_{\text{encoder}}W^K, V=X_{\text{encoder}}W^V Q=XdecoderWQ,K=XencoderWK,V=XencoderWV

标准 GPT 没有Cross-Attention，只有 Masked Self-Attention。只有编码器-解码器模型（如 T5）才有Cross-Attention。

四、多头注意力机制

（一）Multi-head Attention

多头注意力是 Transformer 的标配组件，也是让注意力机制真正强大的关键设计。

多头注意力就是：把 Q、K、V 拆成 h 份，每一份独立做一次注意力，最后把结果拼接起来。换句话说，不是做一次注意力，而是并行做 h 次，每次关注不同的模式。

注意： 多头注意力不是把输入切成 h 份，而是每个头都有独立的投影矩阵，都看完整输入，只是投影到不同子空间。

多头注意力的数学机制：

MultiHead ( Q , K , V ) = Concat ( head 1 , . . . , head h ) W O \text{MultiHead}(Q,K,V)=\text{Concat}(\text{head}_1,...,\text{head}_h)W^O MultiHead(Q,K,V)=Concat(head1,...,headh)WO

其中每个头：

head i = Attention ( Q W i Q , K W i K , V W i V ) \text{head}_i=\text{Attention}(QW_i^Q,KW_i^K,VW_i^V) headi=Attention(QWiQ,KWiK,VWiV)

计算流程：多头注意力的每个头维度是 1/h，总计算量不变。

多头注意力的变体：

类型 Q K、V KV Cache 性能
MHA （Multi-head Attention） h 份独立 h 份独立 最大 最好
GQA（Grouped-Query Attention） h 份独立 g 份共享 中等 接近 MHA
MQA （Multi-Query Attention） h 份独立 1 份共享 最小 略降

（二）MQA （Multi-Query Attention）

MQA（Multi-Query Attention）就是：所有注意力头共享同一份 K 和 V，只有 Q 保持多头独立。

为什么需要MQA： 标准多头注意力 MHA 的显存瓶颈

推理时，GPT 类模型逐词生成，需要缓存历史 K、V 避免重复计算。这就是 KV Cache。标准 MHA 的 KV Cache 大小：

Size = 2 × h × d k × seq_len × batch × layers × bytes \text{Size}=2 \times h \times d_k \times \text{seq\_len} \times \text{batch} \times \text{layers} \times \text{bytes} Size=2×h×dk×seq_len×batch×layers×bytes

MQA 的核心思想： 所有头共享同一份 K、V，只有 Q 保持多头。KV Cache 缩小到原来的 1 / h 1/h 1/h。

为什么 Q 不共享：

因为 Q 代表"当前词想问什么"。每个头应该能问不同的问题（多视角），这是多头注意力的核心价值。而 K、V 是"被查询的信息"，共享它们只是减少了存储，不直接破坏多视角提问的能力。

（三）GQA（Grouped-Query Attention）

GQA 是当前主流大模型（LLaMA 2/3、Mistral、Qwen 等）实际采用的注意力变体。

GQA 的核心思想是把 h 个头分成 g 组，每组共享一份 K、V。Q 仍然每个头独立。

  • 当 g = h g=h g=h：退化为 MHA（每组只有一个头）；
  • 当 g = 1 g=1 g=1：退化为 MQA（所有头共享一份）。

五、FlashAttention

FlashAttention 是算法层面的革命，和 GQA（结构层面）解决的是不同维度的瓶颈。GQA 省的是显存容量（KV Cache），FlashAttention 省的是显存带宽（HBM 读写）。

（一）GPU 内存层级

GPU 的存储不是"一块内存"，而是有层级的（以 A100 为例）：

层级 容量 带宽 特点
SRAM（片上缓存） ~20 MB ~19 TB/s 极快，但极小
HBM（显存） 40~80 GB ~1.5 TB/s 很大，但慢 10 倍以上

标准注意力的问题在于：把巨大的中间矩阵反复读写 HBM 。以序列长度 N = 4096 N=4096 N=4096 为例：

  1. 计算 S = Q K T S=QK^T S=QKT，得到 4096 × 4096 4096 \times 4096 4096×4096 的分数矩阵；
  2. 把这个矩阵写到 HBM （约 64 MB，FP16）；
  3. 从 HBM 读回它，做 softmax；
  4. 再把 softmax 结果写回 HBM；
  5. 从 HBM 读回，乘以 V。

这一来一回，中间那个 N × N N \times N N×N 矩阵被读写了好几次。当 N N N 很大时，这个中间矩阵可能几百 MB，HBM 带宽就成了真正的瓶颈。

关键洞察： 算力（FLOPs）不是瓶颈，HBM 的读写（IO）才是。 FlashAttention 之前，很多"近似注意力"方法试图减少 FLOPs，结果 wall-clock 速度反而没提升，就是因为没抓住 IO 这个真瓶颈。

（二）FlashAttention 的核心思想：IO 感知 + 分块

FlashAttention 的核心思想： 既然 SRAM 很快，那就把计算搬到 SRAM 里做；既然 HBM 是瓶颈，那就尽量别读写它。具体通过三个技术实现：

  • Tiling（分块）：不让大矩阵"落地"
    • 把Q、K、V切成小块（tile），每次只把一个Q块和一个K块加载到SRAM里，在SRAM里算这一小块的注意力，然后立刻处理，绝不把中间矩阵写回HBM；
    • 用双重循环遍历所有块（外层遍历 K/V 块，内层遍历 Q 块），把大矩阵的"物化"彻底避免了。
  • Online Softmax（在线 softmax）：分块算归一化
    • 维护两个"运行中的统计量"------当前块的最大值和指数和，每处理一个新块，就动态更新这两个量，并用一个缩放因子来校正之前已经算出的部分输出；
    • 分块计算的输出，和完整 softmax 的数学结果完全一致。 没有近似，没有精度损失。
  • Recomputation（重计算）：不存中间矩阵
    • 反向传播需要注意力权重矩阵来算梯度。标准做法是前向时把它存在 HBM 里，反向时读出来；
    • FlashAttention 的选择是：前向时只存 softmax 的归一化因子（很小），反向时用这些因子快速"重算"出需要的注意力权重。

六、Transformer架构中的注意力机制

Transformer 原始论文《Attention is All You Need》提出的是编码器-解码器架构：

但现代大模型大多只取其中一部分：

模型 架构 用的注意力
BERT 纯编码器 Self-Attention（无掩码）
GPT / LLaMA / Qwen 纯解码器 Masked Self-Attention
T5 / 原始 Transformer 编码器-解码器 三种都用
ViT 纯编码器 Self-Attention（无掩码）

（一）编码器中的注意力：Self-Attention

编码器每一层只有一个注意力模块：多头自注意力

  • Q、K、V都来自上一层输出；
  • 无掩码：每个词可以看到全序列所有词，包括左右两边；
  • 双向建模：适合理解任务（如分类、问答、NER）。

（二）解码器中的注意力：两种

解码器每层有两个注意力模块：多头掩码自注意力、多头交叉注意力

  • 多头掩码自注意力：
    • 位置：解码器层的第一个子层；
    • Q、K、V 都来自解码器当前序列；
    • 有因果掩码：每个位置只能看自己和左边。
  • 多头交叉注意力
    • 位置：解码器层的第二个子层；
    • Q 来自解码器（上一个子层的输出），K、V 来自编码器输出；
    • 无掩码。

（三）GPT 只有一种注意力

GPT 是纯解码器架构，没有编码器，所以：

  • 没有 Cross-Attention（因为没有编码器输出可查）；
  • 只有 Masked Self-Attention。

GPT的每一层：

对于翻译、问答这类需要"源信息"的任务，GPT把所有信息都拼成一个序列。源句和目标句在同一个序列里，用 Masked Self-Attention 统一处理。这样就不需要 Cross-Attention 了。这是 GPT 架构简化的关键洞察：用"拼接序列 + 因果掩码"替代"编码器 + 交叉注意力"。

七、Pytorch代码

（一）单头自注意力

python 复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        # Q、K、V 三个独立投影矩阵
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        # x: (batch, seq_len, d_model)
        Q = self.W_q(x)  # (batch, seq_len, d_model)
        K = self.W_k(x)  # (batch, seq_len, d_model)
        V = self.W_v(x)  # (batch, seq_len, d_model)
        
        # 1. 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_model ** 0.5)
        # scores: (batch, seq_len, seq_len)
        
        # 2. softmax 归一化
        attn = F.softmax(scores, dim=-1)  # 每行和为 1
        
        # 3. 加权求和 V
        out = torch.matmul(attn, V)  # (batch, seq_len, d_model)
        
        return out, attn


# ===== 测试 =====
if __name__ == "__main__":
    batch_size, seq_len, d_model = 2, 4, 8
    x = torch.randn(batch_size, seq_len, d_model)
    
    attn_layer = SelfAttention(d_model)
    out, attn = attn_layer(x)
    
    print("输入形状:", x.shape)        # (2, 4, 8)
    print("输出形状:", out.shape)      # (2, 4, 8)
    print("注意力矩阵:", attn.shape)   # (2, 4, 4)
    print("每行和:", attn.sum(dim=-1)) # 全为 1

（二）多头自注意力（工程实现）

实际大模型里的写法。核心是维度拆分 + 并行计算。

python 复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 每个头的维度
        
        # 一次性投影所有头（比循环高效）
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)  # 输出投影
    
    def forward(self, x, mask=None):
        # x: (batch, seq_len, d_model)
        batch_size, seq_len, _ = x.size()
        
        # 1. 线性投影 + 拆分成多头
        # (batch, seq_len, d_model) -> (batch, seq_len, h, d_k) -> (batch, h, seq_len, d_k)
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        # 现在 Q、K、V: (batch, h, seq_len, d_k)
        
        # 2. 缩放点积注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        # scores: (batch, h, seq_len, seq_len)
        
        # 3. 加掩码（如果有）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # 4. softmax
        attn = F.softmax(scores, dim=-1)
        # attn: (batch, h, seq_len, seq_len)
        
        # 5. 加权求和
        out = torch.matmul(attn, V)
        # out: (batch, h, seq_len, d_k)
        
        # 6. 拼接多头
        # (batch, h, seq_len, d_k) -> (batch, seq_len, h, d_k) -> (batch, seq_len, d_model)
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        
        # 7. 输出投影
        out = self.W_o(out)
        # out: (batch, seq_len, d_model)
        
        return out, attn


# ===== 测试 =====
if __name__ == "__main__":
    batch_size, seq_len, d_model, num_heads = 2, 6, 12, 3
    x = torch.randn(batch_size, seq_len, d_model)
    
    mha = MultiHeadSelfAttention(d_model, num_heads)
    out, attn = mha(x)
    
    print("输入:", x.shape)          # (2, 6, 12)
    print("输出:", out.shape)        # (2, 6, 12)
    print("注意力:", attn.shape)     # (2, 3, 6, 6) ------ 每个头一个矩阵

（三）掩码多头自注意力（GPT 用的）

python 复制代码 
class MaskedMultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        batch_size, seq_len, _ = x.size()
        
        # 投影 + 拆头
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        # 缩放点积
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        # (batch, h, seq_len, seq_len)
        
        # ===== 关键：因果掩码 =====
        # 生成下三角矩阵（含对角线）
        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device)).bool()
        # causal_mask: (seq_len, seq_len)
        # [[1, 0, 0, 0],
        #  [1, 1, 0, 0],
        #  [1, 1, 1, 0],
        #  [1, 1, 1, 1]]
        
        # 把上三角（未来位置）设为 -inf
        scores = scores.masked_fill(~causal_mask, float('-inf'))
        
        # softmax：-inf 变成 0
        attn = F.softmax(scores, dim=-1)
        
        # 加权求和
        out = torch.matmul(attn, V)
        
        # 拼接 + 输出投影
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        out = self.W_o(out)
        
        return out, attn


# ===== 测试：验证掩码效果 =====
if __name__ == "__main__":
    batch_size, seq_len, d_model, num_heads = 1, 4, 8, 2
    x = torch.randn(batch_size, seq_len, d_model)
    
    masked_mha = MaskedMultiHeadSelfAttention(d_model, num_heads)
    out, attn = masked_mha(x)
    
    print("注意力矩阵（第 1 个头）:")
    print(attn[0, 0].detach().numpy().round(3))
    # 输出应该是下三角结构：
    # [[1.    0.    0.    0.   ]
    #  [0.5   0.5   0.    0.   ]
    #  [0.33  0.33  0.33  0.   ]
    #  [0.25  0.25  0.25  0.25 ]]

（四）自注意力合并版本

python 复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F


class SelfAttention(nn.Module):
    """单头自注意力"""
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
    
    def forward(self, x):
        Q, K, V = self.W_q(x), self.W_k(x), self.W_v(x)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_model ** 0.5)
        attn = F.softmax(scores, dim=-1)
        return torch.matmul(attn, V), attn


class MultiHeadSelfAttention(nn.Module):
    """多头自注意力"""
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x, mask=None):
        batch, seq_len, _ = x.size()
        
        Q = self.W_q(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        return self.W_o(out), attn


class MaskedMultiHeadSelfAttention(MultiHeadSelfAttention):
    """掩码多头自注意力（GPT 核心）"""
    def forward(self, x):
        seq_len = x.size(1)
        # 构造因果掩码
        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device)).bool()
        # 扩展维度以匹配 scores: (batch, h, seq_len, seq_len)
        causal_mask = causal_mask.view(1, 1, seq_len, seq_len)
        return super().forward(x, mask=causal_mask)


# ===== 对比测试 =====
if __name__ == "__main__":
    torch.manual_seed(42)
    batch, seq_len, d_model, num_heads = 1, 4, 8, 2
    x = torch.randn(batch, seq_len, d_model)
    
    print("=" * 50)
    print("1. 单头自注意力")
    sa = SelfAttention(d_model)
    out, attn = sa(x)
    print(f"输出: {out.shape}, 注意力: {attn.shape}")
    print("注意力矩阵:\n", attn[0].detach().numpy().round(3))
    
    print("\n" + "=" * 50)
    print("2. 多头自注意力")
    mha = MultiHeadSelfAttention(d_model, num_heads)
    out, attn = mha(x)
    print(f"输出: {out.shape}, 注意力: {attn.shape}")
    print("第 0 个头:\n", attn[0, 0].detach().numpy().round(3))
    print("第 1 个头:\n", attn[0, 1].detach().numpy().round(3))
    
    print("\n" + "=" * 50)
    print("3. 掩码多头自注意力")
    masked_mha = MaskedMultiHeadSelfAttention(d_model, num_heads)
    out, attn = masked_mha(x)
    print(f"输出: {out.shape}, 注意力: {attn.shape}")
    print("第 0 个头（注意下三角结构）:\n", attn[0, 0].detach().numpy().round(3))

（五）单头交叉注意力

python 复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F

class CrossAttention(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model
        # Q 的投影作用于解码器输入
        self.W_q = nn.Linear(d_model, d_model)
        # K、V 的投影作用于编码器输出
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
    
    def forward(self, x_decoder, x_encoder):
        # x_decoder: (batch, m, d_model) ------ 查询来源
        # x_encoder: (batch, n, d_model) ------ 被查询来源
        Q = self.W_q(x_decoder)  # (batch, m, d_model)
        K = self.W_k(x_encoder)  # (batch, n, d_model)
        V = self.W_v(x_encoder)  # (batch, n, d_model)
        
        # 1. 计算分数：注意形状是 (m, n)，不是方阵
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_model ** 0.5)
        # scores: (batch, m, n)
        
        # 2. softmax
        attn = F.softmax(scores, dim=-1)  # (batch, m, n)
        
        # 3. 加权求和
        out = torch.matmul(attn, V)  # (batch, m, d_model)
        
        return out, attn


# ===== 测试 =====
if __name__ == "__main__":
    batch, m, n, d_model = 2, 3, 5, 8
    x_decoder = torch.randn(batch, m, d_model)  # 解码器：长度 3
    x_encoder = torch.randn(batch, n, d_model)  # 编码器：长度 5
    
    ca = CrossAttention(d_model)
    out, attn = ca(x_decoder, x_encoder)
    
    print("解码器输入:", x_decoder.shape)  # (2, 3, 8)
    print("编码器输入:", x_encoder.shape)  # (2, 5, 8)
    print("输出:", out.shape)              # (2, 3, 8) ------ 长度跟解码器一致
    print("注意力:", attn.shape)           # (2, 3, 5) ------ m × n 矩阵
    print("每行和:", attn.sum(dim=-1))     # 全为 1

（六）多头交叉注意力（工程实现）

python 复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadCrossAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x_decoder, x_encoder, mask=None):
        # x_decoder: (batch, m, d_model)
        # x_encoder: (batch, n, d_model)
        batch, m, _ = x_decoder.size()
        n = x_encoder.size(1)
        
        # 1. 投影 + 拆多头
        # Q 来自解码器
        Q = self.W_q(x_decoder).view(batch, m, self.num_heads, self.d_k).transpose(1, 2)
        # K、V 来自编码器
        K = self.W_k(x_encoder).view(batch, n, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x_encoder).view(batch, n, self.num_heads, self.d_k).transpose(1, 2)
        # Q: (batch, h, m, d_k)
        # K、V: (batch, h, n, d_k)
        
        # 2. 缩放点积
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        # scores: (batch, h, m, n)
        
        # 3. 掩码（可选，通常不需要）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # 4. softmax
        attn = F.softmax(scores, dim=-1)  # (batch, h, m, n)
        
        # 5. 加权求和
        out = torch.matmul(attn, V)  # (batch, h, m, d_k)
        
        # 6. 拼接多头
        out = out.transpose(1, 2).contiguous().view(batch, m, self.d_model)
        
        # 7. 输出投影
        out = self.W_o(out)  # (batch, m, d_model)
        
        return out, attn


# ===== 测试 =====
if __name__ == "__main__":
    batch, m, n, d_model, num_heads = 2, 4, 6, 12, 3
    x_decoder = torch.randn(batch, m, d_model)
    x_encoder = torch.randn(batch, n, d_model)
    
    mhca = MultiHeadCrossAttention(d_model, num_heads)
    out, attn = mhca(x_decoder, x_encoder)
    
    print("解码器输入:", x_decoder.shape)  # (2, 4, 12)
    print("编码器输入:", x_encoder.shape)  # (2, 6, 12)
    print("输出:", out.shape)              # (2, 4, 12)
    print("注意力:", attn.shape)           # (2, 3, 4, 6) ------ 每个头一个 m × n 矩阵

（七）完整解码器层

python 复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    """通用多头注意力：可用于自注意力或交叉注意力"""
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, query, key, value, mask=None):
        batch, m, _ = query.size()
        n = key.size(1)
        
        Q = self.W_q(query).view(batch, m, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch, n, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch, n, self.num_heads, self.d_k).transpose(1, 2)
        
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)
        
        out = out.transpose(1, 2).contiguous().view(batch, m, self.d_model)
        return self.W_o(out), attn


class PositionwiseFFN(nn.Module):
    """前馈网络"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
    
    def forward(self, x):
        return self.fc2(F.relu(self.fc1(x)))


class DecoderLayer(nn.Module):
    """Transformer 解码器层：三种组件协作"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        # 1. 掩码自注意力
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        # 2. 交叉注意力
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        # 3. FFN
        self.ffn = PositionwiseFFN(d_model, d_ff)
        
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, encoder_out, src_mask=None, tgt_mask=None):
        # ---- 1. 掩码自注意力 ----
        # Q、K、V 都来自解码器 x，带因果掩码
        attn_out, _ = self.self_attn(x, x, x, mask=tgt_mask)
        x = self.norm1(x + self.dropout(attn_out))
        
        # ---- 2. 交叉注意力 ----
        # Q 来自解码器 x，K、V 来自编码器 encoder_out
        cross_out, cross_attn = self.cross_attn(x, encoder_out, encoder_out, mask=src_mask)
        x = self.norm2(x + self.dropout(cross_out))
        
        # ---- 3. FFN ----
        ffn_out = self.ffn(x)
        x = self.norm3(x + self.dropout(ffn_out))
        
        return x, cross_attn


# ===== 测试完整解码器层 =====
if __name__ == "__main__":
    batch, tgt_len, src_len = 2, 4, 6
    d_model, num_heads, d_ff = 16, 4, 64
    
    # 模拟输入
    x = torch.randn(batch, tgt_len, d_model)           # 解码器输入
    encoder_out = torch.randn(batch, src_len, d_model) # 编码器输出
    
    # 构造因果掩码
    tgt_mask = torch.tril(torch.ones(tgt_len, tgt_len)).bool()
    tgt_mask = tgt_mask.view(1, 1, tgt_len, tgt_len)
    
    decoder_layer = DecoderLayer(d_model, num_heads, d_ff)
    out, cross_attn = decoder_layer(x, encoder_out, tgt_mask=tgt_mask)
    
    print("解码器输入:", x.shape)            # (2, 4, 16)
    print("编码器输出:", encoder_out.shape)  # (2, 6, 16)
    print("解码器输出:", out.shape)          # (2, 4, 16)
    print("交叉注意力:", cross_attn.shape)   # (2, 4, 4, 6) ------ h=4, m=4, n=6
相关推荐
国际学术会议-杨老师1 小时前
2027年人文发展、城市管理与数字化社会国际会议（HUCD 2027）
大数据·人工智能·人文发展
田里的水稻1 小时前
PyTorch_torchvision---数据集下载\TensorBoard可视化
人工智能·pytorch·python·深度学习·机器学习
科技每日热闻1 小时前
电气设备综合维保服务商怎么选？从预防性维护到数字化运维，看懂五个关键能力——企业选择维保合作伙伴，不能只看“坏了能不能修”，更要看能否覆盖资产与系统的长期运营
大数据·运维·人工智能
Dotrust东信创智1 小时前
HIL台架上的“无人夜班“：东信创智用AI重塑汽车软件测试全链路——测试不再是项目的瓶颈，而是交付的底气
人工智能·汽车
电子制造自留地1 小时前
AI服务器高多层PCB设计：阻抗控制、层间对位与孔金属化工程要点
大数据·网络·人工智能·科技·pcb工艺
QYR-分析1 小时前
汽车车窗调节器市场研究：成熟零部件赛道，模块化集成成为核心方向
人工智能·汽车
星云低代码开发平台1 小时前
ERP 与 CRM 客户信用核对：业务键、时点和权限验收
大数据·数据库·人工智能
TF男孩2 小时前
老AI人面对MCP、Skills等新名词
人工智能·程序员
AI Coder 码农王2 小时前
AgentScope Java v2 第2课，ReActAgent 的循环引擎
人工智能
热门推荐
01「全语录+省流版」邵艾伦 × 孙宇晨：聊风口、聊个人IP、聊年轻人如何抓住 AI 时代的机会02国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）03历年考研数学一、数学二、数学三真题试卷及答案PDF04GitHub 镜像站点05几个好用的ip纯净度检测网站06蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接07Telegram APK 下载安装完整指南 — 2026年最新08Windows 10/11 如何退出微软账户并改用本地账户登录（含找不到选项的解决方法）09苹果iPA游戏软件资源下载网站10微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新