文章目录
- 前言
- 一、提出问题
- 二、模型结构
- 1. Pixel-focused Attention(像素聚焦注意力机制)
- 2. Aggregating Diverse Attentions in a Single Mixer Query embedding
- 3. ConvGLU
- 4. 整体模型架构
- 三、方法论
- 1. PFA设计原理
- 2. Aggregating Diverse Attention原理
- 3. GLU原理
- 4. 长度缩放余弦注意力
- 5. Position bias
- 四、实验
- 总结
前言
近几年,Vision Transformer(ViT)已成为各种CV任务的热门Backbone架构,也算是Transformer在CV领域的经典架构。注意力机制可以有效利用全局信息,获取每个像素间的相关性。但也由于这一特性,其平方复杂度和高内存消耗而面临挑战,这限制了其在高分辨率图像特征上的应用。本文主要介绍最新Backbone TransNeXt的改进和理论。
paper:https://arxiv.org/pdf/2311.17132.pdf
github:https://github.com/DaiShiResearch/TransNeXt
一、提出问题
为了减轻self-Attn中过多的计算成本,典型方法是局部注意力,限制了注意力在特征图上的一个窗口内。但是由于这种方式会牺牲全局感受野的能力,该方法通常需要交替堆叠与不同类型的Token混合器以实现跨窗口信息交换。
另一种典型的方法是空间下采样注意力中的Key和Value(例如,池化,网格采样)。由于它牺牲了Query对特征图的细粒度感知,这种方法也存在一定的局限性。
二、模型结构
在了解作者的研究思路之前,先来看看模型做了哪些改进:
1. Pixel-focused Attention(像素聚焦注意力机制)
整个模块的流程如下:
1)输入图像首先经过两个并行的运算产生K和V,其中蓝色部分通过一个滑动窗口产生向量K和V;黄色部分通过池化获得4个值,构成K和V。
2)两个部分的K分别与细粒度的Q进行注意力运算得到两个注意力图
3)添加位置编码后进行concate,共同经过softmax后分开与各自的V相乘,最后相加并经过线性变换获得最终输出。
2. Aggregating Diverse Attentions in a Single Mixer Query embedding
可以明显发现,相比上述模块,多了一个青色部分和一个额外的query embedding。embedding和Q相加,青色部分额外设置了一个可学习的token作为K和Q做注意力运算并和蓝色部分的Q*K进行融合。
3. ConvGLU
这部分结构就很简单了,不做详细说明
4. 整体模型架构
模型整体架构还是保留了原有的ViT架构。
TransNeXt采用与PVTv2相同的四阶段分层Backbone网络和重叠块嵌入。在阶段4中, 由于特征图大小已减少到 H/32 × W / 32 , 特征池化模块无法正常工作。作者采用了一种修改后的多头自注意力 (MHSA) 版本, 它应用Query嵌入和长度缩放余弦注意力。这与 PVTv2在第四阶段使用MHSA一致。对于1-4阶段的通道混合器, 作者使用卷积GLU与 GELU激活。扩展比也遵循PVTv2的 [ 8,8,4,4 ] 设置。
三、方法论
了解模型的改动之后,再来分析下作者的研究思路和改动原因(个人认为先看模型改进,再看思路更容易理解一些)
1. PFA设计原理
作者基于生物视觉特征进行设计,生物视觉对视觉焦点附近特征的分辨率更高,而对远处特征的分辨率较低。此外,眼球移动时,生物视觉这种特性在图像中任何位置的像素上都保持一致,这意味着像素级的平移等价性。
图中的红色星星表示Q的位置,黑色部分表示Q不能感知到的区域,可以发现传统局部窗口注意力机制没有全局视野,而如果将池化后的图像作为K和V,Q便可以以粗粒度感知全局信息(如图中第三个图),如果再引入局部窗口的K和V,便可以模拟生物视觉中的眼球追踪情况——局部窗口类似于焦点具有更高分辨率,远处特征从池化KV进行感知。
以上就是PFA的设计原理!!!
2. Aggregating Diverse Attention原理
在该部分模型层面设计了额外可学习的Q和K,原理如下:
Query embeding:将可学习的QueryToken集成到Backbone网络的注意力机制中(不同于传统的QKV注意力,它不使用输入中的Query,而是根据当前任务学习一个Query来执行交叉注意力),作者将这种方法归类为可学习的Key-Value(LKV)注意力,并与QKV注意力相平行。作者发现,将传统QKV注意力中所有QueryToken的Query嵌入(QE)添加起来,可以实现与Query无关的信息聚合效果,且额外开销可以忽略不计。
Positional attention: 信息聚合的另一种替代方法是使用一组可学习的Key,这些Key与来自输入的Query相互作用以获得注意力权重,即Query-Learnable-Value(QLV)注意力。这种方法与传统的QKV注意力不同,因为它破坏了Key和Value之间的一对一对应关系,导致为当前Query学习更多的隐式相对位置信息(额外学习一些特征)。因此,它通常与滑动窗口相结合在视觉任务中使用。与静态的亲和矩阵(如卷积或相对位置偏差)不同,这种通过数据驱动建模方法生成的亲和矩阵考虑了当前Query的影响,并可以基于它动态适应。
3. GLU原理
ViT时代的门控通道注意力:以前的工作,由Squeeze-and-Excitation(SE)机制代表,首先将通道注意力引入计算机视觉领域,该机制使用一个带有激活函数的分枝来控制网络输出。在门控通道注意力中,门控分支具有比value分支更大的决策权,最终决定相应的输出元素是否为零。从这个角度来看,SE机制解决了CNN结构中感受野不足的问题。然而,在ViT时代,全局感受野不再是稀缺的。各种全局Token混合器由自注意力表示,已经比全局平均池化实现了更高的全局信息聚合质量。这使得SE机制使用的全局池化方法表现出一些缺陷,例如该方法使特征图上的所有Token共享相同的门控信号,导致其通道注意力缺乏灵活性并过于粗糙。尽管如此,ViT结构缺乏通道注意力。最近的研究发现,将SE机制集成到通道混合器中可以有效提高模型鲁棒性(如图中第四个结构图)。
最近的研究表明,将3×3的零填充卷积引入到视觉Transformer中可以被视为一种条件位置编码(CPE),它有效地捕获了零填充带来的位置信息(图中第五个结构图)。
全新的Channel Mixer设计:
Gated Linear Unit (GLU) :GLU已在各种 NLP任务中显示出比多层感知机(MLP)更好的性能。GLU由两个线性投影组成,其中一个是通过激活函数激活的元素乘法。与SE机制不同,每个Token的gating信号都来自Token本身,并且其感受野大小不超过value分支的感受野。
作者发现,在GLU的gating分支的激活函数之前,简单地添加一个最小形式的卷积,可以使它的结构符合基于最近邻特征的门控通道注意力的设计概念,并将其转换为门控通道注意力机制。
特征分析:在卷积GLU(ConvGLU)中的每个Token都具有基于其最近精细特征的独特gating信号,这解决了全局平均池化在SE机制中的过于粗糙的问题。它也满足了某些没有位置编码设计的ViT模型,这些模型需要由卷积提供的位置信息
此外,这种设计的值分支仍保持与MLP和GLU相同的深度,使其易于反向传播。
4. 长度缩放余弦注意力
与缩放点积注意力不同,长度缩放余弦注意力使用了余弦相似度,已被观察到可以生成更适中的注意力权重,并有效地增强了大型视觉模型的训练稳定性。长度缩放余弦注意力通常将一个额外的可学习系数乘到Query和Key的余弦相似度结果上,使得注意力机制可以有效地忽略不重要的Token。
最近的研究发现,随着输入序列长度的增加,注意力输出的置信度会降低。因此,注意力机制的缩放因子应该与输入序列长度相关。有文章研究了当Query和Key被近似为具有大小 d \sqrt{d} d 的向量,缩放点积注意力熵的估计:
关于变量的定义,文章描述如下:
5. Position bias
这部作者其实主要使用了log-CPB,我其实看的不是很明白,大概意思是对于两条路径分别用log-CPB进行了测试,结果是其中一条路径效果不太好,作者在附录里也详细探讨了实验。log-CPB出自Swin-V2模型,这篇文章我暂时没看过,后续也了解一下。很感兴趣的同学可以去读读Swin-V2。
四、实验
总结
作者基于生物模仿视杆细胞视觉设计的Token混合器聚合注意力和一种具有门控通道注意力的通道混合器卷积GLU。另外作者也在新结构中设计了很多最新技术,将它们结合起来,提出TransNeXt。在预训练模型上达到了SOTA性能。此外,作者还提供了一个CUDA实现,在训练中实现了高达103.4%的加速,在推理中实现了60.5%的加速。
看论文可以发现作者做了大量实验的,公式也较为复杂,其实本文核心的公式就是聚焦注意力了,如下: