2. 注意力机制详解：Self-Attention原理、Multi-Head Attention、位置编码

好，咱们今天来啃一块硬骨头——注意力机制。说实话，我刚入行那会儿，看Transformer论文里的注意力公式，第一反应是「这什么鬼？」。但后来在项目中亲手搭过几版模型后，才真正体会到它的精妙。你想想看，让模型学会「该看哪里」，这本身就是件很酷的事。

2.1 Self-Attention：让模型学会「看哪里」

Self-Attention的核心思想，说白了就是让序列里的每个元素，都去跟其他所有元素「打个招呼」，然后根据亲疏远近重新表达自己。我习惯用一个比喻来理解：你在一个嘈杂的派对上，想听清朋友说话，你会自动忽略远处的噪音，把注意力集中在朋友身上。Self-Attention干的就是这个活。

具体怎么算？三个关键角色：Query、Key、Value。

• Query（查询）：当前元素想问的问题：「我跟谁有关系？」
• Key（键）：其他元素给出的标签：「我是谁，我有什么特征」
• Value（值）：其他元素携带的实际信息：「我身上带着的内容」

计算过程分三步走：

1. 算相似度：用Query跟所有Key做点积，得到注意力分数。分数越高，说明越「投缘」。
2. 归一化：用Softmax把分数变成概率分布，确保所有注意力权重加起来等于1。
3. 加权求和：用归一化后的权重去加权Value，得到最终的输出。

公式长这样：

Attention(Q, K, V) = softmax(Q × K^T / √d_k) × V

这里有个细节——为什么要除以√d_k？我在项目中踩过这个坑。如果不除，当向量维度d_k很大时，点积的结果会非常大，Softmax的梯度会变得极其小，模型根本学不动。除以√d_k相当于把方差拉回1附近，训练稳定很多。

2.2 Multi-Head Attention：多几个角度看问题

一个注意力头只能学到一种「关系模式」。但现实中，词与词之间的关系是多元的——比如「苹果」这个词，既跟「水果」有语义关系，又跟「公司」有品牌关系。一个头搞不定怎么办？那就多搞几个头。

Multi-Head Attention的做法很简单：

• 把Q、K、V分别投影到h个不同的子空间（h是头的数量）
• 在每个子空间里独立做Self-Attention
• 把h个头的输出拼接起来，再做一个线性变换

公式：

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) × W_O
其中 head_i = Attention(Q × W_Q_i, K × W_K_i, V × W_V_i)

我建议头数一般设8或16。太少学不到多样化的关系，太多又容易过拟合。我记得有一次调参，头数从8改成32，效果反而下降了——嗯，不是越多越好。

2.3 位置编码：给模型装上「空间感」

Self-Attention有个「致命」缺陷——它不关心顺序。你想想看，「我打你」和「你打我」，词袋是一样的，但意思完全相反。没有位置信息，模型就是个「词袋模型」。

怎么解决？给每个位置打上一个「位置标签」。有三种主流做法：

2.3.1 绝对位置编码

最经典的做法，来自原始Transformer论文。用正弦和余弦函数生成固定编码：

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

为什么用三角函数？因为不同频率的正余弦可以表示不同粒度的位置关系——低频编码捕捉全局位置，高频编码捕捉局部位置。我个人觉得这个设计很优雅，但有个问题：它只能表示「绝对位置」，无法建模「相对距离」。

2.3.2 相对位置编码

绝对位置编码的局限在于，模型很难学到「词A和词B相距多远」这种关系。相对位置编码直接建模两个token之间的偏移量。

具体做法是在计算注意力分数时，额外加上一个相对位置偏置：

score(i, j) = (Q_i × K_j^T) / √d + b_{i-j}

这里的b_{i-j}是可学习的参数，表示位置差为(i-j)时的偏置。我在做文本生成任务时，换成相对位置编码后，长距离依赖的建模能力明显提升。

2.3.3 旋转位置编码（RoPE）

这是目前最流行的方案，很多大模型（比如LLaMA、ChatGLM）都在用。它的思路很巧妙——不显式地加位置向量，而是把位置信息「旋转」进Q和K里。

具体来说，对每个位置pos，构造一个旋转矩阵R_pos，然后：

Q'_i = R_pos × Q_i
K'_j = R_pos × K_j

这样计算Q'_i × K'_j^T时，结果天然包含了位置差(pos_i - pos_j)的信息。好处很明显：

• 相对位置信息自然融入注意力计算
• 序列长度可以外推（训练时没见过那么长的，推理时也能用）
• 不会增加额外参数

2.4 代码实现：从零搭建注意力机制

光说不练假把式。咱们直接上代码，我习惯用PyTorch实现，清晰易懂。

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

class SelfAttention(nn.Module):
    def __init__(self, d_model, d_k, d_v):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_k)
        self.W_K = nn.Linear(d_model, d_k)
        self.W_V = nn.Linear(d_model, d_v)
        self.scale = math.sqrt(d_k)
    
    def forward(self, x):
        # x: [batch, seq_len, d_model]
        Q = self.W_Q(x)  # [batch, seq_len, d_k]
        K = self.W_K(x)
        V = self.W_V(x)
        
        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        attn_weights = F.softmax(scores, dim=-1)
        
        # 加权求和
        output = torch.matmul(attn_weights, V)
        return output, attn_weights

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_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, seq_len, _ = x.shape
        
        # 线性变换并拆分成多头
        Q = self.W_Q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        
        # 计算注意力
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn_weights = F.softmax(scores, dim=-1)
        context = torch.matmul(attn_weights, V)
        
        # 拼接多头输出
        context = context.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        output = self.W_O(context)
        return output

2.5 位置编码实现

class SinusoidalPositionEncoding(nn.Module):
    """绝对位置编码：正弦+余弦"""
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                           (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        return x + self.pe[:, :x.size(1), :]

class RotaryPositionEncoding(nn.Module):
    """旋转位置编码（RoPE）"""
    def __init__(self, d_model, max_len=2048):
        super().__init__()
        self.d_model = d_model
        # 预计算旋转矩阵
        theta = 10000 ** (-torch.arange(0, d_model, 2).float() / d_model)
        position = torch.arange(max_len).float().unsqueeze(1)
        cos_vals = torch.cos(position * theta).unsqueeze(0).unsqueeze(2)
        sin_vals = torch.sin(position * theta).unsqueeze(0).unsqueeze(2)
        self.register_buffer('cos_vals', cos_vals)
        self.register_buffer('sin_vals', sin_vals)
    
    def forward(self, q, k, pos_ids=None):
        # q, k: [batch, n_heads, seq_len, d_k]
        seq_len = q.size(2)
        cos = self.cos_vals[:, :seq_len, :]  # [1, seq_len, d_k//2]
        sin = self.sin_vals[:, :seq_len, :]
        
        # 将q和k分成两半，分别旋转
        q1, q2 = q.chunk(2, dim=-1)
        k1, k2 = k.chunk(2, dim=-1)
        
        q_rotated = torch.cat([q1 * cos - q2 * sin, q1 * sin + q2 * cos], dim=-1)
        k_rotated = torch.cat([k1 * cos - k2 * sin, k1 * sin + k2 * cos], dim=-1)
        
        return q_rotated, k_rotated

好了，注意力机制的核心内容就这些。从原理到代码，咱们都过了一遍。下一章咱们会把这些模块拼起来，搭一个完整的Transformer模型。到时候你会发现，理解了注意力机制，Transformer剩下的部分其实都是「锦上添花」。