论文信息
题目：UU-Mamba: Uncertainty-aware U-Mamba for Cardiac Image Segmentation
UU-Mamba：用于心脏图像分割的不确定性感知U-Mamba
作者：Ting Yu Tsai，Li Lin，Shu Hu，Ming-Ching Chang，Hongtu Zhu，Xin Wang
源码：https://github.com/tiffany9056/UU-Mamba
论文创新点
- 不确定性感知损失函数：文章提出了一种新的不确定性感知损失函数，该函数整合了基于区域、基于分布和基于像素的损失，以提高分割精度和鲁棒性。这种设计允许模型优先考虑自信的预测并减少噪声数据的影响。
- Sharpness-Aware Minimization (SAM) 优化器：作者将SAM优化器集成到模型中，该优化器通过寻找损失景观中的平坦最小值来增强泛化能力，减少过拟合，并提高模型对输入数据小扰动的鲁棒性。
- UU-Mamba模型架构：文章介绍了UU-Mamba模型，这是U-Mamba模型的一个扩展，通过结合不确定性感知损失和SAM优化器，提高了心脏MRI分割的准确性和可靠性。
摘要
生物医学图像分割对于准确识别和分析医学成像中的解剖结构至关重要，特别是在心脏磁共振成像（MRI）中。手动分割劳动密集、耗时且容易出错，突显了自动化方法的需求。然而，当前的机器学习方法面临过拟合和数据需求等挑战。为了解决这些问题，作者提出了一个新的UU-Mamba模型，将U-Mamba模型与Sharpness-Aware Minimization（SAM）优化器和不确定性感知损失函数集成在一起。SAM通过在损失景观中定位平坦的最小值来增强泛化能力，从而减少过拟合。不确定性感知损失结合了基于区域、基于分布和基于像素的损失设计，以提高分割精度和鲁棒性。我们的方法在ACDC心脏数据集上进行了评估，超越了包括TransUNet、Swin-Unet、nnUNet和nnFormer在内的最先进模型。我们的方法在心脏MRI分割中实现了Dice相似系数（DSC）和均方误差（MSE）得分，证明了其有效性。关键词
心脏图像分割，Mamba，ACDC数据集，不确定性感知损失，锐度感知最小化。关键词
方法
图2显示了我们提出的UU-Mamba架构的概览以及训练改进。§ III-A描述了Mamba块和U-Mamba网络，详细说明了状态空间模型及其在捕获长距离依赖性中的作用。§ III-B提出了我们提出的不确定性感知损失，解释了它如何结合多个损失函数以增强模型性能和鲁棒性。§ III-C描述了锐度感知最小化优化，突出了其在实现损失景观中的平坦最小值以改善泛化和减少过拟合方面的好处。A. Mamba块和U-Mamba网络
U-Mamba网络整合了Mamba块和U-Net的优势，以增强全局上下文理解和提高医学图像分割的准确性。Mamba块为选择性结构化状态空间序列模型（S6）而设计，有效地处理长距离依赖性和序列信息，使其非常适合医学图像分割任务。状态空间模型（SSM）以输入状态向量、控制输入、过程噪声、状态转移矩阵和控制输入矩阵来描述系统，其基本形式由以下公式表示：对于观测矩阵、直通矩阵和观测噪声，观测由以下公式给出：S6架构通过选择性注意力机制和结构化参数化增强了传统的状态空间模型。具体来说，选择性注意力机制可以表示为：其中、和是从状态向量派生的查询、键和值矩阵；是键向量的维度。这种选择性注意力机制专注于相关的输入序列部分，并捕获复杂的依赖性。在Mamba块中整合S6对于顺序医学图像处理至关重要，例如心脏MRI分割，捕获时间动态和结构是必不可少的。作者的方法专门针对每张图像的分割任务，将每个图像视为独立实体，这意味着状态转移和观测矩阵（A、C等）应用于单个图像。U-Mamba利用Mamba的线性扩展优势，改善了CNN的长距离依赖性建模，而没有像ViT和SwinTransformer这样的Transformer中使用的自注意力机制的高计算负担。图2显示了U-Mamba块，它由两个顺序残差块后跟一个Mamba块组成。每个块包括卷积层、实例归一化和Leaky ReLU激活。图像特征被展平、转置、归一化，并通过具有两个并行分支的Mamba块进行处理：一个带有SSM层，一个没有。这些特征通过哈达玛德积合并，投影回并转置到原始形状。完整的U-Mamba网络架构包括一个编码器，这些块捕获局部特征和长距离依赖性，以及一个解码器，由残差块和转置卷积组成，用于详细的局部信息和分辨率恢复。跳跃连接将编码器的层次特征链接到解码器，最终的解码器特征通过一个1×1×1卷积层和一个Softmax层传递以产生最终的分割概率图。B. 不确定性感知损失
将不确定性整合到损失函数中涉及根据每个数据点的估计不确定性对损失函数的不同元素进行加权。这种方法允许模型优先从可靠的实例中学习，并最小化潜在错误或模糊数据的影响。Kendall和Gal引入了同质和异质不确定性的使用，以调整损失函数。同质不确定性在所有数据点之间是恒定的，而异质不确定性在实例之间变化。这种技术使模型能够根据这些不确定性调整其学习过程，通过关注自信的预测并减少模糊预测的影响，增强了韧性和精度。因此，它优化了在各种数据集上的训练过程，并提高了整体性能。为了增强分割精度，作者实现了一个不确定性感知损失函数，它结合了以下基于区域、基于分布和基于像素的损失，利用它们的互补优势。- Dice系数（DC）损失：这种基于区域的度量强调预测和真实区域之间的重叠，保持分割区域的形状和边界的准确性。
- 交叉熵（CE）损失：这种基于分布的损失确保了单个像素的准确分类，提高了分类精度。
- Focal损失：这种基于像素的损失通过赋予难以分类的实例更大的重要性来解决类别不平衡问题，增强了模型处理复杂场景的能力。
给定预测的分割和相应的真实掩码，让表示预测概率和表示真实标签。Dice相似系数（DSC）衡量预测分割和真实之间的重叠。它定义为：DSC的范围从0到1，DSC为1表示预测和真实之间完全重叠，DSC为0表示没有重叠。为了在训练分割模型时将这个度量用作损失函数，我们定义Dice系数（DC）损失为：DC损失旨在通过最大化DSC来最小化预测分割和真实之间的差异。通过最小化DC损失，模型被训练以产生与真实更重叠的分割，从而提高分割精度。交叉熵（CE）损失定义如下标准熵计算：Focal损失旨在解决类别不平衡问题，通过更加关注难以分类的样本：其中是聚焦参数，默认为2。我们通过将这些损失整合到一个不确定性感知框架中来定义不确定性感知损失：其中是各个损失组件的数量，代表每个单独的损失组件（如DC、CE和focal损失），是可学习的参数，根据不确定性调整每个损失组件的贡献。这些可学习参数在训练过程中被优化以最小化整体损失。将不确定性引入损失计算允许模型动态分配每个损失组件的权重。这种方法有效地平衡了全局和局部精度，同时减轻了类别不平衡的影响。因此，模型对噪声或模糊数据更具韧性，导致整体分割性能的提高。C. 锐度感知最小化优化
我们的方法涉及使用锐度感知最小化（SAM）优化来提高U-Mamba模型在分割心脏MRI图像（如ACDC数据集）中的性能。SAM优化使损失景观平坦，从而提高模型的泛化潜力。传统的优化技术旨在找到损失景观中的最低点，但这些点可能是陡峭的，导致在新数据上泛化性能差。相比之下，SAM识别更平滑的最小值，这些是在参数空间中模型性能更一致且不易受干扰的区域。使用SAM的动机在于其减轻过拟合的能力，这是医学图像分割中的一个常见挑战。传统的优化技术通常导致损失景观中的窄谷，导致在新数据上性能不佳。然而，SAM针对的是更平坦的最小值，这些与更好的泛化性能相关联。这对于处理像ACDC这样的复杂和多样化数据集特别有益，心脏MRI图像的变异性如果不妥善管理，可能会导致过拟合。SAM优化涉及一个迭代的两步过程。首先，对于每个小批量，参数被调整以最大化损失的方向。接下来，模型参数被更新以最小化最高改变的损失。这种扰动的目的是找到位于损失景观更平坦区域的模型参数，这些通常对应于更好的泛化和对输入数据小变化的鲁棒性。具体来说，让代表模型参数，代表损失函数，代表训练数据集。让代表扰动，控制周围邻域的大小。扰动被添加到参数以探索由范数约束定义的的局部邻域内的损失景观。SAM优化数学上被表述为：第一步涉及在内找到最大化损失的扰动。这一步识别了的局部邻域中最坏情况的方向，确保模型参数被调整到损失景观不陡峭的区域。在第二步中，模型参数被更新以最小化这个最坏情况扰动位置的损失，使参数对扰动鲁棒。通过迭代应用这两个步骤，SAM确保优化器寻找对扰动鲁棒的参数配置，提高泛化和性能。这种方法帮助模型在损失景观中找到具有更好泛化的更平坦的最小值。实验
声明
本文内容为论文学习收获分享，受限于知识能力，本文对原文的理解可能存在偏差，最终内容以原论文为准。本文信息旨在传播和学术交流，其内容由作者负责，不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题，请及时与我们联系，我们将在第一时间回复并处理。下载1：OpenCV-Contrib扩展模块中文版教程
在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。交流群
欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三 + 上海交大 + 视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~
![]()
![]()
