从零理解CAS卷积加性自注意力机制结构图：原理与实现详解

1次阅读

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

自注意力机制（Self-Attention）在深度学习领域扮演着越来越重要的角色，尤其是在自然语言处理（NLP）和计算机视觉（CV）任务中。传统的自注意力机制通过计算输入序列中各个位置之间的相关性，来捕捉长距离依赖关系。然而，传统的自注意力机制在高维数据（如图像）上计算复杂度较高，容易导致内存和计算资源的浪费。

CAS（Convolutional Additive Self-Attention）卷积加性自注意力机制是一种结合了卷积操作和自注意力机制的创新方法。它通过卷积操作降低计算复杂度，同时保留了自注意力机制捕捉全局依赖的能力。CAS 机制的核心思想是将卷积的局部感受野与自注意力的全局信息结合起来，从而在计算效率和模型性能之间取得平衡。

传统的自注意力机制通过计算查询（Query）、键（Key）和值（Value）之间的点积来获取注意力权重。其计算复杂度为 O(n²)，其中 n 是输入序列的长度。对于高维数据（如图像），这种计算方式会导致巨大的内存和计算开销。

CAS 机制通过以下方式优化了传统自注意力机制：

卷积操作 ：在计算注意力权重之前，先对输入进行卷积操作，减少特征维度，从而降低计算复杂度。
加性注意力 ：使用加性注意力（Additive Attention）代替点积注意力，进一步减少计算量。
局部与全局结合 ：通过卷积操作捕捉局部特征，再通过自注意力机制捕捉全局依赖，实现更高效的特征提取。

以下是一个简化的 CAS 机制结构图：

 输入 -> 卷积层 -> 加性注意力 -> 输出

以下是 CAS 机制的完整 PyTorch 实现代码，包含详细注释：

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

class CASAttention(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super(CASAttention, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=kernel_size//2)
        self.query = nn.Linear(out_channels, out_channels)
        self.key = nn.Linear(out_channels, out_channels)
        self.value = nn.Linear(out_channels, out_channels)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        # 卷积操作降低维度
        x_conv = self.conv(x)
        b, c, h, w = x_conv.shape
        x_flat = x_conv.view(b, c, -1).permute(0, 2, 1)  # (b, h*w, c)

        # 计算查询、键和值
        q = self.query(x_flat)  # (b, h*w, c)
        k = self.key(x_flat)    # (b, h*w, c)
        v = self.value(x_flat)  # (b, h*w, c)

        # 加性注意力
        attention = torch.tanh(q.unsqueeze(2) + k.unsqueeze(1))  # (b, h*w, h*w, c)
        attention = self.softmax(attention.sum(dim=-1))          # (b, h*w, h*w)

        # 加权求和
        out = torch.bmm(attention, v)  # (b, h*w, c)
        out = out.permute(0, 2, 1).view(b, c, h, w)

        return out