MiM-ISTD: Mamba-in-Mamba for Efficient Infrared Small Target Detection

摘要

https://web3.arxiv.org/pdf/2403.02148
近年来，随着基础模型的发展，红外小目标检测（ISTD）取得了显著进展。具体来说，将卷积神经网络（CNNs）与转换器（transformers）相结合的模型能够成功地提取局部和全局特征。然而，转换器的缺点也被继承了下来，即序列长度的二次计算复杂度。受最近具有线性复杂度的远距离建模基础模型Mamba的启发，本文探讨了该状态空间模型在ISTD任务中的有效性和效率方面的潜力。然而，由于缺乏对检测小目标至关重要的局部特征的充分利用，直接应用Mamba模型只能达到次优性能。相反，我们为高效ISTD定制了一种嵌套结构，即Mamba中的Mamba（MiM-ISTD）。它由外部和内部Mamba块组成，以巧妙地捕获全局和局部特征。具体来说，我们将局部图像块视为“视觉句子”，并使用外部Mamba来探索全局信息。然后，我们将每个视觉句子分解为子块作为“视觉单词”，并使用内部Mamba以可忽略的计算成本进一步探索视觉句子中单词之间的局部信息。通过聚合视觉单词和视觉句子特征，我们的MiM-ISTD能够有效地探索全局和局部信息。在NUAA-SIRST和IRSTD-1k上的实验表明，我们方法具有更高的准确性和效率。具体来说，当在$2048\times 2048$图像上进行测试时，MiM-ISTD比最先进的方法快$8$倍，并减少了$62.2\%$的GPU内存使用，克服了高分辨率红外图像的计算和内存限制。

索引词-Mamba中的Mamba，状态空间模型，红外小目标检测

一、引言

红外小目标检测（ISTD）已广泛应用于遥感和军事跟踪系统中。它是一项二元分割任务，旨在将小目标像素从背景中分割出来。这项任务具有挑战性，因为目标非常小，现有方法很容易错过它们或将其与其他背景干扰混淆。

现有的ISTD方法可以分为传统方法和基于深度学习的方法。在早期阶段，传统方法[1]-[9]占据主导地位。然而，这些方法依赖于先验知识和手工制作的特征，当应用于不符合其假设的图像时，准确性有限。

近年来，基于深度学习的方法显著提高了ISTD的性能，其中大多数是基于卷积神经网络（CNN）的方法[10]、[12]-[18]。然而，基于CNN的方法的缺点是它们专注于局部特征，却牺牲了全局上下文。全局上下文对于ISTD也很重要，因为在红外图像中，背景像素和小目标在许多情况下看起来非常相似，仅通过局部特征无法区分，这很容易导致漏检。作为解决方案，一些结合了ViT和CNN的混合方法[11]、[19]-[24]被提出，以依赖ViT建模长距离依赖的能力。然而，这些方法通常由于ViT的二次计算复杂度而承受更重的计算负担。尽管某些工作[19]采用了线性ViT，但其准确性仍然低于具有二次复杂度的设计。考虑到红外遥感领域不乏高分辨率图像（例如，由高分辨率红外军事传感器生成的图像），当分辨率增大时，这种效率缺陷将被放大，从而阻碍实时ISTD。如何在保持高ISTD准确性的同时缓解效率低下是我们主要关注的问题。

最近，状态空间模型（SSMs）在研究人员中引起了越来越多的关注。Mamba [25] 是首个提出的基于SSMs构建的基础模型，与在各种长序列建模任务中蓬勃发展的Transformer相比，它取得了令人瞩目的性能，同时保持了线性复杂度。在很短的时间内，Mamba已在各个领域取得了成功[26]-[29]，并被认为有可能成为继Transformer之后的下一代基础模型。然而，当我们直接将视觉Mamba [28]应用于ISTD时，尽管模型效率令人印象深刻，但检测精度并不高。原因是在ISTD中，目标通常非常小，与其他主要涉及标准大小目标的视觉任务相比，需要更加关注局部特征。不幸的是，Mamba并不擅长捕捉这些关键的局部特征。
我们旨在提出一种基于Mamba的ISTD编码器，以解决局部性缺陷的同时保持卓越的模型效率。为此，我们受到TNT [30]的启发，该模型以极小的计算和内存成本增加为代价，有效地对局部结构特征进行建模，并提出了一种新颖的“Mamba中套Mamba”（MiM-ISTD）架构，以实现更高效的ISTD。为了提升Mamba的特征表示能力，我们将输入图像划分为多个块作为“视觉句子”，并进一步将它们细分为子块作为“视觉单词”。我们使用一个外部Mamba块来提取视觉句子的特征，并进一步利用内部Mamba块来挖掘较小视觉单词的局部特征。特别是，我们使用共享网络独立计算每个视觉句子中视觉单词之间的特征和关系，以确保增加的参数和浮点运算次数（FLOPs）最小。然后，将视觉单词特征重新整合回它们各自的句子中。通过这种方式，MiM-ISTD使我们能够以更精细的粒度提取视觉信息。我们在图1中展示了，随着红外图像分辨率的增加，MiM-ISTD在GPU内存使用量和推理时间方面相较于其他方法展现出卓越的效率。总体而言，MiM-ISTD与其他最先进（SOTA）方法相比，在精度和效率之间取得了最显著的平衡。

