深入解析Transformer中的多头自注意力机制：原理与实现

Transformer模型自2017年由Vaswani等人提出以来，已经成为自然语言处理（NLP）领域的一个里程碑。其核心机制之一——多头自注意力（Multi-Head Attention），为处理序列数据提供了前所未有的灵活性和表达能力。本文将详细解释Transformer中的多头自注意力机制是如何工作的，并提供代码示例。

1. Transformer模型简介

Transformer模型完全基于注意力机制，摒弃了传统的循环神经网络（RNN）结构，这使得模型能够并行处理序列数据，大大提高了训练效率。Transformer模型的关键组件包括编码器（Encoder）、解码器（Decoder）以及它们内部的多头自注意力机制。

2. 自注意力机制

自注意力机制的核心思想是，序列中每个元素都与其他所有元素相关，并且这种关系是通过注意力权重来表示的。自注意力机制可以捕捉序列内部的长距离依赖关系。

3. 多头自注意力的工作原理

多头自注意力是自注意力机制的扩展，它将输入分割成多个“头”，每个头学习输入的不同部分表示，然后将这些表示合并起来，以捕获信息的不同方面。

3.1 计算注意力权重

对于序列中的每个元素，多头自注意力首先计算其与序列中所有元素的关系（即注意力权重）。这通常通过以下公式完成：

[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V ]

其中，( Q )、( K )、( V ) 分别是查询（Query）、键（Key）和值（Value）矩阵，( d_k ) 是键的维度。

3.2 分割成多头

多头自注意力将查询、键和值线性投影到多个不同的空间，然后并行地计算每个头的注意力输出：

[ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O ]

每个头的输出都被拼接起来，并通过一个线性层进行投影，以整合不同头的信息。

4. 代码实现

以下是使用Python和PyTorch库实现多头自注意力机制的示例代码：

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

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

        assert self.head_dim * heads == embed_size, "Embed size needs to be divisible by heads"

        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # Split the embedding into self.heads different pieces
        values = values.reshape(N, value_len, self.heads, self.head_dim)
        keys = keys.reshape(N, key_len, self.heads, self.head_dim)
        queries = query.reshape(N, query_len, self.heads, self.head_dim)

        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(queries)

        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(N, query_len, self.heads * self.head_dim)
        out = self.fc_out(out)
        return out

# Example usage
embed_size = 256
heads = 8
attention_layer = MultiHeadAttention(embed_size, heads)

values = torch.rand(10, 50, embed_size)  # (batch_size, seq_len, embed_size)
keys = torch.rand(10, 50, embed_size)
queries = torch.rand(10, 20, embed_size)  # (batch_size, seq_len, embed_size)
mask = None  # Optional mask for padded sequences

output = attention_layer(values, keys, queries, mask)