滑动窗口的魔法：SW-MSA如何用拼图策略突破Transformer的视野限制-育师

滑动窗口的魔法：SW-MSA如何用拼图策略突破Transformer的视野限制

在计算机视觉领域，Transformer架构正经历着一场静默的革命。传统自注意力机制虽然能够捕捉全局依赖关系，但其平方级的计算复杂度却成为处理高分辨率图像的瓶颈。想象一下，当我们需要处理一张1024x1024像素的图像时，传统的自注意力机制需要计算超过100万对像素点之间的关系——这不仅消耗巨大计算资源，也使得模型难以训练和部署。

1. 窗口自注意力：从全局到局部的优雅妥协

窗口多头自注意力机制（W-MSA）的提出，犹如为Transformer戴上了一副"局部眼镜"。它将特征图划分为不重叠的M×M窗口，仅在每个窗口内部计算自注意力。这种设计带来了三个显著优势：

计算复杂度从O(n²)降至线性：对于h×w的特征图和窗口大小M，计算量从4hwC²+2(hw)²C降至4hwC²+2M²hwC
硬件友好性：规整的窗口划分更适合GPU并行计算
局部特征聚焦：强制模型先学习局部特征间的关联

def window_partition(x, window_size): """将输入特征图划分为不重叠窗口 参数: x: (B, H, W, C)格式的输入张量 window_size: 窗口大小(M) 返回: 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

然而，纯粹的W-MSA存在明显缺陷——窗口间的信息隔离。就像一群人被分隔在不同房间，每个房间内部可以自由交流，但房间之间却完全隔绝。这种设计虽然降低了计算量，却牺牲了模型捕捉长距离依赖的能力。

2. 滑动窗口的智慧：拼图解全局

滑动窗口多头自注意力（SW-MSA）的提出，巧妙地解决了窗口间的信息流通问题。其核心思想可以类比为拼图游戏：

窗口位移：将特征图在高度和宽度方向上各滑动⌊M/2⌋个像素
窗口重组：通过循环位移将边缘区域重新组合成完整窗口
掩码机制：确保不相邻的区域在计算注意力时不会产生关联

def calculate_mask(self, x_size): """计算SW-MSA所需的注意力掩码 参数: x_size: 输入特征图尺寸(H,W) 返回: attn_mask: (num_windows, Wh*Ww, Wh*Ww) """ H, W = x_size img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 h_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) w_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.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, self.window_size) # nW, ws, ws, 1 mask_windows = mask_windows.view(-1, self.window_size * self.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

这种设计的美妙之处在于：

计算效率：保持了W-MSA的线性计算复杂度
全局感知：通过连续的W-MSA和SW-MSA层，信息可以在不同窗口间流动
硬件友好：仍然保持规整的窗口计算模式

3. 相对位置编码：空间关系的优雅表达

在局部窗口中，精确的位置信息尤为重要。SW-MSA引入了相对位置偏置（Relative Position Bias），为注意力得分添加了与位置相关的偏置项：

$$ Attention(Q,K,V) = Softmax(\frac{QK^T}{\sqrt{d}} + B)V $$

其中B是基于相对位置的可学习参数。这种设计比绝对位置编码更适合视觉任务，因为它：

保持平移等变性
更好地建模局部空间关系
在窗口滑动时保持一致性

# 相对位置编码表初始化 self.relative_position_bias_table = nn.Parameter( torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 前向传播中使用 relative_position_bias = self.relative_position_bias_table[ self.relative_position_index.view(-1)].view( self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH