Transformer 注意力机制图解与 PyTorch 手写实现

一、注意力在解决什么问题

RNN 时代，序列靠隐状态传递信息，长距离依赖会衰减。注意力机制让每个位置直接看到所有位置，一步到位。

二、QKV 是什么

把注意力类比数据库查询：

Query（查询）：当前位置想找什么
Key（键）：每个位置能提供的"标签"
Value（值）：每个位置实际携带的信息

流程：用 Q 和所有 K 算相似度 → softmax 归一化成权重 → 对 V 加权求和。

三、缩放点积注意力

Python

import torch
import torch.nn as nn
import math


def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    # (batch, heads, seq_len, d_k)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    attn = torch.softmax(scores, dim=-1)
    return torch.matmul(attn, v), attn

为什么要除以 sqrt(d_k)：当 d_k 很大时，点积方差变大，softmax 会进入梯度极小的饱和区，导致训练困难。

四、多头注意力

Python

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        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)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # 线性变换并拆成多个头
        q = self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 注意力计算
        out, attn = scaled_dot_product_attention(q, k, v, mask)

        # 拼回
        out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.w_o(out), attn

五、完整 Transformer 层

Python

class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, num_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 残差 + LayerNorm
        attn_out, attn = self.attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))
        ffn_out = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_out))
        return x, attn

六、测试

Python

block = TransformerBlock()
x = torch.randn(2, 10, 512)   # batch=2, seq_len=10, d_model=512
out, attn = block(x)
print(out.shape)   # torch.Size([2, 10, 512])
print(attn.shape)  # torch.Size([2, 8, 10, 10])

七、复杂度分析

| 模型 | 每层复杂度 | 序列长度上限 | | --- | --- | --- | | RNN | O(n·d²) | 长序列梯度消失 | | Self-Attention | O(n²·d) | 显存受限，通常 2K-128K | | Linear Attention | O(n·d²) | 理论无限 |

小结

注意力的本质是带权重的加权平均，权重由内容相似度决定。理解 QKV 之后，后面的 FlashAttention、MQA、GQA 都是在这个框架上做工程优化。

一、注意力在解决什么问题

二、QKV 是什么

三、缩放点积注意力

四、多头注意力

五、完整 Transformer 层

六、测试

七、复杂度分析

小结

评论（0）