华中科技 & 地平线 提出 通用视觉 Backbone 网络 ， FLOPs减少 5.2倍，速度提高4.8倍 ！

近期，具有线性复杂度的序列建模网络在多种计算机视觉任务上实现了与视觉Transformer相似的性能，同时使用的FLOPs和内存更少。然而，它们在实际运行速度上的优势并不显著。 为了解决这个问题，作者引入了针对视觉的的门控线性注意力（GLA），利用其卓越的硬件感知能力和效率。作者提出了方向性门控来通过双向建模捕获一维全局上下文，并提出了一个二维门控局部性注入，以自适应地将二维局部细节注入到一维全局上下文中。 作者的硬件感知实现进一步将前向和后向扫描合并为一个单一 Kernel ，提高了并行性并减少了内存成本和延迟。所提出的模型ViG在ImageNet和下游任务上在准确度、参数和FLOPs之间提供了良好的折中，优于流行的基于Transformer和CNN的模型。 值得注意的是，ViG-S在仅使用27%的参数和20%的FLOPs的情况下与DeiT-B的准确度相匹配，在图像上运行速度提高了。在分辨率下，ViG-T使用的FLOPs减少了5.2倍，节省了90%的GPU内存，运行速度提高了4.8倍，并且比DeiT-T的top-1准确度高出20.7%。这些结果使ViG成为一个高效且可扩展的视觉表示学习方法。

1 Introduction

视觉Transformer（ViT）[18]通过将自然语言处理（NLP）中的高级序列建模层Transformer [91]引入计算机视觉领域，从而实现了视觉表示学习的革命性变化。它在各种视觉任务中都取得了巨大成功，作为一个多功能的 Backbone 网络。

然而，Transformer 中 softmax 注意力的二次复杂度对于其在高分辨率图像上的应用提出了重大挑战。许多研究[58; 17; 105]通过借鉴卷积网络的成功，例如限制局部窗口内的注意力计算，来克服这一限制。尽管这种方法实现了与常规CNN相似的线性复杂度，但在捕捉全局上下文方面存在不足。

这引发了一个关键问题：作者能否设计一个基本块，将Transformer 和 CNN的优点结合起来，提供全局感受野和线性复杂度？

来自NLP领域的线性时间序列建模方法[28; 71; 106; 74; 85; 42]的最新发展为此问题提供了一个有希望的解决方案。这些方法与RNN类似，通过将所有历史输入压缩为一个固定大小的状态，然后根据这个压缩状态关注当前输入，而不同于[91]关注所有历史状态的Transformers。为了进一步解决显式循环形式的墙钟效率限制，Mamba [28]，RWKV [71]，和GLA [106]引入了硬件感知实现，与高度优化的Transformers相比，展现出更高的准确性和效率。特别是Mamba已被广泛采用并适应于视觉任务。尽管这些方法在224分辨率上取得了卓越的分类准确度，并在高分辨率的FLOPs和内存方面表现出令人印象的效率，但它们在较低和更常见的分辨率下的墙钟效率通常与ViT/DeiT [18, 88]相当甚至更低。

这个限制激发作者探索将卓越的硬件高效GLA [106, 107]应用于推动线性复杂度视觉序列学习的效率和准确度边界，使其更具实践性，与已建立的高度优化的Transformers和CNNs相竞争。与基于SISO（单输入单输出）状态空间模型[29]的Mamba不同，GLA源自线性注意力[42]，它通过使用线性核近似标准注意力中的softmax，具有简单且对硬件友好的矩阵乘形式。为了进一步增强表现力，GLA引入了一种新颖的数据相关门控机制，以自适应地控制压缩状态遗忘率。

然而，传统的GLA模型是为单向建模设计的，存在时间依赖性。这使得当它应用于视觉任务时，其在NLP中的全局感知能力并不是真正的全局。尽管一些后续工作[38, 104, 40]在视觉领域的Mamba中探索了除了双向建模1D视觉序列之外的其他扫描方向，以近似传统注意力与视觉标记在任何方向和位置交互的能力，但这些方法由于频繁的非顺序内存访问模式而存在显著的内存低效问题，这使得它们在墙钟效率方面不太实用。作者坚持使用双向建模，因为其简单和内存友好的访问模式。