我们的贡献可以总结为以下三个方面：

• 据我们所知，我们是第一个成功将Mamba应用于ISTD的研究，为未来高效且强大的基于Mamba的方法的进步提供了新的基准和有价值的见解。
• 为了将Mamba应用于ISTD领域，我们定制了一种Mamba中套Mamba（MiM-ISTD）结构，以确保在充分提取局部和全局信息的同时，保持较高的效率。
• 在两个公开的ISTD数据集（NUAA-SIRST和IRSTD-1k）上进行的实验证明了我们方法的卓越准确性和效率。具体而言，MiM-ISTD在推理过程中，对于每张$2048\times 2048$的图像，相比最先进（SOTA）方法实现了8倍的速度提升，同时减少了62.2%的GPU内存使用量。

二、相关工作

A. 红外小目标检测网络

一般而言，当前的ISTD网络可以分为两大类：基于卷积神经网络（CNN）的网络和混合网络。基于CNN的网络主要侧重于局部特征提取。Dai等人[16]提出了非对称上下文调制（ACM）以实现跨层信息交换，从而提升ISTD性能。他们还设计了AlcNet[14]，其中包括一个局部注意力模块和一个跨层融合模块，以保留小目标的局部特征。Wang等人提出了MDvsFA[17]，将生成对抗网络（GAN）应用于ISTD，并在漏检和误报之间取得了平衡。BAUENet[13]将不确定性引入ISTD，实现了边界感知分割。DNANet[10]逐步与高低级特征进行交互。Dim2Clear[31]通过探索图像增强的思想，将ISTD视为图像细节重建任务。FC3-Net[32]探索了ISTD任务中的特征补偿和跨级相关性。ISNet[18]设计了一个简单的受泰勒有限差分启发的块和一个双方向注意力聚合模块来检测目标。最近，第一个用于ISTD的扩散模型DCFR-Net[33]已被提出。然而，其准确性低于当前的最先进方法[11]，并且与大多数基于深度学习的ISTD技术相比，其计算效率也大大滞后。

仅依赖局部特征进行红外小目标检测（ISTD）可能会导致小目标漏检，因为小目标可能会融入相似的背景中而变得难以区分。因此，混合方法通过将视觉转换器（ViT）与卷积神经网络（CNN）相结合，用全局上下文信息补充局部细节。IRSTFormer[34]采用分层ViT来建模长距离依赖关系，但对挖掘局部细节的重视不足。ABMNet[19]在CNN和线性ViT结构设计中采用了常微分方程（ODE）方法来进行ISTD。IAANet[20]将简单CNN的局部块输出与原始转换器进行拼接，但在特征提取方面有限，特别是在低对比度场景下。RKformer[24]应用龙格-库塔方法构建耦合的CNN-Transformer块，以突出红外小目标并抑制背景干扰。TCI-Former[11]通过模拟热传导过程来提取小目标特征，并取得了最先进（SOTA）的结果。然而，由于使用了ViT，这些混合方法中的大多数都面临着二次计算复杂度的问题。尽管有些工作[19]采用了线性ViT设计，但其检测精度无法与具有二次复杂度ViT的其他工作相提并论。

为了提高网络效率同时保持高精度，我们从Mamba[25]和TNT[30]中汲取灵感，为ISTD提出了一种Mamba-in-Mamba架构，其中包含了一个Mamba-in-Mamba（MiM）分层编码器，该编码器实现了高效的局部和全局特征提取。

B. 在视觉任务中的Mamba

最近，状态空间序列模型（SSMs）[26]在有效处理长序列建模方面显示出了巨大的潜力，为视觉任务中处理长距离依赖关系提供了一种替代方案。与转换器相比，SSMs更加高效，因为它们随序列长度的增加而线性扩展，并保留了优越的建模长距离依赖关系的能力。一些最新的研究已经证明了Mamba在视觉任务中的有效性[35]。例如，Vim[36]提出了一种具有双向Mamba块的通用视觉主干网络。VMamba[28]提出了一种基于Mamba的分层视觉主干网络和交叉扫描模块，以解决一维序列与二维图像表示之间的差异导致的方向敏感性问题。

值得注意的是，Mamba在医学图像分割领域的应用最为广泛。U-Mamba[27]、Vm-unet[37]、Mamba-unet[38]和SegMamba[39]分别基于nnUNet[40]、Swin-UNet[41]、VMamba[28]和Swin-UNETR[42]提出了具有Mamba块的特定任务架构。P-Mamba[29]将概率运动（PM）扩散与Mamba相结合，在保留目标边缘细节的同时有效地去除了背景噪声。Swin-UMamba[43]验证了基于ImageNet的预训练对于基于Mamba的医学图像分割网络的重要性。Vivim[44]将Mamba引入医学视频对象分割中。
由于这些Mamba模型在各种视觉任务中取得了令人鼓舞的结果，我们打算研究Mamba是否也能像给其他视觉任务带来进步一样，为红外小目标检测（ISTD）带来进步。然而，当我们直接将视觉Mamba块应用于ISTD时，尽管模型效率很高，但检测精度并不显著。原因是ISTD任务中的目标非常小，这要求比其他主要处理常见大小目标的视觉任务更加关注局部特征，而原始的视觉Mamba块无法很好地探索这些局部特征。为了解决这一缺陷，我们为ISTD提出了Mamba-in-Mamba（MiM-ISTD），它同时接收视觉句子和视觉单词流，并将每个流发送到相应的视觉Mamba块，以高效地获得局部和全局特征。

三、提出的方法

在本节中，我们将详细描述我们提出的用于高效ISTD的Mamba-in-Mamba，并详细分析计算复杂度。

A. 预备知识

1） 状态空间模型：状态空间模型（SSMs）通常被用作线性时不变系统，该系统通过中间隐状态$h(t)\in {\mathbb{R}}^N$将一维输入刺激$x(t)\in {\mathbb{R}}^L$转换为输出$y(t)\in {\mathbb{R}}^L$。在数学上，SSMs通常由一组参数来描述，包括状态转移矩阵$A\in {\mathbb{C}}^{N\times N}$、投影参数$B,C\in {\mathbb{C}}^N$，以及跳跃连接$D\in {\mathbb{C}}^1$。

h ′ ( t ) = A h ( t ) + B x ( t ) y ( t ) = C h ( t ) + D x ( t )

$$\begin{array}{l} h^{\prime}(t)=A h(t)+B x(t) \\ y(t)=C h(t)+D x(t) \end{array}$$ h′(t)=Ah(t)+Bx(t)y(t)=Ch(t)+Dx(t)​

2） 离散化：由于状态空间模型（SSMs）具有连续时间的特性，当它们被应用于深度学习场景时，会面临重大挑战。为了解决这个问题，需要将常微分方程（ODE）转换为离散函数。考虑到输入$x_k\in {\mathbb{R}}^{L\times D}$，它是信号流中长度为$L$的采样向量（遵循[45]），可以引入时间尺度参数$\Delta$，将连续参数$A$和$B$转换为它们的离散对应物$\bar{A}$和$\bar{B}$，这遵循零阶保持（ZOH）规则。因此，（1）可以离散化为如下形式：

h k = A ˉ h k − 1 + B ˉ x k , y ( t ) = C h ( t ) + D x ( t ) A ˉ = e Δ A , B ˉ = ( e Δ A − I ) A − 1 B , C ˉ = C

$$\begin{array}{l} h_{k}=\bar{A} h_{k-1}+\bar{B} x_{k}, \quad y(t)=C h(t)+D x(t) \\ \bar{A}=e^{\Delta A}, \quad \bar{B}=\left(e^{\Delta A}-I\right) A^{-1} B, \quad \bar{C}=C \end{array}$$ hk​=Aˉhk−1​+Bˉxk​,y(t)=Ch(t)+Dx(t)Aˉ=eΔA,Bˉ=(eΔA−I)A−1B,Cˉ=C​

其中，$B,C\in {\mathbb{R}}^{D\times N}$且$\Delta \in {\mathbb{R}}^D$。在实践中，我们使用一阶泰勒级数来改进$\bar{B}$的近似值：

$\bar{B}=\left(e^{\Delta A}-I\right)A^{-1}B=(\Delta A)(\Delta A{)}^{-1}\Delta B=\Delta B$

3） 二维选择性扫描：专为一维数据设计的模型无法有效捕获二维空间信息，因此不适合直接将Mamba应用于视觉任务。在[28]中提出的二维选择性扫描（SS2D）可以解决这一问题。SS2D的概述如图（4）所示。SS2D以四个不同方向排列图像块，生成四个独立的序列。这种四向扫描策略确保了特征图中的每个元素都能从不同方向整合来自所有其他位置的信息，从而在不增加线性计算复杂度的情况下创建全局感受野。然后，对每个生成的特征序列使用选择性扫描空间状态序列模型（S6）进行处理，之后再合并以重构二维特征图。给定输入特征$z$，SS2D的输出特征$\bar{z}$可以表示为：

z i = expand ⁡ ( z , i ) z ˉ i = S6 ⁡ ( z i ) z ˉ = merge ⁡ ( z ˉ 1 , z ˉ 2 , z ˉ 3 , z ˉ 4 )

$$\begin{array}{l} z_{i}=\operatorname{expand}(z, i) \\ \bar{z}_{i}=\operatorname{S6}\left(z_{i}\right) \\ \bar{z}=\operatorname{merge}\left(\bar{z}_{1}, \bar{z}_{2}, \bar{z}_{3}, \bar{z}_{4}\right) \end{array}$$ zi​=expand(z,i)zˉi​=S6(zi​)zˉ=merge(zˉ1​,zˉ2​,zˉ3​,zˉ4​)​

其中， i ∈ { 1 , 2 , 3 , 4 } i \in\{1,2,3,4\} i∈{1,2,3,4}代表四个扫描方向之一。expand()和merge()分别指[28]中的扫描扩展和扫描合并操作。图（4）中的S6块允许一维数组中的每个元素通过压缩的隐藏状态与任何先前扫描的样本进行交互。为了更全面地了解S6，[28]提供了深入的解释。
B. 用于高效图像显著目标检测的Mamba中嵌套Mamba

现有的图像显著目标检测（ISTD）方法主要使用卷积神经网络（CNNs）或CNNs与视觉Transformer（ViTs）的混合模型。后者在建模长距离依赖关系方面弥补了前者的不足，但面临着二次计算复杂度的挑战。最近，Mamba模型被提出，以其卓越的模型效率、较低的GPU内存使用量和更好的长距离依赖建模能力而闻名，并已成功应用于视觉领域。因此，我们探索了Mamba是否也能用于提升ISTD的性能。然而，当我们直接将视觉Mamba块应用于ISTD时，准确率并不十分令人印象深刻，因为对于检测小目标至关重要的局部特征没有得到充分探索。为了弥补这一不足，我们提出了一种Mamba中嵌套Mamba（MiM-ISTD）架构，如图3所示，该架构能够在保证卓越模型效率的同时，学习图像中的全局和局部信息。

给定一个二维红外图像$\mathcal{X}\in {\mathbb{R}}^{H\times W\times 3}$，我们将其均匀划分为$n$个块，形成$\mathcal{X}=\left[X^1,X^2,\dots ,X^n\right]$，其中每个块位于${\mathbb{R}}^{n\times p\times p\times 3}$，$(p,p)$表示每个块的分辨率。在MiM-ISTD中，我们将这些块视为代表图像的视觉句子。进一步地，每个块被分割成$m$个更小的子块，使得每个视觉句子成为一系列视觉单词：

$X^i=\left[x^{i,1},x^{i,2},\dots ,x^{i,m}\right],j=1,2,\dots ,m$

其中，$x^{i,j}\in {\mathbb{R}}^{s\times s\times 3}$是第$i$个视觉句子$X^i$的第$j$个视觉单词，$(s,s)$是子块的空间大小。

由于我们的MiM采用了分层编码器结构，因此视觉句子和单词的空间形状是不固定的，并且会随着网络层的加深而逐渐减小。

1）卷积主干（Convolutional Stem）：我们构建了一个卷积主干，其中使用了一叠$3\times 3$的卷积，以在第一阶段生成视觉单词$\in {\mathbb{R}}^{\frac{H}{2}\times \frac{W}{2}\times C}$和视觉句子$\in {\mathbb{R}}^{\frac{H}{8}\times \frac{W}{8}\times D}$，其中$C$是视觉单词的维度，$D$是视觉句子的维度。每个视觉单词对应于原始图像中的一个$2\times 2$像素区域，而每个视觉句子由$4\times 4$个视觉单词组成。与ViTs不同，由于视觉Mamba块的因果性质[28]，我们没有向视觉单词和句子添加位置嵌入偏差。

