超越VMamba | JamMa突破特征匹配瓶颈,用50%参数量实现对Transformer的性能反超

大模型向量数据库机器学习

点击下方卡片,关注

「集智书童」

公众号

点击加入👉

「集智书童」

交流群

picture.image

picture.image

picture.image

picture.image

picture.image

picture.image

导读

现有最先进的特征匹配器利用Transformer捕获长距离依赖关系,但受限于高空间复杂度,导致训练需求大和推理延迟高。在性能和效率之间取得更好的平衡仍然是特征匹配领域的一个挑战。受Mamba的线性复杂度

的启发,作者提出了一种基于Mamba的超轻量级匹配器,命名为JamMa,它可以在单个GPU上收敛,并在推理中实现了令人印象深刻的性能-效率平衡。为了释放Mamba在特征匹配中的潜力,作者提出了联合Mamba(JEGO),它采用了一种扫描-合并策略,能够实现以下功能:(1)对两幅图像进行联合扫描以实现高频相互交互,(2)通过跳过步骤进行高效扫描以减少序列长度,(3)全局感受野,(4)全向特征表示。凭借这些特性,JEGO策略在特征匹配任务中显著优于VMamba和EVMamba提出的扫描-合并策略。与基于注意力的Sparse和半密集匹配器相比,JamMa在性能和效率之间表现出更优的平衡,以低于50%的参数和FLOPs实现了更好的性能。

项目页面:https://leoluxxx.github.io/JamMa-page/

  1. 引言

特征匹配,旨在建立两张图像之间准确的对应关系,是包括运动结构(SfM)[34]和同时定位与建图(SLAM)[12]在内的各种关键任务的基础。然而,特征匹配面临着来自视点、光照和尺度变化的挑战。此外,对实时性能有要求的应⽤对匹配过程的效率施加了巨大的压力。

特征匹配方法通常分为三类:Sparse、半密集和密集方法。Sparse方法[4, 25, 33, 35],与关键点检测器[8, 23]协同工作,旨在建立两组关键点之间的对应关系。因此,当检测器在具有挑战性的场景(如无纹理场景)中表现不佳时,将严重影响Sparse匹配器。相比之下,半密集和密集方法[5, 9, 11, 36]直接在网格点或所有像素之间建立对应关系,增强了在具有挑战性的场景中的鲁棒性。然而,半密集和密集匹配器通常依赖于Transformer[16, 40]来建模密集特征中的长距离依赖关系,导致在处理高分辨率图像时复杂度极高,不切实际。在性能和效率之间实现最佳平衡仍然是在特征匹配中的一项关键挑战。

为了在具有挑战性的场景中实现高效率和鲁棒性,作者探讨了使用Mamba加速半密集范式(semi-dense paradigm)的潜力。Mamba能够以线性复杂度捕捉长距离依赖关系,并在自然语言处理(NLP)领域取得了显著的成果。然而,由于Mamba最初被设计为NLP的序列模型,作者识别出三个关键挑战,这些挑战阻碍了其在特征匹配方面的适用性,并提出了一种联合Mamba与JEGO(Joint Mamba with JEGO)策略来有效解决这些问题。

(1) 缺乏相互交互。Mamba旨在建模单一特征集,即内部交互,而特征匹配本质上需要两组特征之间的交互,即相互交互。因此,需要一种新的扫描策略来有效地建立两张图像之间的跨视图依赖关系。为此,作者探索了两种候选策略:联合扫描和顺序扫描,如图4所示。作者的研究发现,联合扫描,通过在两张图像之间交替扫描以促进高频相互交互,显著优于顺序扫描,在顺序扫描中,每张图像都是单独且按顺序扫描的。考虑到相互交互对于特征匹配的必要性,联合扫描被选为JEGO策略的核心组件,这也导致了“联合Mamba”这一术语的产生。

picture.image

(2) 单向性。Mamba 是单向的,因为它通过一维正向扫描生成序列。然而,在二维图像中,特征在四个基本方向上表现出关系,需要开发新的扫描策略。如图2所示,Vim [46] 和 VMamba [21] 分别提出了双向和四向扫描,但代价是总序列长度增加了2倍和4倍。EVMamba [27] 通过扫描过程中的跳过操作保持了与原始特征数N相等的总序列长度,但未能实现全向性。提出的JEGO策略将跳扫与四向扫描相结合,实现了全向性和总序列长度为N。

picture.image