为了进一步利用视觉数据中固有的方向敏感性（即，来自不同视觉方向的信息在重要性上存在显著差异），作者提出了一种双向GLA（BiGLA）层，通过设计一种方向感知的门控机制来自适应地选择不同方向的全局上下文。这个设计在前向和后向方向之间共享大多数参数。尽管所提出的设计在1D视觉序列上实现了全局上下文，但它仍然无法捕捉到视觉数据的2D特性。

为了解决这个问题，作者提出了一个2D门控局部信息注入机制，以自适应地将卷积提取的2D局部信息压缩到序列建模BiGLA层提取的1D全局上下文中。此外，所提出的参数高效的双向设计使作者能够将双向扫描合并为一个单一 Kernel ，增强了实现的硬件意识，并减少了由额外方向引入的内存成本和延迟。

本文的主要贡献可以总结如下：

• 作者提出了ViG，一个通用的视觉 Backbone 网络，它将门控线性注意力（GLA）的线性复杂性与高效视觉序列学习所需的硬件感知能力相结合。ViG解决了以往基于Transformer和CNN方法局限性，将两者的优势结合，提供了一个高效且可扩展的解决方案。
• 为了更好地适应视觉任务，作者提出了三种关键设计，且开销最小：双向门控线性注意力机制，以捕捉视觉序列的全局1D上下文；方向性门控机制，以自适应地从不同方向选择全局上下文；以及2D门控局部性注入，将2D局部信息整合到1D全局上下文中。作者进一步提供了一个硬件感知的实现，将前向和后向扫描合并为一个单一 Kernel ，提高了并行性并减少了内存成本和延迟。
• 如Fig. 1所示，在 ImageNet 上，与现有的非层次化和层次化模型相比，作者的模型在准确性和参数方面取得了优越性能。对于下游密集预测任务，ViG在各种模型尺寸上的性能优于 ViT 和 VRWKV ，且计算成本更低。ViG的墙钟效率超过了线性复杂性的视觉序列学习方法 ，与已建立且高度优化的 ConvNeXt 和 SwinTransformer 相匹配。

2 Related Work

ViT [18]证明了通过引入NLP中的Transformer ，视觉表示学习可以以序列的方式进行。许多后续工作专注于在不改变softmax注意力的前提下提高ViT的效率和性能。最近，另一系列工作显示，二次softmax注意力可以被先进的类似RNN、线性时间序列建模方法所替代。Vision Mamba 基于线性时间序列建模Mamba块[28]，并通过引入额外的后向SSM层进行构建。VMamba 将交叉扫描引入Mamba并构建分层架构。LocalMamba 优化扫描方向以利用视觉的局部先验。Zigma 和PlainMamba 以之字形方式引入多个扫描方向。许多工作探索了Mamba在各种视觉任务中的有效性。相比之下，VisionRWKV [19]放弃了Mamba块，而是采用了来自NLP的线性复杂度RWKV块以用于视觉任务。

3 Preliminary

在本节中，作者介绍了从标准softmax注意力到高级门控线性注意力（GLA）的演变，为了简单起见，省略了多头机制。

图1：(a) 非层次结构体系结构 [19; 88; 117; 66] 和 (b) 层次结构体系结构 [56; 59; 58; 66; 77; 32] 在ImageNet-1K上的性能比较。作者提出的非层次结构的ViG和层次结构的ViG-H与流行的模型在参数和准确性方面表现出更优越的性能。特别是，提出的基本ViG块以线性复杂度实现全局感受野，而CNN [32; 77; 59]，普通softmax注意力 [88] 和基于窗口注意力 [58] 的块则不能。

4 Method

Overall Architecture

ViG Block

作为基本块的ViG块包括：1)一个能以线性复杂度方式利用图像的1D全局上下文的长期BiGLA层；2)一个能捕捉图像的2D局部细节的短期深度卷积层；3)一个可以自适应地结合全局和局部信息的门控机制；4)一个如同[89, 85, 106]中的SwiGLU FFN层用于通道混合。

4.2.1 Global Bidirectional Gated Linear Attention

在本工作中，作者采用双向建模，因其固有的简洁性和节省内存的访问模式。这种设计的核心是方向感知的门控机制，特别是遗忘门，它精细地控制着信息流。这对于位于目标边界处的标记尤为重要，因为来自不同方向的信息在重要性上差异显著。为了有效地利用这种方向敏感性，作者引入了双向门控线性注意力（BiGLA）层，如图3所示。该层通过共享除遗忘门之外的所有参数来实现参数高效，遗忘门针对每个方向进行了定制：

