Transformer模型中的自注意力机制（Self-Attention）是实现序列元素间依赖关系建模的核心组件。它通过计算序列中每个元素与其他所有元素之间的相关性权重，动态地为每个位置生成一个包含全局上下文信息的加权表示。下面将结合其工作原理、数学公式、具体步骤以及代码示例进行详细说明。 ### 一、自注意力机制的核心思想 自注意力机制允许模型在处理序列时，为序列中的每个位置（如单词或时间步）分配不同的“注意力”权重，这些权重决定了在生成当前位置的表示时，应“关注”序列中其他位置的多少信息。其核心优势在于能够直接捕获序列内部的长距离依赖关系，而无需像RNN那样依赖递归计算。 ### 二、自注意力机制的计算步骤 自注意力机制的计算过程主要涉及**查询（Query）**、**键（Key）** 和**值（Value）** 三个向量，通常通过线性变换从输入序列中派生。具体步骤如下： 1. **线性变换生成Q、K、V** 对于输入序列 \( X \in \mathbb{R}^{n \times d_{\text{model}}} \)（n为序列长度，d_model为模型维度），通过三个不同的权重矩阵 \( W^Q, W^K, W^V \) 分别进行线性变换，得到查询矩阵 \( Q \)、键矩阵 \( K \) 和值矩阵 \( V \)： \[ Q = X W^Q, \quad K = X W^K, \quad V = X W^V \] 其中，\( W^Q, W^K, W^V \in \mathbb{R}^{d_{\text{model}} \times d_k} \)（通常 \( d_k = d_v \)）。 2. **计算注意力分数** 通过计算查询向量与所有键向量的点积，得到未归一化的注意力分数矩阵： \[ \text{Scores} = Q K^T \] 点积结果反映了查询与键之间的相似度，相似度越高表示当前位置应更关注该键对应的值。 3. **缩放与Softmax归一化** 为防止点积结果过大导致梯度不稳定，将分数除以 \( \sqrt{d_k} \) 进行缩放，然后应用Softmax函数将分数转换为概率分布（注意力权重）： \[ \text{Attention Weights} = \text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right) \] 这一步确保了所有权重之和为1，且每个位置的权重反映了其相对重要性。 4. **加权求和得到输出** 将注意力权重与值矩阵 \( V \) 相乘，得到每个位置的加权输出： \[ \text{Output} = \text{Attention Weights} \cdot V \] 输出矩阵的每一行是原始序列中对应位置的上下文感知表示，它融合了序列中所有位置的信息。 ### 三、多头注意力（Multi-Head Attention） 为了增强模型捕捉不同子空间信息的能力，Transformer采用了**多头注意力**机制。具体而言： - 将Q、K、V分别投影到h个不同的子空间（即“头”），在每个头上独立执行上述自注意力计算。 - 将h个头的输出拼接起来，再通过一个线性变换得到最终输出。 **多头注意力的优势**： - 允许模型同时关注序列的不同方面（如语法、语义、指代关系等）。 - 通过并行计算提高效率。 ### 四、代码示例：自注意力机制的PyTorch实现 以下是一个简化的自注意力层实现，包含缩放点积注意力和多头注意力： ```python import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductAttention(nn.Module): """缩放点积注意力实现 [ref_1]""" def __init__(self, d_k): super().__init__() self.d_k = d_k def forward(self, Q, K, V, mask=None): # 计算注意力分数：Q * K^T / sqrt(d_k) [ref_1] scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32)) # 可选：应用掩码（如因果掩码） if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) # Softmax归一化得到注意力权重 [ref_1] attention_weights = F.softmax(scores, dim=-1) # 加权求和 output = torch.matmul(attention_weights, V) return output, attention_weights class MultiHeadAttention(nn.Module): """多头注意力实现 [ref_1]""" def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads == 0, "d_model必须能被num_heads整除" self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads # 线性变换层：生成Q、K、V及最终输出 self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) self.attention = ScaledDotProductAttention(self.d_k) def split_heads(self, x): """将输入张量拆分为多个头""" batch_size, seq_len, d_model = x.size() return x.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) def forward(self, Q, K, V, mask=None): batch_size = Q.size(0) # 线性变换并拆分多头 [ref_1] Q = self.split_heads(self.W_q(Q)) K = self.split_heads(self.W_k(K)) V = self.split_heads(self.W_v(V)) # 计算缩放点积注意力（每个头独立） attention_output, attention_weights = self.attention(Q, K, V, mask) # 合并多头输出 attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # 最终线性变换 output = self.W_o(attention_output) return output, attention_weights # 示例使用 if __name__ == "__main__": d_model = 512 num_heads = 8 seq_len = 10 batch_size = 4 # 模拟输入序列 x = torch.randn(batch_size, seq_len, d_model) # 初始化多头注意力层 mha = MultiHeadAttention(d_model, num_heads) # 前向传播 output, attn_weights = mha(x, x, x) # 自注意力：Q=K=V print(f"输入形状: {x.shape}") print(f"输出形状: {output.shape}") print(f"注意力权重形状: {attn_weights.shape}") ``` ### 五、自注意力机制如何建模依赖关系 1. **全局上下文感知**：每个位置的输出都是序列中所有位置的加权组合，因此任何两个位置间的依赖关系都能被直接捕获，无论它们相距多远[ref_1]。 2. **动态权重分配**：注意力权重是数据驱动的，模型根据当前输入动态决定哪些位置更重要。例如在句子“The animal didn't cross the street because it was too tired”中，模型可以通过自注意力机制学习到“it”应更关注“animal”[ref_1]。 3. **并行计算与长序列处理**：与RNN的序列化处理不同，自注意力可以并行计算所有位置的表示，极大提升了长序列的处理效率[ref_1]。 ### 六、自注意力机制的应用与变体 | 注意力类型 | 计算方式 | 应用场景 | 特点 | |------------|---------|----------|------| | **自注意力** | Q、K、V均来自同一序列 | Transformer编码器 | 捕获序列内部依赖关系[ref_1] | | **交叉注意力** | Q来自解码器，K、V来自编码器 | Transformer解码器 | 实现编码器-解码器间的信息传递[ref_1] | | **因果注意力** | 带掩码的自注意力 | 自回归生成任务 | 防止未来信息泄露，保证自回归性[ref_1] | ### 七、总结 自注意力机制通过查询-键-值的三元组计算，实现了序列元素间依赖关系的直接建模。其核心在于： - **可并行化计算**，适合处理长序列。 - **动态权重分配**，能够捕捉复杂的上下文依赖。 - **多头扩展**，增强了模型在不同表示子空间中的建模能力。 这些特性使Transformer在自然语言处理、时间序列预测等任务中表现出色，成为当前深度学习的基石架构之一[ref_1][ref_2][ref_6]。在实际应用中，还需结合位置编码、残差连接和层归一化等技术，进一步提升模型的训练稳定性和表达能力。