看完这一篇终于把 Transformer 中的注意力机制搞懂了！！建议收藏起来学！！

最新推荐文章于 2025-04-22 15:53:33 发布

AI大模型-派大星

最新推荐文章于 2025-04-22 15:53:33 发布

阅读量1k

点赞数 7

文章标签： transformer 深度学习人工智能开发语言 AI大模型语言模型 agi

本文链接：https://blog.youkuaiyun.com/2401_85327249/article/details/146414389

版权

前言

其核心思想是使用注意力机制来处理输入序列中的每个元素与其他元素之间的关系，从而捕捉全局的依赖关系。

Transformer 中的注意力机制有几种形式，包括自注意力机制、多头自注意力机制、掩蔽自注意力机制和交叉（编码器-解码器）自注意力机制。

自注意力机制

自注意力机制是 Transformer 中的关键机制，它通过计算输入序列中每个元素与其他元素的相关性，来捕捉序列中的全局依赖关系。

其基本原理是让每个输入元素（Token）与整个序列的所有其他元素进行交互，计算其对整个序列的注意力权重，并根据这些权重来生成新的表示。

计算步骤如下所示

线性变换

首先，将输入 X 通过三个不同的线性变换，生成查询 Q、键 K 和值 V 矩阵。

其中，是权重矩阵。
计算注意力权重

通过计算查询 Q 和键 K 之间的相似度，得到的结果是一个矩阵，表示了每个位置的查询向量与其他位置键向量的相似度。
缩放

为了避免数值过大，通常将相似度值除以，其中是键向量的维度。
归一化 (Softmax)

对相似度分数应用 Softmax 函数，将其转换为权重
加权和

最后，将得到的权重与值矩阵 V 进行加权求和，得到每个位置的最终表示。

这样，每个位置的输出就是一个包含了整个序列上下文信息的向量表示。

自注意力的优点

全局信息捕捉：通过自注意力机制，模型能够在每个位置考虑到其他位置的信息，从而能够捕捉到长期依赖关系。
并行计算：与 RNN 等传统方法不同，自注意力机制不依赖于序列的顺序，因此可以进行并行化计算。

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super(SelfAttention, self).__init__()
        self.embed_dim = embed_dim
        
        # 定义可训练参数
        self.W_Q = nn.Linear(embed_dim, embed_dim)
        self.W_K = nn.Linear(embed_dim, embed_dim)
        self.W_V = nn.Linear(embed_dim, embed_dim)
    
    def forward(self, x):
        """
        :param x: 输入序列 (batch_size, seq_len, embed_dim)
        :return: 自注意力后的输出 (batch_size, seq_len, embed_dim)
        """
        Q = self.W_Q(x)  # (batch_size, seq_len, embed_dim)
        K = self.W_K(x)  # (batch_size, seq_len, embed_dim)
        V = self.W_V(x)  # (batch_size, seq_len, embed_dim)

        # 计算注意力分数
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.embed_dim, dtype=torch.float32))
        attn_probs = torch.softmax(attn_scores, dim=-1)  # 归一化

        # 计算加权和
        attn_output = torch.matmul(attn_probs, V)  # (batch_size, seq_len, embed_dim)
        return attn_output

# 测试
batch_size, seq_len, embed_dim = 2, 5, 8
x = torch.rand(batch_size, seq_len, embed_dim)
self_attention = SelfAttention(embed_dim)
output = self_attention(x)
print("Self-Attention Output Shape:", output.shape)

多头自注意力机制

多头自注意力机制是对自注意力机制的扩展，它通过多个注意力头学习不同的特征子空间，以增强模型的表达能力。

具体来说，它将查询、键和值的线性变换拆分成多个头（多个子空间），然后每个头可以在不同的子空间中学习不同的注意力模式。最终，将所有头的输出拼接起来并通过线性变换得到最终的输出。

多头自注意力的计算过程如下。

将输入映射到多个头，每个头独立执行自注意力计算
拼接与线性变换

将所有注意力头的输出拼接起来，并通过线性变换得到最终输出。

其中，h 是头的数量，是输出的线性变换矩阵。

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads"
        
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads  # 每个头的维度
        
        # 线性变换矩阵
        self.W_Q = nn.Linear(embed_dim, embed_dim)
        self.W_K = nn.Linear(embed_dim, embed_dim)
        self.W_V = nn.Linear(embed_dim, embed_dim)
        self.W_O = nn.Linear(embed_dim, embed_dim)  # 最终输出变换

    def forward(self, x):
        batch_size, seq_len, embed_dim = x.shape
        
        # 计算 Q, K, V
        Q = self.W_Q(x)  # (batch_size, seq_len, embed_dim)
        K = self.W_K(x)  # (batch_size, seq_len, embed_dim)
        V = self.W_V(x)  # (batch_size, seq_len, embed_dim)

        # 变形为多头形式: (batch_size, num_heads, seq_len, head_dim)
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
        attn_probs = torch.softmax(attn_scores, dim=-1)

        # 计算加权和
        attn_output = torch.matmul(attn_probs, V)  # (batch_size, num_heads, seq_len, head_dim)

        # 恢复形状: 先transpose再reshape
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, embed_dim)

        # 通过最终的线性变换
        output = self.W_O(attn_output)  # (batch_size, seq_len, embed_dim)
        return output

# 测试
num_heads = 4
multihead_attention = MultiHeadAttention(embed_dim, num_heads)
output = multihead_attention(x)
print("Multi-Head Attention Output Shape:", output.shape)