BERT多头自注意力机制代码实现详解：从理论到实践

1次阅读

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

自注意力机制（Self-Attention）是 Transformer 架构的核心组件，也是 BERT 等预训练模型成功的关键。它通过计算输入序列中每个元素与其他元素的相关性，动态地学习上下文表示。相比传统的 RNN 和 CNN，自注意力机制能够更好地捕捉长距离依赖关系，并且具有天然的并行计算优势。

Query（Q）：表示当前需要关注的元素，用于与其他元素的 Key 进行匹配。
Key（K）：表示其他元素的标识，用于与 Query 计算相似度。
Value（V）：包含实际的信息内容，根据相似度权重进行加权求和。

多头注意力（Multi-Head Attention）是将输入线性投影到多个子空间，每个子空间独立计算注意力，最后将结果拼接。这种方式可以让模型同时关注不同位置的不同特征。

以下是一个完整的 BERT 多头自注意力机制的 PyTorch 实现，包含详细注释：

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # 线性变换层，用于生成 Q、K、V
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)

        # 输出层和 dropout
        self.out_linear = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size = x.size(0)

        # 生成 Q、K、V
        q = self.q_linear(x)
        k = self.k_linear(x)
        v = self.v_linear(x)

        # 分割多头
        q = q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # 应用 mask（如有）if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 计算注意力权重
        attention = torch.softmax(scores, dim=-1)
        attention = self.dropout(attention)

        # 加权求和
        context = torch.matmul(attention, v)

        # 合并多头
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)

        # 输出
        output = self.out_linear(context)

        return output, attention

# 示例用法
if __name__ == "__main__":
    embed_dim = 512
    num_heads = 8
    batch_size = 2
    seq_len = 10

    # 随机生成输入
    x = torch.randn(batch_size, seq_len, embed_dim)

    # 创建多头注意力层
    mha = MultiHeadAttention(embed_dim, num_heads)

    # 前向传播
    output, attention_weights = mha(x)

    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"注意力权重形状: {attention_weights.shape}")