Transformer架构中的(s)w-msa多头自注意力机制详解与性能优化实践

1次阅读

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

传统多头注意力 (MHA) 在处理长度为 $n$ 的序列时，由于需要计算所有 query-key 对的相似度，其时间和空间复杂度均为 $O(n^2)$。当序列长度达到 2048 甚至更长时（如文档级 NLP 任务或高分辨率图像处理），显存消耗会呈现平方级增长，导致训练过程难以进行。

滑动窗口多头自注意力(s)w-msa 通过将序列划分为固定大小的窗口（通常 128-256 个 token），仅在窗口内部计算注意力，将复杂度降低到 $O(n)$。这种设计既保留了局部上下文的建模能力，又通过窗口滑动机制（shifted window）实现了跨窗口的信息交互。

方法类型	计算复杂度	内存占用	适用场景
全局 MHA	O(n²)	极高	短序列任务
稀疏 Attention	O(n√n)	高	特定模式的长序列
(s)w-msa	O(n)	中等	通用长序列任务

测试环境：A100-40GB GPU，batch_size=32

序列长度	全局 MHA 显存	(s)w-msa 显存(窗口 =128)
512	12.1GB	5.3GB
1024	OOM	8.7GB
2048	OOM	14.2GB

import einops

def window_partition(x, window_size):
    """
    Args:
        x: (B, L, C)
        window_size: int
    Returns:
        windows: (num_windows*B, window_size, C)
    """
    B, L, C = x.shape
    x = einops.rearrange(
        x, 
        'b (nw ws) c -> (b nw) ws c', 
        ws=window_size
    )
    return x

window_reverse = lambda windows, window_size, L: einops.rearrange(
    windows,
    '(b nw) ws c -> b (nw ws) c',
    ws=window_size
)

每个注意力头引入可学习的相对位置偏置矩阵 $B \in \mathbb{R}^{(2M-1)×(2M-1)}$（M 为窗口大小）：

$$Attention = Softmax(\frac{QK^T}{\sqrt{d_k}} + B)V$$

# 初始化偏置矩阵
self.relative_position_bias_table = nn.Parameter(torch.zeros((2*window_size-1)**2, num_heads)
)

# 计算相对位置索引
coords = torch.stack(torch.meshgrid([torch.arange(window_size)]*2
))
relative_coords = coords[:, None, :] - coords[:, :, None]  # 2, M, M
relative_coords += window_size - 1  # 转换为非负

通过周期性向左上方移动窗口，使得相邻层的窗口布局不同，实现跨窗口信息流动：

def shifted_window_attention(x, window_size, shift_size):
    if shift_size > 0:
        # 周期性位移
        x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))

    # 常规窗口注意力计算
    attn_mask = create_mask(window_size, shift_size)
    ...

    if shift_size > 0:
        # 位移还原
        x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
    return x

在 WMT14 英德翻译任务上的测试结果：