什么是多头自注意力机制
时间: 2025-05-15 15:46:21 浏览: 15
### 多头自注意力机制的原理及作用
多头自注意力机制(Multi-Head Self-Attention Mechanism)是现代深度学习领域中一种核心技术,广泛应用于自然语言处理和其他序列建模任务。它的设计目标在于增强模型对不同子空间特征的学习能力,并提升计算效率。
#### 1. 原理概述
多头自注意力机制的核心思想是对输入向量进行多次线性变换,分别生成多个不同的表示子空间,在这些子空间上独立执行缩放点积注意力操作[^3]。具体来说:
- 输入向量被投影到三个矩阵 \( Q \) (Query),\( K \) (Key),和 \( V \) (Value) 中。
- 这些矩阵进一步划分为多个头部(Heads),每个头部对应一组新的参数权重。
- 对于每一个头部,按照缩放点积注意力公式计算加权后的上下文表示:
\[
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V
\]
其中 \( d_k \) 是 Key 的维度大小,用于防止梯度消失或爆炸问题[^1]。
完成所有头部运算之后,将各头部的结果拼接起来并通过另一个可训练的线性层得到最终输出[^2]:
```python
import torch.nn as nn
class MultiHeadSelfAttention(nn.Module):
def __init__(self, embed_size, heads):
super(MultiHeadSelfAttention, self).__init__()
assert embed_size % heads == 0
self.heads = heads
self.embed_size = embed_size
self.depth = embed_size // heads
self.values_linear = nn.Linear(embed_size, embed_size)
self.keys_linear = nn.Linear(embed_size, embed_size)
self.queries_linear = nn.Linear(embed_size, embed_size)
self.fc_out = nn.Linear(embed_size, embed_size)
def forward(self, values, keys, queries, mask=None):
N = queries.shape[0]
value_len, key_len, query_len = values.shape[1], keys.shape[1], queries.shape[1]
# Split into multiple heads
values = self.values_linear(values).view(N, value_len, self.heads, self.depth)
keys = self.keys_linear(keys).view(N, key_len, self.heads, self.depth)
queries = self.queries_linear(queries).view(N, query_len, self.heads, self.depth)
energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys]) / (self.depth ** 0.5)
if mask is not None:
energy = energy.masked_fill(mask == 0, float("-1e20"))
attention_scores = torch.softmax(energy, dim=3)
out = torch.einsum("nhql,nlhd->nqhd", [attention_scores, values]).reshape(
N, query_len, self.heads * self.depth
)
return self.fc_out(out)
```
#### 2. 主要作用
多头自注意力机制的作用主要体现在以下几个方面:
- 提升表达能力:通过并行化方式捕捉多种类型的依赖关系,从而更好地理解复杂的数据结构。
- 改善收敛速度:相比单一注意力头的设计,多头架构可以加速优化过程,减少过拟合风险。
- 更强鲁棒性:即使某些特定方向上的信号较弱或者存在噪声干扰,其他方向仍能提供有效补充信息。
---
阅读全文
相关推荐


