(3) 因果性。Mamba 作为因果模型运行,允许序列中的每个特征只能感知到前面的特征。因此,只有序列的末尾E可以访问全局感受野。许多研究已经验证了对于视觉模型,包括视觉Mamba模型,大感受野的重要性。如图2所示,Vim [46] 和 VMamba [21] 通过反向扫描(虚线)实现了全局感受野,将总序列长度增加到2N和4N。EVMamba [27] 保持高效率,但只正向扫描序列,这限制了全局感受野仅限于右下角。在JEGO策略中,四个方向的扫描被安排以确保平衡的感受野,如图2所示,小感受野特征始终与较大感受野的特征相邻。然后使用一个称为聚合器的CNN来将来自四个方向的全局信息聚合到具有小感受野的特征中。作者发现,在平衡感受野上的简单聚合器具有惊人的有效影响,直观地确保每个特征都是全局的和全向的,如图6所示。

picture.image

总结而言,本文提出了以下贡献:

  • • 作者提出了一种基于联合Mamba和JEGO策略的超轻量级半密集特征匹配器——JamMa。
  • • 提出了一种名为JEGO的扫描合并策略,该策略能够高效地联合扫描跨视图特征,并生成全局全向特征。
  • • 在平衡感受野上的高频互交互作用和局部聚合器产生了出人意料的鲁棒特征,同时计算开销最小。定量和定性实验展示了JamMa在性能效率平衡方面的优异表现,并突出了JEGO策略的关键作用。
  1. 相关工作

2.1 局部特征匹配

Sparse匹配包括如SIFT [23]、SURF [2]、BRIEF [3] 和 ORB [31] 等已知的定制方法。随着深度学习的兴起,研究重点转向基于学习的方法。R2D2 [29] 和 SuperPoint [8] 等模型使用卷积神经网络 (CNN) 来提高鲁棒性。SuperGlue [33] 首次使用 Transformer [40] 对Sparse特征之间的全局关系进行建模,通过自注意力和交叉注意力丰富特征。最近的基于注意力的Sparse方法 [4, 20, 25, 35] 通常解决计算效率低下的问题,但仍然严重依赖于检测器的性能。半密集匹配方法由 LoFTR [36] 领先提出,它提出了一种从粗到细的范式,在网格之间建立粗匹配,然后使用精细特征调整匹配点。ASpanFormer [5] 提出了一种新的替代线性注意力 [16] 的方法,该方法根据不确定性自适应地调整注意力的范围。MatchFormer [41] 放弃了基于 CNN 的特征提取器,并提出了一种纯注意力匹配器以提高鲁棒性。半密集方法在无纹理场景中表现出色,因为它们匹配图像上分布的网格。然而,大量网格点也引入了更大的计算负担,尤其是在匹配器利用注意力机制建立长距离依赖关系时。

密集匹配方法直接估计密集光流以在两幅图像之间为每个像素建立对应关系。一些方法[7, 18, 30]利用4D相关体积来预测密集光流。最近,DKM[9]和RoMa[11]引入了核化的全局匹配器和嵌入解码器来初步预测粗略光流,随后使用堆叠特征图进行细化。虽然DKM和RoMa在准确性和鲁棒性方面优于Sparse和半密集方法,但它们需要显著更多的参数和更长的运行时间。

2.2. 状态空间模型

状态空间模型(SSM)最初是为了建模连续线性时不变系统[43]而设计的,其中输入信号

被映射到其输出信号

作为

是SSM参数。为了将SSM应用于离散系统,例如序列到序列的任务,S4 [15] 提出使用零阶保持规则将SSM参数转换为它们的离散对应参数。然而,S4在所有时间步之间共享参数,这严重限制了其表示能力。Mamba [14] 提出了S6,该模型使SSM参数依赖于输入序列。这种数据依赖性显著提升了Mamba的性能,使其与Transformer相当。同时,Mamba高效的主要原因在于,SSM中的重复计算可以通过全局卷积核一次性完成,该核可以通过SSM参数进行预计算。

在Mamba取得成功之后,该框架在计算机视觉任务中的应用出现了激增。视觉Mamba模型通常将2D特征图扫描成1D序列,以Mamba对其进行建模,然后将1D序列合并回2D特征图。Vim[46]和VMamba[21]分别提出了双向扫描和四向扫描,将原本的因果Mamba扩展以实现全局感受野。EVMamba[27]提出了跳过扫描以减少序列长度,从而显著提高效率。上述视觉Mamba模型主要关注单图像任务,例如图像分类。相比之下,作者提出了联合Mamba,该模型旨在联合建模两张图像以进行特征匹配。

  1. 研究方法

3.1. JamMa概述

