Transformer架构中2D多头自注意力计算流程图的实现与优化

1次阅读

没有评论

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

自注意力机制是 Transformer 架构的核心组件，它通过对输入序列中不同位置的关系进行建模，实现了对长距离依赖的捕捉。然而，传统的自注意力计算存在两个主要问题：

计算效率低下：自注意力计算的时间复杂度为 O(n^2)，其中 n 是序列长度。对于长序列，计算量会急剧增加，导致训练和推理速度变慢。
内存占用过高：自注意力计算需要存储大量的中间结果，尤其是多头注意力机制中，每个头都需要独立的计算和存储，这进一步加剧了内存压力。

针对上述问题，常见的优化方案包括原始实现、分块计算和并行计算。以下是它们的优缺点对比：

原始实现：
优点：实现简单，易于理解。
缺点：计算效率低，内存占用高。
分块计算：
优点：通过将大矩阵分块处理，减少内存占用。
缺点：增加了计算复杂度，可能引入额外的开销。
并行计算：
优点：利用多线程或多 GPU 加速计算，显著提升效率。
缺点：实现复杂，需要处理线程同步问题。

综合考虑，我们选择 矩阵分块和并行计算 的组合方案，既能降低内存占用，又能提升计算效率。

在多头自注意力中，输入序列通过线性变换生成查询（Q）、键（K）和值（V）矩阵。具体步骤如下：

将输入序列 X 分别与权重矩阵 W_Q、W_K、W_V 相乘，得到 Q、K、V。
将 Q、K、V 按头数分割成多个子矩阵，每个子矩阵对应一个注意力头。

对于每个注意力头，计算注意力权重：

计算 Q 和 K 的点积，得到注意力分数。
对注意力分数进行缩放（除以 sqrt(d_k)），其中 d_k 是键向量的维度。
对缩放后的分数应用 softmax 函数，得到注意力权重。

将每个头的注意力权重与 V 相乘，得到加权后的值。
将所有头的加权值拼接起来，通过线性变换得到最终输出。

以下是基于 PyTorch 的实现代码：

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

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

        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch_size, seq_len, d_model = x.size()

        # 生成 Q, K, V
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)

        # 分头
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn_weights = F.softmax(scores, dim=-1)

        # 加权求和
        output = torch.matmul(attn_weights, V)

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        output = self.W_O(output)

        return output

我们对比了原始实现和优化后的实现在不同序列长度下的性能：