Hybrid Mamba for Few-Shot Segmentation
作者: Qianxiong Xu, Xuanyi Liu, Lanyun Zhu, Guosheng Lin, Cheng Long, Ziyue Li, Rui Zhao
分类: cs.CV
发布日期: 2024-09-29
备注: This paper is accepted by NIPS'24
🔗 代码/项目: GITHUB
💡 一句话要点
提出混合Mamba网络(HMNet)用于解决小样本分割中支持信息利用不足的问题。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 小样本分割 Mamba模型 序列建模 支持重述 查询截取 混合网络 计算机视觉
📋 核心要点
- 现有小样本分割方法依赖交叉注意力融合支持信息，但计算复杂度高，且易忽略序列依赖。
- 提出混合Mamba网络(HMNet)，通过支持重述和查询截取机制，增强支持信息的利用。
- 在标准数据集上进行实验，HMNet 显著提升了小样本分割的性能，验证了方法的有效性。
📝 摘要（中文）
许多小样本分割(FSS)方法使用交叉注意力来融合支持前景(FG)到查询特征中，但存在二次复杂度问题。最近的Mamba模型能够很好地捕获序列内部依赖关系，且复杂度仅为线性。因此，本文旨在设计一种类交叉注意力的Mamba模型来捕获FSS的序列间依赖关系。一个简单的想法是扫描支持特征，选择性地将其压缩到隐藏状态中，然后将该隐藏状态用作初始隐藏状态来顺序扫描查询特征。然而，这种方法存在(1)支持信息遗忘问题：在扫描查询特征时，查询特征也会逐渐被压缩，导致隐藏状态中的支持特征不断减少，许多查询像素无法融合足够的支撑特征；(2)类内差距问题：查询前景本质上更类似于自身，而不是支持前景，即查询可能更倾向于融合来自隐藏状态的自身特征，而非支持特征，然而FSS的成功依赖于对支持信息的有效利用。为了解决这些问题，本文设计了一种混合Mamba网络(HMNet)，包括(1)支持重述Mamba，用于在扫描查询时定期重述支持特征，从而使隐藏状态始终包含丰富的支持信息；(2)查询截取Mamba，用于禁止查询像素之间的相互作用，并鼓励它们从隐藏状态中融合更多的支持特征。因此，支持信息得到更好的利用，从而获得更好的性能。在两个公共基准上进行了大量实验，表明了HMNet的优越性。
🔬 方法详解
问题定义：小样本分割(FSS)旨在仅利用少量带标注的样本（支持集）来分割新的未见过的图像（查询集）。现有方法，特别是基于交叉注意力的方法，在融合支持集信息到查询集特征时，计算复杂度高（二次复杂度），且容易忽略序列内部的依赖关系。此外，简单地将支持集信息压缩到隐藏状态中，再顺序扫描查询集特征，会导致支持信息遗忘和类内差距问题，即查询特征会逐渐覆盖支持特征，且查询前景更倾向于融合自身特征而非支持特征。
核心思路：本文的核心思路是设计一种混合Mamba网络(HMNet)，该网络能够有效地利用支持集信息，同时避免高计算复杂度和信息遗忘问题。HMNet通过引入支持重述和查询截取机制，确保隐藏状态始终包含丰富的支持信息，并鼓励查询特征融合更多的支持特征。
技术框架：HMNet的整体框架包括特征提取模块（通常是预训练的卷积神经网络）、支持重述Mamba模块和查询截取Mamba模块。首先，使用特征提取模块提取支持集和查询集的特征。然后，支持重述Mamba模块周期性地将支持集特征重新注入到隐藏状态中，以防止信息遗忘。最后，查询截取Mamba模块禁止查询像素之间的相互作用，促使它们更多地融合来自隐藏状态的支持特征。
关键创新：HMNet的关键创新在于两个方面：(1) 支持重述Mamba：通过周期性地重述支持特征，解决了支持信息遗忘的问题，确保隐藏状态始终包含丰富的支持信息。(2) 查询截取Mamba：通过禁止查询像素之间的相互作用，鼓励查询特征更多地融合来自隐藏状态的支持特征，解决了类内差距问题。与传统的交叉注意力方法相比，HMNet具有线性复杂度，更适合处理高分辨率图像。
关键设计：支持重述Mamba模块的关键设计在于重述的频率和方式。论文中可能采用了固定的重述间隔，或者根据隐藏状态的变化动态调整重述频率。查询截取Mamba模块的关键设计在于如何有效地禁止查询像素之间的相互作用，可能通过mask或者特殊的激活函数来实现。具体的损失函数可能包括交叉熵损失和Dice损失等，用于优化分割结果。
🖼️ 关键图片
📊 实验亮点
HMNet 在两个公开数据集上进行了实验，结果表明 HMNet 显著优于现有的 FSS 方法。具体来说，HMNet 在 PASCAL-5i 数据集和 COCO-20i 数据集上都取得了 state-of-the-art 的性能。实验结果验证了 HMNet 在解决支持信息遗忘和类内差距问题方面的有效性，以及其在线性复杂度下实现高性能的潜力。
🎯 应用场景
该研究成果可应用于医学图像分析、遥感图像解译、自动驾驶等领域。在医学图像分析中，可以利用少量标注的病灶图像来分割新的病理图像，辅助医生进行诊断。在遥感图像解译中，可以利用少量标注的地物图像来分割新的遥感图像，用于土地利用监测和灾害评估。在自动驾驶中，可以利用少量标注的道路场景图像来分割新的道路场景图像，提高车辆的感知能力。
📄 摘要（原文）
Many few-shot segmentation (FSS) methods use cross attention to fuse support foreground (FG) into query features, regardless of the quadratic complexity. A recent advance Mamba can also well capture intra-sequence dependencies, yet the complexity is only linear. Hence, we aim to devise a cross (attention-like) Mamba to capture inter-sequence dependencies for FSS. A simple idea is to scan on support features to selectively compress them into the hidden state, which is then used as the initial hidden state to sequentially scan query features. Nevertheless, it suffers from (1) support forgetting issue: query features will also gradually be compressed when scanning on them, so the support features in hidden state keep reducing, and many query pixels cannot fuse sufficient support features; (2) intra-class gap issue: query FG is essentially more similar to itself rather than to support FG, i.e., query may prefer not to fuse support features but their own ones from the hidden state, yet the success of FSS relies on the effective use of support information. To tackle them, we design a hybrid Mamba network (HMNet), including (1) a support recapped Mamba to periodically recap the support features when scanning query, so the hidden state can always contain rich support information; (2) a query intercepted Mamba to forbid the mutual interactions among query pixels, and encourage them to fuse more support features from the hidden state. Consequently, the support information is better utilized, leading to better performance. Extensive experiments have been conducted on two public benchmarks, showing the superiority of HMNet. The code is available at https://github.com/Sam1224/HMNet.