深入理解CA注意力机制与多头注意机制：从原理到PyTorch实战

1次阅读

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

传统注意力机制（如 Self-Attention）在自然语言处理（NLP）和计算机视觉（CV）任务中存在两个主要问题：

长序列处理效率低下 ：计算复杂度为 O(n²)，当序列长度增加时，内存和计算资源消耗急剧上升。
空间信息捕捉不足 ：在 CV 任务中，传统注意力机制难以有效建模像素间的空间位置关系。

CA（Coordinate Attention）和多头注意力机制针对这些问题提出了改进方案。CA 注意力通过分解空间注意力到水平和垂直两个方向，有效降低了计算复杂度；多头注意力机制则通过并行计算多个注意力头，增强了模型的表达能力。

机制类型	计算复杂度	适用场景	计算效率
Self-Attention	O(n²)	通用序列建模	低
CA Attention	O(n)	空间敏感任务	高
Multi-Head	O(n²/k)	需要多视角建模任务	中等

CA 注意力由两个关键组件构成：

通道注意力 ：通过全局平均池化捕获通道间的依赖关系
位置注意力 ：分解为水平方向和垂直方向的位置编码

数学表达为：

$$CA(X) = Conv([P_h(X); P_v(X)]) \otimes C(X)$$

其中 $P_h$ 和 $P_v$ 分别表示水平和垂直位置注意力，$C$ 表示通道注意力。

多头注意力的核心思想是将输入投影到多个子空间并行计算注意力：

将 query/key/value 线性投影到 h 个不同的子空间
在每个子空间独立计算注意力
合并所有头的输出

公式表达：

$$MultiHead(Q,K,V) = Concat(head_1,…,head_h)W^O$$

其中每个头的计算为：

$$head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)$$

import torch
import torch.nn as nn

class CAAttention(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        # 通道注意力
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.channel_attention = nn.Sequential(nn.Linear(channel, channel // reduction),
            nn.ReLU(),
            nn.Linear(channel // reduction, channel)
        )

        # 位置注意力
        self.conv_h = nn.Conv2d(channel, channel, kernel_size=1)
        self.conv_v = nn.Conv2d(channel, channel, kernel_size=1)

    def forward(self, x):
        b, c, h, w = x.size()

        # 通道注意力
        y = self.avg_pool(x).view(b, c)  # [b, c]
        y = self.channel_attention(y).view(b, c, 1, 1)  # [b, c, 1, 1]

        # 水平位置注意力
        x_h = torch.mean(x, dim=3, keepdim=True)  # [b, c, h, 1]
        x_h = self.conv_h(x_h)  # [b, c, h, 1]

        # 垂直位置注意力
        x_v = torch.mean(x, dim=2, keepdim=True)  # [b, c, 1, w]
        x_v = self.conv_v(x_v)  # [b, c, 1, w]

        # 合并位置信息
        y = y * x_h.expand(-1, -1, -1, w) * x_v.expand(-1, -1, h, -1)

        return x * y.sigmoid()

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_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, q, k, v, mask=None):
        batch_size = q.size(0)

        # 线性投影并分头 [batch_size, seq_len, d_model] -> [batch_size, seq_len, n_heads, d_k]
        q = self.W_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        k = self.W_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        v = self.W_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 计算注意力分数 [batch_size, n_heads, seq_len_q, seq_len_k]
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)

        # 应用 mask（如有）if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 计算注意力权重
        attn_weights = torch.softmax(scores, dim=-1)

        # 加权求和 [batch_size, n_heads, seq_len_q, d_k]
        output = torch.matmul(attn_weights, v)

        # 合并多头 [batch_size, seq_len_q, n_heads * d_k]
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 最终投影
        return self.W_o(output)

使用 torch.profiler 分析多头注意力的内存使用情况：

with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
    profile_memory=True
) as prof:
    output = mha(q, k, v)

print(prof.key_averages().table(sort_by="self_cuda_memory_usage"))

在 PyTorch 中启用混合精度训练：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    output = model(inputs)
    loss = criterion(output, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()