FasterViT | 英伟达提出分层注意力，达成SOTA性能-吞吐量均衡

AIWalker

发布于 2023-09-01 10:15:43

3660

发布于 2023-09-01 10:15:43

文章被收录于专栏：AIWalker

https://arxiv.org/abs/2306.06189

https://github.com/NVlabs/FasterViT

出发点

本文旨在面向主流硬件设备(如GPU)研发具有高吞吐量的骨干架构。当前主流硬件设备包含多个CUDA与Tensor核计算单元，它需要进行频繁的数据迁移进行计算，可能会受到数据移动带宽影响。因此，以下两种类型算子需要进行精心平衡以最大化吞吐量：

受计算量约束的算子称之为math-limited，笔者将其称之为计算约束算子;
受访存迁移约束的算子称之为memory-limited，笔者将其称之为访存约束算子。

在分层架构中，特征分辨率会随推理进行而逐渐缩小，如从输入的

224 \times 224

到最后一个stage的

7 \times 7

。

网络的起始阶段特征分辨率比较大(如

112\times 112 \times 64

)，而通道数较少，此时会受访存迁移约束。在这个阶段，计算密集型的算子更为适合，如稠密卷积，而非depth-wise卷积；类似的，非矩阵操作形式的算子(如非线性、池化、BN)属于访存约束算子，应当尽可能少的使用。

在网络的后半部分更倾向于计算约束算子。以网络第三阶段为例，特征分辨率为

14 \times 14

，这为更强表达能力的算子(如LN、SE、自注意力等对吞吐量影响较小)提供了操作空间。

基于上述考量，本文提出了一种新型的混合架构FasterViT，它取得了更优的性能-吞吐量均衡，达成了新的Pareto前沿，可参考下图。

本文方案

上图给出了本文所提FasterViT架构示意图，从中可以看到:

在网络的前半部分，特征分辨率比较大，推理效率主要受访存约束，故它仅采用了卷积类操作以充分利用稠密卷积；
在网络的后半部分，特征分辨率比较小，推理效率会受计算约束，故它采用了attention等操作提升表达能力。

Stem 该部分通过两个连续

3\times 3

卷积(stride为2)将输入图像

\mathbf{x} \in \mathbb{R}^{H \times W \times 3}

转换为重叠块，输出维度为

的嵌入向量(embedding)。此外，每个卷积后接BN与ReLU激活函数。

DownBlock FasterViT延续了分层结构，相邻stage之间通过下采样模块进行空域分辨率收缩。下采样模块由LN与

stride=2

的

3\times 3

卷积构成。

ConvBlock Stage1与Stage2由残差模块构成，定义如下：

\hat{\mathbf{x}} = \text{GELU}(\text{BN}\text{Conv}_{3 \times 3}(\mathbf{x})) \\ \mathbf{x} = \text{BN}(\text{Conv}_{3 \times 3}(\hat{\mathbf{x}})) + \mathbf{x}

HierarchicalAttention 上图给出了本文所提新型窗口注意力示意图，通过如下操作，本文构建了分层注意力(Hierarchical Attention)的概念。

它以SwinT所引入的局部窗口为出发点，
引入了一种CT(Carrier Tokens)进行完整局部窗口信息聚合，第一个注意力模块作用于CTs以进行全局信息建模；
将局部窗口Token与CTs进行拼接并执行自注意力处理以促进局部与全局信息交互。

假设给定输入特征

\mathbf{x} \in \mathbb{R}^{H \times W \times d}

(为简单起见，这里假设

H=W

)，我们首先将输入特征拆分为

n \times n

局部窗口，表示如下：

\hat{\mathbf{x}} = \text{Split}_{k \times k}(\mathbf{x})

本文方案的关键在于Carrier Tokens(CTs)的构建，它有助于以更低的消耗保持注意力机制。首先，我们通过池化操作初始化CTs，描述如下：

\hat{\mathbf{x}}_c = \text{Conv}_{3\times3}(\mathbf{x}) \\ \hat{\mathbf{x}}_{ct} = \text{AvgPool}_{H^2\rightarrow n^2L}(\hat{\mathbf{x}})

注1: 这里的

\text{Conv}_{3\times3}

表示高效位置编码。注2: 上述CT初始化步骤仅在每个stage执行一次，每个Stage有自己独一的CTs。

在每个HAT模块内部，CTs执行如下注意力处理：

\hat{\mathbf{x}}_{ct} = \hat{\mathbf{x}}_{ct} + \gamma_1 \cdot \text{MHSA}(\text{LN}(\hat{\mathbf{x}}_{ct})) \\ \hat{\mathbf{x}}_{ct} = \hat{\mathbf{x}}_{ct} + \gamma_2 \cdot \text{MLP}_{d \rightarrow 4d \rightarrow d}(\text{LN}(\hat{\mathbf{x}}_{ct}))

接下来，为进行长短程空域信息建模，我们在

\hat{\mathbf{x}}_1, \hat{\mathbf{x}}_{ct, 1}

之间进行信息交互，处理流程如下：

\hat{\mathbf{x}}_{w} = \text{Concat}(\hat{\mathbf{x}}_1, \hat{\mathbf{x}}_{ct,1}) \\ \hat{\mathbf{x}}_{w} = \hat{\mathbf{x}}_{w} + \gamma_1 \cdot \text{MHSA}(\text{LN}(\hat{\mathbf{x}}_{w})) \\ \hat{\mathbf{x}}_{w} = \hat{\mathbf{x}}_{w} + \gamma_2 \cdot \text{MLP}_{d \rightarrow 4d \rightarrow d}(\text{LN}(\hat{\mathbf{x}}_{w}))

最后，上述所得Tokens将被拆分为拆分以用于后续分层注意力模块中：

\hat{\mathbf{x}}_1, \hat{\mathbf{x}}_{ct,1} = \text{Split}(\hat{\mathbf{x}}_w)

在每个stage，上述步骤将迭代执行多次。为促进长短程交互，在每个stage的尾部参考EdgeViT进行了全局信息传播。每个stage的输出计算如下：

\mathbf{x} = \text{Upsample}_{n^2L \rightarrow H^2}(\hat{\mathbf{x}}_{ct,1}) + \text{Merge}_{n^2k^2 \rightarrow H^2}(\hat{\mathbf{x}}_1)

此外，对于窗口Tokens与CTs，参考SwinV2注入了绝对位置偏置，而位置编码方式的灵活性促使其可以用于任务输入分辨率。