从零理解(s)w-msa多头自注意力：原理剖析与PyTorch实战

1次阅读

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

传统多头自注意力（MSA）在处理长度为 $n$ 的序列时，需要计算所有位置对的注意力分数，导致内存复杂度高达 $O(n^2)$。这在处理图像（如将 224×224 图像展平为 50176 长度序列）或长文档时会立刻耗尽显存。

(s)w-msa（shifted window multi-head self-attention）通过两个关键设计解决该问题：

窗口划分：将特征图划分为 $M \times M$ 的非重叠窗口，每个窗口内独立计算注意力，复杂度降为 $O(M^2 \times n/M^2)=O(n)$
移位窗口 ：通过循环移位(cyclic shift) 使相邻窗口间能交换信息，避免完全隔离

对于每个窗口内的 $M \times M$patch，计算：
$$\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}} + B)V$$
其中 $B$ 是相对位置偏置，通常用可学习的 $\hat{B} \in \mathbb{R}^{(2M-1)\times(2M-1)}$ 通过索引映射得到。

关键步骤是处理循环移位后不同子窗口的拼接关系：

对特征图进行 $\lfloor M/2 \rfloor$ 的循环移位
计算注意力时，对来自不同原始窗口的 patch 添加 $-100$ 的掩码值
数学表达为：
$$\text{Mask}(i,j) = \begin{cases}
0 & \text{若}i,j\text{同属原窗口} \
-\infty & \text{否则}
\end{cases}$$

def window_partition(x: torch.Tensor, window_size: int):
    """
    输入: (B, H, W, C)
    输出: (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 cyclic_shift(x: torch.Tensor, shift_size: int):
    """
    输入: (B, H, W, C)
    输出: 移位后的同形状张量
    """
    shifted_x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
    return shifted_x

def create_mask(H: int, W: int, window_size: int, shift_size: int):
    """
    生成移位后的注意力掩码
    返回: (1, num_patches, num_patches)
    """
    img_mask = torch.zeros((1, H, W, 1))
    h_slices = [slice(0, -window_size),
                slice(-window_size, -shift_size),
                slice(-shift_size, None)]
    w_slices = [slice(0, -window_size),
                slice(-window_size, -shift_size),
                slice(-shift_size, None)]

    cnt = 0
    for h in h_slices:
        for w in w_slices:
            img_mask[:, h, w, :] = cnt
            cnt += 1

    mask_windows = window_partition(img_mask, window_size)
    mask_windows = mask_windows.view(-1, window_size * window_size)
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
    return attn_mask