JamMa的整体架构如图3所示。作者使用CNN编码器提取粗特征

和细特征

。随后,粗特征通过联合Mamba进行处理,即JEGO扫描

Mamba

JEGO合并,以执行内部和相互交互,感知两张图像的全局和全方位信息。最后,作者利用粗到细匹配(C2F)模块生成匹配结果。具体来说,作者建立了双向粗匹配,并在细匹配和亚像素细化模块中调整粗匹配点[38]。

picture.image

3.2 局部特征提取

作者采用ConvNeXt V2 [44]作为编码器,从图像

中提取局部粗特征

和精细特征

作者观察到ConvNeXt的设计在轻量级设置中表现出色。具体来说,一个基于ConvNeXt的编码器,参数量为0.65M,足以支持JamMa实现具有竞争力的性能。

3.3 联合曼巴

联合Mamba,包括JEGO扫描、Mamba和JEGO合并,被提出用于高效地建立粗特征

的全局依赖关系,以实现鲁棒的粗匹配。

3.3.1 JEGO 扫描

联合扫描。现有的Mamba模型是为单图像任务设计的,而特征匹配需要两组特征之间的交互。例如,基于注意力的匹配器利用自注意力和交叉注意力进行内部和相互交互。因此,需要一种新的扫描策略来实现内部和相互交互。

作者首先将粗略特征

在水平和垂直方向上进行拼接。

基于

,作者探讨了两种候选扫描策略:顺序扫描和联合扫描,如图4所示。在顺序扫描中,在移动到下一张图像之前,先扫描一张图像的所有特征。相比之下,联合扫描在两张图像之间交替进行,以实现高频相互交互,这对跨视图依赖性有益。如表4所示,联合扫描在特征匹配方面比顺序扫描显著更有效。高效四向扫描。除了提出的联合扫描外,JEGO扫描还结合了EVMamba [27]中引入的跳过扫描策略,该策略以间隔扫描特征而不是彻底扫描以减少序列长度。然而,EVMamba只执行正向扫描,导致全局感受野仅限于图像的右下角,缺乏全向性。通过如图2所示的四个方向上安排序列的起始点

和终止点

,JEGO扫描实现了平衡的感受野和全向性。

picture.image

通过在

上进行水平扫描和在

上进行垂直扫描,生成了四个方向上的长度为

的序列

分别代表右侧、左侧、上侧和下侧扫描的序列。然后,这些序列通过四个Mamba块独立处理,以建立长距离依赖关系。

操作

表示对每个通道进行切片,从高度

和宽度

开始,每隔

步跳过一次。每个序列的起始点

的排列方式为

3.3.2 JEGO 合并

在各自序列中建立长程依赖关系后,将序列合并成二维特征图对于视觉任务至关重要。针对此问题,作者提出了JEGO合并策略,以生成两张图像的全局全方位特征图。

首先,作者将四个序列中的特征恢复到其原始扫描位置,从而得到联合特征图

然后,作者将联合特征图

在空间上分割为图像

的特征

。水平和垂直特征相加以获得合并特征图

尽管

在宏观层面上是全向的,并表现出平衡的接受域,但每个单独的特征仍然感知单向信息,并在微观层面上具有局部接受域,除了四个序列的端点 E。因此,作者提出使用门控卷积单元作为局部信息聚合器,在

窗口中聚合信息,将来自不同方向和接受域的信息巩固向中心。

表示具有3个核大小的二维卷积,而

是门控信号,它根据输入确定哪些信息被过滤。聚合器确保

在微观层面上具有全局性和全方位性,这对于JamMa的性能至关重要。

3.4 粗到细匹配

基于聚合的粗特征FA、B和精细特征

,作者在XoFTR [38]中采用了粗到细的匹配模型来生成匹配,如图5所示。

picture.image

粗匹配。两集合粗特征之间的粗相似度矩阵

计算如下:

其中

是一个温度参数,

表示内积。然后作者对相似度矩阵分别进行行 Softmax

和列 Softmax

操作,从而获得匹配概率矩阵

作者可以通过

以多对一的方式分配

[24],同样地,可以通过

分配给

,从而得到两组匹配

。具体来说,作者根据以下两个标准建立

: (1) 在

中的行最大值或

中的列最大值,以及 (2) 置信度大于阈值

。然后,将

的并集视为粗匹配

,它比通过双软最大化操作得到的单对一匹配更鲁棒。

精细匹配。对于每个粗匹配,从精细特征

中裁剪出两个

的特征窗口

,其中

是粗匹配的数量。然后,

