Diffusion Mamba：用于CT到MRI转换的Mamba扩散模型

• • 提出背景
  • • 拆解
    • 左侧：整体框架
    • 中间：Mamba块的细节
    • 右侧：螺旋扫描的细节

 

---

提出背景

论文：https://arxiv.org/pdf/2406.15910

代码：https://github.com/wongzbb/DiffMa-Diffusion-Mamba

使用图像生成模型，尤其是扩散模型和生成对抗网络（GANs），将计算机断层扫描（CT）图像转换为磁共振成像（MRI）图像的可能性。

这种转换尝试在不增加额外成本的情况下扩展诊断检查的范围。CT成像在成本效益和速度方面表现优异，尤其适用于骨骼和钙化组织，但在软组织的成像上存在限制，而MRI则提供了更高清晰度的软组织图像，但成本和时间上的需求较高。

Mamba模型，作为一种在SSMs中引入时变参数以提高并行处理效率的技术，已经被扩展应用到视觉和医学图像生成中。

这项技术特别适合处理大规模和高维度数据，但在处理图像数据时，直接应用Mamba可能会影响图像的空间完整性。因此，引入了新的扫描技术，如螺旋扫描模块，以保持图像结构的连续性和完整性。

此外，这项研究还提出了通过视觉嵌入器预训练和跨序列监督来优化重要区域的权重分配，从而更精确地重建MRI图像。

通过这种方法，研究人员能够在保持参数数量不变的情况下，比传统的基于CNNs和ViTs的方法取得更好的性能。

拆解

扩散模型架构Diffusion Mamba（简称DiffMa），该模型专为MRI图像的扩散模型训练而设计。

这张图描绘了Diffusion Mamba（DiffMa）框架的结构细节，特别强调了用于MRI图像生成的扩散模型的各个组件和工作流程。

左侧：整体框架

• 嵌入层：首先，噪声潜在MRI图像通过嵌入层进行处理，这一层处理输入的图像，将其分解成多个块，并且为每个块添加条件标签（如m、c、t等）。
• Mamba块序列：这些输入块随后被一系列Mamba块处理，这些块通过长跳连接和线性处理层次地处理这些数据。
• 重排和线性层：处理后的数据被重新排列，并通过线性层进行进一步的处理，最终生成预测的噪声图像。

中间：Mamba块的细节

• 条件处理：输入令牌被带有条件信息的多层感知机（MLP）处理，这有助于模型根据给定的医学条件进行更精确的生成。
• 螺旋扫描：采用螺旋扫描技术来处理输入数据，此技术通过对数据进行螺旋形扫描来强调结构信息，增强生成图像的质量。
• 层标准化和比例变换：在螺旋扫描后，数据通过层标准化进行标准化，并进行必要的比例和位移调整。

右侧：螺旋扫描的细节

• 前向扫描和1D卷积：展示了螺旋扫描内部的细节，包括前向扫描和1D卷积层，这有助于处理和识别图像序列中的关键特征。
• 层标准化：最后，通过层标准化确保数据的一致性和标准化。

这个图形展示了DiffMa模型在处理医学图像，尤其是MRI图像时的复杂性和多层次处理结构，突出了在模型中使用的先进技术，如螺旋扫描和条件化的多层感知机。

这种结构旨在提高生成图像的精确度和相关性，特别是在处理具有复杂内部结构的医学图像时。

DiffMa使用视觉Mamba处理图像块序列。

特别关注螺旋扫描和跨序列注意力两个组件。

1. 子解法1：使用螺旋扫描模块
   • 功能：在处理图像块时保持空间连续性。
   • 原因：传统的扫描模式可能无法有效考虑空间连续性，螺旋扫描可以通过特定的排序模式保持图像块之间的结构关系，从而保持空间信息的完整性。

这张图显示了2D图像螺旋扫描的八种不同模式。

每个模式包括两种方向的扫描方式，这些模式在处理图像数据时被用来以不同的顺序遍历图像的像素或区块。

