Transformer模型中自注意力机制的计算复杂度分析

技术背景

在自然语言处理领域，传统的基于循环神经网络（RNN）的序列编码层在处理长序列时存在效率和性能问题。Google Research提出的Transformer模型，通过自注意力机制完全替代了传统的RNN层，为机器翻译等任务带来了显著的性能提升。在Transformer论文的表1中，作者比较了不同序列编码层的计算复杂度，并指出当序列长度n小于向量表示的维度d时，自注意力层比RNN层更快。然而，实际的计算复杂度似乎与论文中的说法存在差异。

实现步骤

设X是自注意力层的输入，其形状为(n, d)，其中n是词向量的数量（对应行数），d是每个词向量的维度。计算自注意力层的输出需要以下步骤（为简化起见，考虑单头自注意力）：

线性变换：将X的行进行线性变换，计算查询矩阵Q、键矩阵K和值矩阵V，每个矩阵的形状均为(n, d)。这通过将X与3个形状为(d, d)的学习矩阵进行后乘来实现，计算复杂度为$O(n d^2)$。
计算层输出：根据论文中的公式1，计算SoftMax(Q Kt / sqrt(d)) V，其中softmax是按行计算的。计算Q Kt的复杂度为$O(n^2 d)$，将结果与V进行后乘的复杂度同样为$O(n^2 d)$。

核心代码

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

import torch
import torch.nn as nn

class SelfAttention(nn.Module):
    def __init__(self, d):
        super(SelfAttention, self).__init__()
        self.W_q = nn.Linear(d, d)
        self.W_k = nn.Linear(d, d)
        self.W_v = nn.Linear(d, d)

    def forward(self, X):
        Q = self.W_q(X)
        K = self.W_k(X)
        V = self.W_v(X)

        d_k = Q.size(-1)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
        attention_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, V)
        return output