得益于基础模型的发展，红外小目标检测（ISTD）算法取得了显著进展。特别是，结合卷积网络和 Transformer 结构的模型能够很好地提取局部和全局特征。同时，它们也继承了基础模型的缺陷，例如， Transformer 的二次计算复杂度，这影响了效率。
受近期一个具有线性复杂度用于长距离建模的基础模型Mamba的启发，本文探索了这种状态空间模型在ISTD中的潜力。然而，直接应用是不合适的，因为对检测小目标至关重要的局部特征无法得到充分利用。相反，作者为高效的ISTD定制了一种Mamba-in-Mamba（MiM-ISTD）结构。例如，作者将局部Patch视为“视觉句子”，并将它们进一步分解为作为“视觉单词”的子Patch以进一步探索局部性。给定视觉句子中每个单词之间的交互将在几乎无计算成本的情况下计算。通过聚合单词和句子特征，MiM-ISTD的表现能力可以得到显著增强。
在NUAA-SIRST和IRSTD-1k上的实验证明了作者方法在准确性和效率方面的优势。具体而言，MiM-ISTD比SOTA快，并且在推理过程中，每个图像的GPU内存使用减少了73.4，克服了在处理高分辨率红外图像时执行基于Mamba理解的计算和内存限制。
1 Introduction
红外小目标检测（ISTD）已经在遥感检测和军事跟踪系统中得到广泛应用。这是一项二元分割任务，旨在从背景中分割出小目标像素点。这项任务具有挑战性，因为目标非常小，以至于现有方法很容易漏检或者将它们与其他背景干扰混淆。
当前ISTD方法可以分为传统方法和基于深度学习的方法。在早期阶段，传统方法占据主导地位。然而，这些方法依赖于先验知识和手工制作特征，导致在应用于不符合其假设的图像时准确性有限。
近年来，基于深度学习的方法显著提升了入侵目标检测（ISTD）的性能，其中大部分方法是基于卷积神经网络（CNN）的。然而，基于CNN的方法的缺点在于，它们对局部特征的强调是以牺牲全局上下文为代价的。全局上下文对ISTD来说很重要，因为背景像素和小型目标常常相似到仅凭局部特征无法区分，这很容易导致检测遗漏。作为一种解决方案，一些将视觉变换（ViT）与CNN结合的混合方法被提出，以依赖ViT对长距离依赖建模的能力。然而，这些方法通常因ViT的二次计算复杂度而承受更重的计算负担。
尽管某些工作采用了线性ViT，但其准确性仍然低于具有二次复杂度的设计。考虑到在红外遥感领域高分辨率图像并不罕见（例如，由高分辨率红外军事传感器产生的图像），这种效率缺陷将被放大，当分辨率变大时，它会阻碍实时ISTD。如何在保持高ISTD准确性的同时减轻这种低效率，是作者的主要关注点。
近期，状态空间模型（SSMs）在研究者中越来越受到关注。Mamba 是首个由SSMs构建的提出的基础模型，并且在各种长序列建模任务中与日益流行的 Transformer 模型相比，保持了线性复杂度的同时取得了有希望的表现。在短时间内，Mamba在各个领域都取得了成功，并被认为是继 Transformer 之后的下一代基础模型。然而，当作者直接将视觉Mamba 应用到ISTD时，尽管模型效率令人印象深刻，但检测准确度并不高。原因是ISTD中的目标通常非常小，与其他主要涉及标准尺寸目标的视觉任务相比，需要更加重视局部特征。不幸的是，Mamba在捕捉这些关键局部特征方面并不擅长。
作者旨在定制一个纯基于Mamba的编码器，以解决局部性缺陷，同时仍然保持卓越的模型效率。为此，作者受到[29]的启发，并提出了一种新颖的Mamba-in-Mamba（MiM-ISTD）架构，以提高ISTD的效率。为了增强Mamba的特征表示能力，作者首先将输入图像划分为几个块，作为“视觉句子”，然后再将它们进一步划分为子块，作为“视觉词汇”。
除了外部Mamba块用于提取视觉句子的特征外，作者还将在MiM块中嵌入一个内部Mamba块，以挖掘较小视觉词汇的局部特征。特别是，每个视觉句子中的视觉词汇特征和关系是使用共享网络独立计算的，以确保添加的参数量和FLOPs最小。然后，将视觉词汇特征合并回各自的句子。通过这种方式，MiM-ISTD使作者能够以更精细的粒度提取视觉信息。
作者在图1中展示，在GPU内存使用和推理时间方面，MiM-ISTD比其他方法表现出更高的效率，尤其是当红外图像的分辨率增加时。同时，在处理各种分辨率的红外图像时，MiM-ISTD仍然可以展现出显著的准确性与效率平衡。
作者的贡献可以概括为三个方面：
- 据作者所知，作者是首个成功将Mamba应用于ISTD的研究者，为未来在高效有力的基于Mamba方法的发展提供了新的基准和一些启示。
- 为了将Mamba应用于ISTD领域，作者设计了一种内外嵌套的Mamba结构（MiM-ISTD），以确保在充分提取局部和全局信息的同时，还能保证更高的效率。
- 在两个公开的ISTD数据集NUAA-SIRST和IRSTD-1k上的实验证明了作者方法在准确性和效率方面的优越性。特别是，MiM-ISTD在推理时，对于每一张的图像，在比现有技术水平（SOTA）快10倍的同时，还将GPU内存使用降低了73.4。
2 Related Work
Infrared Small Target Detection Networks
总的来说，当前的ISTD网络可以分为两类：基于CNN的和混合型网络。基于CNN的网络主要关注局部特征提取。Dai等人提出了一种非对称上下文调制（ACM）来实现跨层信息交换，以提高ISTD性能。他们还设计了AlcNet，包括一个局部注意力模块和一个跨层融合模块，以保留小目标局部特征。Wang等人提出了MDvsFA，该方法将生成对抗网络（GAN）应用于ISTD，并在漏检与误报之间取得了折中。
BAUENet将不确定性引入到ISTD中，实现了边界感知分割。DNANet逐步与高低级特征进行交互。Dim2Clear通过探索图像增强思想，将ISTD视为图像细节重建任务。FC3-Net探索了特征补偿和跨 Level 相关性用于ISTD任务。ISNet设计了一个受简单Taylor有限差分启发的块和一个双方向注意力聚合模块来检测目标。
仅依赖局部特征进行红外小目标检测（ISTD）可能会导致小型目标的检测遗漏，这些目标可能会融入相似背景中变得无法区分。因此，混合方法通过结合视觉 Transformer （ViT）和卷积神经网络（CNN）来补充局部细节和全局上下文。IRSTFormer采用分层ViT来建模长距离依赖，但对挖掘局部细节的强调不足。ABMNet 在CNN和线性ViT结构设计中都采用了常微分方程（ODE）方法进行红外小目标检测。
IAANet 将简单CNN的局部Patch输出与原始 Transformer 连接起来，但在低对比度场景中特征提取能力有限。RKformer 应用Runge-Kutta方法构建耦合的CNN-Transformer块以突出红外小目标并抑制背景干扰。然而，由于使用了ViT，这些混合方法大多数都遭受二次计算复杂度的困扰。尽管有些工作采用了线性ViT设计，但其检测精度无法与其他具有二次复杂度ViT的工作相媲美。
为了在保持高准确度的同时提高网络效率，作者受到了Mamba 和TNT两者的启发，为ISTD提出了一个Mamba-in-Mamba架构，其中包含一个纯Mamba-in-Mamba金字塔编码器，实现了高效的地方和全局特征提取。
Mamba in Vision Tasks
近期，状态空间序列模型（SSMs）在高效处理长序列建模方面显示出潜力，为视觉任务中解决长距离依赖提供了另一种选择。与 Transformer 相比，SSMs由于与序列长度成线性扩展而更为高效，并且在建模长距离依赖方面保持卓越的能力。几项最新研究已证明了Mamba在视觉任务中的有效性。例如，Vim 提出了一个带有双向Mamba块的通用视觉主干网络。VMamba 提出了一个基于分层Mamba的视觉主干网络以及一个跨扫描模块，以解决由于1D序列与2D图像表示之间的差异引起的方向敏感性问题。
值得注意的是，Mamba已经在医疗图像分割领域得到广泛应用。U-Mamba，Vm-unet，Mamba-unet和SegMamba分别基于nnUNet，Swin-UNet，VMamba和Swin-UNETR提出了针对特定任务的架构，并采用了Mamba模块。P-Mamba将PM扩散与Mamba结合，在有效去除背景噪声的同时保留目标边缘细节。Swin-UMamba验证了基于ImageNet的预训练对于基于Mamba网络的医疗图像分割的重要性。Weak-Mamba-UNet和Semi-Mamba-UNet分别用于弱监督和半监督的医疗图像分割。Vivim将Mamba应用于医疗视频目标分割。
鉴于这些Mamba模型在各项视觉任务中取得了鼓舞人心的成果，作者打算研究Mamba是否也能为ISTD带来像其他视觉任务那样的进步。然而，当作者直接将视觉Mamba模块应用于ISTD时，尽管模型效率很高，但检测准确度并不显著。原因是ISTD任务中的目标非常小，这需要比主要针对常见大小目标的其它视觉任务更多地关注局部特征，而原始的视觉Mamba模块并不能很好地探索这些局部特征。为了弥补这一缺陷，作者提出了针对ISTD的Mamba-in-Mamba（MiM-ISTD）方法，该方法同时考虑视觉句子和视觉词的流，并将每个流分别发送到相应的视觉Mamba模块，以高效地获取局部和全局特征。
3 Method
作者在图3中展示了Mamba-in-Mamba（MiM）的整体架构。它主要由以下部分组成：
- 一个卷积以生成块和子块，
- 一个MiM金字塔编码器，用于在不同尺度上提取全局和局部特征，
- 一个带有几个上采样块的解码器，用于预测分割 Mask 。
Preliminaries
3.1.1 State Space Models.
状态空间模型（SSMs）通常被用作线性时不变系统，它将一维输入刺激 通过中间隐含状态 转换为一个输出 。
在数学术语中，SSMs 通常由线性常微分方程（ODEs）（方程1）来描述，其中系统由一组参数定义，包括状态转移矩阵 、投影参数 和跳跃连接 。
3.1.2 Discretization.
状态空间模型（SSMs）在应用于深度学习场景时，由于其连续时间性质，带来了重大挑战。为了解决这一问题，至关重要的是通过离散化过程将常微分方程（ODE）转换为离散函数。这种转换对于同步模型与输入数据中体现的底层信号的采样率至关重要，从而便于进行计算效率高的操作。
考虑到输入数据中的，这是信号流程中的采样向量，长度为，遵循[17]，可以将时间尺度参数引入到连续参数和中，根据零阶保持（ZOH）规则，转换为它们的离散对应物和。因此，方程1可以离散化如下：
在公式中， 且 。在实际操作中，遵循[12]的方法，作者使用一阶泰勒级数来细化的近似：
3.1.3 2D Selective Scan.
二维视觉数据与一维语言序列之间的区别使得不能直接将Mamba应用于视觉任务。例如，虽然在视觉任务中2D空间信息至关重要，但在1D序列建模中并非主要关注点，这导致了受限的接收域，忽略了与未扫描的图像块之间可能存在的关系。在[27]中提出的2D选择性扫描（SS2D）可以解决这个问题。
SS2D的概述如图2所示。SS2D以四个不同的方向排列图像块，生成四个独立的序列。四向扫描策略确保特征图中的每个元素从不同方向的所有其他位置整合信息，从而在不增加线性计算复杂性的情况下创建全局接收域。然后使用选择性扫描空间状态序列模型（S6）处理每个生成的特征序列。将特征序列组合起来重构2D特征图。给定输入特征，则SS2D的输出特征可以表示为：
其中 表示四种扫描方向之一。 和 分别指代文献[27]中的扫描扩展和扫描合并操作。方程式4中的S6块是SS2D的核心SSM算子。它使得1D数组中的每个元素都能通过一个压缩的隐藏状态与之前扫描过的任何样本进行交互。为了更全面地理解S6，文献[27]提供了深入的说明。
Mamba-in-Mamba for Efficient ISTD
当前ISTD方法主要采用卷积神经网络（CNN）或卷积神经网络和视觉 Transformer （ViTs）的混合。后者补偿了前者在建模长距离依赖方面的不足，但受到了二次计算复杂性的困扰。最近，提出了Mamba模型。它以卓越的模型效率、较低的GPU内存使用和更好的长距离依赖建模而闻名，并已应用于视觉领域。
因此，作者尝试探索Mamba是否也能应用于提升ISTD的性能。然而，当作者直接将视觉Mamba模块应用于ISTD时，准确度并不十分令人印象深刻，因为对于检测小目标至关重要的局部特征探索得较少。为了解决这个问题，作者提出了一种Mamba-in-Mamba（MiM-ISTD）架构，如图3所示，以在保证卓越模型效率的同时学习图像中的全局和局部信息。
给定一个2D红外图像 ，它被均匀地划分为 个块以形成 ，其中每个块在 中， 表示每个块的分辨率。在MiM-ISTD中，作者将这些块视为表示图像的视觉句子。此外，每个块被分割成 个更小的子块，使得每个视觉句子成为视觉词汇的序列：
其中 是第 个视觉句子中的第 个视觉词， 是子块的空间大小。由于作者的MiM采用了金字塔编码器结构，视觉句子和词汇的空间形状是不固定的，并且随着网络层级的加深将会逐渐减小。
3.2.1 Convolutional Stem.
作者构建了一个卷积干细胞，其中使用了一堆卷积，以在第一阶段生成视觉词和视觉句子，其中是视觉词的维度，是视觉句子的维度。每个视觉词对应于原始图像中一个像素的区域，而每个视觉句子由个视觉词组成。与ViTs不同，由于视觉mamba块[27]的因果性质，作者没有向视觉词和句子中添加位置嵌入偏差。
3.2.2 MiM Pyramid Encoder.
3.2.3 Decoder Structure.
与在编码器中用于降采样的 Patch 合并操作相比，作者在解码器中使用了 Patch 扩展层来进行升采样。在每一个解码器阶段，作者将升采样后的编码器输出与扩展的解码器特征相结合，并将整合后的特征发送到基本的ResNet模块中。最后，这些特征通过一个全连接网络 Head 和一个插值操作来得到最终的 Mask 预测。
4 Experiments
Experimental Settings
4.1.1 Datasets.
作者选择NUAA-SIRST和IRSTD-1k作为训练、验证和测试的基准数据集。NUAA-SIRST包含了427幅不同尺寸的红外图像，而IRSTD-1k则由1,000幅真实尺寸为的红外图像组成。IRSTD-1k是一个更具挑战性的ISTD数据集，包含了更丰富的场景。对于每个数据集，作者使用80的图像作为训练集，20作为测试集。
4.1.2 Evaluation Metrics.
4.1.3 Optimization.
算法是在Pytorch中实现的，使用自适应梯度（AdaGrad）作为优化器，初始学习率设为0.05，权重衰减系数设为0.0004。训练使用了2块NVIDIA A6000 GPU，批量大小设置为32。损失函数采用了Dice损失。在SIRST和IRSTD-1k上的训练分别进行了800个和600个周期。
Quantitative Comparison
4.2.1 Accuracy Comparison.
作者选择了一些现有的SOTA ISTD方法进行比较。如表1所示，在两个数据集上，作者的MiM-ISTD在像素级指标和目标级指标方面通常都能取得最佳性能。
对于像素级指标（），MiM-ISTD由于进一步集成了内部Mamba块进行局部特征建模而取得了最佳性能。这种方式使得一些可能被外部Mamba块忽略的可区分细节能够得到注意，这提高了检测的准确性。
对于目标 Level 的指标（），在和之间达到平衡是具有挑战性的，因为这两个指标是相互排斥的：较高的也会增加较高的的可能性。从表1中，作者可以看到与其他方法相比，MiM-ISTD通常能够获得最佳的目标 Level 指标结果，尤其是在NUAA-SIRST测试集中检测到所有小目标（）。这一结果表明，作者的MiM-ISTD可以学习更好的表示来克服漏检和误报。
4.2.2 Efficiency Comparison.
作者还比较了不同方法在参数数量（M）、浮点运算数（G）、推理时间（秒）以及训练期间在红外图像数据集上的GPU内存使用（M），如表2所示。与其他方法（ACM除外）相比，作者的MiM-ISTD通常具有明显更少的参数、GFLOPs、推理时间和内存使用。这是因为作者采用了高效的Mamba结构，并使用共享网络来计算视觉句子中每个视觉单词之间的关系，因此增加的参数和GFLOPs可以忽略不计。
即使与最轻量级的ISTD模型ACM相比，作者的MiM-ISTD仍然具有更少的GFLOPs和显著更高的准确度。此外，在更高分辨率的红外图像场景中，作者在图1中比较了MiM-ISTD在推理时间和GPU内存使用上的优势将进一步扩大，同时始终保持卓越的准确度。总的来说，MiM-ISTD在效率和准确度之间达到了最佳的平衡。
Visualization
4.3.1 Visualization of Mask Results.
不同方法的近距离视觉结果展示在图5中，其中当前方法或多或少存在检测不完整和漏检的问题。与其它前沿技术（SOTAs）相比，MiM-ISTD更好地减少了这些情况，更完整地检测到所有小目标形状。这是因为集成内部Mamba块有助于作者的网络进一步挖掘更多的局部特征，这促进了小目标的更精细检测。这一点也可以通过将MiM-ISTD的视觉结果与“无内部Mamba”的视觉结果进行比较来证明，放弃内部Mamba会导致更差的检测性能。
4.3.2 Visualization of Feature Maps.
作者可视化TCI-Former和MiM-ISTD学习到的特征，以进一步理解所提出方法的效果。特征图是通过根据它们的空间位置 Reshape 块嵌入形成的。第1、2、3和4阶段的特征图输出显示在图6（a）中，其中为每个输出随机采样了9个特征图通道。在MiM-ISTD中，与TCI-Former相比，在更深层次上更好地保留了局部信息。
此外，MiM-ISTD在各个通道之间的特征一致性也比TCI-Former要高，这意味着MiM-ISTD提取的特征更专注于目标。这些优势归功于引入内部Mamba块以进一步建模局部特征。此外，作者还使用t-SNE[47]可视化了第3阶段特征输出的所有64个通道，以证明作者的分析，如图6（b）所示。作者观察到MiM-ISTD的特征更集中，展现出更强的辨别力。
Ablation study
4.4.1 Ablation study of Each Module.
消融研究各个模块的结果展示在表3中。作者的 Baseline 模型（编号#1）使用视觉Mamba编码器作为作者的编码器，具有递减的空间分辨率，其中没有使用内部Mamba块。
作者还用标准卷积、批量归一化和激活操作（编号#2）替换了MiM-ISTD中的内部Mamba块，以检验作者内部Mamba块的效果。作者发现，作者当前的设置带来了最佳的结果，因为卷积通常关注局部信息，而内部Mamba块能够捕捉复杂的空间关系并整合多方向信息，从而对小型目标进行更为精细的感知和识别。
4.4.2 Ablation Study of the Granularity of Patch Division.
为了探索每个视觉词中包含的信息的影响，作者调整了子块的表现粒度，使得每个视觉词对应于原始图像中一个像素区域（编号#1），与作者当前设置的进行比较。作者还改变了视觉句子的初始空间形状，从变为（编号#2），以消除视觉句子粒度的影响。作者采用了表现最佳配置作为作者当前的设置。
5 Conclusion
在本文中，作者提出了Mamba-in-Mamba（MiM-ISTD）结构以实现高效的ISTD。作者将图像均匀地划分为块（类似于视觉句子），并将每个块进一步划分为多个较小的子块（视觉词汇）。作者设计了一个基于纯Mamba的MiM金字塔编码器，该编码器包含堆叠的MiM块。每个MiM块包含一个外部的Mamba块来处理句子嵌入，以及一个内部的Mamba块来建模词嵌入之间的关系。
经过线性投影后，视觉词汇嵌入被加到视觉句子嵌入中。与现有的SOTAs相比，MiM-ISTD可以实现更高效地建模局部和全局信息。在两个数据集上的实验验证了MiM-ISTD的有效性。
参考
[1].MiM-ISTD: Mamba-in-Mamba for Efficient Infrared Small Target Detection.
作者：AI视界引擎
文章来源：AI视界引擎
推荐阅读
- 编译入门那些事儿（5）：Jump Table 优化介绍
- 南开提出 Range-View | 激光雷达技术新进展在自动驾驶等多任务中的应用
- IEEE 754 Floating NaN在深度学习计算中的行为
- 手机端智能体Mobile-Agent开源，人人都可以拥有手机端超级助理
更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
