【ViT系列】TransNeXt: Robust Foveal Visual Perception for Vision Transformers

本文主要是介绍【ViT系列】TransNeXt: Robust Foveal Visual Perception for Vision Transformers，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

论文链接：https://arxiv.org/pdf/2311.17132.pdf

代码链接：https://github.com/DaiShiResearch/TransNeXt

一、摘要

1、引入了Pixel-focused Attention（PFA），它采用双路径设计。在一个路径中，每个查询对其最近邻特征具有细粒度的注意力，而在另一个路径中，每个查询对空间下采样特征具有粗粒度的注意力，从而实现全局感知。（模拟生物中央凹视觉和连续的眼球运动）
2、将查询嵌入和位置注意机制纳入到像素关注中，提出了聚合像素关注，简称为聚合注意。
3、提出了一种名为Convolutional GLU的新型通道混合器Mixer。它弥合了GLU和SE机制之间的差距，使每个token能够根据其最近邻图像特征获得通道注意力，从而增强了局部建模能力和模型鲁棒性。

二、实验结果介绍

TransNeXt-Tiny仅在ImageNet-1K上预训练，实现了84.0%的ImageNet准确率，超过了ConvNeXt-B。在COCO目标检测中，它使用DINO检测头获得了55.1的盒子mAP，优于在384×2分辨率上预训练的ConvNeXt-L高出1.7。我们的TransNeXt-Small/Base在仅仅5个时期的384×2分辨率微调后，实现了86.0%/86.2%的ImageNet准确率，超过了之前的最先进MaxViT-Base在30个时期微调的0.3%/0.5%。此外，在384×2分辨率上对极具挑战性的ImageNet-A测试集进行评估时，我们的TransNeXt-Small/Base模型实现了58.3%/61.6%的令人印象深刻的top-1准确率，明显优于ConvNeXt-L的7.6%/10.9%，为ImageNet-1K监督模型设定了新的鲁棒性基准。

三、关键问题

1、最初为语言建模开发的视觉Transformer的Transformer编码器设计在下游计算机视觉任务中存在固有限制。具体而言，自注意力中全局亲和矩阵的计算由于其二次复杂度和高内存消耗而构成挑战，限制了其在高分辨率图像特征上的应用。

2、现有的研究提出了大量稀疏注意力机制。

a.一种代表性方法是局部注意力[41]，它将注意力限制在特征图上的一个窗口内。然而，由于有限的感受野，这种方法通常需要与不同类型的Token Mixer交替堆叠，以实现跨窗口信息交换。
b.另一种代表性方法是对注意力的键和值进行空间下采样（如池化[61,62,64]，网格采样[57]）。由于这种方法牺牲了查询对特征图的细粒度感知，它也存在一定的局限性。最近的研究[7,57]交替堆叠了空间下采样注意力和局部注意力，取得了令人瞩目的性能结果。然而，最近的研究[12,59]和实验[28]表明，具有残差块[19]的深度网络表现出类似于更浅网络的集成效果，这表明通过堆叠块实现的跨层信息交换可能并不像预期的那样有效。

3、由于深度退化效应，许多高效的ViT模型无法通过堆叠形成足够的信息混合。即使通过深层堆叠，它们的窗口分区的痕迹总是形成不自然的伪影，如图2所示。

四、具体理论

1、Aggregated Pixel-focused Attention 像素聚焦注意力

像素聚焦注意力的设计分为双通道的设计、Activate and Pool和填充掩码Padding mask。

双通道设计：包括以查询为中心的滑动窗口注意力和池化注意力。此外，为了在两个注意力路径之间引入耦合，计算两个路径的查询-键相似性结果的重要性，并在同一个 softmax中计算结果，然后再分别作用到中心和非中心区域。这导致了细粒度和粗粒度特征之间的竞争，将像素聚焦注意力转化为多尺度注意机制。

原文描述如下图所示：
在这里插入图片描述

Activate and Pool：空间维度的下采样使用平均池化算子，但它会显著丢失信息。因此，在特征图池化之前使用单层神经网络进行投影和激活，以提前压缩和提取有用信息，从而提高下采样后的信息压缩率。在池化后使用层归一化来规范化输出，以确保 $X$ 和 $\sigma(X)$ 的方差一致性。所提出的下采样算子，称为“Activate and Pool”，可以用以下方程表示：

$\sigma(X) = LayerNorm(AvgPool(GELU(Linear(X)))),$

用“Activate and Pool”机制替换了PVTv2-li [62]中的下采样模块，并为CIFAR-100上的消融实验设计了一个2M大小的模型。该模块将PVTv2-li的top-1准确率从68.1%提高到70.4%。

