单头注意力.py

# 导入PyTorch库及其神经网络模块
import torch
import torch.nn as nn

# 定义自注意力机制类（版本2）
class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
        """
        初始化函数
        d_in: 输入特征维度（词嵌入向量长度）
        d_out: 输出特征维度（QKV向量空间维度）
        qkv_bias: 是否启用偏置项
        """
        super().__init__()
        # 定义三个独立的线性变换层，用于生成查询向量(Query)、键向量(Key)和值向量(Value)
        # W_query将输入从d_in维映射到d_out维：Q = XW_query^T + b
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        # 同理生成Key向量：K = XW_key^T + b
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        # 同理生成Value向量：V = XW_value^T + b
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

    def forward(self, x):
        """
        前向传播函数
        x: 输入张量形状为(seq_len, d_in)，seq_len为序列长度
        """
        # 通过线性变换生成Key向量 (seq_len, d_out)
        keys = self.W_key(x)
        # 生成Query向量 (seq_len, d_out)
        queries = self.W_query(x)
        # 生成Value向量 (seq_len, d_out)
        values = self.W_value(x)
        
        # 计算注意力分数矩阵：QK^T (seq_len, seq_len)
        # 每个元素表示query_i和key_j的点积相似度
        attn_scores = queries @ keys.T  # 矩阵乘法：(n,d) @ (d,m) => (n,m)
        
        # 计算注意力权重（应用softmax归一化）
        # 缩放因子：除以√d_k（key的维度）防止内积过大导致梯度消失
        # softmax(QK^T/√d_k) => (seq_len, seq_len)
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5,  # 分母为√d_out
            dim=-1  # 在最后一个维度进行softmax
        )
        
        # 计算上下文向量：注意力权重乘以Value向量
        # (seq_len, seq_len) @ (seq_len, d_out) => (seq_len, d_out)
        # 每个位置的输出向量是所有Value向量的加权平均
        context_vec = attn_weights @ values
        
        return context_vec

# 定义输入数据：6个词的嵌入向量（词表征）
inputs = torch.tensor(
    [[0.43, 0.15, 0.89], # 第1个词 "Your" 的嵌入向量
    [0.55, 0.87, 0.66],  # 第2个词 "journey" 的嵌入向量
    [0.57, 0.85, 0.64],  # 第3个词 "starts" 的嵌入向量
    [0.22, 0.58, 0.33],  # 第4个词 "with" 的嵌入向量
    [0.77, 0.25, 0.10],  # 第5个词 "one" 的嵌入向量
    [0.05, 0.80, 0.55]]  # 第6个词 "step" 的嵌入向量
)

# 获取输入维度d_in（词向量长度=3）
d_in = inputs.shape[-1]

# 设置输出维度d_out（QKV空间维度=2）
d_out = 2

# 设置随机种子以保证结果可复现
torch.manual_seed(789)

# 实例化自注意力模型
sa_v2 = SelfAttention_v2(d_in, d_out)

# 前向传播计算并打印结果
print(sa_v2(inputs))

更新日志

2025/6/26 11:30

查看所有更新日志

• dfb81-update于 2025/6/26
• c2341-update于 2025/6/16
• dc6b2-update于 2025/5/23

版权所有

版权归属：NateHHX

许可证：署名 4.0 国际 (CC-BY-4.0)