【魔改Mamba系列】混合Mamba算法在少样本分割中的应用

科技   2024-11-15 10:05   中国香港  

点击下方深度学习爱好者”,选择加p"星标"或“置顶

论文信息

题目:Hybrid Mamba for Few-Shot Segmentation

混合Mamba算法在少样本分割中的应用

作者:Qianxiong Xu, Xuanyi Liu, Lanyun Zhu, Guosheng Lin, Cheng Long, Ziyue Li, Rui Zhao

源码:https://github.com/Sam1224/HMNet

论文创新点

  1. 引入Mamba模型于少样本分割(FSS):作者首次将Mamba模型应用于FSS领域,以捕获序列间的依赖性。这种方法有效地利用了Mamba的线性复杂度特性,解决了传统注意力机制中的二次复杂度问题。
  2. 混合Mamba网络(HMNet)的设计:提出了混合Mamba网络(HMNet),该网络包含自Mamba块(SMB)和混合Mamba块(HMB),其中HMB进一步由支持重述Mamba(SRM)和查询截断Mamba(QIM)组成。这种结构旨在更有效地融合查询和支持特征,提高FSS的性能。
  3. 支持重述Mamba(SRM):为了解决支持遗忘问题,作者设计了支持重述Mamba(SRM)。SRM通过定期回顾支持特征,确保隐藏状态始终包含丰富的支持信息,从而提高了支持信息的利用率。
  4. 查询截断Mamba(QIM):针对类内差距问题,作者提出了查询截断Mamba(QIM)。QIM通过拦截查询像素之间的相互作用,鼓励查询像素更多地融合来自支持特征的信息,而不是仅仅依赖于自身的特征,从而增强了FSS的效果。

摘要

许多少样本分割(FSS)方法使用交叉注意力将支持前景(FG)融合到查询特征中,而不考虑二次复杂度。最近的进步Mamba也能很好地捕获序列内依赖性,但复杂度仅为线性。因此,我们的目标是设计一个交叉(类注意力)Mamba来捕获FSS的序列间依赖性。一个简单的方法是扫描支持特征,有选择地将它们压缩到隐藏状态中,然后将其作为初始隐藏状态用于顺序扫描查询特征。然而,它存在两个问题:(1)支持遗忘问题:在扫描查询特征时,支持特征也会逐渐被压缩,因此隐藏状态中的支持特征会不断减少,许多查询像素无法融合足够的支持特征;(2)类内差距问题:查询前景本质上更类似于自身,而不是支持前景,即查询可能不愿意融合来自隐藏状态的支持特征,而是更倾向于它们自己的特征,但FSS的成功依赖于有效使用支持信息。为了解决这些问题,我们设计了一个混合Mamba网络(HMNet),包括(1)一个支持重述Mamba,以定期回顾扫描查询时的支持特征,使隐藏状态始终包含丰富的支持信息;(2)一个查询截断Mamba,禁止查询像素之间的相互交互,并鼓励它们从隐藏状态中融合更多的支持特征。因此,支持信息得到了更好的利用,从而提高了性能。在两个公共基准测试上进行了广泛的实验,展示了HMNet的优越性。

关键词

少样本分割、交叉注意力、Mamba、序列间依赖性

4. 方法

4.1 重新审视Mamba

Mamba[5]的本质是一个结构化的状态空间模型(SSM),它起源于连续系统,即线性时不变(LTI)系统,将1D序列从x(t)映射到y(t)。这种映射是通过中间隐藏状态h(t)和线性常微分方程(ODEs)[25]实现的,如下所示:
其中A是演化参数,B和C表示两个投影参数。然后,SSM采用零阶保持(ZOH)策略将连续系统转换为离散系统:
其中Δ表示时间尺度参数,方程1可以重写为:
如Mamba[5]所述,离散参数是恒定的动态,因此它们不能有效地将信息压缩到隐藏状态并从上下文中融合正确的信息,导致无法捕获长距离依赖性。为此,Mamba提出了为SSM配备选择机制(称为选择性SSM),使参数B、C和Δ输入依赖,被验证能够捕获复杂的相关性。

4.2 混合Mamba网络(HMNet)

