深入解析BERT多头注意力机制：从原理到实践

1次阅读

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

在自然语言处理（NLP）领域，注意力机制（Attention Mechanism）已经成为现代神经网络架构的核心组件之一。它最早由 Google 团队在 2017 年提出的 Transformer 模型中引入，随后在 BERT、GPT 等预训练模型中大放异彩。注意力机制的核心思想是让模型能够 ” 关注 ” 输入序列中与当前任务最相关的部分，而不是像传统的 RNN 或 CNN 那样对所有输入一视同仁。

这种机制特别适合处理序列数据，因为它可以动态地为每个位置的输出分配不同的权重，从而捕捉长距离依赖关系。在机器翻译、文本摘要等任务中，注意力机制已经证明了其强大的能力。

多头注意力（Multi-Head Attention）是注意力机制的进阶版本，也是 BERT 模型的关键组成部分。它的基本思想是将注意力机制 ” 复制 ” 多份，每份都从不同的角度（即不同的表示子空间）学习输入序列的特征，最后将这些不同的视角综合起来。

首先，输入序列经过三个不同的线性变换得到查询（Q）、键（K）和值（V）矩阵。
然后，计算每个头的注意力分数：
$$Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V$$
将多个头的输出拼接起来，再经过一个线性变换得到最终输出。

这里的 $d_k$ 是键向量的维度，$\sqrt{d_k}$ 的缩放是为了防止点积变得过大而导致 softmax 函数的梯度消失。

下面是一个简单的多头注意力层的 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__()
        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads"

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

        # 定义线性变换层
        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):
        batch_size, seq_length, embed_dim = x.size()

        # 生成 Q,K,V 矩阵
        qkv = self.qkv_proj(x)
        qkv = qkv.reshape(batch_size, seq_length, self.num_heads, 3*self.head_dim)
        qkv = qkv.permute(0, 2, 1, 3)  # [batch, heads, seq_len, 3*head_dim]
        q, k, v = qkv.chunk(3, dim=-1)  # 分成 Q,K,V

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

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

        # 计算输出
        output = torch.matmul(attention, v)
        output = output.permute(0, 2, 1, 3)  # [batch, seq_len, heads, head_dim]
        output = output.reshape(batch_size, seq_length, embed_dim)

        # 最后的线性变换
        output = self.out_proj(output)
        return output

多头注意力机制虽然强大，但也带来了显著的计算开销，主要有以下几个方面需要考虑：