2）MiM分层编码器：我们MiM的核心部分是其具有不同令牌数量的四阶段分层编码器，如图3所示。在四个阶段中，视觉单词的空间形状分别设置为$\frac{H}{2}\times \frac{W}{2}$、$\frac{H}{4}\times \frac{W}{4}$、$\frac{H}{8}\times \frac{W}{8}$和$\frac{H}{16}\times \frac{W}{16}$。视觉句子的空间形状分别设置为$\frac{H}{8}\times \frac{W}{8}$、$\frac{H}{16}\times \frac{W}{16}$、$\frac{H}{32}\times \frac{W}{32}$和$\frac{H}{64}\times \frac{W}{64}$。我们采用[41]中的块合并作为下采样操作。每个阶段由多个MiM块组成，这些块同时处理单词级和句子级特征。视觉单词$x^{i,j}$通过线性投影映射到一系列单词嵌入$w^{i,j}$中，表示为：

$W^i=\left[w^{i,1},w^{i,2},\dots ,w^{i,m}\right],w^{i,j}=FC\left(\mathrm{Vec}\left(x^{i,j}\right)\right)$

其中，$w^{i,j}\in {\mathbb{R}}^c$是第$i$个视觉句子的第$j$个单词嵌入，$c$是单词嵌入的维度，$W^i$是第$i$个视觉句子的单词嵌入集合，$\mathrm{Vec}()$表示向量化操作。
MiM块处理两种不同的数据流：一种遍历视觉句子，另一种管理每个句子中的视觉单词。对于单词嵌入，应利用视觉单词之间的关系，如下所示：

W l i ′ = W l − 1 i + Mamba ⁡ ( L N ( W l − 1 i ) ) W l i = W l i ′ + Conv ⁡ F F N ( L N ( W l i ′ ) )

$$\begin{array}{l} W_{l}^{i^{\prime}}=W_{l-1}^{i}+\operatorname{Mamba}\left(L N\left(W_{l-1}^{i}\right)\right) \\ W_{l}^{i}=W_{l}^{i^{\prime}}+\operatorname{Conv}_{F F N}\left(L N\left(W_{l}^{i^{\prime}}\right)\right) \end{array}$$ Wli′​=Wl−1i​+Mamba(LN(Wl−1i​))Wli​=Wli′​+ConvFFN​(LN(Wli′​))​

其中，$l$表示从1到$L$的一系列块中第$l$个块的索引，$L$表示此类块的总数。第一个单词嵌入输入$W_0^i$是(6)中的$W^i$。$W_l^{i^{\prime }}$表示中间特征。所有转换后的单词嵌入用${\mathcal{W}}_l=\left[W_l^1,W_l^2,\dots ,W_l^n\right]$表示。这可以看作是一个内部Mamba块，如图4(a)所示。在此过程中，通过计算每两个视觉单词之间的相互作用来构建视觉单词之间的关系。例如，在包含小目标的补丁中，表示目标的单词将与其他与目标相关的单词具有更强的关系，而与背景部分的交互则较少。

在句子层面，我们生成句子嵌入记忆作为句子级表示序列的存储，初始化为零：${\mathcal{S}}_0=\left[S_0^1,S_0^2,\dots ,S_0^n\right]\in {\mathbb{R}}^{n\times d}$。在每一层中，单词嵌入序列通过线性投影映射到句子嵌入的域中，并随后整合到句子嵌入中：

$S_{l-1}^i=S_{l-1}^i+FC\left(\mathrm{Vec}\left(W_l^i\right)\right),l=1,2,\dots ,L$

通过这种方式，句子嵌入可以通过单词级特征进行增强。我们使用另一个称为外部Mamba块的Mamba块来转换句子嵌入：

S l ′ = S l − 1 + Mamba ⁡ ( L N ( S l − 1 ) ) S l = S l ′ + Conv ⁡ F F N ( L N ( S l ′ ) )

$$\begin{aligned} \mathcal{S}_{l}^{\prime} & =\mathcal{S}_{l-1}+\operatorname{Mamba}\left(L N\left(\mathcal{S}_{l-1}\right)\right) \\ \mathcal{S}_{l} & =\mathcal{S}_{l}^{\prime}+\operatorname{Conv}_{F F N}\left(L N\left(\mathcal{S}_{l}^{\prime}\right)\right) \end{aligned}$$ Sl′​Sl​​=Sl−1​+Mamba(LN(Sl−1​))=Sl′​+ConvFFN​(LN(Sl′​))​