如图2所示,我们提出了混合Mamba网络(HMNet)以将高效的Mamba与FSS结合起来。按照现有的FSS方法[16, 40, 48, 54],查询图像IQ和支持图像IS被送入预训练的骨干网络,如VGG16[37]或ResNet50[9],以获得中级查询和支持特征。然后,它们被送入一些交替出现的自Mamba块(SMB)和混合Mamba块(HMB)进行特征增强。具体来说,SMB旨在对查询和支持特征进行建模,以捕获序列内相关性,而HMB(第4.2.1节)旨在将足够的支持FG特征融合到查询FG特征中。特别是,HMB进一步包含一个支持重述Mamba(SRM)和一个查询截断Mamba(QIM),以减轻上述支持遗忘和类内差距问题(第1节)。最后,增强的查询特征被送入解码器[54]以获得预测
请注意,我们的贡献是改进交叉Mamba以捕获支持-查询序列间依赖性,而不是改进自Mamba。因此,我们采用VMamba[25]构建SMB,其详细信息显示在图6(附录D.1)中。2D图像特征将被重塑为4个序列,根据不同的扫描方向。然后,这些序列被单独的Mamba(具有不同参数)扫描以增强特征。最后,4个序列被重新形状为4个特征,它们被求和以获得输出特征。

4.2.1 混合Mamba块(HMB)

如图3所示,HMB的细节如下。为了便于说明,我们在描述变量形状时省略了批量大小和隐藏维度。
特征准备。尽管查询和支持图像共享相同的FG类别,它们通常具有不同的BG类别,因此通常的做法是遮蔽支持背景[16, 40],即支持特征是稀疏的。一些方法[31, 48]已经验证了对支持特征进行一定程度的下采样不会降低性能,但可以节省大量内存。因此,我们采用相同的策略获得。然后,我们制作两个副本的查询特征:(1)对于SRM,我们保持原始粒度进行密集分割,但将特征分割成补丁,其中每个补丁的大小与下采样的支持特征相同;(2)对于QIM,我们下采样查询特征为,因此其粒度将与一致,以更好地融合特征。
其中是下采样/分割比例,本文中经验设置为4。
SRM。如图3(1)所示,SRM旨在在扫描查询时定期回顾支持特征。我们首先将查询补丁重塑为4个查询序列,具有不同的扫描方向。以为例:
其中表示重塑的查询补丁(内部补丁像素扫描方向与补丁扫描方向相同),表示补丁连接。然后,我们重复并回顾次,并将其重塑为4个支持序列
接下来,我们将支持和查询序列沿像素维度连接:
其中表示对支持和查询补丁。支持放在查询前面,因为我们的目标是将前者传播到后者。此后,被展平成1D序列,我们使用4组Mamba参数扫描序列:
其中是增强的查询特征,从序列中断开连接并重塑。是从序列头部取出的重塑支持特征,我们不考虑其他位置的支持特征,因为它们已经与查询特征混合。是在扫描第一个支持特征后从第i个序列获得的隐藏状态。
QIM。为了鼓励查询特征融合更多的支持特征,我们进一步设计了QIM(图3(2)),以拦截查询像素之间的相互交互,并将隐藏状态并行传播到每个查询像素。因此,不需要根据不同方向将查询像素重塑为多个序列。我们直接将展平成1D序列,并取SRM获得的平均隐藏状态(纯支持特征)作为QIM的初始隐藏状态:
然后,方程3可以直接用矩阵乘法计算并重写为:
其中C, 构成一组Mamba参数。方程10可以解释为方程3的特殊情况,其中每个序列的长度为1。为了便于更好地学习参数,我们共享SRM的参数与QIM,这些参数是随着长序列学习的。
特征集成。最后,我们使用求和融合获得混合Mamba的输出特征:
其中Up(·)表示将特征上采样到原始大小。

4.2.2 计算复杂度

我们遵循Vim[68]来分析计算复杂度。给定一个输入序列,注意力[41]、Mamba[5](1方向)和VMamba[25](4方向)的复杂度为:
其中表示序列的长度,D是隐藏维度,N是一个小常数(例如,16),表示隐藏状态的大小。在混合Mamba中,SRM将查询和支持序列组装成4个(4方向,更长的序列),而QIM处理展平的$F_{Q_{\downarrow}} \in \mathbb{R}^{M/\alpha^2 \times D}(1方向,更短的序列),因此总复杂度为:
其中是下采样比率(例如,4),复杂度仍然是与序列长度线性关系。

5. 实验

声明

本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与我们联系,我们将在第一时间回复并处理。
下载1:Pytoch常用函数手册
在「深度学习爱好者」公众号后台回复:Pytorch常用函数手册,即可下载全网第一份Pytorch常用函数手册,涵盖Tensors介绍、基础函数介绍、数据处理函数、优化函数、CUDA编程、多线程处理等十四章章内容。
下载2:Python视觉实战项目52讲
小白学视觉公众号后台回复:Python视觉实战项目即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。

交流群

欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉驶、计算摄影、检测、分割、识别、医学影像、GAN算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~

深度学习爱好者
分享机器学习、深度学习和Python等知识与原理,每天分享深度学习与计算机视觉领域的经典和最新的论文总结,带读者一起跟踪前言科技!
 最新文章