Transformer注意力机制与FFN层深度解析

1. Transformer注意力机制深度解析

在自然语言处理领域，Transformer模型彻底改变了序列建模的范式。作为其核心组件，注意力机制通过动态权重分配实现了对输入序列的上下文感知处理。让我们从一个实际案例开始理解这个机制的价值：当模型处理句子"The animal didn't cross the street because it was too tired"时，传统的RNN结构难以准确判断"it"指代的是"animal"还是"street"，而注意力机制通过计算词与词之间的关联度，可以自动学习到"it"与"animal"的强相关性。

1.1 单头注意力机制详解

单头注意力是理解整个机制的基础，其计算过程可以分为三个关键阶段：

计算阶段分解：

1. 查询-键匹配阶段：通过点积计算(Q·K^T)建立词与词之间的关联强度
2. 缩放阶段：除以√dk进行数值稳定
3. 权重分配阶段：应用softmax生成概率分布
4. 值加权阶段：用注意力权重对V矩阵进行加权求和

python复制# 单头注意力计算示例代码
def scaled_dot_product_attention(Q, K, V):
    d_k = K.size(-1)  # 获取key的维度
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    attention_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, V)
    return output

维度缩放(√dk)的数学原理：
假设Q和K的每个元素是独立同分布的随机变量，均值为0，方差为1，那么点积Q·K^T的每个元素的方差就是dk。通过除以√dk，我们将方差重新缩放为1，防止梯度在反向传播时变得过大或过小。具体推导如下：

Var(∑qiki) = ∑Var(qiki) = dk·Var(qi)Var(ki) = dk (当Var(qi)=Var(ki)=1时)

实际训练中的现象观察：

• 未缩放时：随着dk增大，softmax输入值的极差可能达到数十甚至上百，导致梯度消失
• 适当缩放后：注意力权重分布更加平衡，模型收敛速度提升约30-50%
• 过度缩放：可能削弱模型捕捉长距离依赖的能力

1.2 多头注意力机制设计哲学

多头注意力的核心创新在于将单一的注意力过程分解为多个并行的"子注意力"，每个头学习不同的关注模式。这种设计源于三个关键认知：

1. 表示空间分解：将高维语义空间分解为多个子空间，允许模型在不同子空间学习不同的关系类型
2. 注意力多样性：强制模型建立多种关联模式，避免单一注意力机制的局限性
3. 组合爆炸优势：通过头的组合产生指数级的关系表示能力

具体实现流程：

1. 线性投影：将原始Q、K、V分别通过h个不同的线性变换矩阵投影
2. 分头计算：将投影后的矩阵分割成h个头
3. 并行注意力：每个头独立计算缩放点积注意力
4. 结果拼接：将所有头的输出拼接起来
5. 最终投影：通过线性层整合多头信息

python复制class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, h=8):
        super().__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        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, Q, K, V):