从零理解Transformer自注意力机制：LLM语义关联的核心实现

共计 1836 个字符，预计需要花费 5 分钟才能阅读完成。

自注意力机制是 Transformer 模型的核心组件，也是现代大型语言模型（LLM）能够捕捉文本长距离依赖关系的关键所在。对初学者来说，理解自注意力机制的工作原理是掌握 Transformer 架构的重要第一步。

1. 自注意力机制的数学原理

自注意力机制的核心思想是通过计算输入序列中每个位置与其他位置的关联程度（注意力分数），动态地为每个位置生成一个加权表示。具体来说，它涉及三个关键矩阵：Query（Q）、Key（K）和 Value（V）。

1. QKV 矩阵运算 ：
2. 输入序列 X（维度 [batch_size, seq_len, d_model]）通过三个不同的线性变换得到 Q、K、V 矩阵
3. 计算公式：Q = XW_Q, K = XW_K, V = XW_V

4. 其中 W_Q、W_K、W_V 是可学习的权重矩阵

5. 注意力分数计算 ：

6. 注意力分数 = softmax(QK^T/√d_k)
7. 这里除以√d_k（key 的维度）是为了防止点积结果过大导致 softmax 梯度消失

8. 得到的注意力权重矩阵维度为 [batch_size, seq_len, seq_len]

9. 上下文向量生成 ：

10. 最终输出 = 注意力权重 * V
11. 维度保持为 [batch_size, seq_len, d_model]

2. PyTorch 代码实现

下面是一个带 mask 的缩放点积注意力函数的实现示例：

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

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Args:
        Q: Query tensor [batch_size, seq_len, d_k]
        K: Key tensor [batch_size, seq_len, d_k]
        V: Value tensor [batch_size, seq_len, d_v]
        mask: Optional mask tensor [batch_size, seq_len, seq_len]
    Returns:
        context: [batch_size, seq_len, d_v]
        attention_weights: [batch_size, seq_len, seq_len]
    """
    d_k = Q.size(-1)
    # 计算注意力分数
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))

    # 应用 mask（如果有）if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    # 计算注意力权重
    attention_weights = F.softmax(scores, dim=-1)

    # 计算上下文向量
    context = torch.matmul(attention_weights, V)

    return context, attention_weights

3. 典型应用场景

1. 处理可变长输入 ：
2. 在实际应用中，输入序列长度往往不一致
3. 需要通过 padding（填充）使所有序列长度相同

4. 使用 padding mask 可以避免注意力机制关注这些填充位置

5. 多头注意力 ：

6. 多头注意力将 Q、K、V 拆分成多个头并行计算
7. 每个头可以关注序列的不同方面

8. 最后将各头的输出拼接并通过线性变换得到最终结果

9. 位置编码 ：

10. 自注意力机制本身没有位置信息
11. 需要通过位置编码（Positional Encoding）为输入添加位置信息
12. 常用的位置编码方式包括正弦 / 余弦函数和可学习的位置嵌入

4. 避坑指南

1. 梯度消失问题 ：
2. 深层 Transformer 容易出现梯度消失

3. 解决方案：在每层后添加 LayerNorm 和残差连接

4. 多头注意力优化 ：

5. 使用矩阵并行计算而非循环计算各头
6. 可以通过 einops 库简化多头注意力的实现

7. 示例：rearrange(x, 'b s (h d) -> b h s d', h=n_heads)

8. 内存优化 ：

9. 长序列会导致注意力矩阵占用大量内存
10. 可以考虑使用内存高效的注意力实现
11. 或者在训练时使用梯度检查点技术

5. 思考题

1. 为什么 Transformer 需要位置编码？如果去掉位置编码会有什么影响？
2. 如何修改上面的代码实现因果注意力（Causal Mask）？这种 mask 在什么场景下特别有用？
3. 在多头注意力中，为什么每个头的维度通常是总维度除以头数？如果保持每个头的维度不变会有什么问题？

通过本文的学习，你应该对 Transformer 自注意力机制有了基本的理解。建议尝试自己实现一个完整的 Transformer 编码器层，这将帮助你更深入地掌握这一重要技术。