Swin Transformer: 基于移位窗口的分层视觉Transformer

摘要

将Transformer从语言领域迁移到视觉领域面临的挑战源于两个领域之间的差异,例如视觉实体尺度变化大,以及图像像素分辨率相较于文本中的单词更高。为解决这些差异,我们提出了一种层次化Transformer,其表示通过移位窗口计算。移位窗口方案通过将自注意力计算限制在非重叠的局部窗口内,同时允许跨窗口连接,从而带来更高的效率。这种层次化架构具有在不同尺度建模的灵活性,并且计算复杂度与图像大小呈线性关系。Swin Transformer的这些特性使其兼容广泛的视觉任务。

1 引入

背景:长期以来,计算机视觉领域的建模主要由卷积神经网络(CNNs)主导。CNN架构通过更大规模[30,76]、更广泛的连接[34]以及更复杂的卷积形式[70,18,84]不断演进,变得日益强大。

另一方面,自然语言处理(NLP)领域的网络架构演进则遵循了不同的路径。当前主流的架构是Transformer,该模型专为序列建模与转换任务设计,其显著特点在于利用注意力机制建模数据中的长程依赖关系。Transformer在语言领域的巨大成功促使研究者探索其在计算机视觉领域的适应性。近期研究表明,Transformer在特定视觉任务上展现出显著潜力,尤其是在图像分类[20]与视觉-语言联合建模[47]方面取得了令人瞩目的成果。

将Transformer在语言领域的高性能迁移至视觉领域面临显著挑战:这些挑战主要源于两种模态之间的差异。其中一个差异涉及尺度问题。与语言Transformer中作为基本处理单元的固定尺度词元不同,视觉元素的尺度可能存在显著变化。现有基于Transformer的模型[64,20]采用固定尺度的令牌处理机制,这一特性难以适应视觉应用需求。另一差异在于图像像素的分辨率远高于文本段落中的词汇。许多视觉任务(如语义分割)需要在像素级别进行密集预测,而Transformer的自注意力计算复杂度与图像尺寸呈平方关系,导致其难以处理高分辨率图像。

