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

1次阅读

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

在自然语言处理（NLP）领域，传统的 RNN 和 CNN 模型在处理长序列时存在明显的局限性。RNN 虽然能够捕捉序列信息，但由于其顺序计算特性，难以并行化处理，导致训练效率低下。而 CNN 虽然可以通过卷积核捕捉局部特征，但对于全局依赖关系的建模能力较弱。

注意力机制的提出，为解决这些问题提供了新的思路。它允许模型在处理每个位置时，动态地关注序列中的其他位置，从而更好地捕捉长距离依赖关系。Transformer 模型正是基于这一思想，通过自注意力机制和多头注意力机制，实现了高效的序列建模。

自注意力机制的核心思想是通过计算序列中每个位置与其他位置的相关性，动态地分配注意力权重。具体来说，自注意力机制通过三个矩阵（Query、Key、Value）来计算注意力分数。

QKV 矩阵 ：输入序列经过线性变换得到 Query（Q）、Key（K）、Value（V）三个矩阵。
缩放点积注意力 ：计算 Q 和 K 的点积，然后除以一个缩放因子（通常是 K 的维度平方根），再经过 softmax 归一化得到注意力权重。
加权求和 ：用注意力权重对 V 进行加权求和，得到最终的输出。

数学公式如下：

Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V

多头注意力机制是自注意力机制的扩展，通过并行计算多个自注意力头，然后将结果拼接起来，再经过线性变换得到最终输出。这种方式可以让模型同时关注不同子空间的信息，从而提升表达能力。

数学公式如下：

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 SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (self.head_dim * heads == embed_size), "Embedding 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

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(MultiHeadAttention, self).__init__()
        self.attention = SelfAttention(embed_size, heads)
        self.norm = nn.LayerNorm(embed_size)
        self.feed_forward = nn.Sequential(nn.Linear(embed_size, embed_size),
            nn.ReLU(),
            nn.Linear(embed_size, embed_size),
        )
        self.dropout = nn.Dropout(0.1)

    def forward(self, value, key, query, mask):
        attention = self.attention(value, key, query, mask)
        x = self.dropout(self.norm(attention + query))
        forward = self.feed_forward(x)
        out = self.dropout(self.norm(forward + x))
        return out

自注意力机制和多头注意力机制的计算复杂度与序列长度的平方成正比，这在处理长序列时会带来显著的计算和内存开销。以下是一些优化策略：