香草多头自注意力机制
时间: 2025-05-03 13:48:14 浏览: 35
### 香草多头自注意力机制的概念与实现
香草多头自注意力机制(Vanilla Multi-Head Self-Attention Mechanism)是一种广泛应用于深度学习中的技术,尤其在自然语言处理领域取得了显著的成功。它通过并行计算多个注意力头来捕捉输入序列的不同子空间特征,从而增强模型的表现力。
#### 多头自注意力机制的核心概念
多头自注意力机制基于单头注意力扩展而来,其核心思想是将输入向量映射到不同的线性变换空间中,并分别计算这些空间内的注意力权重[^1]。具体来说,对于给定的查询(Query)、键(Key)和值(Value),每个多头会独立地执行以下操作:
1. **线性投影**:每个头部会对 Query、Key 和 Value 进行一次线性变换,得到对应的 Q'、K' 和 V'。
2. **缩放点积注意力**:利用缩放点积公式 \( \text{Attention}(Q', K', V') = \text{softmax}(\frac{Q'(K'^T)}{\sqrt{d_k}})V' \),其中 \( d_k \) 是 Key 的维度大小。
3. **拼接与最终投影**:各头部的结果会被拼接在一起并通过另一个线性层进行转换,形成最终输出。
这种设计允许网络关注不同位置的信息组合方式,在诸如 Transformer 架构中有广泛应用案例。
#### 实现代码示例
以下是使用 PyTorch 编写的简单版本 Vanilla Multi-Head Self-Attention Module 的实现:
```python
import torch
import torch.nn as nn
import math
class MultiHeadSelfAttention(nn.Module):
def __init__(self, embed_dim, num_heads, dropout=0.1):
super(MultiHeadSelfAttention, self).__init__()
assert embed_dim % num_heads == 0
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
self.o_proj = nn.Linear(embed_dim, embed_dim)
self.dropout_layer = nn.Dropout(dropout)
def forward(self, x):
batch_size, seq_len, _ = x.size()
qkv = self.qkv_proj(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
query, key, value = qkv.unbind(dim=0)
attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(self.head_dim)
attention_probs = torch.softmax(attention_scores, dim=-1)
attention_probs = self.dropout_layer(attention_probs)
context_vector = (attention_probs @ value).transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
output = self.o_proj(context_vector)
return output
```
上述代码定义了一个标准的多头自注意模块,支持任意嵌入维度以及指定数量的头数配置。
#### 应用场景拓展
除了 NLP 中的经典应用外,该机制也被引入计算机视觉任务里,比如用于图像分类或者单目深度估计等问题解决上[^2]。某些研究还提出了时间解耦与时空分离型自我专注组件结合常规前馈神经网路构成两种独特变压器结构形式[^3]。
阅读全文
相关推荐





