点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
精确对齐多模态图像与固有特征差异是一项关键挑战，在可变形图像配准中具有重要意义。传统基于学习的配准方法通常将配准网络视为黑箱，不具有可解释性。
一个核心洞察是，在跨模态之间分离对齐特征和非对齐特征可以带来好处。然而，对于图像配准任务中的突出方法，如卷积神经网络，它们往往无法通过其局部感受野来捕捉长程依赖性。
当给定的图像对由于缺乏有效地学习长程依赖性和对应关系而导致较大错位时，这些方法通常会失败。
在本论文中，作者提出了一种名为MambaReg的新颖Mamba基础架构，以解决这些挑战。通过作者提出的几个子模块，MambaReg可以有效地将负责配准的跨模态无关特征与依赖模态的非对齐特征进行分离。
通过选择性地关注相关特征，作者的网络能够适当地捕捉多模态图像之间的相关性，从而实现聚焦变形场预测和精确图像对齐。
Mamba基础架构无缝地将卷积层的局部特征提取力量与Mamba的长程依赖性建模能力集成在一起。在公共非刚性RGB-IR图像数据集上的实验表明，作者的方法在配准准确性和变形场平滑性方面均优于现有方法。
1 Introduction
图像配准是计算机视觉领域的一个基本任务。多模态图像配准（MIR）旨在将来自不同模态的两张图像在空间上进行对齐，这两张图像从不同的视角捕获了相同的场景，如[1] RGB图像提供了强度和纹理细节，而红外图像则提供了在RGB光谱中不可见的热性质信息。这种对齐对于后续的多模态任务，如图像融合[2]和分割[3]至关重要。然而，这些多模态图像是由不同的传感器获取的，它们通常会发生错位，并且包含需要和对齐的和不需对齐的特征，需要进一步关注。
经典的配准方法主要涉及迭代相似度度量指标（如互信息[4]、相关系数[5]和结构特征描述符[6,7]）的优化，以在多模态之间最大化相似度。与基于深度学习的模型相比，这些经典方法耗时且性能较差。近年来，随着深度学习的成功，许多基于学习的MIR方法出现，并广泛分为两类。
一类训练深度网络以优化多模态注册的某些相似度度量。例如，de Vos等人[10]提出了一种双卷积神经网络，用于最小化多模态图像之间的归一化相关系数度量，以进行注册。
然而，这些方法高度依赖于相似度度量的有效性。另一类采用图像到图像的转换，将MIR简化为单模态注册任务。Arar等人[14]采用GAN将一种模态转换为另一种模态，并在同一模态内执行注册。然而，训练GAN具有挑战性，可能引入人工噪声，从而阻碍注册。
这些方法将网络视为黑箱，无法提供更强的解释性和不同模态之间的关联性。一个核心的洞察是，将负责对齐和不受对齐影响的特征从不同模态中解耦是有利于多模态配准的。
这提出了一个有趣的问题：如何构建能够识别对齐和非对齐不同特征的配准模型？
同时，图像配准任务中的一些主要方法，如卷积神经网络，由于其局部感受野无法捕捉到长程依赖性，因此在其给出的图像对中存在较大的错位时，往往无法有效地学习长程依赖性和对应关系，从而导致方法失效。现有方法通常要求图像对大致对齐。尽管视觉变形器可以通过其注意力机制利用长程依赖性，但它们的高二次复杂度和大规模数据的依赖性使其难以实现。
因此，另一个问题是：如何使配准模型在不消耗过多资源的情况下利用长程依赖性？
为了解决上述问题，作者提出了一种基于Mamba的多模态图像配准网络，名为MambaReg，用于变形多模态配准任务。作者提出的MambaReg由两个交互分支组成，通过作者提出的模块将不同模态的配准和非配准特征解耦。为了利用长期依赖性，作者将Mamba的[16]长序列建模能力集成到作者的模型中。总之，本文作者的贡献如下：
作者提出了一种基于Mamba的多模态图像注册网络MambaReg，用于可变形多模态注册任务。据作者所知，提出的MambaReg是第一个基于Mamba的无监督图像注册框架。
作者提出了不同的模块，将负责配准的独立于模态的特征与依赖模态且非对齐的特征进行解耦。
为了应对该领域中的数据稀缺问题，作者使用公开的MSU-PID数据集重构了一个植物RGB-IR配准数据集。
大量实验表明，MambaReg在非网格RGB-IR图像数据集上优于现有方法，并实现了最先进的性能。
2 Related Works
多模态图像配准（Multi-modal Image Registration，MIR）旨在将来自不同模态的图像对齐，这比单模态配准更具挑战性，因为不同模态之间的对应关系复杂。为了处理刚性多模态配准问题，已经探索了多种基于深度学习的策略。Nguyen等人[17]开发了一种用于对具有照明变化的跨模态图像进行对齐的快速深度卷积神经网络，通过最小化光度损失。Pielawski等人[18]将对比学习引入到多模态配准，将其转移到单模态任务。Shao等人[19]使用多尺度 Transformer 网络来确定跨分辨率图像之间的对应关系，并实现高达分辨率的准确全同构估计。
传统变形多模态配准方法优化相似度指标，如互信息（MI）[20]，以及独立于模态的局部描述符（MIND）[7]。然而，优化相似度指标的过程通常是迭代和耗时的。为了应对这一挑战，Wu等人[21]提出了第一个深度网络，用于学习应用特定的相似度度量。Wang等人[22]通过学习编码器-解码器模块绕过域转换，生成独立于模态的特征，然后将这些特征输入空间 Transformer 网络（STN）[23]以学习变形变换。
最近，基于生成对抗网络（GAN）的图像到图像转换方法在变形多模态配准中展现了最先进的性能。Qin等人[12]的基于解耦图像表示的无监督多模态图像配准网络（UMDIR）。Arar等人[14]提出了一种结合STN和转换网络的方法，用于RGB/深度和RGB/NIR多模态配准，该方法采用一种新颖的训练方案，交替进行两种不同的空间转换流程。然而，训练GAN具有挑战性，图像转换可能引入有害的人工噪声，影响配准。Li等人[25]将图像的不同模态嵌入到基于解耦表示的共同特征空间中。最近，Dent等人[26]提出了一种引入学习的卷积稀疏编码（LCSC）到多模态配准的方法，以实现更好的可解释性。
Mamba 状态空间序列模型（SSMs）[27]是一种系统，它将一维函数或序列映射，可以表示为以下线性常微分方程（ODE）:
在此，作者考虑状态矩阵 和 是其参数， 表示隐性状态。SSMs 提供了几种理想特性，例如每个时间步的线性计算复杂度和并行计算以实现高效的训练。尽管 SSMs 具有每个时间步的线性计算复杂度和并行计算能力，但它们通常具有高内存需求，并且容易产生消失梯度，这限制了它们在序列建模中的应用。
为了解决这些问题，Structured State Space Sequence Models (S4) [28] 通过在状态矩阵 A 上施加结构形式并利用高阶多项式投影算子（HIPPO）[29] 进行初始化，从而构建具有丰富能力和高效长程推理能力的深度序列模型。
作为一种新的网络结构，S4 在 Long Range Arena Benchmark [31] 上超过了 Transformers [30]。最近，Mamba [16] 进一步改进了 SSMs 在离散数据建模（如文本和基因组）方面的性能。
Mamba 引入了一个输入相关的选择机制，与传统的 SSMs 不同，从而实现了从输入中进行高效的信息过滤。此外，Mamba 采用了一种硬件感知的算法，在序列长度上线性缩放，以在扫描时计算模型。Mamba 在现代硬件上的速度比以前的方法更快。此外，Mamba 架构将 SSM 块与线性层合并，明显更简单，并在各种长序列领域实现了最先进的性能。
3 Proposed Method
作者提出的MambaReg网络架构如图1所示。为了实现多模态可变形注册，MambaReg包括三个模块，分别是基于模态相关特征提取器（MDFE）、基于模态无关特征提取器（MIFE）和基于Mamba的多模态注册模块（M3RM）。在多模态图像注册（MIR）中，并非所有特征都为注册过程做出积极贡献；有些特征甚至由于不同模态之间信息内容的不一致而降低了注册的有效性。关键洞察是分离出负责注册过程的模态无关（MI）特征与模态相关（MD）特征，并仅在MI特征上进行注册。
为了分离MI和MD特征，作者利用两个MDFE模块分别处理来自两个不同模态的移动图像Ix和固定图像Iy。MDFE模块的目标是从每个模态中提取出不参与后续图像注册过程的MD特征。因此，MI特征MIx和MIy可以通过从Ix和Iy中减去MDx和MDy得到。所提出的M3RM模块将拼接的MI特征MIx和MIy作为输入，预测注册场φ。
MIFE模块的目标是进一步编码移动图像MIx的模态无关特征，并生成其在两个模态中的重构表示Rx和Ry。为了获取移动模态无关特征，通过应用空间变换网络（STN）[23]对Rx进行变形得到。
固定图像的MD特征MDy随后被添加，以得到重构的固定图像Iy。在MambaReg的训练阶段，移动图像和固定图像的分割标签保持未见，以确保网络以无监督方式学习。接下来，作者将介绍三个模块的设计和训练策略。
### Modality-Dependent Feature Extractor(MDFE)
在多模态图像配准中，并非所有特征都对配准产生积极影响；相反，由于不同模态之间的信息内容不匹配，一些特征甚至可能降低配准的有效性。因此，MDFE模块旨在提取依赖于图像模态和的特征，这些特征与图像配准无关，可以从原始图像中减去以获得干净的信息。
如图1所示，每个模态都有一个专门的MDFE：一个用于移动图像模态，另一个用于固定图像模态。移动图像和固定图像分别表示为和。移动图像的MDFE模块的设计旨在从移动图像中提取模态的依赖特征。
首先，经过卷积操作产生初步的特征表示，然后通过一个软阈值层来增强显著特征。接着，一系列基于 Mamba 的 learnable convolutional sparse coding (MLCSC) 块被应用来生成 Modality-Dependent Convolutional Sparse Representations (MDCSRs)。
每个MLCSC块都包含学习卷积稀疏编码（LCSC）算法[32]和双向Mamba（Bi-Mamba）层[33]，允许在图像数据中捕获长期依赖关系。在Bi-Mamba层中的选择性SSM有助于将模式相关的特征与模式无关的特征区分开来，从而生成更好的MDCSRs。
然后，将MDCSRs与模式部门滤波器（MD Filters）卷积，MD Filters是卷积字典滤波器，用于捕获模式相关的特征。对于固定图像的MDFE模块，它将固定图像作为输入，并采用与模式中相同的相同方式提取模式相关的特征。
Modality-Invariant Feature Extractor(MIFE)
语态不变特征提取器（MIFE）模块旨在生成重建的语态无关表示，这些表示捕捉了不同语态之间共享的关键特征，从而促进有效的多模态图像配准。如图1所示，MIFE架构反映了MDFE模块，包括一系列MLCSC块和语态特定的滤波器。然而，MIFE不是在原始输入图像上运行，而是在通过MDFE的输出获得的语态无关表示上运行。
通过一系列MLCSC块，其中包含LCSC算法和双向Mamba层，MIFE捕获了长程依赖性，并从语态相关的特征中区分出语态无关的特征。在双向Mamba层的选择性SSM有助于这一过程，使得可以生成语态不变的卷积稀疏表示（MICSRs）。随后，这些MICSRs与两个语态无关的滤波器（分别为滤波器和滤波器）卷积，以从各自语态中重构语态无关的表示和。
重构的表示 和 ，通过 MIFE 模块获得，包含输入模态之间共享的必要不变特征，从而通过关注共同信息内容而不是特定模态特性，促进有效的多模态图像配准。
Mamba-based Multi-Modal Registration Module(M3RM)
如图1所示，基于Mamba的多模态注册模块（M3RM）从两个模态中获取MI特征和，用于预测多模态注册任务中的变形场。M3RM模块包含一个U-Net [34]架构和U-Net编码器最后一层的Mamba块。这种设计旨在将卷积层的局部特征提取能力与Mamba在提取长期依赖关系中至关重要的特征的能力相结合，从而实现多模态注册。
Two-stage Training Strategy
遵循InMIR[26]的研究，作者采用两阶段的训练策略，通过预训练AG-Net以更好地将MI特征从MD特征中解耦。如图2所示，AG-Net包括2个MDFE模块和1个MIFE模块。MDFE模块从两个模态中提取MD特征。然后，通过从原始输入图像中减去相应的MI特征来计算MI特征。
给定MI特征作为输入，MIFE模块通过一系列MLCSC块和模态特定滤波器来重构模态无关的表示。
由于AG-Net是在完全对齐的图像上进行训练的，作者只需将MD特征添加到重构表示上，即可得到重构的移动图像。
AG-Net 是由以下损失函数进行预训练的：
由于 已经完全对齐，所以在第二阶段训练中，作者通过在 MambaReg 中的总损失函数中的指导损失，使用预训练的 AG-Net 来指导 MI 特征提取。
Binary ROI Mask
为了使模型在训练和推理过程中更专注于感兴趣区域（ROI），作者为MSU-PID数据集[35]生成了ROI Mask ，该 Mask 在无监督的情况下突出显示感兴趣的目标，同时抑制背景区域。如图3所示，ROI Mask 是通过一系列数字图像处理技术获得的。这个 Mask 是基于Otsu二值化方法[36]生成的，该方法自动确定图像中目标和背景像素的优化阈值。
Otsu方法假设图像包含两个像素类别（前景和背景），并计算最小化两个类别之间的类内方差的最优阈值。通过应用这个阈值，作者得到一个二进制 Mask ，其中属于感兴趣目标的像素被赋予1的值，背景像素被赋予0的值。然后，作者使用矩形结构元素对二进制 Mask 进行膨胀，以连接附近区域并放大它们，确保包括像植物的分支和茎等显著但较小的特征在前景中。最后，作者使用预设的像素大小过滤掉小而可能无关的组成部分，以丢弃噪声和次生艺术。这一步骤旨在丢弃噪声和次生艺术。在下一节中，作者将讨论作者如何将ROI Mask 集成到端到端训练过程中。
### Loss Function
MambaReg 的总体损失函数，用 表示，是一个由四个组成部分加权求和：ROI Mask 加权相似损失 、平滑损失 、引导损失 和重建损失 。
3.6.1 1) ROI Mask Weighted Similarity Loss
为了实现多模态变形图像配准的更好对齐，作者利用均方误差（MSE）来度量变形的图像与真实图像的相似性。如3.5节所述，在植物生长的早期阶段，植物可能是一个图像中的小物体。因此，生成的ROI Mask 被用作相似性损失的权重，以鼓励模型更关注感兴趣区域而不是背景。作者定义ROI Mask 加权相似性损失如下：
，其中是通过使用预测的变换场通过STN[23]对进行变换得到的。
为了鼓励平滑变形，作者采用了VoxelMorph [8]中的平滑损失，以约束预测的变形场上的较大梯度，如下所示：
在变形场中的像素p表示。
指导损失：如第3.4节所述，作者利用预训练的AG-Net提取的MI特征作为监督信息，用于指导MambaReg-Net的MI特征提取，从而提高配准精度。指导损失定义为利用均方误差（MSE）损失来鼓励来自AG-Net和MambaReg-Net的相似MI特征，这遵循了InMIR [26]的设定。
重建损失：如图1所示，MambaReg将运动和固定图像分别重构为和。为了在这样一个重建过程中实现一致性，重建损失被定义为，
整体损失：MambaReg的整体损失函数是上述四个组成部分损失函数的加权求和：
其中，、、 和 是每个损失组成部分的权重。
4 Experimental Results
Dataset
基于Cruz等人收集的植物表型多模态影像数据库（MSU-PID）[35]，作者重新生成了一个非刚性RGB-IR植物影像数据集，用于变形多模态图像配准任务。原始MSU-PID数据集包括9天（从上午9点到晚上11点）捕获的16棵拟南芥植物的影像，每天产生15帧。其中，每天有4帧由手工标注的实例分割标签。每张图像包含的实例数量从4到14个不等，反映了植物的自然生长过程。
然而，直接使用MSU-PID数据集进行多模态图像配准面临几个挑战：
(1) 同一植物和时间戳的RGB-IR对未预对齐；
(2) 植物生长导致新叶子的出现，从而在运动和固定图像对之间实例之间的差异；
(3) 容器中存在显著的背景干扰。
为了应对这些问题，作者按照以下步骤重建了数据集。首先，作者使用每个植物最后一个时刻的分割标签来生成植物特定的边界框。利用这些边界框，作者裁剪了RGB和红外图像及其相应的标签。这使得RGB-IR图像对的中心对齐，可以专注于植物区域而不是周围的容器。
然后，从裁剪的数据中，作者选择了同一植物在不同时间戳的RGB-IR图像对，允许大的变形，同时确保两种模式都包含相同数量的实例。在无监督图像配准的训练集中，有12,275个图像对没有标注，而测试集中有1,633个图像对带有标注，可用于计算Dice得分[37]。最后，作者从12,275个图像对中随机选择了300对用于训练，从1,633个图像对中选择了900对用于测试。此外，如第3.5节所述，作者还为作者的重建数据集中的所有RGB和红外图像生成了ROI二值 Mask 。作者的重建数据集将公开发布。
Implementation Details
该网络使用128x128大小的图像进行训练。MDFE和MIFE模块中的MLCSC块数量设置为4。在损失函数中，，，和分别设置为100，10，25和10。总共的训练轮次为1000。作者使用Adam优化器以基本学习率1e-4训练网络。在训练过程中，作者使用PolyLR调度器根据多项式衰减策略调整学习率。作者在单个NVIDIA RTX A6000 GPU上进行实验。模型初始化时使用随机权重，随机种子设置为3407。所有实现都使用PyTorch深度学习框架。
Evaluation Metrics
对于变形注册，作者采用Dice分数[37]，均方误差（MSE），归一化相关系数（NCC）和SSIM来衡量注册精度。Dice分数，NCC和SSIM的较大值以及MSE的较小值表示图像注册的更好性能。
由于作者在重建的MSU-PID数据集中，叶子大小不同，因此根据GT中每个标签像素的数量对Dice分数的计算进行加权。计算公式如下：
总共有N个标签，和分别是固定图像和变形图像中第i个标签的像素集。加权Dice评分提供了一种更平衡的测量注册性能的方法。
Quantitative and Qualitative Results
作者将作者提出的MambaReg方法与最先进的方法InMIR[26]进行了比较。图5可视化了由MambaReg和其他 Baseline 模型预测的移动图像及其对应的变形场。定性比较显示MambaReg在注册大变形图像对方面表现出色。与最先进结果的定量比较如图1所示。
对于Dice，NCC和SSIM等指标，越高越好。对于MSE等指标，越低越好。表1说明，作者的模型在Dice，NCC和SSIM指标上的数量最高，而在MSE指标上的数量最低，这表明作者模型的性能最好。作者可以得出结论，作者的提出的算法在所有指标上实现了最佳性能，验证了作者提出的模型的有效性。
### Ablation Study
基础模型。所提出的MambaReg具有6个基础模型（即，B1，B2，B3，B4，B5，B6），如表2所示。所有基础模型都包含三个模块，MDFE，MIFE和M3RM。主要的区别在于是否包含不同的Mamba模型，例如1d Mamba或1d Bi-Mamba。
(i) B1不包含任何基于Mamba的层。
(ii) B2不包含基于Mamba的层，除非涉及1d Bi-Mamba的MDFE。
(iii) 在B3中，MDFE和MIFE涉及1d Bi-Mamba，但M3RM不包含任何Mamba-base层。
(iv) 在B4中，所有模块，包括MDFE，MIFE和M3RM都涉及1d Mamba。
(v) 在B5中，所有模块，包括MDFE，MIFE和M3RM都涉及1d Bi-Mamba。
(vi) B6是作者的完整模型，名为MambaReg。所有模块，包括MDFE，MIFE和M3RM都涉及1d Bi-Mamba。同时，B6引入了二进制ROI Mask ，以更注重感兴趣的区域。
消融分析。 表2中显示了消融研究的结果。当作者将1d Bi-Mamba引入MDFE、MIFE和M3RM时，与B1相比，模型在DSC上的改进分别为1.4%、2.4%和4.0%。结果证实了作者将1d Bi-Mamba层引入MDFE、MIFE和M3RM模块的有效性。与将1d Mamba引入三个模块的B4相比，B5在DSC上的改进为1.5%。结果证实了1d Bi-Mamba的有效性，这比1d Mamba更好。当作者引入二进制ROI Mask 以更关注感兴趣区域时，与B5相比，B6在DSC上的改进为0.7%。结果证实了ROI Mask 的引入是有效的。因此，结果证明了作者所提出的MambaReg的有效性。
5 Conclusion
MambaReg，提出的基于Mamba的架构用于无监督变形多模态图像注册，通过解耦模态相关和模态无关的特征以及利用图像内的长程依赖，已经证明其优越性。
专门模块，即模态相关特征提取器和解耦模态无关特征提取器，使得相关特征的聚焦提取和处理成为可能。在重建的MSU-PID数据集上的大量实验表明，MambaReg在注册精度方面达到了最先进水平。
[0]. MambaReg: Mamba-Based Disentangled Convolutional Sparse Coding for Unsupervised Deformable Multi-Modal Image Registration.
点击上方卡片，关注 「AI视界引擎」 公众号