从原理到实践：深入解析Transformer中的自注意力机制与移位窗口优化

1次阅读

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

Transformer 模型中的自注意力机制（Self-Attention）因其强大的序列建模能力而广受欢迎。然而，原始的自注意力机制存在一个显著的缺陷：计算复杂度。具体来说，自注意力机制的计算复杂度为 O(N^2)，其中 N 是输入序列的长度。这意味着当处理长序列时，计算资源的需求会急剧增加，导致训练和推理速度变慢，内存占用过高。

此外，常规的窗口注意力（Window Attention）虽然通过将输入序列划分为多个局部窗口来降低计算复杂度，但这种方法在跨窗口信息交互上存在局限性。窗口之间的信息流动受限，可能影响模型对全局依赖关系的捕捉能力。

移位窗口机制（Shifted Window Multi-Head Self-Attention, SW-MSA）是一种优化窗口注意力的方法。其核心思想是通过周期性移位窗口来实现跨窗口的信息交互。具体来说，SW-MSA 在连续的 Transformer 层中交替使用常规窗口划分和移位窗口划分。

常规窗口划分 ：将输入序列划分为不重叠的局部窗口，每个窗口内部进行自注意力计算。
移位窗口划分 ：将窗口进行周期性移位（例如，向右下方移动半个窗口大小），使得相邻窗口的部分区域重叠，从而促进跨窗口的信息流动。

这种交替使用的方式既保持了计算效率，又增强了模型捕捉长距离依赖的能力。

大核注意力（Large Kernel Attention, LKA）通过使用大卷积核来捕捉长程依赖关系。LKA 的核心在于将大卷积核分解为多个小卷积核的组合，例如将 7 ×7 卷积分解为 3 ×3 深度卷积（Depth-wise Convolution）和 3 ×3 空洞卷积（Dilated Convolution）。这种分解方式既能减少计算量，又能保持大卷积核的感受野。

数学上，LKA 可以表示为：

$$
\text{LKA}(X) = \text{DW-Conv}{3\times3}(\text{Dilated-Conv}(X))
$$

其中，DW-Conv 表示深度卷积，Dilated-Conv 表示空洞卷积。通过这种分解，LKA 能够有效地捕捉长距离依赖，同时保持较低的计算复杂度。

以下是移位窗口划分和还原的 PyTorch 代码实现：

import torch
import torch.nn as nn

def window_partition(x, window_size):
    """
    将输入划分为不重叠的窗口
    Args:
        x: (B, H, W, C)
        window_size: 窗口大小
    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows

def window_reverse(windows, window_size, H, W):
    """
    将窗口还原为原始特征图
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size: 窗口大小
        H: 特征图高度
        W: 特征图宽度
    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x

以下是 LKA 的 PyTorch 代码实现：

class LargeKernelAttention(nn.Module):
    def __init__(self, dim, kernel_size=7):
        super().__init__()
        self.dim = dim
        self.kernel_size = kernel_size

        # 深度卷积
        self.dw_conv = nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim)
        # 空洞卷积
        self.dilated_conv = nn.Conv2d(dim, dim, kernel_size=3, padding=2, dilation=2, groups=dim)

    def forward(self, x):
        """
        Args:
            x: (B, C, H, W)
        Returns:
            out: (B, C, H, W)
        """
        out = self.dw_conv(x)
        out = self.dilated_conv(out)
        return out