《HIRI-ViT: Scaling Vision Transformer with High Resolution Inputs》解读

期刊：TPAMI

年份：2024

摘要

视觉Transformer(ViT)和卷积神经网络(CNN)的混合深度模型已经成为一类强大的视觉任务骨干。扩大这种混合主干网的输入分辨率自然会增强模型的能力，但不可避免地要承受二次扩展的沉重计算成本。相反，我们提出了一种具有高分辨率输入(即HIRI-ViT)的新型混合主干网，将流行的四级ViT升级为针对高分辨率输入量身定制的五级ViT。HIRI-ViT是建立在将典型的CNN操作以经济高效的方式分解为两个并行CNN分支的开创性思想之上的。一个高分辨率分支直接将主要高分辨率特征作为输入，但使用较少的卷积操作。另一个低分辨率分支首先执行下采样，然后在这些低分辨率特征上使用更多的卷积操作。在识别任务(ImageNet-1K数据集)和密集预测任务(COCO和ADE20K数据集)上的实验都证明了iri - vit的优越性。更值得注意的是，在相当的计算成本(~ 5.0 GFLOPs)下，iri - vit在输入为448×448的ImageNet上实现了迄今为止公布的最佳Top-1准确率为84.3%，在输入为224×224的情况下绝对提高了iFormer-S的83.4%。

Introduction

研究问题：论文的核心问题是，是否存在一种有原则的方法来扩展具有高分辨率输入的CNN+ViT混合骨干网络，同时保持可比的计算开销。

HIRI-ViT的提出：为了解决这个问题，作者提出了HIRI-ViT，这是一个为高分辨率输入量身定制的新型五阶段ViT混合骨干网络。HIRI-ViT通过在早期阶段使用双分支构建块来平衡性能和计算成本。

主要贡献：论文的主要贡献是提出了一种新的方法来扩展具有高分辨率输入的CNN+ViT混合骨干网络，同时保持有利的计算成本。此外，作者还展示了通过一系列实验来证明HIRI-ViT在多个视觉任务（如图像识别、目标检测和实例/语义分割）中的优越性。

在ImageNet-1K数据集上，Swin Transformer和具有不同输入分辨率的HIRI-ViT的性能和计算成本(即GFLOPs和推理时间)

Preliminaries

多阶段Vision Transformer (M-ViT)：

传统M-ViT的基本结构，通常由一个stem层加上四个stage组成，类似于卷积神经网络的结构。
输入图像首先通过stem层被分割成多个小块（patches），每个小块被视为一个“视觉标记”（visual token）。

Stem层：

两种Stem层的设计：ViT-stem和Conv-stem。ViT-stem通常是一个步长为4的卷积层，用于将输入图像划分为小块。Conv-stem则由多个3x3卷积层堆叠而成，有助于稳定网络优化过程并提高性能。

多头自注意力（Multi-head Self-attention, MHA）：

将输入特征线性变换为queries、keys和values，然后分解到不同的head中，并计算注意力权重。

前馈网络（Feed-Forward Network, FFN）：

原始FFN的结构，包括两个全连接层和中间的非线性激活函数。
通过添加额外的卷积操作来升级FFN，从而引入2D归纳偏差，形成卷积前馈网络（Convolutional Feed-Forward Network, CFFN）。

HIRI-VIT

将典型的四阶段M-ViT升级为一个新的五阶段ViT家族，其中包含早期阶段的双分支构建块，将单分支CNN操作分解为两个并行的CNN分支。这种方法导致为高分辨率输入量身定制的有利计算开销。

关键组件：

HR-stem：高分辨率Stem层，通过两个并行CNN分支处理输入图像。
HR-block：高分辨率块，替代了第一阶段和第二阶段的Transformer块，以高效编码高分辨率输入。
Inverted Residual Downsampling (IRDS)：一种新的下采样方法，使用两个并行分支进行空间下采样。
Normalization：在不同阶段使用批量归一化（BN）和层归一化（LN）以提高性能和加快推理速度。
EMA Distillation：一种新的指数移动平均蒸馏策略，用于训练HIRI-ViT，提高学生网络的训练效果。

3.1 High Resolution Stem

Stem层是Vision Transformer的初始部分，负责将输入图像分割成小块（patches），这些小块作为"视觉标记"（visual tokens）输入到后续的Transformer层。

