深入解析8头自注意力机制的两层Transformer：从原理到实现

1次阅读

没有评论

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

Transformer 模型自 2017 年由 Vaswani 等人提出后，已成为自然语言处理领域的基石架构。其核心创新在于完全依赖自注意力机制（Self-Attention）来建模序列关系，摒弃了传统的循环神经网络结构。

自注意力机制的本质是计算序列中每个元素与其他元素的关联权重。给定输入序列 $X \in \mathbb{R}^{n \times d}$（n 为序列长度，d 为特征维度），其计算过程可分为三步：

通过可学习的权重矩阵 $W_Q, W_K, W_V$ 分别生成查询（Query）、键（Key）和值（Value）矩阵
计算注意力分数 $\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V$
其中 $\sqrt{d_k}$ 的缩放操作是为了防止点积结果过大导致 softmax 梯度消失

多头注意力（Multi-Head Attention）是标准自注意力的扩展版本，其核心思想是将注意力机制并行执行多次（本例中为 8 次）。具体优势体现在：

并行捕获不同特征：每个注意力头可以关注输入序列的不同方面（如语法结构、语义关系等）
提高模型容量：通过增加参数量使模型能学习更复杂的模式
实验验证优势：在机器翻译任务中，8 头注意力比单头注意力 BLEU 值平均提升 2 - 3 个点

关键的计算差异在于：

输入特征被分割到 8 个头的子空间（假设原始维度 d =512，则每个头处理 64 维特征）
各头独立计算注意力后拼接结果：$\text{MultiHead}(Q,K,V) = \text{Concat}(head_1,…,head_8)W^O$
最终通过输出矩阵 $W^O$ 融合各头信息

以下是完整的实现代码（要求 PyTorch 1.8+）：

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        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, x):
        batch_size = x.size(0)

        # 线性变换并分头
        Q = self.W_Q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1,2)
        K = self.W_K(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1,2)
        V = self.W_V(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1,2)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = nn.Softmax(dim=-1)(scores)

        # 加权求和
        context = torch.matmul(attn, V)
        context = context.transpose(1,2).contiguous().view(batch_size, -1, self.num_heads*self.d_k)

        return self.W_O(context)

class TransformerLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, ff_dim=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 自注意力子层
        attn_output = self.self_attn(x)
        x = self.norm1(x + self.dropout(attn_output))

        # 前馈子层
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output))

        return x

class TwoLayerTransformer(nn.Module):
    def __init__(self, vocab_size=10000, d_model=512, num_heads=8):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([TransformerLayer(d_model, num_heads) 
            for _ in range(2)
        ])

    def forward(self, x):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x)
        return x

实际部署时需注意以下关键点：