Swin Transformer 详解:窗口注意力如何成为新一代视觉 backbone

摘要

ViT 把图像切成 patch 直接做全局自注意力,效果虽好,但序列长度随分辨率平方增长,高分辨率场景根本算不动,而且全程单一分辨率的特征对检测、分割这类密集预测任务并不友好。Swin Transformer(微软亚洲研究院,2021)用两个设计回答了这个问题:窗口内注意力控制计算量,移位窗口(Shifted Window)补上窗口间的信息交互,再配合分层下采样搭出金字塔特征。这篇文章沿着数据在 Swin 里的流动顺序,从 Patch Embedding 一路讲到 Patch Merging,把每一步的张量形状拆开来看,最后结合官方代码仓库的配置和训练设定,聊聊它为什么能同时当分类 backbone 和检测、分割的底座。

一、背景:ViT 的两个麻烦

ViT(Vision Transformer)证明了纯 Transformer 架构在图像分类上完全能打,但它的做法很"直白":把图像切成 16×16 的 patch,每个 patch 拉平成一个 token,然后所有 token 之间算全局自注意力。这套做法在 224×224 的分类任务上问题不大,可一旦搬到视觉任务的真实场景,两个麻烦立刻浮现。

第一个麻烦是计算量。注意力机制的计算复杂度是序列长度的平方,而图像里像素点太多了------想要更精细的特征就得切更碎的 patch,序列就变得很长,长序列上算注意力慢得难以接受。检测、分割任务动辄用 800×1333 的输入,照 ViT 的切法会得到上千个 token,平方复杂度直接爆炸。能不能用窗口和分层的形式替代长序列全局计算?这正是 Swin 的本质出发点。

第二个麻烦是多尺度特征。CNN 里我们习惯了感受野逐层扩大的过程:浅层看纹理边缘,深层看语义主体,检测和分割的neck(如 FPN)就是靠这种金字塔结构吃饭的。Transformer 这边该怎么体现"感受野"这个概念?Swin 给出的答案是分层------像 CNN 一样逐阶段下采样,让每个 stage 输出不同分辨率的特征图。

二、整体架构:一座四层的金字塔

【图1:Swin Transformer 整体架构】图中从左到右是完整的数据通路:输入 H×W×3 的图像先经过 Patch Partition 和 Linear Embedding,变成 H/4×W/4×C 的 token 序列;随后是四个 Stage,每两个 Stage 之间夹一个 Patch Merging 做下采样,特征图尺寸从 H/4 一路降到 H/32,通道数则按 C、2C、4C、8C 翻倍。每个 Stage 内部堆叠若干 Swin Transformer Block,数量分别标注为 ×2、×2、×6、×2。关键看点是这张图同时体现了"分层"和"注意力计算在 Block 内部完成"两件事------分层负责多尺度,Block 负责改进的注意力计算,这也是 Swin 相比 ViT 的全部结构增量。

整体流程可以概括成两句话:先得到各 patch 特征构建序列,再分层计算 attention(也就是一个逐步下采样的过程)。其中 Block 是最核心的部分,对 attention 的计算方式做了改进------具体怎么改进,就是后面 W-MSA 和 SW-MSA 的故事。

三、从图像到序列:Patch Embedding

第一步和 ViT 类似,把图像变成序列。以 224×224×3 的输入为例:

复制代码
输入:(224, 224, 3)
卷积:Conv2d(3, 96, kernel_size=(4, 4), stride=(4, 4))
输出:(3136, 96)

用一个 4×4、步长 4 的卷积一次完成切 patch 和线性映射,3136 = (224/4) × (224/4),即每个 4×4 的 patch 变成一个 96 维向量。序列长度 3136,如果直接在这上面做全局注意力,计算量已经相当可观了------Swin 的策略不是缩短序列,而是改变注意力的作用范围。

值得一提的是,这里也可以按需求更改卷积参数(比如更大的 patch 尺寸换取更短的序列),官方后续的 384 分辨率版本就是同样的 patch 大小、更大的输入图。