对于ViT风格的模型，作者将 Patch 大小设置为16，并堆叠12个ViG块。然后，通过直接调整嵌入维度，作者获得了3种不同大小的模型变体（ViG-T，ViG-S和ViG-B），它们的参数与DeiT-T，S和B相似。对于分层模型，作者也提出了3种变体（ViG-H-T，ViG-H-S和ViG-H-B），遵循SwinTransformer的设计。作者将 Patch 大小设置为4，并在不同的阶段应用作者提出的ViG块。详细的架构信息在附录B中提供。

Efficient Implementation

模型的实际效率不仅取决于理论上的FLOPs，而且在很大程度上由实现的硬件感知度决定。这意味着实现的模型应当：

1. 理解内存层次结构；
2. 利用专门的计算单元（GPU上的张量核心进行快速矩阵乘法）；
3. 具有高度的并行性。得益于GLA的硬件感知实现，ViG中最耗计算的部分可以用矩阵乘法表示，并在更快的SRAM上执行，而不是在较慢的高带宽内存（HBM）上，这通过利用张量核心并减少HBM的I/O开销，带来了卓越的墙面时间效率。

硬件感知的双向设计。鉴于将2D图像以1D展平序列表示的多向性，基于双向建模的开创性工作Vim选择分别调用两个顺序核来处理正向和反向，这种做法在并行性方面效率低下。在这项工作中，作者提出了一种硬件感知的双向设计，通过将方程（6）的正向和反向融合到一个核中，以实现更高的并行性。此外，由于BiGLA层的参数高效设计，作者可以减少在高带宽内存（HBM）中实现反向视觉序列的物质化，从而节省内存成本。

5 Experiment

作者进行了大量的实验来验证作者提出模型的有效性。作者在ImageNet [14]上展示了主要结果，并将它们与各种其他模型进行了比较。此外，作者在下游密集预测任务上对作者的模型进行了基准测试，包括在COCO [52]数据集上的目标检测以及在ADE20K [116]上的语义分割。

Image Classification

设置. 作者在ImageNet-1K数据集上进行了分类实验的训练，这是一个广泛使用的大规模图像分类基准。为了与之前的工作公平比较，作者主要遵循DeiT和SwinTransformer [88, 58]的训练和评估设置。具体来说，所有模型都是从零开始训练300个周期。图像在训练和评估时被裁剪为。更多细节在附录C中提供。

与非分层架构的比较. 表1将ViG与基于不同序列建模层的普通非分层架构进行了比较，包括Transformer、SSM和线性RNN。结果显示，所提出的ViG在各种模型尺寸上，在参数和准确度之间实现了优越的权衡，如图1（a）所示。值得注意的是，ViG-S的参数数量与DeiT-S几乎相同，并且其top-1准确度显著高出1.9%，与DeiT-B的性能相当（仅低0.1%），但参数数量减少了3.7倍，FLOPs减少了5倍，吞吐量提高了2倍。此外，ViG-B达到了82.6%的top-1准确度，比DeiT-B高出0.8%，比WRWKV-B高出0.6%，比S4ND-ViT-B高出2.2%。

在实际吞吐量方面，ViG超过了其他线性复杂度的序列建模方法，特别是在小型和微型模型尺寸上，比Vim-T快1.8倍，比Vim-S快1.6倍。与具有多方向扫描的对标方法相比，由于双向建模的记忆友好访问模式，ViG-T比LocalVim-T快5.2倍，比PlainMamba-L1快4.6倍。ViG仅比DeiT稍慢，这可以归因于Transformer的高度并行性和输入小型图像时的低FLOPs。如图4所示，随着图像尺寸的增加，ViG实现了指数级的优势。

与分层架构的比较. 表1将ViG-H与分层架构进行了比较，包括先进的基于CNN的RegNet和ConvNeXt，以及基于Transformer的SwinTransformer、基于SSM的S4ND和VAMmba。由于所提出的VGLA块的线性复杂度，ViG-H实现了与基于窗口的Transformer和CNN相似的FLOPs，同时具有全局感受野的优势。如图1（b）所示，ViG-H在不同模型尺寸上取得了最佳的top-1准确度。

