论文: https://arxiv.org/pdf/2111.09883.pdf
代码: GitHub - microsoft/Swin-Transformer: This is an official implementation for "Swin Transformer: Hierarchical Vision Transformer using Shifted Windows".
在Swin Transformer的基础上,提出了以下的改进措施:
1)post normalization technique and scaled cosine attention提升大型视觉模型的稳定性;
2)log-spaced continuous position bias 进行低分辨率预训练模型向高分辨率模型迁移。
3)implementation details 大幅节省GPU显存占用以使得大视觉模型训练变得可行。
基于上述技术与自监督预训练,训练了一个包含 3 billion 参数量的SwinTransformer模型并将其迁移到不同的高分辨率输入的下游任务上,取得了SOTA性能。
Swin Transformer V2 在四个具有代表性的基准上刷新纪录:在 ImageNet-V2 图像分类任务上 top-1 准确率为 84.0%,COCO 目标检测任务为 63.1 / 54.4 box / mask mAP,ADE20K 语义分割为 59.9 mIoU,Kinetics-400 视频动作分类的 top-1 准确率为 86.8%。
Swin Transformer是一种通用的视觉骨干模型,在不同的视觉任务(包含图像分类、目标检测以及语义分割)上均取得了极强性能。Swin Transformer的主要思想:为常规Transformer Encoder架构引入了几个重要的视觉信号先验信息 ,包含分层、局部以及平移不变形。基础Transformer单元提供了强建模能力,视觉先验信息使其可以处理不同视觉任务。
原始的SwinTransformer采用了常规的预归一化技术
原始Swin Transformer引入了一个额外参数偏置,公式如下:
其中,B是每个head的相对位置偏置,对于dense recognition任务非常重要。当进行不同分辨率模型迁移时,常规方案是对该偏置矩阵进行双三次插值。
在对SwinTransformer进行容量与窗口分辨率的scaling过程中,我们发现以下两个问题:
在大型模型中,跨层激活函数输出的幅值的差异变得更大。激活值是逐层累积的,因此深层的幅值明显大于浅层的幅值。当我们将原来的 Swin Transformer 模型从小模型放大到大模型时,深层的 activation 值急剧增加。最高和最低幅值之间的差异达到了10^4。
当我们进一步扩展到一个巨大的规模 (658M 参数) 时,它不能完成训练,如下图所示。
许多下游视觉任务需要高分辨率的图像或窗口,预训练模型时是在低分辨率下进行的,而 fine-tuning 是在高分辨率下进行的。针对分辨率不同的问题传统的做法是把位置编码进行双线性插值 ,这种做法是次优的。如下图3所示是不同位置编码方式性能的比较,当我们直接在较大的图像分辨率和窗口大小测试预训练的 Imagenet-1k 模型 (分辨率256×256,window siez=8×8) 时,发现精度显著下降。
使用post normalization technique and scaled cosine attention解决不稳定问题;
使用log-spaced continuous position bias解决跨窗口分辨率迁移问题。
原始SwinTransformer的预归一化配置下,每个残差模块的输出激活值与主分支直接合并,导致主分支在更深层的幅值越来越大,进而导致训练不稳定。
为缓解该问题,我们提出了Post Normalization:每个残差模块的输出先进行归一化再与主分支进行合并,因此主分支的幅值不会逐层累积。从上面的Figure2可以看到:使用Post Normalization的模型激活幅值更温和。
在最大的模型中,我们每6个Transformer模块额外引入一个LN单元以进一步稳定训练。
在原始自注意力计算过程中,像素对的相似性通过query与key的点积计算。我们发现:在大模型中,某些模块与head的注意力图会被少量像素对主导 。为缓解该问题,我们提出了Scaled Cosine Attention(SCA),公式如下:
不同于直接对偏置参数直接优化,连续位置偏置方法不是直接优化参数偏差,而是在相对坐标上引入一个小的元(meta)网络:
注:g是一个很小的网络,比如2层MLP中间加一个ReLU。元网络对任意相对坐标生成偏置参数,因而可以自然地进行任意窗口尺寸的迁移。在推理阶段,每个相对位置的偏置可以预先计算并保存,按照原始方式进行推理。
当变化很大的窗口大小之间迁移时,有较大比例的相对坐标范围需要外插。为缓解该问题,我们采用了log-spaced坐标:
△x和△y是线性空间的坐标
通过对数空间坐标,在进行块分辨率迁移时,所需的外插比例会更小。比如,将8*8 window预训练模型向16*16 window 迁移时,输入坐标范围从[-7,7]*[-7,7]调整为[-15,15]*[-15,15],外插比例为8/7=1.14倍。而采用对数空间坐标,输入坐标范围从[-2.079,2.079]*[-2.079,2.079]调整为[-2.773,2.773]*[-2.773,2.773],外插比例0.33倍。表1 给出了不同位置偏置下的迁移性能对比,可以看到:当向更大窗口尺寸迁移时,对数空间连续位置偏置性能最佳 。
大分辨率输入与大容量模型存在的另一个问题是GPU显存占用不可接受问题 。我们采用了以下实现改善该问题:
(来自论文:Zero: Memory optimizations toward training trillion parameter models)
(来自论文:Training deep nets with sublinear memory cost)
更大的模型需要更多的数据(data hungry)。为解决该问题,之前的大模型训练通过采用额外的数据或者自监督预训练。我们对这两种策略进行了组合。
通过上述训练方案,训练了一个具有3 B参数量的SwinTransformer模型并在多个基准数据集上取得了SOTA性能。
我们保持与SwinTransformer相同的stage、block以及通道配置得到了四个版本的SwinTransformerV2:
我们进一步对SwinV2进行更大尺寸缩放得到了658M与3B参数模型:(每6层额外增加一个LN)
Swin Transformer V2 在四个具有代表性的基准上刷新纪录:在 ImageNet-V2 图像分类任务上 top-1 准确率为 84.0%,COCO 目标检测任务为 63.1 / 54.4 box / mask mAP,ADE20K 语义分割为 59.9 mIoU,Kinetics-400 视频动作分类的 top-1 准确率为 86.8%。
SwinV2-G 实验设置:预训练采用 192×192 图像分辨率,以节省训练成本,实验采用 2-step 预训练方法:首先,在 ImageNet-22K-ext 数据集上使用自监督方法对模型进行 20epoch 的预训练。其次,在 ImageNet-1K V1 和 V2 分类任务上,继续将模型进行 30epoch 预训练。
表 2 将 SwinV2-G 模型与之前在 ImageNet-1K V1 和 V2 分类任务上的最大 / 最佳视觉模型进行了比较。SwinV2-G 是之前所有密集(dense)视觉模型中最大的。它在 ImageNet V2 基准测试中达到了 84.0% 的 top-1 准确率,比之前最好的 ViT-G (83.3%) 高 0.7%。但是,SwinV2-G 在 ImageNet-1K V1 上的准确率比 CoAtNet-7 略低(90.17% 比 90.88%)。
表 3 将 SwinV2-G 模型与之前在 COCO 目标检测和实例分割任务上取得最佳性能模型进行了比较。SwinV2-G 在 COCO test-dev 上实现了 63.1/54.4 box/max AP,比 SoftTeacher(61.3/53.0) 提高了 + 1.8/1.4。这表明扩展视觉模型有利于目标检测任务中的密集视觉识别任务。
下表 4 将 SwinV2-G 模型与之前在 ADE20K 语义分割基准上的 SOTA 结果进行了比较。Swin-V2-G 在 ADE20K val 集上实现了 59.9 mIoU,比之前的 SOTA 结果(BEiT)58.4 高了 1.5。这表明扩展视觉模型有益于像素级视觉识别任务。在测试时使用更大的窗口大小还可以带来 +0.2 的增益,这可能归功于有效的 Log-spaced CPB 方法。
下表 5 将 SwinV2-G 模型与之前在 Kinetics-400 动作分类基准上的 SOTA 结果进行了比较。可以看到,Video-SwinV2-G 实现了 86.8% 的 top-1 准确率,比之前的 SOTA (TokenLearner)85.4% 高出 +1.4%。这表明扩展视觉模型也有益于视频识别任务。在这种场景下,在测试时使用更大的窗口大小也可以带来额外增益 ( +0.2% ),这也要归功于有效的 Log-spaced CPB 方法。
post-norm 和缩放余弦注意力的消融实验:下表 6 展示了 post-norm 和缩放余弦注意力方法应用于原始 Swin Transformer 方法的性能表现。可以看到,这两种方法都提高了 Swin-Tiny、Swin-Small 和 Swin-Base size 的准确率,整体提升分别为 +0.2%、+0.4% 和 +0.5%,表明它们对更大的模型更有益。