深入解析2D多头自注意力计算流程图：从原理到高效实现

1次阅读

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

自注意力机制（Self-Attention）是现代深度学习模型中的核心组件之一，尤其在 Transformer 架构中发挥了关键作用。它通过计算输入序列中每个元素与其他元素的关系权重，动态地聚合上下文信息，从而捕捉长距离依赖关系。2D 多头自注意力（2D Multi-Head Self-Attention）是自注意力机制的一种扩展形式，广泛应用于计算机视觉、自然语言处理和多模态任务中。

基本概念 ：自注意力机制的核心思想是通过计算查询（Query）、键（Key）和值（Value）之间的相似度，生成注意力分数，再根据这些分数加权聚合值向量。多头自注意力则通过将输入拆分为多个子空间（头），分别计算注意力，最后合并结果，增强了模型的表达能力。
重要性 ：在 Transformer 模型中，自注意力机制替代了传统的循环神经网络（RNN）和卷积神经网络（CNN），解决了长序列建模中的梯度消失和计算效率问题。2D 多头自注意力进一步扩展了这一能力，适用于图像、视频等二维数据。

尽管自注意力机制表现优异，但其传统实现方式存在一些显著问题：

计算复杂度高 ：自注意力机制的计算复杂度与序列长度的平方成正比（O(n²)），对于长序列或高分辨率图像，计算开销极大。
内存消耗大 ：存储注意力分数矩阵需要大量内存，尤其是在多头情况下，内存占用会成倍增加。
并行化难度 ：传统实现中，注意力分数的计算和聚合可能涉及复杂的矩阵操作，难以高效并行化。

这些问题限制了自注意力机制在大规模数据上的应用，因此需要优化实现方案。

2D 多头自注意力的计算流程可以分为以下几个关键步骤：

输入投影 ：将输入数据（如图像特征图）通过线性变换生成查询（Q）、键（K）和值（V）矩阵。
多头拆分 ：将 Q、K、V 矩阵按头的数量拆分，每个头独立计算注意力。
注意力分数计算 ：计算每个头的注意力分数，通常使用缩放点积注意力（Scaled Dot-Product Attention）。
注意力权重聚合 ：对注意力分数进行 Softmax 归一化，并加权聚合值矩阵。
多头合并 ：将所有头的输出拼接，并通过线性变换生成最终输出。

以下是计算流程图的伪代码表示：

# 输入: 特征图 x (B, C, H, W)
# 输出: 自注意力结果 out (B, C, H, W)

# 1. 投影生成 Q, K, V
q = linear_q(x)  # (B, C, H, W)
k = linear_k(x)  # (B, C, H, W)
v = linear_v(x)  # (B, C, H, W)

# 2. 多头拆分
q = split_heads(q, num_heads)  # (B, num_heads, C//num_heads, H*W)
k = split_heads(k, num_heads)  # (B, num_heads, C//num_heads, H*W)
v = split_heads(v, num_heads)  # (B, num_heads, C//num_heads, H*W)

# 3. 计算注意力分数
attn_scores = torch.matmul(q, k.transpose(-2, -1)) / sqrt(d_k)  # (B, num_heads, H*W, H*W)
attn_weights = F.softmax(attn_scores, dim=-1)

# 4. 加权聚合
out = torch.matmul(attn_weights, v)  # (B, num_heads, H*W, C//num_heads)

# 5. 多头合并
out = merge_heads(out)  # (B, C, H, W)
out = linear_out(out)  # 最终输出

以下是一个完整的 2D 多头自注意力实现，包含详细注释：

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

class MultiHead2DSelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # 线性变换层
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, C, H, W = x.shape
        x = x.flatten(2).transpose(1, 2)  # (B, H*W, C)

        # 生成 Q, K, V
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        # 多头拆分
        q = q.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn_weights = F.softmax(attn_scores, dim=-1)

        # 加权聚合
        out = torch.matmul(attn_weights, v)
        out = out.transpose(1, 2).contiguous().view(B, -1, self.embed_dim)

        # 输出投影
        out = self.out_proj(out)
        out = out.transpose(1, 2).view(B, C, H, W)
        return out

为了提升 2D 多头自注意力的计算效率，可以采用以下优化技巧：