Transformer 架构 - 注意力机制代码实现及测试

1 课程概览

本课实现自注意力机制代码。自注意力机制中 Q=K=V，计算公式为 softmax(Q·K^T / √d_k) · V。步骤：①定义查询张量特征维度 d_k；②计算原始注意力分数 scores = Q·K^T / √d_k；③掩码处理（可选）；④softmax 归一化；⑤dropout；⑥乘以 V 得到最终结果。

2 核心概念与定义

自注意力机制：Q=K=V 的注意力机制。
Q（Query）：查询张量，形状 (batch, seq_len, d_model)。
K（Key）：键张量，形状同 Q。
V（Value）：值张量，形状同 Q。
d_k：Q 的特征维度（最后一个维度的大小）。
scores：原始注意力分数。
mask：掩码张量，用于遮盖某些位置。
dropout：随机失活，防止过拟合。

3 模型与算法详解

自注意力机制计算流程

1. 定义 d_k = query.size(-1)
2. 计算 scores = Q·K^T / √d_k
3. 掩码处理（可选）：scores.masked_fill(mask == 0, -1e9)
4. softmax 归一化：p_attn = softmax(scores)
5. dropout：p_attn = dropout(p_attn)
6. 最终结果：output = p_attn · V

形状变化

步骤	形状	说明
Q, K, V	(2, 4, 512)	batch=2, seq_len=4, d_model=512
d_k	512	Q 的最后一个维度
Q·K^T	(2, 4, 4)	注意力分数矩阵
scores	(2, 4, 4)	缩放后
p_attn	(2, 4, 4)	softmax 后
output	(2, 4, 512)	最终结果

K 的转置

K.transpose(-2, -1) 交换最后两个维度。

原始 K：(2, 4, 512)
转置后：(2, 512, 4)
用负数索引：-2 和 -1，适用于任意维度

4 数学原理与推导

自注意力机制公式

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

步骤分解

计算注意力分数： $$\text{scores} = \frac{QK^T}{\sqrt{d_k}}$$
掩码处理（可选）： $$\text{scores}_i = \begin{cases} \text{scores}_i, & \text{if } \text{mask}_i \neq 0 \ -\infty, & \text{if } \text{mask}_i = 0 \end{cases}$$
softmax 归一化： $$\text{p_attn} = \text{softmax}(\text{scores})$$
最终结果： $$\text{output} = \text{p_attn} \cdot V$$

为什么要除以 √d_k？

防止点积结果过大，导致 softmax 梯度消失。

5 代码示例

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


def attention(query, key, value, mask=None, dropout=None):
    """
    自注意力机制计算
    :param query: 查询张量 (batch, seq_len, d_model)
    :param key: 键张量 (batch, seq_len, d_model)
    :param value: 值张量 (batch, seq_len, d_model)
    :param mask: 掩码张量（可选）
    :param dropout: dropout 层（可选）
    :return: (注意力输出, 注意力权重)
    """
    # 1. 定义查询张量特征维度 d_k
    d_k = query.size(-1)
    
    # 2. 计算原始注意力分数 scores = Q·K^T / √d_k
    # K.transpose(-2, -1): 交换最后两个维度
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 3. 掩码处理（可选）
    if mask is not None:
        # 将 mask=0 的位置替换为 -1e9
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # 4. softmax 归一化
    p_attn = F.softmax(scores, dim=-1)
    
    # 5. dropout（可选）
    if dropout is not None:
        p_attn = dropout(p_attn)
    
    # 6. 乘以 V 得到最终结果
    output = torch.matmul(p_attn, value)
    
    return output, p_attn


# 测试代码
def dm04_test_attention():
    """测试自注意力机制"""
    # 参数设置
    batch = 2
    seq_len = 4
    d_model = 512
    
    # 创建 Q, K, V（自注意力机制中 Q=K=V）
    query = torch.randn(batch, seq_len, d_model)
    key = torch.randn(batch, seq_len, d_model)
    value = torch.randn(batch, seq_len, d_model)
    
    # 创建 dropout 层
    dropout = torch.nn.Dropout(0.1)
    
    # 计算注意力
    output, p_attn = attention(query, key, value, dropout=dropout)
    
    print(f"Q 形状: {query.shape}")
    print(f"K 形状: {key.shape}")
    print(f"V 形状: {value.shape}")
    print(f"d_k: {d_model}")
    print(f"注意力权重形状: {p_attn.shape}")
    print(f"输出形状: {output.shape}")
    
    return output


if __name__ == "__main__":
    dm04_test_attention()

代码说明

代码	说明
`query.size(-1)`	获取最后一个维度的大小（d_k）
`key.transpose(-2, -1)`	交换最后两个维度
`torch.matmul(query, key.transpose(-2, -1))`	Q·K^T
`math.sqrt(d_k)`	√d_k
`scores.masked_fill(mask == 0, -1e9)`	掩码处理
`F.softmax(scores, dim=-1)`	softmax 归一化
`torch.matmul(p_attn, value)`	乘以 V