四、W-MSA:把注意力关进窗口里

拿到序列后,Swin 的第一板斧是窗口划分。以第一个 Stage 的特征图(56, 56, 96)为例:

复制代码
window_partition:
输入:(56, 56, 96)
窗口大小 7 → 划分成 8×8 = 64 个窗口
输出:(64, 7, 7, 96) → 展平后 (64, 49, 96)

划分前注意力的基本单位是序列里的单个 token,划分后单位变成了窗口------共 64 个窗口,每个窗口内部 7×7=49 个 token。窗口之间互不打扰,各算各的自注意力。这一下就把平方复杂度的"序列长度"从 3136 压到了 49,而且 64 个窗口可以并行计算。

窗口内部的标准多头自注意力(W-MSA,Window Multi-head Self-Attention)的张量形状值得仔细看一遍:

复制代码
qkv 放在一起:(3, 64, 3, 49, 32)
              ↑   ↑  ↑   ↑   ↑
           qkv三矩阵 窗口数 heads 49个token 每head特征 96/3=32

attention 得分:(64, 3, 49, 49)

每个头都会得出自己窗口内的注意力矩阵。qkv 三个矩阵提前拼在一起做一次线性投影,是官方实现里省事的写法。另外官方实现在 QK^T 的结果上还加了相对位置偏置(relative position bias),因为窗口内注意力本身不带任何位置信息,这个偏置让模型能感知 token 之间的相对空间关系。

注意力算完得到新特征 (64, 49, 96)------64 个窗口、每窗 49 个点、每点 96 维。接下来 window_reverse 通过 reshape 把它还原回 (56, 56, 96):尺寸和输入一致,但每个位置已经融合过窗口内的注意力信息。partition 和 reverse 就是一对可逆的整形操作,核心代码整理如下:

python 复制代码
def window_partition(x, window_size):
    # x: (B, H, W, 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()
    windows = windows.view(-1, window_size, window_size, C)
    return windows  # (num_windows*B, 7, 7, C)