图中每种扫描方式的详细说明如下：

1. 左上至右下：两行四列的图表展示了螺旋形扫描的不同方式。

   • 第一行从左到右显示了螺旋向内和向外的顺时针方向。
   • 第二行从左到右显示了逆时针方向的螺旋向内和向外。

2. 每种模式的用途：

   • 顺时针和逆时针螺旋：两种基本方向，顺时针通常用于常规的数据扫描，而逆时针则可能用于特殊情况或提供不同的数据视角。
   • 向内和向外：螺旋可以从外围向中心移动，或者从中心向外扩展，这两种方式可以根据数据的特性和需要突出的图像区域来选择使用。

3. 实际应用：

   • 每个块在处理图像时选择其中一个扫描模式，这有助于模型更好地理解和解析图像的结构和空间关系。
   • 不同的螺旋扫描模式可以用来强调图像的不同特性，例如，从外向内的扫描可能更关注中心特征，而从内向外的扫描则可能更注重边缘特征。

这些螺旋扫描模式为图像分析提供了灵活的方法，使得图像处理模型能够根据不同的应用场景选择最合适的数据遍历方式，从而优化性能和结果。

通过这样的技术，可以在保留关键结构信息的同时，增加处理图像数据的维度和深度。

2. 子解法2：采用跨序列注意力机制

   • 功能：增强模型对于不同图像块序列中的细微差异的识别能力。
   • 原因：在扩散模型中，理解和利用图像块间的关系对于生成高质量的MRI图像至关重要。跨序列注意力通过对CT图像块生成的软掩模和嵌入进行处理，使模型能够更加精确地关注图像中的关键区域。

3. 子解法3：利用变分自编码器（VAE）生成噪声潜在空间表示

   • 功能：为MRI图像提供一个初始的噪声潜在空间表示，作为DiffMa模型的输入。
   • 原因：噪声潜在空间表示可以模拟实际MRI图像在成像过程中可能出现的随机变异，为模型训练提供更加真实的数据基础。

4. 子解法4：使用视觉编码器和视觉嵌入器对CT图像进行处理

   • 功能：将CT图像转换为适合模型处理的嵌入向量。
   • 原因：通过专门的视觉编码器，如BioMedCLIP，和预训练的视觉嵌入器处理CT图像，可以精确地捕捉CT图像与MRI图像之间的差异，为后续的扩散过程提供关键信息。

这张图展示了视觉嵌入器（Vision Embedder）的框架，这是用于处理CT图像数据的神经网络架构，旨在生成图像嵌入和软掩膜，以供进一步的图像处理和分析使用。

1. 潜在CT：

   • 原始的CT图像数据作为输入，显示在图像左下角。

2. 卷积层和展平（Conv2d & Flatten）：

   • 图像数据首先经过二维卷积层（Conv2d）处理，这有助于提取图像的空间特征。
   • 卷积后的数据通过展平操作转换成一维数组，以便进行进一步处理。

3. 池化层（MaxPool2d & AvgPool2d）：

   • 使用最大池化（MaxPool2d）和平均池化（AvgPool2d）技术对卷积后的特征进行下采样，这有助于减少计算量并保留关键信息。

4. 多层感知机（MLP）：

   • 池化后的数据送入多层感知机（MLP），MLP能够进一步处理数据，学习更高级的特征表示。

5. 标准化层和加权组合（Layer Norm, ⊗ & σ）：

   • 在数据进入MLP之前，先经过层标准化（Layer Norm），保证数据在不同的批次间具有相同的分布。
   • 数据通过一个特殊的操作（⊗ & σ），可能表示了一个加权或标准化步骤，用于调整和优化特征权重。

6. 输出：CT嵌入和软掩膜（CT Embeddings & Soft Masks）：

   • 最终，生成的CT嵌入和软掩膜被输出。CT嵌入为后续处理提供了密集的特征表示，而软掩膜可能用于指导或优化图像生成过程中的注意力机制。