填充掩码：在滑动窗口路径中，位于特征图边缘的像素不可避免地与边界外的零填充计算相似性。为了防止这些零相似性影响softmax操作，使用填充掩码将这些结果设置为 $-\inf$ 。

2、Aggregating Diverse Attentions in a Single Mixer 聚合多样化注意力在单个混合器中

(1) 查询嵌入
与传统的QKV注意力不同，因为它不使用输入的查询，而是学习由当前任务定义的查询来执行交叉注意力。因此，文中将这种方法归类为可学习键值（LKV）注意力，与QKV注意力相对应。**在传统的QKV注意力中为所有查询标记添加一个可学习的查询嵌入（QE）**可以实现类似的信息聚合效果，而额外开销可以忽略不计。我们只需要修改公式1如下：
$S_{(i,j)∼ρ(i,j)} = (Q_{(i,j)} + QE)K^{T}_{ρ(i,j)},$

$S_{(i,j)∼σ(X)} = (Q_{(i,j)} + QE)K^{T}_{σ(X)}.$

(2) 位置注意力
查询可学习值（QLV）注意力（使用一组可学习的键与源自输入的查询交互以获得注意力权重），相对于静态相对位置偏差表现出更大的鲁棒性，并且可以进一步增强局部建模能力。在每个注意力头中引入一组可学习标记 $R^{d \times k^{2}}$ ，允许这些标记与查询交互以获得额外的动态位置偏差，并将其添加到 $A_{(i,j)∼ρ(i,j)}$ 。仅使用这种增强需要额外的计算开销 $HW k^{2} C$ ，修改方程4如下：
$(A_{(i,j)∼ρ(i,j)} + Q_{(i,j)}T) V_{ρ(i,j)} + A_{(i,j)∼\sigma(X)} V_{\sigma(X)}.$

3、Overcoming Multi-scale Image Input 克服多尺度图像输入

(1) 长度缩放余弦注意力
缩放余弦注意力通常会将一个额外的可学习系数 $\lambda$ 乘以查询和键的余弦相似性结果，从而使注意力机制能够有效地忽略不重要的标记。最近的研究发现，随着输入序列长度的增加，注意力输出的置信度会降低。因此，注意力机制的缩放因子应与输入序列的长度相关联。进一步提出，注意力的设计应该表现出熵不变性，以促进更好地泛化到未知长度。提供了一个在查询和键被近似为具有 $\sqrt{d}$ 大小的向量时，缩放点积注意力在序列长度 $n$ 时熵的估计。公式8如下：
$H i \approx l o g n - 0.24 λ d + O (1) .$
对于余弦相似性，将查询和键定义为分别沿其头维度应用L2归一化的Q和K，它们的大小均为1。为了保持熵不变性并忽略常数项，我们设置 $λ≈\frac{logn}{0.24}$ 。鉴于公式8仅仅是一个估计，设置 $\lambda = \tau logn$ ，其中 $\tau$ 是一个可学习变量，为每个注意力头初始化为1/0.24。我们提出长度缩放余弦注意力如下：
$softmax(\tau log N * \dot Q \dot K^{T}) V,$
这里，N表示每个查询与之交互的有效键的计数，不包括被掩码标记的计数。特别是，在应用于Transformer Decoder时，被因果掩码掩盖的未来标记不应计入N。在像素关注的背景下，N被计算为 $∥\sigma(X)∥ - ∥\mu(i,j)∥$ ，其中 $\mu(i,j)$ 代表位置(i,j)处的填充掩码标记集合。

(2) 克服多尺度图像输入的位置偏差
为进一步增强针对多尺度图像输入的像素关注力的外推能力，文中采用不同的方法在两个路径上计算 $B (i, j) \sim ρ (i, j)$ 和 $B(i,j)∼\sigma(X)$ 。在池化特征路径上，使用对数间隔连续位置偏差（log-CPB）[41]，一个包含ReLU [47]的2层MLP来从 $Q (i, j)$ 和 $\sigma(X)$ 之间的空间相对坐标 $Δ(i,j)∼\sigma(X)$ 计算 $B(i,j)∼\sigma(X)$ 。在滑动窗口路径上，直接使用可学习的 $B (i, j) \sim ρ (i, j)$ 。一方面，这是因为滑动窗口的大小是固定的，不需要通过log-CPB对未知相对位置偏差进行外推，从而节省计算资源。另一方面，作者观察到使用log-CPB计算 $B (i, j) \sim ρ (i, j)$ 会导致性能下降。文中认为这是因为 $Δ (i, j) \sim σ (X)$ 表示细粒度标记和粗粒度标记之间的空间相对坐标，而 $Δ (i, j) \sim ρ (i, j)$ 表示细粒度标记之间的空间相对坐标，它们的数值含义不同。