尺度:某个维度的“大小”或“范围”。NLP 中输入的本质是离散的语义符号,模型接收的输入是 “词” 或 “子词”(如 “猫”、 “##ing”)。这些是离散的、有明确语义的符号。尺度指依赖关系的距离。cv 中模型接收到的输入是像素,单个像素什么都代表不了,一个“语义实体”在图像中没有固定的“令牌边界”。一辆车可能由几百个像素组成,形状不规则。模型必须学会将一组连续的、特定模式的像素归类为“车”。既要“发现/定义”什么是基本部件(多尺度实体),又要理解部件间的关系。

对于视觉的下游任务比如检测、分割,有多尺寸的特征是至关重要的。

令牌:将图像分割成的固定大小的,每个块就是一个视觉令牌。

固定尺度的令牌不适合的原因大物体被分割成很多块,模型需要将这多个令牌信息重新组合才能理解,组合过程困难,且可能丢失全局信息;小物体可能只占一块不到一半的面积,令牌的特征被背景污染,细节丢失,模型难以识别。

图1:为解决这些问题,我们提出了一种名为Swin Transformer的通用视觉Transformer骨干网络,该网络能够构建层次化特征图,并实现与图像尺寸呈线性关系的计算复杂度。如图1(a)所示,Swin Transformer通过从小尺寸图像块(灰色轮廓标注)开始,在更深的Transformer层中逐步合并相邻图像块(patch merging),构建出层次化特征表示。借助这种层次化特征图,Swin Transformer可以便捷地集成特征金字塔网络(FPN)或U-Net等先进密集预测技术。线性计算复杂度的实现得益于在图像划分的非重叠局部窗口(红色轮廓标注)内计算自注意力(对于视觉来说,全局去算自注意力实际有些浪费资源)每个窗口包含固定数量的图像块(图像面积放大x倍,窗口数也放大x倍,达到线性),从而使计算复杂度与图像尺寸呈线性关系。相较于先前产生单一分辨率特征图且具有平方计算复杂度的Transformer架构[20],这些优势使Swin Transformer能够成为适用于多种视觉任务的通用骨干网络。

ViT:处理的都是单一尺寸、且低分辨率特征图。并因全局计算自注意力而具有与输入图像尺寸成平方关系的计算复杂度。

生成多尺度的特征:卷积:池化操作;swin :patch merging 类似池化,把相邻的小patch合成大patch,大patch可以看到小patch的特征,从而得到多尺寸的特征图。

图2:连续自注意力层之间移动窗口划分的方式,如图2所示。这种窗口偏移策略将前一层的窗口连接起来,为它们之间提供了重要的联系,从而显著增强了模型的建模能力(参见表4)。该策略在实际延迟方面也表现出高效性:一个窗口内的所有查询补丁共享同一组键集合¹,这有助于硬件中的内存访问。相比之下,早期基于滑动窗口的自注意力方法[33, 50]由于不同查询像素对应不同的键集合²,在通用硬件上延迟较高。我们的实验表明,所提出的偏移窗口方法比滑动窗口方法延迟低得多,同时在建模能力上表现相近(参见表5和表6)。偏移窗口方法也被证明对全MLP架构[61]有益。

查询补丁(patch):当前要计算注意力的像素/块

键集合:注意力计算中作为key的像素/块集合

硬件上的高效:创造了规整的数据局部性,让GPU能高速、批量处理数据,避免了频繁、低效的内存访问。

两个核心技术

Shift window(图2):主要用于建立同分辨率下的跨窗口连接,每层都要进行。实现原本被分隔开的两个小块现在可能出现在同一个新窗口内,可以直接计算注意力。从而达到局部注意力的计算量(O(N)),获得类似全局注意力的信息流通能力。

patch merging(图1):主要用于构建特征金字塔,分辨率由于下采样减半,空间细节永久性消失。负责从细节到语义、局部到全局的纵向信息抽象。

层次化建模:指的是网络逐层、渐进地构建特征,形成分辨率由高到低、语义由浅入深的“特征金字塔”。

为什么层次化建模可以解决多尺度的问题:

【1】感受野的逐层扩大

【2】特征的逐层抽象:分辨率降低的过程其实也是信息压缩和重组的过程。(细节到部件到物件到高级语义)

2 相关工作

CNN,自注意力,CNN中的自注意力机制,transformer,ViT

3 方法

3.1 总体架构

Swin Transformer架构的概览(图3左):图中以微型版本(SwinT)为例进行说明。

该架构首先通过图像块分割模块(1. patch partition)将输入的RGB图像分割为不重叠的块,类似于ViT的处理方式。每个块被视为一个token,其特征由原始像素RGB值拼接而成。在我们的实现中,采用4×4的块尺寸,因此每个块的特征维度为4×4×3=48。

随后通过线性嵌入层(2. linear embedding)将这些原始值特征映射到任意维度(记为C)

将48维的原始像素块,通过线性变换矩阵转变为C维向量,各个维度代表更抽象、更有意义的特征,也为后续所有Swin Transformer块的处理提供了统一且规整的输入维度,将RGB颜色通道的信息与空间信息进行交叉组合(建立通道关联

而从线性层出来的序列长度是56×56=3136,([3136,c]的矩阵,但在内存上是三位特征图,二维位置信息存在)太长了无法接受。于是swin transformer就引入了基于窗口的自注意力计算,每个窗口里面只有7×7=49个patch,可以接受。

在这些块标记上应用多个经过自注意力(基于窗口)计算改进的Transformer块(3. Swin Transformer block)。这些Transformer块保持标记数量不变(H/4 × W/4),与线性嵌入层共同构成“阶段1”。

为生成层次化表征(为了构建多尺度的信息),随着网络加深,通过块合并层(4. patch merging)逐步减少token数量(降低分辨率)。首个块合并层将每组2×2相邻块的特征进行拼接,并对拼接后的4C维特征应用线性层(为了和卷积神经网络对齐)。这一操作使标记数量减少至原来的1/4(分辨率下采样2倍),同时将输出维度设置为2C。  

随后应用Swin Transformer块进行特征变换,并保持分辨率为H/8 × W/8。这一包含块合并与特征变换的首个模块称为“阶段2”。该过程重复两次,分别构成“阶段3”和“阶段4”,输出分辨率依次为H/16 × W/16和H/32 × W/32。这些阶段共同生成层次化表征,其特征图分辨率与典型卷积网络(如VGG[52]和ResNet[30])保持一致。因此,所提出的架构能够便捷地替代现有各类视觉任务方法中的骨干网络。

最后也是用全局平均池化(分类任务)。根据下游任务也可以接检测头等等。

Swin Transformer模块(图3右):通过将Transformer块中的标准多头自注意力(MSA)模块替换为基于移位窗口的模块(详见第3.2节),同时保持其他层级不变,从而构建而成。如图3(b)所示,一个Swin Transformer模块包含一个基于移位窗口的MSA模块,其后接一个两层MLP(中间采用GELU非线性激活函数)。在每个MSA模块和每个MLP之前均应用层归一化(LN),并在每个模块后添加残差连接。

3.2 基于移位窗口的自注意力机制

本研究提出在局部窗口内计算自注意力。窗口以非重叠方式均匀划分图像。假设每个窗口包含M×M个图像块,在包含h×w个图像块的输入图像上,全局多头自注意力模块与基于窗口的多头自注意力模块的计算复杂度分别为:  

Ω(MSA) = 4hwC² + 2(hw)²C, (1)  

Ω(W-MSA) = 4hwC² + 2M²hwC, (2)  

M=7,hw较大

先做一次基于窗口的自注意力,再做一次基于移动窗口(3 patch)的自注意力,达到窗口间的互相通信

移动窗口掩码机制(图4):为了避免窗口数增加、各个窗口间数量不相等,会进行循环位移,但是拼接后的元素原来不相邻的部分不应该去做自注意力,比如天空在草地下面。 于是进行掩码。算完多头自注意力后再将循环位移还原回去,保持图片的相对位置(语义信息)。

相对位置编码:相对值->索引(归正)->训练表、查表

3.3 架构变体(改C和transformer block数量)

4 实验

4.1 在ImageNet-1K数据集上的图像分类

4.2 在COCO数据集上进行目标检测

4.3 在ADE20K数据集上进行语义分割

4.4 消融实验

用了移动窗口和相对位置编码,在分类任务上提升并不太明显,更大的帮助在目标检测和语义分割上。

5 结论

本文提出了一种名为Swin Transformer的新型视觉Transformer,它能够生成层次化的特征表示,并且在计算复杂度上与输入图像大小呈线性关系。Swin Transformer在COCO目标检测和ADE20K语义分割任务上实现了最先进的性能,显著超越了先前的最佳方法。我们希望Swin Transformer在各种视觉问题上的强大表现能够促进视觉与语言信号的统一建模。作为Swin Transformer的关键组成部分,基于移位窗口的自注意力机制在视觉问题上被证明既有效又高效,我们也期待在自然语言处理领域进一步探索其应用潜力。

附录A

Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