CAS卷积加性自注意力机制结构图解析与高效实现方案

1次阅读

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

传统自注意力机制在长序列处理时存在 O(n²) 的计算复杂度问题，这导致其在处理长文本或高分辨率图像时显存消耗急剧增加。同时，纯卷积网络虽然计算高效，但在全局建模能力上存在明显局限。CAS 机制通过结合卷积的局部性和自注意力的全局性，提供了一种高效的混合架构。

CAS 结构由三个核心组件构成：

卷积特征提取层 ：采用多尺度卷积核捕获局部特征
加性注意力计算层 ：使用加性注意力公式计算全局关系
门控融合层 ：动态调整局部和全局特征的融合比例

简易 ASCII 结构图示意：

Input
  │
  ▼
[Conv Feature Extraction]
  │
  ▼
[Additive Attention]
  │
  ▼
[Gated Fusion]
  │
  ▼
Output

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

class CASLayer(nn.Module):
    def __init__(self, dim, heads, kernel_size=3):
        super().__init__()
        # 卷积特征提取
        self.conv = nn.Sequential(nn.Unfold(kernel_size, padding=kernel_size//2),
            nn.Linear(kernel_size**2 * dim, dim)
        )
        # 加性注意力参数
        self.W_q = nn.Linear(dim, dim)
        self.W_k = nn.Linear(dim, dim)
        # 门控融合
        self.gate = nn.Linear(2*dim, dim)

    def forward(self, x):
        b, n, d = x.shape
        # 1. 卷积特征提取
        conv_feat = self.conv(x.transpose(1,2)).transpose(1,2)

        # 2. 加性注意力计算
        q = self.W_q(x)  # [b,n,d]
        k = self.W_k(x)  # [b,n,d]
        # 计算注意力分数 a_ij = ReLU(W_q*q_i + W_k*k_j)
        attn = F.relu(q.unsqueeze(2) + k.unsqueeze(1))  # [b,n,n,d]
        attn = attn.mean(-1)  # 降维到 [b,n,n]
        attn = F.softmax(attn, dim=-1)

        # 3. 门控融合
        global_feat = torch.bmm(attn, x)
        gate = torch.sigmoid(self.gate(torch.cat([conv_feat, global_feat], -1)))
        return gate * conv_feat + (1-gate) * global_feat

在 SQuAD 1.1 数据集上的测试结果：