三种Stem对比：ViT-stem通常是将其划分为patches，然后对patches进行一个步长为4的卷积层，而Conv-stem由多个3x3卷积层堆叠而成。

图3 (a) 典型视觉Transformer中的 ViT-stem， (b) Conv-stem，(c) HIRI-ViT 中的 HR-stem。

为应对随着输入分辨率的增加，使用Conv-stem的计算成本（GFLOPs）显著增加这一问题，HR-Stem将单分支Conv-stem改造为两个并行的CNN分支来设计：高分辨率分支和低分辨率分支。高分辨率分支直接处理高分辨率特征，使用较少的卷积操作；低分辨率分支首先进行下采样，然后使用更多的卷积操作。

高分辨率分支：高分辨率分支包含一个轻量级的深度卷积（DWConv），后跟一个步长为2的3x3卷积，用于直接处理高分辨率特征。

低分辨率分支：低分辨率分支首先使用步长为2的3x3卷积进行下采样，然后应用两个卷积层（3x3和1x1卷积）来引入归纳偏差。

评价：HR-stem设计不仅保留了Conv-stem的高模型容量，同时在高分辨率输入下显著降低了计算成本。

3.2 High Resolution Block

为了解决在混合骨干网络中的前两个阶段输入分辨率较大，传统的Transformer块的计算成本相对较高这个问题，提出了高分辨率块（HR-block），这是一种新的构建块，用于替换第一阶段和第二阶段的Transformer块。

高分辨率分支：高分辨率分支设计为轻量级深度卷积，用于捕获高分辨率输入的粗略信息。

低分辨率分支：

低分辨率分支首先使用步长深度卷积进行下采样，随后应用前馈操作处理低分辨率特征图。
前馈操作包括两个全连接层，中间有一个非线性激活函数。即FFN=Conv 1x1 +Hard wish+Conv 1x1

3.3 Inverted Residual Downsampling

在标准的多阶段Vision Transformer（如Swin Transformer）中，通常使用步长卷积来进行下采样，这在计算上可能不够高效。为了解决这个问题，作者提出了IRDS，这是一种受卷积神经网络（CNN）启发的下采样方法，具有两个并行分支。

图5 (a) 典型 Vision Transformer 中的下采样 (DS) ，(b)IRDS-a , (c) IRDS-b

IRDS-a：对于前两个阶段的高分辨率输入，作者采用了IRDS-a。IRDS-a首先使用一个步长为2的3x3卷积来扩展维度并减少空间大小，然后使用1x1卷积来缩小通道维度。

IRDS-b：对于后续阶段的下采样，作者使用了IRDS-b。IRDS-b类似于反转残差块，但区别在于只在第一个卷积后应用归一化和激活操作。

3.4 Normalization of Block

归一化的重要性：归一化是深度学习中用于稳定训练过程和提高模型性能的关键技术。它可以减少内部协变量偏移（Internal Covariate Shift），加速收敛，并提高模型的泛化能力。

批量归一化（Batch Normalization, BN）：

在传统的卷积神经网络（CNN）中，BN通常用于每个卷积层之后，以归一化层的输出。
BN可以加速训练过程，提高模型的稳定性，并在一定程度上起到正则化的作用。

层归一化（Layer Normalization, LN）：

LN与BN不同，它对每个样本的每个特征通道进行归一化，而不是对每个批次的每个样本进行归一化。因此对对训练批量大小（例如对象检测和语义分割）较小的密集预测任务更友好。
LN 导致高分辨率输入的计算成本更高。

HIRI-ViT中的归一化策略：

在高分辨率输入的前三个阶段，使用BN而不是LN，因为BN对于具有高分辨率特征的卷积操作更为合适。
在低分辨率输入的最后两个阶段，使用LN，因为LN对于小批量训练更为有效，并且可以提高模型在密集预测任务中的性能。

3.5 EMA Distillation

指数移动平均（Exponential Moving Average，EMA）是一种时间序列预测方法，通过给最近的数据更高的权重来计算平均值。在深度学习中，EMA通常用于平滑模型参数，提高训练稳定性。

传统EMA：教师网络（teacher network）的参数根据学生网络（student network）的参数进行更新，但这种信息传递是单向的，可能导致次优解。