外部Mamba块可以建模句子嵌入之间的关系。MiM块的输入和输出是视觉单词嵌入和句子嵌入，因此MiM分层编码器可以定义为：

${\mathcal{W}}_l,{\mathcal{S}}_l=\mathrm{MiM}\left({\mathcal{W}}_{l-1},{\mathcal{S}}_{l-1}\right),l=1,2,\dots ,L$

在我们的MiM块中，视觉Mamba块[28]被用于内部和外部Mamba块中，每个块后面都耦合了一个卷积前馈网络，如图4(b)所示。内部Mamba块通过建模视觉单词之间的关系来进行局部特征提取，而外部Mamba块则通过建模视觉句子之间的关系来捕获全局特征。卷积前馈网络负责用更精细的局部细节来增强这些特征。我们默认将每个阶段的块数$L_1,L_2,L_3,L_4$设置为2,2,2,2。通过堆叠8层的MiM块，我们构建了MiM分层编码器。

我们将每个阶段的视觉句子嵌入输出馈送到相应的解码器阶段。四阶段编码器输出的空间形状为$\frac{H}{8}\times \frac{W}{8}$、$\frac{H}{16}\times \frac{W}{16}$、$\frac{H}{32}\times \frac{W}{32}$和$\frac{H}{64}\times \frac{W}{64}$，这与典型骨干网络的特征图尺度不同。为了解决这种差异，这些输出在到达解码器之前需要通过直接的上采样模块。该模块包括一个步长为2的$2\times 2$转置卷积，后面跟着批量归一化[46]、GeLU[47]、步长为1的$3\times 3$卷积、另一个批量归一化和GeLU。因此，我们的MiM分层编码器可以生成相对于输入图像步长为4、8、16和32像素的特征图。
3）解码器结构：与编码器中使用的用于下采样的补丁合并[41]操作不同，我们在解码器中使用补丁扩展层[41]进行上采样。在每个解码器阶段，我们将上采样的编码器输出与扩展的解码器特征进行整合，并将整合后的特征发送到基本ResNet块。最后，特征通过全连接网络头部和插值操作得到最终的掩码预测。

4）复杂度分析：给定一个视觉序列 $T\in {\mathbb{R}}^{1\times n\times d}$，SSM的计算复杂度为 $3nEN+nEN=96nd+32nd=128nd$ [36]，其中 $N$ 是SSM的维度，默认设置为16，且 $E=2d$。相比之下，自注意力的计算复杂度为 $4n{d}^2+2n^{2}d$。显然，自注意力对序列长度 $n$ 表现出二次复杂度，而SSM是线性的。这种计算效率使得我们的MiM-ISTD相比其他基于二次Transformer的模型（如IAANet[20]和RKformer[24]）更具可扩展性。

我们提出的每个MiM块都包含16个内部Mamba块和一个外部Mamba块，每个Mamba块都包含核心的视觉Mamba部分和一个卷积前馈网络。为了简化计算，我们在计算浮点运算次数（FLOPs）时仅考虑最关键的视觉Mamba部分。这个核心部分包括一个SSM和三个线性层。我们将其FLOPs计算为$128nd+3n{d}^2$。因此，一个MiM块中所有内部Mamba块和一个外部Mamba块的FLOPs可以分别计算为$128mnc+3mn{c}^2$和$128nd+3n{d}^2$，其中$m$是一个视觉句子中的视觉单词数量，$c$是单词嵌入维度。因此，MiM块的总FLOPs $FLOP{s}_{MiM}$为$128mnc+128nd+3mn{c}^2+3n{d}^2$，仍然保持线性复杂度。而标准Transformer块的FLOPs $FLOP{s}_{ST}$为$2nd(6d+n)$[30]。考虑到在高分辨率红外图像中，$c\ll d$且$m\ll n$，$FLOP{s}_{MiM}$与$FLOP{s}_{ST}$的比率接近0，这意味着我们的MiM块在后续实验中展示了优越的精度-效率平衡的同时，仅引入了微不足道的FLOP增加。

四、实验

A. 实验设置

