医图论文 ICLR‘25 | UD-MAMBA：一种用于医学图像分割的像素级不确定性驱动的Mamba模型

最新推荐文章于 2025-04-07 15:52:16 发布

程序员笑武

最新推荐文章于 2025-04-07 15:52:16 发布

阅读量1.4k

点赞数 7

文章标签：计算机视觉人工智能 python 开发语言知识图谱机器学习

本文链接：https://blog.csdn.net/m0_59164304/article/details/144110340

版权

论文信息

题目：UD-MAMBA: A Pixel-Level Uncertainty-Driven Mamba Model for Medical Image Segmentation

UD-MAMBA：一种像素级不确定性驱动的Mamba模型用于医学图像分割

论文创新点

作者引入了一种新颖的基于通道不确定性的像素级选择性扫描方法，解决了传统基于位置的顺序扫描方法的局限性。
作者引入了可学习参数来平衡不同扫描方向上的特征重要性，并采用余弦一致性损失来对齐前向和后向扫描结果，提高了特征一致性。
在三个医学成像数据集上的广泛实验表明，UD-Mamba有效地识别了模糊区域，与现有的基于Mamba的方法相比，取得了更可靠的分割结果。

摘要

近期进展强调了Mamba框架，这是一个状态空间模型（SSMs），以其在线性计算复杂度下捕捉长距离依赖关系而闻名。尽管Mamba在医学图像分割中展现出了竞争性能，但由于传统基于位置的扫描方法的间歇性特性以及医学图像中常见的复杂、模糊边界，它在建模局部特征方面遇到了困难。为了克服这些挑战，作者提出了不确定性驱动的Mamba（UD-Mamba），它通过将通道不确定性纳入扫描机制来重新定义像素顺序扫描过程。UD-Mamba引入了两种关键扫描技术：顺序扫描，它通过逐行扫描优先处理高不确定性区域；以及跳过扫描，它以固定间隔垂直处理列，从高到低或从低到高不确定性。顺序扫描有效地聚集高不确定性区域，如边界和前景对象，以提高分割精度，而跳过扫描增强了背景和前景区域之间的交互，允许及时整合背景信息以支持更准确的前景推断。认识到从某些区域扫描到不确定区域的优势，作者引入了四个可学习参数来平衡从不同扫描方法提取的特征的重要性。此外，采用余弦一致性损失来减轻在扫描过程中从不确定区域过渡到确定区域时的缺点。作者的方法在三个不同的医学成像数据集上验证了其稳健的分割性能，涉及病理学、皮肤病变和心脏任务。

3 方法

在本节中，作者首先介绍与状态空间模型（SSMs）相关的概念。接下来，作者提供了所提出的UD-Mamba架构的全面概述，并在图3中说明了整体框架。最后，作者详细阐述了UD-Mamba的关键组成部分，详细说明了不确定性驱动的选择性扫描块（UD-SSB）的操作工作流程和派生的优化策略。

3.2 UD-MAMBA

UD-Mamba架构利用了一个简洁而强大的UNet框架，其基本层为不确定性驱动（UD）块。如图3所示，设计包括三个关键组件：一个将输入图像转换为一系列补丁的补丁嵌入层，用于后续处理；一个由UD块组成的编码器-解码器结构，用于捕获和整合不同尺度的局部和全局特征；以及一个基于解码特征产生最终像素级分割输出的分割头。编码器-解码器配置通过跳跃连接得到增强，这些连接有助于整合多尺度特征表示。在UD块内，不确定性驱动的选择性扫描块（UD-SSB）作为关键元素。这种架构选择增强了跨层的信息传播，最终提高了分割精度。

3.3 不确定性驱动的选择性扫描块

为了解决传统状态空间模型（SSMs）如Mamba在有效建模局部特征方面的局限性，作者提出了一种基于像素级不确定性的扫描方法。这种方法与传统的像素顺序扫描机制不同，它利用像素级的不确定性来通知扫描序列。如图4 I所示，作者的不确定性驱动的选择性扫描块（UD-SSB）引入了五个关键组件：通道不确定性计算、基于不确定性的排序、扫描扩展操作、S6块处理和恢复操作。

给定一个输入特征张量，其中、、和分别表示批量大小、通道数、高度和宽度。作者提出以下方法：