(3) 聚合注意力
通过应用前述多样的注意力聚合方法和技术来增强多尺度输入的外推能力，提出了像素关注力的增强版本，称为聚合像素关注力，简称为聚合注意力（AA）。

在这里插入图片描述

4、特征分析

(1) 计算复杂度
在这里插入图片描述
当池化大小 $H_{p} \times W_{p}$ 设置为与输入大小无关的值时，Ω(PFA)和Ω(AA)都与输入序列的长度成线性关系。这意味着PFA和AA都可以以线性复杂度模式进行推断。

(2) 最佳精度-效率权衡
通过实证研究，作者观察到滑动窗口的大小对模型性能影响微乎其微。因此，采用最小形式的3×3滑动窗口来捕获接近视觉焦点的特征，显著减少了计算和内存消耗。将这归因于存在池化特征路径，使每个查询具有全局感受野，从而大大减少了扩展滑动窗口大小以扩展感受野的需求。详细的消融研究结果和讨论可在附录中找到。

5、Convolutional GLU

(1) 门控线性单元（GLU）[11,51]是一种通道混合器，已被证明在各种自然语言处理任务中优于多层感知器（MLP）。GLU由两个线性投影组成，这两个投影进行逐元素相乘，其中一个投影由门控函数激活。与SE机制不同，GLU的每个标记的门控信号来自标记本身，并且没有比值分支更大的感受野。

(2) 在GLU的门控分支的激活函数之前简单地添加一个最小形式的3×3深度卷积可以使其结构符合门控通道注意力的设计概念，并将其转换为基于最近邻特征的门控通道注意力机制。文中将这种方法命名为卷积GLU，如图4所示。
在这里插入图片描述

(3) **卷积GLU（ConvGLU）**中的每个标记都具有一个独特的门控信号，基于其最近的细粒度特征。这解决了SE机制中全局平均池化过于粗粒度的缺点。它还满足了一些ViT模型的需求，这些模型没有位置编码设计，需要深度卷积提供的位置信息。此外，该设计的值分支仍然保持与MLP和GLU相同的深度，使其易于反向传播。当将参数体积与扩展比率为R和卷积核大小为k×k的卷积前馈（ConvFFN）[62]保持一致时，ConvGLU的计算复杂度为 $C^{2} + \frac{2}{3} RHW Ck^{2}$ ，小于ConvFFN的 $2RHW C^{2} + RHW Ck^{2}$ 。这些属性使ConvGLU成为一种简单但更强大的混合器，满足ViTs的多样化需求。

6、Architecture Design of TransNeXt

为了确保后续消融实验的一致性，TransNeXt采用了与PVTv2相同的四阶段分层主干和重叠补丁嵌入。聚合注意力在第1-3阶段的池化特征大小也设置为 $H /32 \times W /32$ ，与PVTv2相同。在第4阶段，由于特征图大小已经减小到 $H /32 \times W /32$ ，特征池化模块无法正常工作。文中采用了一种修改版的多头自注意力（MHSA），应用了查询嵌入和长度缩放的余弦注意力。这与PVTv2在第四阶段使用MHSA的做法一致。对于第1-4阶段的通道混合器，使用带有GELU激活的卷积GLU。扩张比例也遵循PVTv2的**[8,8,4,4]设置。为了确保与典型MLP参数的一致性，卷积GLU的隐藏维度是设置值的2.3倍。此外，将头维度设置为24**，以便在通道维度上能够被3整除。

五、实验

1、ImageNet-1K classification
在这里插入图片描述

2、Object detection, instance segmentation and Semantic segmentation
在这里插入图片描述

3、Multi-scale Inference
在这里插入图片描述

4、A roadmap from PVT to TransNeXt
在这里插入图片描述

5、Attention Map
在这里插入图片描述

6、Architecture Setting
在这里插入图片描述

更多实验结果看附录。提供了一个CUDA实现，可在训练中实现高达103.4%的加速和推断中的60.5%的加速。更详细的实验数据和讨论见附录。

六、总结

文中提出了一种基于仿生中央视觉设计的Token Mixer，即聚合注意力，以及一种带有门控通道注意力的通道混合器，即卷积GLU。将它们结合起来提出了一种强大且高度稳健的视觉模型TransNeXt，该模型在分类、检测和分割等各种视觉任务中实现了最先进的性能。TransNeXt在多尺度推断中的出色表现突显了其在解决深度退化问题方面相对于大内核策略的优势。

这篇关于【ViT系列】TransNeXt: Robust Foveal Visual Perception for Vision Transformers的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！