
©作者 | 娄蒙、俞益洲
单位 | 香港大学计算机系
香港大学计算机系俞益洲教授及其研究团队开发了一种新型的通用视觉 backbone—TransXNet,该模型同时融合了全局和动态性。其创新之处在于引入了一个即插即用的模块,名为 Dual Dynamic Token Mixer(D-Mixer)。
与以往的卷积和 self-attention 混合模块或全局与局部信息融合模块不同,D-Mixer 专注于动态性的捕捉。具体来说,D-Mixer 能够利用网络上一层的 feature maps 信息,生成用于提取当前 feature maps 的全局和局部动态权重,从而显著增强模型对多尺度特征的动态表征能力。
为了验证 TransXNet 的性能,研究团队分别在 ImageNet-1K、COCO 2017 和 ADE20K 数据集上进行了广泛的实验,结果表明,TransXNet 作为一个通用视觉模型,展现出了卓越的性能。

论文标题:
TransXNet: Learning both global and local dynamics with a dual dynamic token mixer for visual recognition
论文地址:
https://arxiv.org/abs/2310.19380
https://ieeexplore.ieee.org/document/10948454
代码地址:
https://github.com/LMMMEng/TransXNet
团队主页:

引言
在本文中,作者深入探讨了 self-attention 和卷积在权重计算方式上的不同,并基于此提出了改进方案。作者指出,self-attention 之所以性能卓越,部分原因在于其权重(即 attention matrix)会根据不同的输入动态变化,这种 input-dependent 的特性使得模型能够更好地捕捉输入数据的特征。
相比之下,传统的卷积操作中的权重是静态的,与输入数据无关,即 input-independent,这限制了其表征能力。据此,作者进一步分析了将 self-attention 和卷积简单结合时可能遇到的问题:
-
表征能力差异:由于卷积缺乏类似 self-attention 的动态性,这导致在卷积-attention混合模块中,self-attention的语义建模能力可能会被削弱,因为卷积部分的静态特性限制了整体模块的表征能力。 -
深层模型的挑战:在深层模型中,self-attention 能够利用先前特征图的全局和局部信息来生成动态的 attention matrix,从而提高模型的性能。然而,卷积核的静态特性使其无法利用这些信息来动态调整,这限制了模型在深层结构中的表现。

方法
为了解决上述问题,作者提出了一个即插即用模块 D-Mixer。该模块能够动态地捕获全局和局部信息,从而使网络兼具大感受野和稳定的归纳偏置。
如图 1(a)所示,对于给定的输入特征图,D-Mixer 首先会将其沿通道拆分为两部分,并分别将这两部分输入到两个不同的模块中,即 Input-dependent Depthwise Convolution(IDConv)以及 Overlapping Spatial Reduction Attention(OSRA)。
输出结果会进一步进行拼接,并且由一个轻量级的 Squeezed Token Enhancer(STE)进行整合。以下为不同模块的具体描述:

Input-dependent Depthwise Convolution (IDConv):如图 1(b)所示,对于任意输入特征图,作者首先使用一个自适应平均池化层聚合空间上下文,并将空间维度压缩为 K×K 大小,其中 K 为待生成的动态卷积核的大小。
紧接着,输出的特征图被送入两个 1×1 卷积层中,并生成一个多组空间注意力图,维度为 A∈G×C×K×K,其中 G、C 分别表示注意力图的组数以及输入特征图的通道大小。并且在 G 维度上使用 softmax 函数对注意力图进行归一化,使其具备自适应的选择能力。
最后,通过将注意力图和一组维度同为 G×C×K×K 的可学习参数进行逐元素相乘来生成动态卷积核,并将其用于输入特征图的特征提取。IDConv 可以根据上一层特征图携带的语义信息来动态地生成卷积核的权重,进而可以与 self-attention 进行深度嵌合。
Overlapping Spatial Reduction Attention (OSRA):如 1 图(c)所示,OSRA 的整体流程和 PVT 中提出的 Spatial Reduction Attention(SRA)类似,都采用了 token-to-region 的全局建模方式。
主要不同之处在于,SRA 在降采样的过程中使用的是 Non-overlapping 的卷积,即卷积核大小和步长大小一致。因此,一些连续的局部信息会在这个过程中被破坏,从而降低 key/value 分量包含的空间结构信息。
据此,OSRA 引入了 Overlapping 卷积来进行降采样,即卷积核的大小大于步长。此外,OSRA 还引入了一个轻量的局部算子(以 3×3 的 depthwise 卷积实现)来进一步提炼 key/value 分量的局部信息。
Squeezed Token Enhancer (STE):为了高效地将 IDConv 和 OSRA 处理后的特征进行融合,本文引入了一个轻量级的特征融合模块 STE。如图1(d)所示,STE 通过一个 3×3 的 depthwise 卷积和两个通道压缩的 1×1 卷积来实现特征融合。值得一提的是,STE 相较于直接使用一个全连接层进行特征融合具有更好的性能以及更优的计算复杂度。
TransXNet 整体架构:如图 2 所示,TransXNet 的基本构成部分主要包含 D-Mixer 和 Multiscale Feed-forward Network(MS-FFN)。其中,MS-FFN 使用了基于不同卷积核大小的 depthwise 卷积来提取多尺度信息。由于 TransXNet 同样沿用了 ResNet 和 Swin Transformer 中的金字塔设计,因此,该网络可以应用到各类视觉任务中。


实验

目标检测:如表 1 所示,在 COCO 2017 数据集上,TransXNet 同样展示出了更优的性能。值得注意的是,TransXNet 在处理小目标时比强调局部性的 Slide-Tranformer 性能更好,充分展现了 TransXNet 利用全局和局部动态性的强大能力。
语义分割:如表 2 所示,TransXNet 在 ADE20K 上也进行了全面的评估,其性能在与一些强大的 Vision Backbones 的比较中脱颖而出,并且有着更优秀的tradeoff。
消融研究:如表 3 所示,TransXNet 对核心模块进行了广泛的消融实验,例如与不同 token mixer 的性能对比以及和不同动态卷积的性能对比等。可以看出,TransXNet 中的核心部件相较于其他算法有着更加出众的性能。除此之外,原文还包含有更多的性能和效率对比及分析。




可视化研究
动态卷积在卷积 – Attention 混合模型中重要性: 如图 4 (a) 所示,在动态卷积的加持下,卷积-Attention 混合模型能够拥有更大的有效感受野(Effective Receptive Field, ERF)。其中,所提出的 IDConv 展现出了最好的性能,表明其可以更好地配合 self-attention 模块来捕获上下文信息。
不同 Vision Backbones 的 ERF 对比: 如图 4 (b) 所示,TransXNet 在具有最大感受野的同时还具备显著的局部敏感度,这是其他模型无法兼备的能力,进一步表明了全局和局部动态性的重要性。

(文:PaperWeekly)