### 如何在 PyTorch 中实现注意力机制 #### 注意力机制的核心概念 注意力机制是一种用于捕捉输入序列中不同部分之间关系的技术。其基本思想是通过加权的方式突出重要特征，从而提高模型性能。以下是几种常见的注意力机制及其特点： 1. **自注意力机制 (Self-Attention)** 自注意力允许模型关注同一序列中的其他位置，以便更好地编码上下文信息[^2]。 2. **多头注意力 (Multi-Head Attention)** 多头注意力通过多个平行的注意力层捕获不同的子空间表示，增强了模型表达能力。 3. **点积注意力 (Dot-Product Attention)** 这种方法通过计算查询向量 \(Q\) 和键向量 \(K\) 的点积来获得注意力权重[^3]。 4. **加性注意力 (Additive Attention)** 或串联注意力 (Concat Attention) 使用非线性函数（如 tanh）处理拼接后的查询和键向量组合。 --- #### Self-Attention 的实现步骤 下面是一个简单的 Self-Attention 实现示例代码，展示了如何利用 PyTorch 计算注意力得分并生成输出。 ```python import torch import torch.nn as nn import math class ScaledDotProductAttention(nn.Module): def __init__(self, d_k): super(ScaledDotProductAttention, self).__init__() self.d_k = d_k def forward(self, Q, K, V, mask=None): # 计算注意力分数：(Query * Key.T) / sqrt(d_k) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # [batch_size, n_heads, seq_len, seq_len] if mask is not None: scores.masked_fill_(mask == 0, -1e9) # 对于被遮挡的位置设置为极小值 attn_weights = torch.softmax(scores, dim=-1) # 归一化得到注意力分布 output = torch.matmul(attn_weights, V) # 将注意力权重应用于 Value 向量 return output, attn_weights ``` 上述代码定义了一个缩放点积注意力模块 `ScaledDotProductAttention`，其中： - 输入张量 \(Q\), \(K\), \(V\) 分别代表查询、键和值； - 输出包括经过加权求和后的结果以及注意力权重矩阵。 --- #### Multi-Head Attention 的实现 为了进一步提升模型表现，可以通过引入多头注意力扩展单头注意力的功能。下面是完整的实现方式： ```python class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_k = d_model // num_heads self.num_heads = num_heads 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.fc_out = nn.Linear(d_model, d_model) self.attention = ScaledDotProductAttention(self.d_k) def split_heads(self, x, batch_size): """将输入分割成多个头部""" x = x.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # [batch_size, num_heads, seq_len, d_k] return x def combine_heads(self, x): """重新合并多个头部的结果""" batch_size, _, seq_len, _ = x.size() x = x.transpose(1, 2).contiguous().view(batch_size, seq_len, -1) # [batch_size, seq_len, d_model] return x def forward(self, Q, K, V, mask=None): batch_size = Q.size(0) # 线性变换获取 Query, Key, Value Q = self.split_heads(self.W_q(Q), batch_size) # [batch_size, num_heads, seq_len, d_k] K = self.split_heads(self.W_k(K), batch_size) V = self.split_heads(self.W_v(V), batch_size) # 执行缩放点积注意力操作 out, attn_weights = self.attention(Q, K, V, mask=mask) # [batch_size, num_heads, seq_len, d_k] # 合并多个头部并将结果传递给全连接层 out = self.combine_heads(out) # [batch_size, seq_len, d_model] out = self.fc_out(out) # [batch_size, seq_len, d_model] return out, attn_weights ``` 此代码实现了多头注意力机制的关键组件，包括分头 (`split_heads`) 和合头 (`combine_heads`) 操作。 --- #### 示例应用 假设我们有一个长度为 5 的序列嵌入维度为 8，则可以如下调用以上类完成一次前向传播： ```python # 参数配置 seq_length = 5 embedding_dim = 8 num_heads = 2 batch_size = 3 # 初始化输入数据 Q = torch.randn((batch_size, seq_length, embedding_dim)) K = torch.randn((batch_size, seq_length, embedding_dim)) V = torch.randn((batch_size, seq_length, embedding_dim)) # 创建 MultiHeadAttention 层实例 multi_head_attention = MultiHeadAttention(embedding_dim, num_heads=num_heads) # 前向传播 output, attention_weights = multi_head_attention(Q, K, V) print("Output shape:", output.shape) # 应该返回 [batch_size, seq_length, embedding_dim] print("Attention weights shape:", attention_weights.shape) # 返回 [batch_size, num_heads, seq_length, seq_length] ``` --- #### 总结 通过上述代码片段可以看出，在 PyTorch 中实现注意力机制并不复杂。无论是基础版还是高级版本都可以灵活调整参数适应具体应用场景需求[^1]。