地平线ViG基于视觉Mamba的通用视觉主干网络

ViG基于视觉Mamba的通用视觉主干网络

Vision Mamba的成功预示着将视觉表征学习转换为线性复杂度视觉序列表征学习具有巨大的潜力。尽管以Vision Mamba为代表的线性视觉序列表征学习结构在高清图像上展示了显著的效率优势，但在更为常见的分辨率，仍略逊于Transformer和CNN。

为了进一步推进线性复杂度视觉序列表征学习的效率，我们结合当今时代计算设备的硬件特性去设计结构，将自然语言序列建模中硬件计算更为友好的门控线性注意力模块GLA引入到视觉表征学习中，进一步设计参数高效的双向建模、长短上下文的动态门控机制和硬件感知的双向算子设计。通过以上设计，我们提出的ViG模型在各种任务上，精度、参数量、效率均优于主流的Transformer和CNN模型。

• 论文链接：

https://arxiv.org/abs/2405.18425

• 本工作已被AAAI 2025收录。

Vision Mamba的成功预示着将视觉表征学习转换为线性复杂度视觉序列表征学习具有巨大的潜力。尽管以Vision Mamba为代表的线性视觉序列表征学习结构在高清图像上展示了显著的效率优势，但在更为常见的分辨率，仍略逊于Transformer和CNN。

为了进一步推进线性复杂度视觉序列表征学习的效率，我们结合当今时代计算设备的硬件特性去设计结构，将自然语言序列建模中对硬件计算更为友好的门控线性注意力模块(Gated Linear Attention, GLA)引入到视觉表征学习中，并且进一步设计参数高效的双向建模，长短上下文的动态门控机制和硬件感知的双向算子设计。通过以上设计，我们提出的ViG模型在各种任务上，精度、参数量、效率均优于主流的Transformer和CNN模型。

背景介绍

Transformer的自注意力（Self-Attention, SA）机制需要对历史所有的输入进行交互，导致计算量随着输入序列长度呈现二次方增长：

而线性注意力(Linear Attention, LA)机制则将注意力机制中的softmax算子替换为简单的矩阵相乘，从而能够通过矩阵乘法的结合律，先进行KV的计算从而将计算量压缩为线性：

  线性门控注意力(Gated Linear Attention, GLA)机制则进一步引入门控机制控制对历史信息的遗忘和更新从而提高了原有线性注意力机制的表征能力：  

方法概述

双向门控注意力机制

GLA原是用于处理一维时序信号如文本，图片作为二维信号，在一维序列表达上具有多向的特点，针对这一特点，我们在原GLA的设计中，通过仅仅引入双向门控设计，便能够极大地提升对视觉信号的空间表征能力。   我们提出的BiGLA算子将前向和反向的视觉序列压缩到固定大小的隐状态和中去。  

硬件感知的双向算子设计

为了进一步提升效率，我们设计了一个硬件感知的双向实现，将BiGLA的前向扫描和后向扫描合并到一个Triton算子中，无需实例化反向序列，只需要维护一个单向视觉序列，便能够进行多向扫描和融合，从而大幅减小显存占用，并提升硬件运行速度。

长短上下文的动态门控制机

BiGLA算子中双向隐状态和是沿着一维视觉序列进行长上下文全局压缩，为了进一步加强对图片二维空间细节的感知，我们引入了短上下文的卷积门控设计：  

结构设计

围绕上述的基础ViG模块我们进一步搭建两个结构变种：类似于Vision Transformer的朴素直筒结构ViG和类似于CNN的层次化金字塔结构ViG-H。

实验结果

上图显示我们提出的ViG和ViG-H在参数量和精度的权衡上打败了先进的基于Transformer和CNN的模型。特别值得注意的是，我们基础ViG模块同时具备全局感受野和线性复杂度，这是之前CNN，基于原始注意力机制Transformer和基于窗口化注意力机制Transformer所达不到的。

上图两个表格也是显示我们在速度上能够和当前先进模型在224x224的图片输入上匹配。通过综合对比我们可以看到，ViG-S 与 DeiT-B精度相当并且速度快了一倍，但只用了其27%的参数和20%的计算量。当分辨率增加到1024时，ViG-T相较于基于Transformer的DeiT-T将计算量缩减5.2倍，GPU显存节省90%，速度快了3.8倍，精度更是高了20.7%。这些结果证实了ViG作为一个高效且可缩放的基础视觉骨干网络的广阔潜力。

下图则是进一步显示了，我们随着分辨率增大而凸显的性能优势，即随着输入图片分辨率的增大，ViG计算量更低，显存占用更少，速度更快，精度更高。

路线图也充分展示了设计的有效性。

我们的双向设计只引入了2%的参数量便得到了11.1精度的提升，同时我们设计的硬件感知实现减少了19%的推理代价和13%的显存占用，使其更加适合低算力场景的部署运行。

总结和展望

本工作提出的 ViG 方法在高效视觉序列建模领域展现了重要的创新性和实际应用价值。通过引入双向门控线性注意力 (BiGLA) 和二维门控局部注入机制，ViG 有效弥补了现有方法在全局感知能力和局部细节捕获上的不足，实现了兼具高效性与准确性的视觉表示学习。这一方法不仅在图像分类、目标检测和语义分割等任务中表现卓越，还显著提升了模型在高分辨率场景下的性能与资源利用效率。

尤其值得注意的是，ViG 的硬件感知优化设计充分结合了计算设备的硬件特性，大幅降低了内存占用和计算成本。通过在多种任务中的验证，ViG 展现了出色的精度、参数效率、显存优化和推理速度，为未来高清视觉信号处理以及基于大型语言模型 (LLM) 的多模态序列表征学习提供了强有力的技术支撑。