通过 MLP-Mixer [37] 进行处理,以轻量级的方式实现相互交互。类似于公式 (9),通过内积计算精细相似度矩阵

,并应用双软最大化(Dual-Softmax)计算精细匹配概率矩阵

进一步,作者应用互最近邻(MNN)准则来建立一对一精细匹配

。由于基于分类的匹配

无法达到亚像素精度,因此采用基于回归的细化方法进一步调整匹配。对于每个精细匹配,从特征窗口中裁剪出

,然后由 MLPMixer 处理。使用 Tanh 激活函数计算偏移量

,然后将这些偏移量加到

上进行调整,从而得到最终的亚像素匹配

3.5 监督

作者的损失函数包含三个部分:粗匹配损失

,细匹配损失

,以及亚像素损失

。真实匹配矩阵

由相机位姿和深度图生成。

粗匹配损失

是通过计算

之间的Focal Loss FL 得到的。

精细匹配损失

是通过计算 Pgt 和

之间的Focal Loss来得到的。

与XoFTR [38]一致,作者实现了对称极线距离函数来计算亚像素精炼损失。

是匹配点

的齐次坐标,而

表示真实基础矩阵。

  1. 实验

4.1 实施细节

所JamMa使用Pytorch [26] 实现。在特征提取方面,作者采用了ConvNeXt V2-N的前两个阶段,该阶段包含

个参数。JamMa由4个Mamba块组成,可以处理四个方向上的序列。粗特征和细特征的通道维度分别为

。粗特征分辨率是

,细特征分辨率是

。温度参数

设置为 0.1,粗匹配阈值

设置为 0.2。JamMa在MegaDepth数据集 [19] 上训练了30个epoch,批大小为2,使用Adam W优化器 [22],并且没有针对其他任务进行微调。在训练过程中,图像被调整大小并填充到

的大小。初始学习率设置为 0.0002,并采用具有1个epoch线性预热阶段的余弦衰减学习率调度器。网络训练在单个NVIDIA 4090 GPU上需要

小时,所有评估也在该GPU上进行。

4.2 相对位姿估计

数据集。作者在MegaDepth数据集[19]上评估了姿态估计的匹配器。测试图像被调整大小并填充至

,与训练大小相同。根据[20]的建议,作者使用LO-RANSAC[17]来估计所有方法的本质矩阵,因为它比传统的RANSAC更鲁棒,从而提高了评估的可靠性。每种方法的详细评估设置见附录。

指标。在阈值

下姿态误差的累积曲线下面积(AUC)被报告。作者还报告了包括参数、FLOPs和运行时间在内的三个效率指标,以及六个效率和性能指标的平均排名。

结果。在14种评估方法中,将平均排名7.5定义为预期的性能-效率平衡,例如,最佳性能与最差效率。如表2所示,JamMa的平均排名为3.5,领先于预期的平衡以及其他匹配器,差距较大。在性能方面,JamMa在评估的半密集和Sparse匹配器中展现出最佳的姿态估计结果,仅略逊于优先考虑准确性的密集匹配器DKM和RoMa。在效率方面,JamMa超过了所有评估的匹配器,除了XFeat和XFeat*。JamMa的参数数量明显少于基于注意力的匹配器,使其成为超轻量级。

picture.image

4.3. 单应性估计

数据集。HPatches数据集[1]包含用于评估单应性估计的平面场景。作者评估了JamMa在108个HPatches序列上的性能,其中包含52个具有光照变化的序列和56个具有视角变化的序列。

度量。遵循LoFTR [36]中提出的评估协议,作者计算了角点点的平均重投影误差,并报告了角点误差的AUC值,分别对应于3、5和10像素的阈值。所有图像都调整大小,使其较小的一维等于480像素。对所有方法应用Vanilla RANSAC [13]来估计单应性。为了公平起见,选择了半密集方法的顶部1000个匹配项。

结果。如表3所示,JamMa在单应性估计方面取得了具有竞争力的性能。与最先进的匹配器ASpanFormer相比,JamMa在阈值10时落后0.2%,但在阈值3时领先0.7%。值得注意的是,JamMa的体积远小于ASpanFormer,其参数仅为后者的36%。作者将这种有利的表现-效率平衡归因于JEGO策略产生的全局全方位表示以及具有线性复杂度的Mamba。

picture.image

4.4. 消融研究

作者对MegaDepth数据集进行了定性和定量消融研究,以全面理解JamMa。表4中的模型在544分辨率下训练了15个epoch,并使用与第4.2节相同的设置进行测试。

