一文搞懂自注意力（PyTorch）

自注意力机制是Transformer模型的核心组件，它允许模型在处理序列数据时，通过计算序列中不同位置元素之间的相关性得分，动态地调整对每个元素的关注程度，从而捕捉到序列内部的复杂依赖关系。

一、自注意力机制 

自注意力机制（Self- Attention）是什么？自注意力机制能够动态地捕捉序列中不同位置元素之间的依赖关系，并根据这些依赖关系生成新的序列表示。

它之所以被称为“自注意力”，是因为它在单一序列中通过计算序列元素之间的相互依赖关系来生成新的特征表示。这与传统的注意力机制有所不同，后者通常涉及两个序列之间的交互。

自注意力机制的查询和键来自同一组元素，即查询和键都是同一序列（如一句话中的词元或同一张图像中的不同patch）的特征，彼此之间进行注意力计算。

二、注意力分数 

如何实现注意力机制？在自注意力机制中，通过缩放点积计算注意力得分，并利用这些得分对值向量进行加权求和，从而实现了自注意力机制，它能够捕捉序列内部元素之间的依赖关系。 

 

import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
    def forward(self, x):
        # x shape: (batch_size, seq_len, embed_dim)
        Q = self.query(x)
        K = self.key(x)
        V = self.value(x)
        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.embed_dim ** 0.5)
        attention = F.softmax(scores, dim=-1)
        # 加权求和
        output = torch.matmul(attention, V)
        return output

 注意力分数是什么？注意力分数用来量化注意力机制中某一部分信息被关注的程度，反映了信息在注意力机制中的重要性。在注意力机制中，模型会根据注意力分数来决定对不同输入信息的关注程度。

Q（Query）、K（Key）、V（Value）通过映射矩阵得到相应的向量，通过计算Q与K的点积相似度并经过softmax归一化得到权重，最后使用这些权重对V进行加权求和得到输出。 

 

Q、K、V计算过程是什么？对于输入序列的每个单词，通过计算其Query与所有单词Key的点积得到注意力分数，经Softmax归一化后得到注意力权重，再用这些权重对Value向量进行加权求和，以得到包含丰富上下文信息的新单词表示。

import torch
def compute_attention_scores(Q, K, V):
    """
    计算自注意力分数
    输入:
        Q: Query矩阵, 形状为 [batch_size, seq_len, d_k]
        K: Key矩阵, 形状为 [batch_size, seq_len, d_k]
        V: Value矩阵, 形状为 [batch_size, seq_len, d_v]
    输出:
        output: 加权后的结果 [batch_size, seq_len, d_v]
        attention_weights: 注意力权重 [batch_size, seq_len, seq_len]
    """
    # 1. 计算Q与K的点积（相似度）
    scores = torch.matmul(Q, K.transpose(-2, -1))  # [batch_size, seq_len, seq_len]
    # 2. 缩放：除以sqrt(d_k)防止梯度消失
    d_k = K.size(-1)
    scores = scores / (d_k ** 0.5)
    # 3. 归一化：Softmax得到注意力权重
    attention_weights = torch.softmax(scores, dim=-1)
    # 4. 用权重对V加权求和
    output = torch.matmul(attention_weights, V)  # [batch_size, seq_len, d_v]
    return output, attention_weights
# 示例测试
if __name__ == "__main__":
    # 生成随机数据
    batch_size = 2
    seq_len = 3
    d_k = 4  # Q和K的维度
    d_v = 6  # V的维度
    Q = torch.randn(batch_size, seq_len, d_k)
    K = torch.randn(batch_size, seq_len, d_k)
    V = torch.randn(batch_size, seq_len, d_v)
    # 计算注意力
    output, weights = compute_attention_scores(Q, K, V)
    print("输入形状:")
    print(f"Q: {Q.shape}, K: {K.shape}, V: {V.shape}\n")
    print("输出形状:")
    print(f"加权结果: {output.shape}")
    print(f"注意力权重: {weights.shape}\n")
    print("注意力权重示例（第一个样本）:")
    print(weights[0].detach().numpy().round(2))

• 生成Q、K、V向量：对于输入序列中的每个单词，都会生成对应的Query（查询）、Key（键）和Value（值）向量。这些向量通常是通过将单词的嵌入向量（Embedding Vector）输入到一个线性变换层得到的。

• 计算Q、K的点积（注意力分数）：计算Query向量与序列中所有单词的Key向量之间的点积，得到一个分数。这个分数反映了Query向量与每个Key向量之间的相似度，即每个单词与当前位置单词的关联程度。

• Softmax函数归一化（注意力权重）：这些分数会经过一个Softmax函数进行归一化，得到每个单词的注意力权重。这些权重表示了在理解当前单词时，应该给予序列中其他单词多大的关注。

• 注意力权重加权求和（加权和向量）：这些注意力权重与对应的Value向量进行加权求和，得到一个加权和向量。这个加权和向量会被用作当前单词的新表示，包含了更丰富的上下文信息。

  ---