### 关于 '十六个头是否比一个头好' 的代码实现
在探讨多头注意力机制（Multi-Head Attention, MHA）的有效性和必要性的过程中，研究者们提出了多个假设并进行了广泛的实验验证。具体到“Are Sixteen Heads Really Better than One？”这一课题上，项目提供了详细的代码实现来支持理论分析。
#### 多头注意力机制简介
多头注意力允许模型共同关注来自不同表示子空间的信息[^1]。这种设计使得每个头部可以专注于输入的不同方面，从而提高整体性能。对于具有16个头的MHA来说，其核心在于如何有效地分配这些不同的视角以捕捉更丰富的特征组合。
#### 代码实现细节
以下是基于Transformer架构中的一个多头自注意层的具体Python实现：
```python
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model=512, num_heads=8):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0
self.d_k = d_model // num_heads
self.num_heads = num_heads
# Linear layers to project queries, keys and values from input embeddings.
self.q_linear = nn.Linear(d_model, d_model)
self.v_linear = nn.Linear(d_model, d_model)
self.k_linear = nn.Linear(d_model, d_model)
# Output linear layer after concatenating all heads together again.
self.out = nn.Linear(d_model, d_model)
def forward(self, q, k, v, mask=None):
batch_size = q.size(0)
# Perform linear operation and split into h heads.
k = self.k_linear(k).view(batch_size, -1, self.num_heads, self.d_k)
q = self.q_linear(q).view(batch_size, -1, self.num_heads, self.d_k)
v = self.v_linear(v).view(batch_size, -1, self.num_heads, self.d_k)
# Transpose to get dimensions [batch_size, num_heads, seq_len, d_k].
k = k.transpose(1, 2)
q = q.transpose(1, 2)
v = v.transpose(1, 2)
# Calculate attention using function we will define next line.
scores = self.attention(q, k, v, self.d_k, mask, batch_size)
# Concatenate last two dimensions back together.
concat = scores.transpose(1, 2).contiguous().view(
batch_size, -1, self.num_heads * self.d_k)
output = self.out(concat)
return output
@staticmethod
def attention(query, key, value, d_k, mask=None, dropout=None):
"""
Scaled Dot Product Attention implementation.
"""
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
p_attn = F.softmax(scores, dim=-1)
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn
```
这段代码展示了如何构建一个多头注意力模块，并将其应用于给定序列上的查询、键和值向量之间计算相似度得分矩阵。值得注意的是，在实际应用中通常还会加入残差连接以及LayerNorm操作来稳定训练过程。
#### 实验与讨论
为了评估16个头相对于单一更大尺寸头的优势与否，研究人员对比了多种配置下的模型效果差异。结果显示增加更多独立的工作单元确实有助于提升某些任务的表现；然而并非总是越多越好——过多数量可能会引入冗余甚至降低泛化能力[^4]。因此，选择合适的头数应当依据具体的场景需求来进行调整优化。
扫一扫
884
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