联合扫描。如图4(1)所示,将联合扫描替换为顺序扫描会导致性能显著下降(-2.3%,-2.6%,-2.6%),这表明联合扫描在特征匹配方面更为有效,因为它促进了两个图像之间的高频相互交互。请注意,VMamba和EVMamba本身缺乏相互交互能力,因此,作者在这两种方法中引入联合扫描进行消融研究(3)和(4)。

聚合器。作者在图6中展示了有效的感受野(ERF),其中JamMa的粗略特征实现了全局感受野,并能从各个方向感知信息。没有聚合器,JamMa将失去之前分析过的全局感受野和全向性。表4(2)中的结果表明,没有聚合器时,性能严重下降(-2.2%,-2.2%,-2%)。附录中报告了更多关于聚合器的消融研究。此外,(2)和(3)之间的比较证实,即使在没有聚合器的情况下,EVMamba中的四向扫描也优于双向扫描。

与基于Mamba的方法比较。如表4(3)(4)所示,作者评估了JEGO策略与EVMamba和VMamba中JamMa。与EVMamba相比,JamMa表现出显著的性能提升(+2.6%、+2.5%、+2.2%)。虽然VMamba如图6所示实现了全局感受野,但这是以4倍的总序列长度为代价的。利用JEGO策略,JamMa层在提供更优性能的同时,将VMamba层的速度提升了3倍,从而实现了9.5%的整体速度提升(68.3ms vs. 61.8ms)。

与基于注意力的方法比较。在(5)纯注意力的情况下,当使用24GB VRAM时会出现内存不足(OOM)错误。与(6)线性注意力相比,JamMa实现了性能提升

,同时速度提升了

倍。JamMa将大多数匹配器 Bottleneck ,即基于注意力的交互,降低到几乎可以忽略不计的水平,显著提高了整体效率。JamMa和(6)的总运行时间为

(总速度提升

),而4个JamMa层和4个线性注意力层包含的参数为1.8M对

(参数减少

)。图7进一步展示了不同分辨率下粗层处理时间。

picture.image

编码器。表5展示了采用XFeat、SuperPoint和ResNet作为编码器的JamMa变体。JamMa在不同编码器下均能提供具有竞争力的结果,其中ConvNeXt在性能与效率平衡方面表现最佳。

picture.image

  1. 讨论

如图9所示,作者对JamMa的运行时进行了成分分析以进一步加速。作者的研究发现,在处理Joint Mamba的粗粒度特征时,粗层不再是效率 Bottleneck ,如基于注意力的方法中观察到的。相反,使用Softmax在大型相似度矩阵上进行的粗粒度匹配已成为新的 Bottleneck ,需要优化。作者探讨了[42]中提出的无Softmax粗粒度匹配,并观察到在不降低性能的情况下速度提升了10ms。请注意,为了公平起见,这种优化在第4节中被排除。

picture.image

如表6所示,作者在ScanNet [6]数据集上进行了零样本姿态估计实验。虽然JamMa在室外场景中取得了具有竞争力的性能,但这并不完全适用于室内场景,作者将其归因于其参数数量减少(比比较方法少50%以下),这可能会限制其泛化能力。请注意,与其它方法相比,JamMa仍然在性能效率平衡方面取得了合理的表现。

picture.image

  1. 结论

本文提出了JamMa,这是一种基于Mamba的超轻量级特征匹配器,其性能具有竞争力。作者探讨了图像对的各种扫描合并策略,这是一个之前未曾研究过的主题,因为现有的视觉Mamba模型主要关注单图像任务。首先,作者发现联合扫描与高频互作用显著优于顺序扫描,类似于自注意力和交叉注意力。其次,将四向扫描与聚合器结合,生成全局和全方位特征,这在性能中起着至关重要的作用。第三,Mamba与跳过扫描结合,使JamMa极为高效,甚至超越了大多数基于注意力的Sparse方法。

实验表明,JamMa在性能和效率之间实现了显著的平衡。

参考

[1]. JamMa: Ultra-lightweight Local Feature Matching with Joint Mamba

picture.image

扫码加入👉

「集智书童」

交流群

(备注:

方向

学校/公司+

昵称

picture.image

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
云原生机器学习系统落地和实践
机器学习在字节跳动有着丰富业务场景:推广搜、CV/NLP/Speech 等。业务规模的不断增大对机器学习系统从用户体验、训练效率、编排调度、资源利用等方面也提出了新的挑战,而 Kubernetes 云原生理念的提出正是为了应对这些挑战。本次分享将主要介绍字节跳动机器学习系统云原生化的落地和实践。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论