1） 数据集：我们选择NUAA-SIRST[16]和IRSTD1k[18]作为训练、验证和测试基准。NUAA-SIRST包含各种大小的427张红外图像，而IRSTD-1k包含1000张尺寸为$512\times 512$的真实红外图像。我们将NUAA-SIRST的所有图像大小调整为$512\times 512$用于训练和测试。IRSTD-1k是一个更具挑战性的ISTD数据集，包含更丰富的场景和1000张红外图像。对于每个数据集，我们使用$80\%$的图像作为训练集，$20\%$的图像作为测试集。

2）评估指标：我们在像素级和对象级评估指标上将我们的方法与其他最先进的（SOTA）方法进行比较。像素级指标包括交并比（IoU）和归一化交并比（nIoU），而对象级指标包括检测概率（$P_d$）和误报率（$F_a$）。

IoU衡量在给定数据集中检测相应对象的准确性。nIoU是IoU的归一化，它可以在红外小目标的结构相似性和像素准确性之间取得更好的平衡。IoU和nIoU的定义如下：

$$IoU=\frac{A_i}{A_u},nIoU=\frac{1}{N}\sum _{i=1}^N\left(\frac{TP[i]}{T[i]+P[i]-TP[i]}\right),$$

其中$A_i$和$A_u$分别是预测和真实值之间交集和并集区域的面积。$N$是样本总数，$TP[.]$是真阳性像素的数量，$T[.]$和$P[.]$分别是真实值和预测的正像素的数量。

$P_d$计算正确预测的目标数量$N_{\text{pred}}$与所有目标$N_{\text{all}}$之间的比例。$F_a$指的是错误预测的目标像素数量$N_{\text{false}}$与红外图像中所有像素$N_{\text{all}}$之间的比例。$P_d$和$F_a$的计算方式如下：

$$P_d=\frac{N_{\text{pred}}}{N_{\text{all}}},F_a=\frac{N_{\text{false}}}{N_{\text{all}}}.$$

预测的正确性取决于预测目标与真实值之间的质心距离是否小于3个像素。
3）优化：该算法在Pytorch中实现，采用自适应梯度（AdaGrad）作为优化器，初始学习率设置为0.06，权重衰减系数设置为0.0004。使用2个NVIDIA A6000 GPU进行训练，批处理大小设置为32。采用Dice损失[48]作为损失函数。在SIRST和IRSTD-1k数据集上的训练分别需要1000个epoch和800个epoch。

B. 定量比较

1）准确性比较：我们选择了一些最先进的红外小目标检测（ISTD）方法进行比较。如表I所示，我们的MiMISTD方法在像素级和对象级指标上，在两个数据集上都普遍达到了最佳性能。

对于像素级指标（IoU, nIoU），我们的方法由于进一步集成了内部曼巴（Inner Mamba）模块进行局部特征建模，因此取得了最佳性能。通过这种方式，一些可能被外部曼巴（Outer Mamba）模块忽略的可区分细节可以被注意到，从而提高了检测准确性。

对于对象级指标（$P_d$，$F_a$），如何在$P_d$和$F_a$之间达到平衡是一个挑战，因为较高的$P_d$也会增加较高$F_a$的可能性。从表I中，我们可以看到我们的方法通常获得了最佳的对象级指标结果，特别是在NUAA-SIRST测试集中检测到了所有小目标（$P_d=100\%$）。该结果表明，我们的MiM-ISTD能够学习更好的表示来克服漏检和误报。

2） 效率比较：我们还比较了不同方法在参数数量（M）、浮点运算次数（GFLOPs）、推理时间（s）和训练时GPU内存使用量（M）方面的效率，这些数据是在$512\times 512$红外图像数据集上获得的，如表II所示。除了ACM之外，与其他方法相比，我们的MiM-ISTD的参数数量、GFLOPs、推理时间和内存使用量都显著较少。这是因为我们采用了高效的曼巴结构，并使用了一个共享网络来计算视觉句子中每个视觉单词之间的关系，因此增加的参数和GFLOPs可忽略不计。即使与最轻量级的ISTD模型ACM相比，我们的MiM-ISTD的GFLOPs也更少，且准确率显著提高。此外，在更高分辨率的红外图像场景中，如图1所示，MiMISTD在保持优越准确率的同时，其在推理时间和GPU内存使用方面的优势将进一步扩大。值得注意的是，与没有内部曼巴模块的版本相比，我们注意到模型复杂度、参数数量和GPU内存使用量略有下降。然而，推理速度保持不变，平均准确率却出现了明显下降。这表明，虽然移除内部曼巴模块后模型变得稍微轻量一些，但这是以牺牲整体准确率为代价的。总的来说，MiM-ISTD达到了最佳的效率-准确率平衡。
C. 可视化

