摘要
这篇论文提出了一种新型的视觉Transformer,名为BiFormer,它采用了双层路由注意力(Bi-Level Routing Attention, BRA)机制。注意力机制是视觉变换器的核心构建模块,能够捕获数据中的长期依赖性。然而,这种能力的代价是计算量大和内存占用高。为了解决这个问题,作者提出了一种动态稀疏注意力机制,通过双层路由来实现计算的灵活分配,并具有内容意识。具体来说,对于一个查询(query),首先在粗粒度区域级别过滤掉不相关的键值对,然后在剩余候选区域的联合上应用细粒度的令牌到令牌的注意力。BiFormer在多项计算机视觉任务上验证了其有效性,包括图像分类、目标检测和语义分割。
拟解决的问题
- 传统注意力机制在计算所有空间位置的成对令牌亲和性时,具有高计算复杂度和内存占用。
- 现有方法通过引入手工设计的稀疏性来缓解这个问题,但这些方法要么使用静态模式,要么在所有查询之间共享采样的键/值子集,可能不是最优的。
创新之处
- 提出了一种新颖的动态稀疏注意力机制,即双层路由注意力(BRA),它通过粗粒度区域级别来过滤不相关的键值对,然后对剩余的路由区域应用细粒度的注意力。
- BRA利用稀疏性节省计算和内存,同时只涉及对GPU友好的密集矩阵乘法。
- 通过BRA作为构建模块,提出了一种新的通用视觉变换器BiFormer,它在查询适应性方面关注一小部分相关令牌,从而在不牺牲性能的情况下提高计算效率。
从传统的全注意力(Vanilla Attention)到所提出的双层路由注意力(Bi-level Routing Attention)的几种变体:
(a) Vanilla Attention(传统全注意力):这是最基本的注意力机制,它在全局范围内操作,计算所有空间位置的成对令牌亲和性,导致高计算复杂度和重内存占用。
(b) Local Attention(局部注意力):局部注意力限制了注意力操作仅在局部窗口内进行,减少了计算量,但可能牺牲了一些全局上下文信息。
(c) Axial Attention(轴向注意力):轴向注意力是一种沿特定轴(通常是x轴或y轴)的稀疏注意力模式,它减少了一些计算量,但仍然保持了一定程度的全局感知能力。
(d) Dilated Attention(扩张注意力):扩张注意力通过在局部区域内使用扩张的窗口模式来计算注意力,这种模式可以捕获更广泛的上下文信息,同时保持计算效率。
(e) Deformable Attention(可变形注意力):可变形注意力允许根据图像内容自适应地调整注意力窗口,通过变形规则网格来实现稀疏注意力,从而更灵活地捕获图像特征。
(f) Bi-level Routing Attention(双层路由注意力):这是本文提出的核心机制,它首先在粗粒度的区域级别上搜索top-k个相关区域,然后在这些区域的联合上应用细粒度的令牌到令牌注意力。这种机制实现了动态的、查询感知的稀疏性,有效平衡了计算效率和性能。
方法论
- 采用四阶段金字塔结构,每个阶段使用不同的BiFormer块数和通道数。
- 在每个BiFormer块中,首先使用3×3深度卷积来隐式编码相对位置信息,然后顺序应用BRA模块和两层MLP模块,分别用于跨位置关系建模和每个位置的嵌入。
Bi-Level Routing Attention (BRA)
1.区域划分和输入投影:将输入特征图 𝑋 划分为 𝑆×𝑆 个非重叠区域,每个区域包含的特征向量,通过线性变换得到查询 𝑄,键 𝐾,和值 𝑉 张量:
2.区域到区域路由:构建一个区域级别的有向图来确定每个区域应该关注的其他区域。首先,通过区域平均得到区域级别的查询和键:
通过Q和K计算区域间的亲和矩阵的邻接矩阵:即是计算两两区域之间的相关性
其中,接着,通过top-k操作保留和每个区域最重要的 𝑘 个连接(最相关的k个),形成路由索引矩阵:
其中,的第 i 行包含第 i 个区域最相关的 k 个索引。
3.令牌到令牌注意力:
利用路由索引矩阵 来应用细粒度的注意力。对于区域 𝑖 中的每个查询令牌,它将关注由索引的 𝑘 个路由区域中的所有键/值对。首先聚集键和值张量:
然后应用注意力机制:
其中,LCE(𝑉) 是一个局部上下文增强项,通常由深度卷积实现。
复杂度比较
复杂度为:
其他方法:
- Vanilla Attention:传统的全注意力机制,计算复杂度为 ,因为它需要对所有空间位置计算成对的令牌亲和性。
- Quasi-Global Axial Attention:一种近似全局的轴向注意力机制,计算复杂度为 ,它通过限制注意力范围来减少计算量,但仍高于BRA。
结论
作者提出的BiFormer通过双层路由注意力机制实现了高效的计算分配,具有动态、查询感知的稀疏模式。在图像分类、目标检测、实例分割和语义分割等多个流行视觉任务上,BiFormer展现出了卓越的性能。尽管BiFormer由于额外的GPU内核启动和内存事务而具有较低的吞吐量,但作者指出,通过工程努力,如GPU内核融合,可以减轻这个问题,并在未来的工作中探索具有硬件意识的高效稀疏注意力和视觉变换器。