BERT多头自注意力机制代码实现与性能优化实战

1次阅读

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

Transformer 架构之所以在 NLP 领域大放异彩，关键在于其核心组件——自注意力机制。这种机制能够让模型在处理序列数据时，动态地关注输入序列中不同位置的信息，从而捕捉长距离依赖关系。BERT 作为 Transformer 的代表模型之一，其强大的表征能力很大程度上得益于多头自注意力机制的设计。

多头自注意力机制的数学表达可以分为以下几个关键步骤：

QKV 计算：
输入序列经过三个不同的线性变换得到查询 (Query)、键(Key) 和值 (Value) 矩阵：
```
Q = XW_Q, K = XW_K, V = XW_V
```
其中 W_Q, W_K, W_V 是可训练的参数矩阵。
缩放点积注意力：
计算注意力权重并进行缩放：
```
Attention(Q,K,V) = softmax(QK^T/√d_k)V
```
这里 d_k 是键向量的维度，缩放因子√d_k 用于防止点积结果过大导致 softmax 梯度消失。
多头拼接：
将多个头的注意力输出拼接后通过线性变换：
```
MultiHead(Q,K,V) = Concat(head_1,...,head_h)W_O
```
其中每个头的计算都是独立的注意力机制。

下面是使用 PyTorch 实现多头自注意力机制的完整代码，包含详细的形状注释：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"

        # 初始化 QKV 和输出投影矩阵
        self.qkv_proj = nn.Linear(embed_dim, 3*embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, mask=None):
        """
        Args:
            x: 输入张量，形状为(batch_size, seq_len, embed_dim)
            mask: 可选，注意力掩码，形状为(batch_size, 1, 1, seq_len)
        Returns:
            输出张量，形状与输入相同
        """
        batch_size, seq_len, embed_dim = x.shape

        # 步骤 1：生成 QKV
        qkv = self.qkv_proj(x)  # (batch_size, seq_len, 3*embed_dim)
        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch_size, num_heads, seq_len, head_dim)
        q, k, v = qkv[0], qkv[1], qkv[2]  # 每个形状都是(batch_size, num_heads, seq_len, head_dim)

        # 步骤 2：计算缩放点积注意力
        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)

        # 步骤 3：应用注意力权重到 V
        output = torch.matmul(attn_weights, v)  # (batch_size, num_heads, seq_len, head_dim)

        # 步骤 4：拼接多头输出
        output = output.transpose(1, 2)  # (batch_size, seq_len, num_heads, head_dim)
        output = output.reshape(batch_size, seq_len, embed_dim)

        # 最终投影
        output = self.out_proj(output)
        return output

原始实现的复杂度主要体现在以下方面：
– 内存占用：存储中间注意力分数矩阵需要 O(batch_sizenum_headsseq_len^2)的空间
– 计算量：注意力分数的计算和 softmax 操作都是 O(seq_len^2)的复杂度

我们可以利用 PyTorch 的 einsum 函数来优化矩阵乘法：

# 替换原来的 matmul 计算
attn_scores = torch.einsum('bhid,bhjd->bhij', q, k) / (self.head_dim ** 0.5)

Flash Attention 是一种新型的注意力计算方式，可以显著减少内存访问次数：

try:
    from flash_attn import flash_attn_qkvpacked_func

    # 替换原有注意力计算
    output = flash_attn_qkvpacked_func(torch.stack([q,k,v], dim=2),
        dropout_p=0.0,
        softmax_scale=1.0/(self.head_dim ** 0.5),
        causal=False
    )
except ImportError:
    # 回退到原始实现
    pass