实际吞吐量的结果表明，ViG-H超过了基于SSM的S4ND和VAMmba，并且与经过良好建立和高度优化的ConvNeXt和SwinTransformer的性能相匹配。

Object Detection

设置. 作者在COCO 2017数据集[52]上进行了目标检测和实例分割实验。作者使用Mask-RCNN[31]作为检测Head，并遵循VRWKV[19]将ViT-Adapter[6]集成到作者的普通ViG模型上。更多细节在附录C中提供。

设置. 作者在ADE20K[116]数据集上进行了语义分割的训练实验。作者使用UperNet[98]作为分割头，并按照VRWKV[19]采用ViT-Adapter[6]来适应作者的普通ViG进行分割。训练细节与VRWKV[19]和Vim[117]相同，详细内容在附录C中。

结果. 如表2所示，对于中等分辨率输入图像，在不同模型大小上，ViG在FLOPs和分割准确率上都优于二次复杂度的基于Transformer的ViT和线性复杂度的基于RNN的VRWKV。例如，在微小尺寸模型中，ViG比ViT在mIoU上高1.2，FLOPs低5G，比VRWKV在mIoU上高0.5，FLOPs低1.7G。在小尺寸模型中，ViG比ViT在mIoU上高1.7，FLOPs低14G，比VRWKV在mIoU上高0.7，FLOPs低6.3G。这些结果证明了与VRWKV相比，ViG在密集像素预测任务上的优越适应性。

Ablation Study

路线图。在表3中，作者展示了如何以最小的成本将因果序列建模的GLA引入到提出的ViG中。"第3行的改进的块嵌入层"意味着作者没有直接应用带有大核的卷积，而是采用了步长为8的卷积后接步长为2的卷积，这仅增加了0.03 M参数，但在几乎不影响推理效率的情况下提高了准确性。ViG-T引入了仅0.11M参数，并且比普通GLA的top-1准确率提高了11.1%，比DeiT提高了5.0%的top-1准确率。通过进一步引入硬件感知的双向实现，作者在效率上大幅提升（提高了21.9%的吞吐量并节省了13.3%的GPU内存），并缩小了与DeiT的差距，即使在低分辨率的图像上也是如此。

ViG的效率。在图4中，作者将ViG-T与Vim-T、VRWKV-T和DeiT-T进行了比较，重点关注理论上的FLOPs、实际的延迟和RTX 4090 GPU上的内存使用情况。作者测试了从到的递增输入图像分辨率的完整模型。得益于线性复杂度的序列建模，提出的架构相对于ViT的优势随着分辨率的增加而增长。当分辨率达到时，ViG-T使用的FLOPs降低了5.2，节省了90%的GPU内存，并且比DeiT-T快4.8。与具有线性复杂度的Vim相比，ViG也展示了1.6更低的FLOPs，节省了26%的GPU内存，并且运行速度快2.2。此外，在FLOPs、延迟和内存方面，ViG超越了VRWKV。

准确率_与_分辨率。在图4（d）中，作者测试了在不同分辨率下，在分辨率上训练的模型。随着分辨率的增加，ViG在准确性方面优于ViT、Vim、VRWKV，甚至基于分层CNN的ResNet50。结果表明，ViG得益于更好的2D感知能力，并在分辨率外推方面表现出优越的泛化能力。

6 Conclusion

在本文中，作者介绍了ViG，一个通用视觉 Backbone 网络，它将门控线性注意力（GLA）引入视觉领域，实现了高效的视觉表示学习。作者的方法通过在保持全局感受野的同时以线性复杂度运行，解决了传统Transformers和CNNs的固有局限性。

作者提出通过双向GLA建模和2D门控局部性注入来有效地捕捉全局上下文和局部细节，从而显著提高准确性。此外，作者的硬件感知实现减少了由额外方向带来的开销，提高了效率。ViG在低分辨率图像分类、中分辨率分割和高分辨率检测的优越实验结果使其成为现有通用视觉 Backbone 网络的非常有竞争力的替代品。

局限性。尽管ViG提出了硬件感知实现以提高效率，如表格3所示，并在高分辨率输入图像上显示出明显的优势，但在小尺寸输入图像上仍然略逊于DeiT。

作者将在未来的工作中进一步优化实现，以提高硬件感知能力。

参考

[1].ViG: Linear-complexity Visual Sequence Learning.