知识蒸馏：与传统的知识蒸馏不同，EMA蒸馏不依赖于大规模预训练网络，而是直接利用教师网络的知识来辅助学生网络的学习。

EMA：一种双向信息交互的EMA蒸馏策略：1）根据学生网络的参数更新教师网络，2）利用教师网络学习到的概率分布来指导学生网络的训练，从而实现更优的解决方案。

蒸馏过程：

使用Cutmix/Mixup等数据增强技术生成新的训练样本。

其中： $\left ( x_{a},y_{a} \right )$ 和 $\left ( x_{b},y_{b} \right )$ 是训练样本，其中 $M\in \left \{ 0,1 \right \}^{H\times W}$ 表示矩形掩码，⊙ 是元素乘法。

将原始样本 $x_{a}$ 和 $x_{b}$ 输入教师网络 $F^{t}$ (去掉最后一个特征图的平均池化操作)，得到概率分布图 $P_{a}$ 和 $P_{b}$ 。
为生成的样本 $\widetilde{x}$ 分配一个混合目标标签

其中AvgPool表示沿空间维度使用平均池化操作.

结合了Cutmix和教师网络生成的两个目标标签

$\overline{y}=\alpha \widetilde{y}+\left ( 1-\alpha \right )\widehat{y}$

在训练期间将混合样本 $\left ( \widetilde{x},\overline{y} \right )$ 输入学生网络 $F^{s}$ 。

3.6 Architecture Details

其中 $H_{i}$ 、 $C_{i}$ 和 $E_{i}$ 分别表示阶段 i 中前馈层的头数、通道维度和扩展率。

表 1 详细说明了我们的 HIRI-ViT 系列的架构。遵循现有 CNN+ViT 混合主干的基本网络配置，我们构建了不同模型大小的 HIRI-ViT 的三种变体，即 HIRIViT-S（小尺寸）、HIRI-ViT-B（基本大小）和 HIRI-ViTL（大尺寸）。具体来说，HIRIViT的整个架构由一个HR-stem层和五个阶段组成。对于具有高分辨率输入的前两个阶段，我们将传统的 Transformer 块替换为我们的轻量级高分辨率块，以避免巨大的计算开销。对于第三阶段，我们仅使用 CFFN 块来处理中间分辨率的特征图。与传统的 Vision Transformer 类似，我们在最后两个阶段使用 Transformer 块，具有低分辨率输入。

Result

HIRI-ViT与ImageNet-1K验证集上最先进的常规骨干的图像分类性能比较，分辨率更高。我们根据 GFLOPs 对主干进行分组，并根据Top-1 准确度对它们进行排名。

表 2 说明了我们的 HIRI-ViT 家族与其他具有较大输入图像大小的最先进的视觉主干之间的比较。对于这种更高分辨率的输入(768 × 768)升级的HIRI-ViT，我们在8个V100 gpu上采用AdamW优化器，动量为0.9，初始学习率为1.0e^−5，权重衰减为1.0e^−8。优化过程包括30个epoch，余弦衰减学习率调度器。同样，对于计算成本相当的每一组，与其他分辨率较高的视觉骨干相比，我们的HIRI-ViT始终获得性能提升。这些结果清楚地验证了我们提出的针对高分辨率输入量身定制的五阶段 ViT 主干的有效性。这种设计新颖地将典型的CNN操作并行分解为高分辨率分支和低分辨率分支，从而在更高分辨率的输入设置下保持良好的计算成本。

Conclusion

在这项工作中，为高分辨率输入设计了一个新的五阶段ViT主干，即HIRI-ViT，它将典型的CNN操作并行分解为高分辨率和低分辨率分支。凭借这种原则性的五阶段和两支路设计，我们的HIRIViT能够以经济高效的方式扩展具有高分辨率输入的Vision Transformer主干。在ImageNet-1K(图像分类)、COCO(目标检测和实例分割)和ADE20K数据集(语义分割)上进行了大量实验，以验证我们的iri -ViT对竞争对手CNN或ViT主干的有效性。
尽管有这些观察结果，仍然存在一些悬而未决的问题。虽然五级结构和两支路设计明显提高了视觉Transformer的扩展效率，但我们观察到当采用六级结构时，性能/计算成本的改善较少。此外，如何扩大高分辨率输入的视频视觉Transformer仍然是一个重大挑战。