深入解析BERT双向自注意力机制：从数学原理到PyTorch实现

1次阅读

没有评论

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

在自然语言处理（NLP）领域，传统 RNN（循环神经网络）曾是处理序列数据的首选。然而，RNN 存在一些固有缺陷，尤其是在处理长序列时：

梯度消失 / 爆炸问题：RNN 在反向传播时，梯度需要通过时间步逐步传递，导致长距离依赖难以学习。
顺序计算限制：RNN 必须按时间步依次计算，无法充分利用现代 GPU 的并行计算能力。
信息瓶颈：RNN 的隐藏状态需要压缩所有历史信息，容易丢失关键细节。

这些问题促使了自注意力机制的诞生，它通过直接建模序列中所有位置的关系，解决了上述痛点。BERT 作为基于 Transformer 的模型，其核心正是双向自注意力机制。

自注意力机制的核心是计算查询（Query）、键（Key）和值（Value）矩阵的交互。以下是关键步骤的数学表达：

QKV 矩阵计算：
[
Q = XW_Q, \quad K = XW_K, \quad V = XW_V
]
其中，(X)是输入序列，(W_Q, W_K, W_V)是可学习的权重矩阵。
缩放点积注意力：
[
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
]
缩放因子 (\sqrt{d_k})（(d_k) 是键的维度）用于防止点积过大导致梯度消失。
多头注意力拼接：
[
\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W_O
]
每个头独立计算注意力后拼接，再通过 (W_O) 投影到输出空间。

以下是工业级优化的 AttentionLayer 类实现，包含关键注释和优化技巧：

import torch
import torch.nn as nn
import torch.nn.functional as F

class AttentionLayer(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__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

        # Linear projections for Q, K, V
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, embed_dim = x.shape
        assert embed_dim == self.embed_dim, "Input embedding dim must match layer embed_dim"

        # Project Q, K, V and split into heads
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot-product attention
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        attn_output = torch.matmul(attn_weights, v)

        # Merge heads and project
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        return self.out_proj(attn_output)