def window_reverse(windows, window_size, H, W):
    # windows: (num_windows*B, 7, 7, C)
    B = int(windows.shape[0] / (H / window_size * W / 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  # (B, H, W, C)

但只算窗口内部有个明显的问题:窗口之间从此没有任何信息交换,模型容易局限在自己的"小领地"里。Swin 的解法是让窗口动起来------这就是和 W-MSA 成对出现的 SW-MSA。

五、SW-MSA:让窗口之间说话

5.1 为什么要 shift

【图2:两个连续的 Swin Transformer Block】图中画了前后串联的两个 Block:前一个用 W-MSA,后一个用 SW-MSA,各自内部都是 LayerNorm → 注意力 → 残差连接 → LayerNorm → MLP → 残差连接的结构,中间用 zl、ẑl 标注了特征传递。关键看点是这两个 Block 必须成对出现------W-MSA 负责窗口内计算,SW-MSA 负责窗口间打通,串在一起才构成 Swin 论文里所说的一个完整 block,单独只用其中一个都会丢掉一半信息。

W-MSA 只算窗口内部,窗口之间的边界区域就失去了交互渠道。Swin 的做法是在相邻的两层之间把窗口划分网格整体挪动半个窗口的位置,让原来不在一个窗口里的 patch 有机会落进同一个窗口。

【图3:Shifted Window 划分示意】图中从左到右四个子图演示了一只猫的特征图如何被重新划分:(a) 是原始特征图,尺寸 56×56;(b) 是常规的 7×7 网格窗口划分(红色格线);© 显示窗口整体向左上偏移 3 个像素后的状态;(d) 是偏移后重新形成的窗口划分,可以看到边缘出现了一圈不完整的窗口,图中标注了 (1)(2) 两种区域。关键看点是偏移之后,相邻层之间同一块空间区域会落入不同的窗口组合,跨窗口的信息交互就发生在这些"错位"的窗口内部。

位移的具体操作就是像素点挪一下位置------cyclic shift(循环移位),把特征图整体滚动几个像素,滚出去的部分从另一侧绕回来:

【图4:Swin 移位注意力的完整流程】图中从左到右是四步:window partition 把特征图切成规则窗口;cyclic shift 之后窗口内容错位,出现几个不完整窗口;中间部分画了两组 masked MSA------对错位后的不规则区域分别做注意力,并用掩码屏蔽无关区域;最后 reverse cyclic shift 把特征移回原位。关键看点是 masked MSA 这一步:Swin 并没有真的按"9 个不均匀窗口"去算注意力,而是用掩码在规则窗口内把跨区域的交互掐掉,这一步是整个 SW-MSA 的工程精髓。

5.2 计算量的坑:4 个窗口变 9 个

移位直接带来了一个数学上的小麻烦。以最小示例来看:原来特征图划分成 4 个规则窗口,移位半格之后,网格错位处会多出一条条不完整的窗口带,4 个窗口变成了 9 个大小不一的窗口,下面这张对比图就是这件事的直观版本:

【图5:移位导致窗口数从 4 变成 9】图中左边是 Layer l 的规则窗口划分,4 个整齐的窗口;右边是 Layer l+1 移位半格后的划分,窗口数变成了 9 个,大小形状各不相同。关键看点是这些不均匀窗口破坏了统一的 7×7 尺寸与数量约定,如果按它们各自计算,注意力的计算量就不再和 W-MSA 对齐,工程上也非常难看,这正是掩码方案要解决的问题。

Swin 的处理方式很聪明:移位之后不再按错位的网格切窗,而是把特征图重新按原来的规则网格划分,只是现在每个规则窗口里可能混着来自不同区域的 patch。下面这组图把这件事拆到了最细的粒度:

【图6:移位后窗口内的区域编号】图中左右各画了一个 4×4 的小特征图:左边是移位后新划分的窗口,按区域编号 0 到 8,其中 4 号是中间最大的完整区域;右边是另一种移位幅度下的窗口划分,能看到一个窗口里同时包含区域 4、5、3 等不同来源的 patch。关键看点是"区域"这个概念------同一区域内的 patch 在移位前后是真正相邻、应该互相注意的,跨区域的 patch 则是被 cyclic shift 硬凑到一个窗口里的"假邻居"。

对混装窗口里的注意力矩阵逐个位置看,会发现只有同区域的位置才应该有交互:

【图7:掩码后的稀疏注意力矩阵】图中左侧是一个窗口内 Q 向量的排列(数字代表所属区域编号),中间是 K 转置后的排列,右侧是 Q·K^T 得到的 4×4 注意力矩阵:只有区域编号相同的位置(两个 5 之间、两个 3 之间)保留了注意力值,其余位置全部置灰屏蔽。关键看点是注意力矩阵的稀疏模式完全由区域编号决定------这正是 attention mask 的作用方式:给不该交互的位置加上一个很大的负数,过 softmax 之后权重趋近于 0,等效于这些位置从来没参与过计算。

总结成规则就是:计算时只需要算自己区域的,其他的都是无关的;mask 的设置只需让对应位置的值为负无穷,softmax 自然把它们的权重压成零。算完之后别忘记把特征图 reverse 回去------循环移位怎么滚出去的,就怎么滚回来,恢复原始的空间对齐。至此,W-MSA 和 SW-MSA 这对组合构成了 Swin Transformer 的核心计算模块。

这个 mask 方案还带来一个额外的工程红利:所有窗口都保持 7×7 的统一尺寸、统一数量,SW-MSA 的计算量和 W-MSA 完全一致,shift 带来的信息交互是"免费"的。

六、Patch Merging:分层的关键

窗口注意力解决了计算量,多尺度还得靠下采样。Swin 的下采样不叫池化,叫 Patch Merging------一个间接的下采样操作:

【图8:Patch Merging 的间隔采样与拼接】图中从左到右:原始特征图(H×W×C)先按行列的奇偶性做间隔采样,拆出 x0、x1、x2、x3 四个 H/2×W/2 的子图,图中标注了各自的切片下标(如 x0 = 0::2, 0::2);四个子图沿通道维拼接成 H/2×W/2×4C 的张量,再排列成长序列。关键看点是"间隔采样"这个动作:它相当于把 2×2 邻域的四个邻居拆到四个不同子图里,空间减半、信息不丢。

具体来说,对 H 和 W 维度按 2 为步长做间隔采样后拼接,得到 (H/2, W/2, 4C):

python 复制代码
x0 = x[..., 0::2, 0::2]   # 偶数行偶数列
x1 = x[..., 1::2, 0::2]   # 奇数行偶数列
x2 = x[..., 0::2, 1::2]   # 偶数行奇数列
x3 = x[..., 1::2, 1::2]   # 奇数行奇数列
x = torch.cat([x0, x1, x2, x3], -1)   # (B, H/2, W/2, 4C)
x = self.norm(x)
x = self.reduction(x)     # Linear: 4C -> 2C

和池化不同,Patch Merging 不丢信息------2×2 邻域的四个点全部保留,只是从空间维搬到了通道维;官方实现在拼接后又接一个线性层把 4C 压回 2C,让通道数维持"每个 stage 翻倍"的节奏。一次下采样后序列长度从 3136 变成 784(56×56 → 28×28),然后继续走 W-MSA + SW-MSA 这两兄弟的流程,各个 stage 就这样串起来。最后根据任务选择合适的 head:分类接全局平均池化 + 线性层,检测和分割则取各 stage 的金字塔特征图接 FPN 之类的结构。

七、从配置文件看全貌

Swin 提供了大、中、小多个版本的模型,可以按算力自由选择合适的规格。四个变体的差异全部浓缩在配置文件里:

模型 EMBED_DIM DEPTHS NUM_HEADS
Swin-Tiny 96 2, 2, 6, 2 3, 6, 12, 24
Swin-Small 96 2, 2, 18, 2 3, 6, 12, 24
Swin-Base 128 2, 2, 18, 2 4, 8, 16, 32
Swin-Large 192 2, 2, 18, 2 6, 12, 24, 48

表里每一列都能和前面的原理对上:EMBED_DIM 是 Patch Embedding 输出的通道数 C;DEPTHS 是四个 Stage 各堆多少个"双 Block"(W-MSA + SW-MSA 算一组,对应图 1 里的 ×2、×2、×6、×2);NUM_HEADS 随通道数翻倍而同步增加,保证每个头的特征维度不至于过大。命名规则也写在配置里,比如 swin_tiny_patch4_window7_224:patch 4×4、窗口 7×7、输入 224------论文里的关键设计参数全在名字上。

训练侧的设定同样值得抄作业。官方默认配置里:优化器 AdamW,学习率 5e-4 配余弦退火,20 个 epoch 预热,权重衰减 0.05,梯度裁剪阈值 5.0;数据增强用 Mixup(0.8)+ CutMix(1.0)+ RandAugment + Random Erasing,标签平滑 0.1;随机深度(drop path)按模型规模从 0.1 到 0.5 递增,层数越深、模型越大,这个正则开得越狠。数据管道还支持把 ImageNet 打包成 zip 直接读(ZIP_MODE + CachedImageFolder),避免海量小文件拖垮 IO。这些细节组合起来,才是 Transformer 系模型在 ImageNet 上稳定收敛的完整配方。

八、实验表现

Swin 发布时在三大任务上同时刷新了纪录。分类上,ImageNet-1K 的 top-1 精度从 Tiny 到 Large 一路走高,用 ImageNet-22K 做预训练后还能再涨一截;检测上以 Swin 为 backbone 的 Cascade Mask R-CNN 在 COCO 上大幅超越当时的 CNN 方案;分割上 UperNet 加 Swin 同样拿下 ADE20K 的最佳成绩。更重要的不是刷分本身,而是它验证了一个判断:视觉任务需要的是"分层 + 高效注意力"的 Transformer,而不是照搬 NLP 的全局注意力。官方把各规格预训练权重全部开源,作为新一代 backbone 可以直接套用到各项下游任务中,这也是它影响力远超论文数字的原因。

九、实现视角的踩坑复盘

对照原理和代码,有几个容易踩的点值得记下来:

  • 窗口划分对尺寸有整除要求。 56×56 除以窗口 7 正好得 8×8 个窗口,但换一个不整除的输入尺寸就会切出残缺窗口。官方实现里对不整除的情况会先 padding 到窗口大小的整数倍,算完再裁回去------自己复现时如果漏了这步,形状错位会报得很隐蔽。
  • W-MSA 和 SW-MSA 必须成对使用。 只用 W-MSA,窗口间永远没有信息流;只用 SW-MSA,每层都在做昂贵的跨区域 mask 计算。论文里的"一个 Block"指的是两个子模块的串联,配置里 DEPTHS 的 2、6 也都是按"对"来数的。
  • mask 与还原是一套组合拳。 shift 之后忘记给跨区域位置加 mask,模型不会报错,注意力会静默地混入"假邻居"信息,精度下降但很难定位原因;同样,算完忘记 reverse shift,特征图和标签的空间对齐就错位了。这两步任何一步漏掉都是 silent bug。
  • attention mask 的索引要提前算好。 9 个区域对应的 mask 矩阵在初始化时构造一次并缓存,训练中反复使用;如果每次 forward 现算,SW-MSA 的开销会明显放大。
  • 移位量是窗口大小的一半。 shift_size = window_size // 2,这个值还要保证小于窗口大小;某些层(分辨率太小的 stage)会取消移位退回 W-MSA,这也是官方实现里根据 H、W 动态判断的细节。

十、总结

回头看,Swin Transformer 的每个设计都在回答 ViT 留下的具体问题:注意力算不动,就用窗口把平方复杂度关进 7×7 的小格子里;窗口之间没交流,就用半个窗口的循环移位加一层掩码,让跨窗口交互以零额外计算量的方式发生;密集预测需要多尺度,就用 Patch Merging 搭出从 H/4 到 H/32 的金字塔。分层 + 移位窗口这两个词,撑起了 Transformer 在视觉领域真正站稳脚跟的那一代 backbone,之后大量视觉工作(包括视频、多模态)都直接沿用了这套分层窗口的思路。理解它最好的方式,就是把张量形状从 Patch Embedding 到 window_reverse 全程推一遍------形状对上了,设计动机自然就清楚了。


相关推荐
落雪snowflake1 小时前
论文精讲:A Language-Informed Loss for Visual Generation中文题目:一种面向视觉生成的语言信息引导损失函数
人工智能·深度学习·计算机视觉
葫三生2 小时前
关于三生原理系列被技术开发社区批量收录或引发策展边界讨论的一种思考
大数据·人工智能·深度学习·算法·机器学习·数据挖掘·开源
持续学习的程序员+12 小时前
跟随一次“沙拉酱瓶入篮”任务,理解 RPent(harness vla) 的完整工作流程
人工智能·深度学习·机器学习
jay神2 小时前
基于 YOLOv8 的无人机识别与检测系统
人工智能·深度学习·yolo·计算机视觉·无人机·课程设计·计算机毕业设计
田里的水稻2 小时前
FA_CNN卷积神经网络四---线性层及其拉平PyTorch
人工智能·pytorch·python·深度学习·神经网络·cnn
碧海银沙音频科技研究院2 小时前
实现炬芯ATS308X的usb cdc 串口IAP 固件升级功能实现
人工智能·嵌入式硬件·深度学习·算法
深蓝学院3 小时前
机器视觉中CNN与Transformer的原理及应用差异
计算机视觉
跟我学机器学习3 小时前
拥有短期记忆的RAG Agent:记忆增强检索实战
人工智能·深度学习·语言模型·数据挖掘·rag
QN1幻化引擎3 小时前
信号场注意力(SFA):把 KV Cache 变成 O(1) 的双通道注意力——0.1B/0.3B 受控 A/B 精度反升 15.5% / 24.3%
pytorch·深度学习·大模型·transformer·注意力机制·推理优化