通道不确定性计算：为了计算每个空间位置在所有通道上的不确定性图，作者定义：

在此背景下，作者使用标准差作为不确定性度量，这一选择在第4.4.2节中的结果得到了验证。具体来说，对于输入特征图，作者计算每个空间位置（h，w）在所有通道上的标准差：

其中表示所有通道在该空间位置的平均值。这个计算捕捉了跨通道的像素级标准差，其中较高的不确定性通常对应于关键区域，如对象边界或前景区域，而较低的不确定性表示背景一致性。通过关注像素级不确定性，我们可以更精确地识别医学图像分割中的关键区域，这在识别病理区域或器官边界时至关重要。

基于不确定性的排序：然后，不确定性图被按降序排序，得到，它将空间位置从高不确定性区域（前景和边界）排序到低不确定性区域（背景）。这允许模型在后续操作中优先处理具有更高复杂性或重要性的区域：

特征图重排：使用排序后的索引，作者重排原始特征图以创建，其中高不确定性区域被密集处理。这种重组为有效的扫描做好准备：

扫描扩展操作：作者在重排的特征图上实现了两种不同的扫描操作：1）顺序扫描（Scanse）：此操作按像素不确定性的降序处理空间位置，这意味着优先处理高不确定性区域，如前景对象和边界。这种方法确保模型密集地建模关键高不确定性区域，允许模型专注于对准确分割至关重要的区域。2）跳过扫描（Scansk）：此操作在不确定性谱上定期间隔选择空间位置，促进背景和前景区域之间的交互。通过及时整合背景信息，跳过扫描有助于保持图像的整体背景结构，同时细化前景的细节，从而实现更平衡的分割结果。结合顺序和跳过扫描使模型能够有效捕获局部和全局特征。

S6块处理：然后扫描的特征被S6块处理：

恢复操作：最后，重排和处理的特征被恢复到它们原始的空间配置，确保输出的空间结构与输入一致。这确保了模型保留了对准确医学图像分割至关重要的位置信息：

3.4 不确定性驱动的选择性扫描优化策略

如图4 II所示，UD-SSB应用四种不同的扫描序列：从高到低不确定性水平的顺序和跳过扫描（y1和y2），以及从低到高不确定性水平的顺序和跳过扫描（y3和y4）。通常，高不确定性区域更可能对应于目标区域和关键边界，而低不确定性区域通常与背景相关。在作为自回归模型运行的Mamba框架中，每个输出都依赖于由先前输入派生的隐藏状态。从低不确定性到高不确定性区域的扫描允许模型首先处理更简单的背景信息，在处理更复杂区域之前积累隐藏状态储备。如图2所示，这种方法优于相反的扫描顺序。因此，为了利用这一属性，作者提出了两种优化策略，以利用从高到低不确定性扫描的好处，同时减轻从低到高不确定性扫描的固有局限性。

3.4.1 不同扫描序列的重新加权

为了优化每个扫描序列的贡献，作者引入了四个可学习的参数（α1，α2，α3，α4），每个参数对应一个扫描方向。这些参数用于增强从高到低不确定性扫描的优势，同时调节每个单独扫描序列的贡献。重新加权机制在数学上定义为：

这种方法确保每个扫描方法根据其在捕获关键图像区域方面的有效性成比例地贡献，更多地强调从更确定到较不确定区域的扫描。

3.4.2 双向扫描之间的一致性约束

为了解决在解码阶段与低到高不确定性扫描相关的局限性并提高整体分割性能，作者在解码器的末尾引入了余弦一致性约束。此约束应用于两个方向（从高到低和从低到高不确定性）执行的顺序和跳过扫描之间。通过将低到高不确定性扫描的结果与高到低不确定性扫描的结果对齐，作者确保了不同扫描方向上的特征表示的一致性。为了保持位置一致性，所有输出在应用恢复操作到后得出。余弦一致性损失定义为：

其中表示前向和后向顺序和跳过扫描之间的平均余弦相似性。通过最大化这种相似性，我们的目标是最小化两种扫描方向之间的差异，从而加强最终分割输出的一致性。

最后，总体损失函数结合了监督损失和余弦一致性损失：

其中表示组合的交叉熵和Dice损失（CeDice损失），是余弦相似性损失，是平衡这两个组成部分的超参数。

4 实验

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述