多头注意力与位置嵌入：从原理到实现的Transformer核心技术解析

1次阅读

共计 2398 个字符，预计需要花费 6 分钟才能阅读完成。

在自然语言处理（NLP）领域，自注意力机制已经成为现代模型的核心组件。传统的 RNN 和 CNN 在处理序列数据时存在长距离依赖问题，而自注意力机制通过直接计算序列中所有位置之间的关系，有效解决了这一难题。Transformer 模型的成功很大程度上归功于其创新的多头注意力和位置嵌入技术，这些技术使得模型能够并行处理序列数据并保留位置信息。

多头注意力是 Transformer 模型的核心创新之一，它允许模型在不同的子空间中学习不同的注意力模式。具体来说：

数学原理 ：给定查询(Q)、键(K) 和值 (V) 矩阵，注意力分数计算为：
```
Attention(Q,K,V) = softmax(QK^T/√d_k)V
```
其中 d_k 是键向量的维度。
多头实现：将 Q、K、V 投影到 h 个不同的子空间，分别计算注意力后拼接结果：
```
MultiHead(Q,K,V) = Concat(head_1,...,head_h)W^O
```
优势分析：
允许模型关注不同位置的子空间
提高了模型的表达能力
并行计算效率高

由于自注意力机制本身不考虑序列顺序，位置嵌入为模型提供了位置信息：

绝对位置编码：
使用正弦和余弦函数的固定模式
公式：PE(pos,2i)=sin(pos/10000^(2i/d_model))
相对位置编码：
关注元素间的相对位置关系
通常通过可学习参数实现

优点：
– 实现简单
– 不需要额外参数
– 可处理任意长度序列

缺点：
– 对长序列的泛化能力有限
– 难以捕捉精细的相对位置关系

优点：
– 更好地建模局部依赖
– 对序列长度变化更鲁棒

缺点：
– 实现复杂
– 增加模型参数
– 可能影响训练稳定性

以下是 PyTorch 实现的核心代码片段：

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        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.W_o = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # 线性变换并分头
        q = self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k)
        k = self.W_k(k).view(batch_size, -1, self.num_heads, self.d_k)
        v = self.W_v(v).view(batch_size, -1, self.num_heads, self.d_k)

        # 转置以获得维度 (batch_size, num_heads, seq_len, d_k)
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 注意力权重
        attn_weights = torch.softmax(scores, dim=-1)

        # 上下文向量
        context = torch.matmul(attn_weights, v)

        # 拼接多头输出
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.num_heads * self.d_k)

        return self.W_o(context)

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()

        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))

        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(1)]