1）掩码结果的可视化：图5展示了不同方法的局部视图的可视化结果，其中现有的方法或多或少都存在检测不完整和漏检的问题。与其他最先进的方法相比，我们的MiM-ISTD更好地减少了这些情况，并更完整地检测到了所有小目标的形状。这是因为集成内部曼巴模块（Inner Mamba block）有助于我们的网络进一步挖掘更多的局部特征，从而促进了小目标的更精细检测。这也可以通过比较MiM-ISTD的可视化结果与“无内部曼巴”的可视化结果来证明，其中放弃内部曼巴会导致更差的检测性能。

2）特征图的可视化：为了更深入地了解所提出方法的效果，我们对TCI-Former和MiM-ISTD学习到的特征进行了可视化。特征图是通过根据空间位置重塑块嵌入来形成的。图6a展示了阶段1、2、3和4的特征图输出，其中每个输出随机采样了9个特征图通道。与TCI-Former相比，在MiMISTD中，深层网络更好地保留了局部信息。此外，与现有的最先进方法TCI-Former[11]相比，MiM-ISTD在每个通道之间的特征一致性更高，这意味着MiM-ISTD提取的特征更加关注于目标。这些优势得益于引入了内部曼巴模块来进一步建模局部特征。此外，我们在图6b中使用t-SNE[52]对阶段3的特征输出的64个通道进行了可视化，以证明我们的分析。我们观察到，与没有内部曼巴模块的模型以及现有的最先进方法[11]相比，MiM-ISTD的特征更加集中。这一观察结果与特征比较得到的结果一致。总的来说，我们的MiM-ISTD表现出更强的判别能力。
D. ROC结果

虽然IoU（交并比）、nIoU（归一化交并比）、$P_d$（检测率）和$F_a$（误检率）在固定阈值下衡量分割性能，但ROC曲线可以在多个不同阈值下提供整体评估。我们还在图7中比较了其他最先进方法的ROC曲线。可以注意到，我们的MiM-ISTD在两个数据集上都表现出整体令人印象深刻的性能，特别是在低误检率区间内，其真正检测率迅速上升，显示出强大的检测能力。此外，在较高的误检率下，MiM-ISTD仍然保持相对较高的真正检测率，表明整体鲁棒性良好。这些结果进一步证明了我们的方法相比其他最先进方法的优越性。

E. 消融实验

1）各模块的消融实验：各模块的消融实验结果如表III所示。我们的基线（编号#1）使用普通的视觉曼巴（Visual Mamba）[28]编码器作为我们的编码器，具有递减的空间分辨率，其中没有使用内部曼巴模块。我们还用标准的卷积、批量归一化和激活操作（编号#2）替换了MiM-ISTD中的内部曼巴模块，以考察内部曼巴模块的效果。我们发现，我们当前的设置带来了最佳结果，因为卷积通常关注局部信息，而内部曼巴模块能够捕获复杂的空间关系并整合多方向信息，从而对小目标有更精细的感知和识别能力。
2）补丁划分粒度的消融实验：

为了探索从卷积主干开始产生的每个视觉单词中包含的信息的影响，我们调整了子补丁的表示粒度，使得每个视觉单词对应于原始图像中的一个$1\times 1$（编号#1）或$4\times 4$（编号#2）像素区域，与我们当前设置的$2\times 2$相对比。我们还固定了视觉单词的$2\times 2$接收域，并改变了视觉句子（visual sentence）的初始空间形状，从$\frac{H}{8}\times \frac{W}{8}$变为$\frac{H}{4}\times \frac{W}{4}$（编号#3）和$\frac{H}{16}\times \frac{W}{16}$（编号#4），以消融视觉句子粒度的影响。我们可以观察到，过大或过小的划分粒度都不能带来最理想的性能。我们采用了在当前设置中表现最佳的配置。

五、结论

在本文中，我们提出了一种用于高效小目标检测的“曼巴中的曼巴”（MiM-ISTD）结构。我们将图像均匀划分为补丁作为视觉句子，并进一步将每个补丁划分为多个更小的子补丁作为视觉单词。我们设计了一个基于纯曼巴的MiM分层编码器，该编码器包含堆叠的MiM块。每个MiM块包含一个外部曼巴块来处理句子嵌入，以及一个内部曼巴块来建模单词嵌入之间的关系。视觉单词嵌入在经过线性投影后被添加到视觉句子嵌入中。实验表明，我们的方法能够更有效地对局部和全局信息进行建模。