10M参数拿下94%Dice：SegDT把医学图像分割搬进轻量时代

• 医学图像分割对于许多医疗任务至关重要，包括疾病诊断和治疗计划。一个关键领域是皮肤病变的分割，这对于诊断皮肤癌和监测患者至关重要。在此背景下，本文介绍了基于扩散 Transformer （DiT）的新型分割模型SegDT。 

• SegDT设计用于在低成本硬件上运行，并集成了整流流，该技术能够在减少推理步骤的同时提高生成质量，同时保持标准扩散模型的可灵活性。

•  yinqing-SegDT_2507在三个基准数据集上进行了评估，并与几种现有工作进行了比较，在保持快速推理速度的同时实现了最先进的结果。这使得所提出的模型对实际医疗应用具有吸引力。

• 这项工作提升了深度学习模型在医学图像分析中的性能和能力，为医疗专业行人提供了更快、更准确的诊断工具。

01 引言

皮肤癌是全局主要的健康问题之一，其早期检测对于提高生存率至关重要。医学图像分割在这一过程中发挥着关键作用，它能够实现病灶边界的精确界定。在自动图像分割的成功方法中，深度学习脱颖而出，彻底改变了该领域，其中卷积神经网络（CNN）如UNet [18] 和 DeepLabV3+ [3] 展现出令人瞩目的性能。然而，CNN在处理长距离依赖关系方面存在困难，这限制了它们对复杂或不规则形状病灶进行分割的能力。作为替代方案，Transformer [20] 受到其在自然语言处理（NLP）和视觉领域成功的启发，通过自注意力机制捕获全局上下文，提供了有吸引力的解决方案，TransUNet [2] 和 Swin-UNet [1] 等模型显示出更高的准确率（尽管它们的高计算成本可能阻碍实际应用）。另一方面，扩散模型最近已成为一种强大的技术，通过迭代去噪过程逐步优化分割，在各种医学图像任务中取得了最先进的结果 [21]。虽然它们非常精确，但其计算成本和较长的推理时间对实际部署构成了挑战。

为解决上述部分局限性，作者提出了SegDT，一种专为低成本GPU高效分割设计的额外小型扩散Transformer（DiT）模型。SegDT集成了整流流，以加速推理同时保持高分割精度。

本工作的主要贡献包括：

1. 提出了一种针对资源受限GPU优化的紧凑型DiT架构；
2. 描述了一种通过减少采样步骤实现高效推理的整流流集成方法；
3. 在三个基准数据集上进行了大量实验，与现有方法相比取得了最先进的性能。

02 相关工作

本节讨论了医学图像分割的相关先前工作，重点关注卷积神经网络（CNNs）、Transformer、混合架构和扩散模型，特别强调与皮肤病变分割相关的方法以及高效架构。

2.1 基于CNN的分割

卷积神经网络（CNNs）多年来一直是医学图像分割领域的主流方法。U-Net [18]凭借其编码器-解码器结构和 Shortcut ，通过实现有限数据下的有效学习，革新了该领域。其成功源于能够捕捉局部和全局上下文信息，这对于精确分割至关重要。

针对U-Net提出了许多变体和扩展，例如DeepLabV3+ [3]，该模型结合卷积来捕获多尺度信息，以及ResUNet++ [11]，该模型利用残差连接提升训练稳定性。然而，由于感受野有限，CNNs在捕捉长距离依赖关系方面往往存在困难，这对于分割复杂或不规则形状的病变（如皮肤癌成像中遇到的病变）可能至关重要。尽管这些方法能取得令人印象深刻的结果，但其局限性促使研究行人引入基于Transformer的方法。

2.2 基于Transformer的分割

Transformer模型最初是为自然语言处理[20]开发的，最近在计算机视觉领域[7]取得了显著成功。其核心优势在于自注意力机制，该机制能够有效捕捉长距离依赖关系和全局上下文。在医学图像分割中，TransUNet[2]是一项开创性工作，结合了Transformer和U-Net的优势。该模型采用Transformer编码器提取全局特征，并使用U-Net解码器保留局部细节。

虽然TransUNet展示了Transformer的潜力，但其计算复杂度可能成为限制因素。Swin-UNet[1]通过引入具有移位窗口的层次化Swin Transformer架构解决了这一问题，实现了自注意力的高效计算。这种方法在局部和全局特征提取之间取得了平衡，以较低的计算成本实现了具有竞争力的性能。然而，基于Transformer的模型仍然可能资源密集，这促使研究行人开发更资源高效的架构。

2.3 混合架构：CNN-Transformer

认识到卷积神经网络（CNN）和Transformer的互补优势，研究行人探索了混合架构。DS-TransUNet [14] 将可变形自注意力机制融入TransUNet框架，使模型能够聚焦于图像中的相关区域。

BRAU-Net++ [13] 结合了卷积特征提取器与基于Transformer的全局推理，在多个基准测试中取得了当前最佳结果。这些混合方法旨在利用CNN的局部特征提取能力和Transformer的全局上下文建模能力。虽然有效，但这些方法通常伴随着复杂度增加和资源需求提升。MobileUNETR [17] 采取了不同策略，专注于效率。

它提出了一种轻量级的基于Transformer的模型，专为移动和边缘设备优化。虽然与更大规模的模型相比牺牲了一些精度，但它实现了更快的推理速度，使其适用于资源受限的环境。

2.4 用于分割的扩散模型

扩散模型近来成为图像生成与分割领域的一种强大工具。它们通过逐步向图像中添加噪声，直至图像变为纯噪声（正向扩散），然后学习逆转这一过程以生成原始图像（逆向扩散）来运作。在医学图像分割中，MedSegDiff [21] 在扩散框架内集成了Transformer，使模型能够捕捉精细的解剖学细节。MedSegDiff-V2 [22] 通过引入多分辨率特征进一步改进了这一方法。尽管这些基于扩散的模型展示了令人印象深刻的性能，但它们的计算成本和较长的推理时间可能成为实时部署的限制因素。

作者的工作通过引入SegDT，一种结合了整流流的扩散Transformer（DiT）模型，来解决这个问题，该模型通过加速推理使其更适合实际应用。与以往通常依赖U-Net架构的扩散模型不同，作者探索了一种更高效的基于Transformer的设计。此外，作者采用整流流方法，以更少的采样步骤实现高质量的分割，显著减少了推理时间。这是作者工作与以往基于扩散的分割方法的关键区别。

03 提出方法

本节详细介绍了SegDT，作者提出的用于医学图像皮肤病变分割的架构。图1展示了SegDT架构，该架构由变分自动编码器（VAE）编码器、DiT和VAE解码器组成。VAE组件基于预训练的用于稳定扩散的微型自动编码器（TAESD）4，因其紧凑的尺寸和计算效率而被选择。TAESD的小型尺寸和快速编码/解码速度通过生成紧凑的潜在表示，压缩宽度和高度8倍，从而提升了SegDT的整体效率。这个更小的潜在空间允许更高效的DiT，减少了计算负担，并进一步加速了分割过程。DiT组件基于DiT-XS（超小型）变体，该变体采用2的块大小，如[16]中所述。DiT在图像生成方面表现出色，使其成为作者分割任务的合适基础。SegDT旨在解决在噪声医学图像中精确病变边界描绘的挑战——这是准确皮肤科诊断和治疗计划的关键要求。此外，模型的紧凑架构和推理过程中整流流的集成能够通过减少所需的推理步骤，实现快速高效的分析。

3.1 整流流高效推理

在前向扩散过程中，一个分割标签  （由其编码的潜在表示   表示）通过在   个时间步长上逐步添加高斯噪声而逐渐扰动。这导致了一系列噪声分割  （以及相应的噪声潜在表示   ）。逆向扩散过程旨在从噪声  （或   ）中重建原始分割 Mask  （或   ），条件是给定对应医学图像的潜在表示   。受高效扩散采样技术 （如去噪扩散隐式模型（DDIMs）［19］中使用的那些技术）的启发，作者采用整流流方法来学习一个改进的、更高效的逆向过程。

yinqing－SegDT＿2507不是直接预测每个时间步长添加的噪声，而是在潜在空间   中学习一个速度场   。这个速度场表示从   移动到   所需的变化方向和幅度，条件为   、时间步长   以及图像的潜在表示   。反向扩散过程随后遵循这个学习到的速度场。作者使用数值积分方法（如欧拉法）来近似反向轨迹。

其中   表示积分步长的大小。在作者的归一化公式中，作者假设连续时间区间为   。当该区间被划分为离散的时间步长   时，每一步对应于   的增量。速度场   通过最小化损失函数进行学习。给定一个带有噪声的潜在表示  （通过向编码后的真实 Mask   添加噪声获得），模型预测速度   。该预测值与从已知添加到   的噪声中推导出的目标速度（用于生成   ）进行比较，以获得   。这种比较常用的指标是均方误差（MSE）。该训练过程使模型能够准确预测逆转扩散过程所需的速度。

学习速度场而非直接预测噪声，可被视为一种形式正则化，促进更平滑且可预测的逆向扩散过程。它还为整合条件信息提供了更大的灵活性，例如由提供的解剖学背景。尽管学习扩散速度场的概念已被探索（例如[15]），但“正则化流”这一特定术语及其精确定义仍在扩散建模领域发展中。作者使用该术语强调了从噪声潜在空间到分割的更直接、高效路径的目标。

3.2 基于Transformer的扩散模型（SegDT）

本节介绍了基于Transformer的图像分割扩散模型SegDT的架构。SegDT采用预训练的变分自编码器（VAE）将图像映射到潜在空间。该潜在表示随后由DiT进行处理。SegDT采用12个DiT模块。

在训练过程中，DiT的输入为真实分割 Mask 的潜在表示。相比之下，在推理过程中，模型通过处理以固定种子初始化的随机采样、噪声潜在向量来生成分割结果。关键在于，VAE解码器组件在训练期间未被使用，因为损失函数直接在潜在空间内计算。这种方法显著减少了训练时间。

SegDT架构包含以下组件：

VAE编码器: 采用预训练的VAE编码器将尺寸为   的输入图像映射到低维潜在空间。在训练过程中，该编码器以相同的方式处理相应的真实标签 Mask 。生成的潜在表示维度为   ，其中   ，空间维度缩小了 8 倍，使得   和   。该潜在表示作为DIT模型的输入。

Patch 嵌入器: 采用共享架构处理输入张量和条件张量的 Patch 嵌入模块被应用。对于输入，它处理来自VAE编码器的潜在表示。在推理过程中，当没有真实 Mask 可用时，它处理随机化的张量作为条件输入。输入张量被展平，然后线性投影到   维的嵌入空间，其中   是来自DIT的token长度的嵌入维度。Patch 数量   由  ＇s确定，其中   是DiT模型中使用的 Patch 大小，  代表总 Patch 数量。

时间步嵌入器: 一个时间步嵌入模块被用于生成一个   维向量（其中   ，与图像块嵌入维度一致），该向量编码了扩散时间步   。此嵌入向量随后被用于调节 DIT 模块，为扩散过程当前阶段提供信息。

DIT模块: DiT模块是SegDT的核心，由一系列DIT模块构成。每个模块接收嵌入的潜在表示   、时间步嵌入t以及条件嵌入   作为输入。这些模块逐步细化潜在表示   。每个DIT模块的详细架构如图 2所示。每个DIT模块包含：

• 自适应层归一化（adaLN）：对激活值进行归一化处理，然后使用参数  （用于缩放）、  （用于平移）以及  （用于缩放）对归一化后的激活值进行调制，以实现平移和缩放。
• 自注意力机制：模型在图像嵌入中建模全局空间关系。
• 跨注意力：整合来自条件嵌入y的上下文指导。
• FFN （FFN）：增强特征表示。FFN隐藏层的大小由m1p_ratio  hidden_size确定。使用近似tanh计算的GELU作为激活函数。
• DropPath正则化：应用DropPath正则化以防止过拟合。

在DIT模块之后，对   的输出应用线性层，将其转换为形状为   的张量。该线性变换将每个图像块嵌入从维度   投影到新的维度  （相当于   ）。随后，该张量被 Reshape 为   ，以与VAE的潜在空间维度   保持一致，但通道深度增加了一倍（   而不是   ）。在这个   的张量中，前   个通道表示预测的噪声，后续   个通道表示预测的方差，主要用于训练期间。

推理时，经过此去噪过程后，模型的输出  （仅前   个通道）被输入到一个预训练的VAE解码器。该解码器随后从这个表示中重建图像，生成一个大小为   的3通道图像。最后，这个重建的图像通过二值化转换为最终的分割 Mask。

04 实验分析

4.1 数据集

ISIC 2016 [10]、2017 [6] 和 2018 [5] 挑战赛被广泛认为是医学图像分割领域的基准测试，特别是在皮肤病变分析方面。这些挑战赛提供了大规模的 Token 数据集，用于从皮肤镜图像中检测和分割黑色素瘤。ISIC 2016 数据集包含 900 张训练图像和 335 张测试图像，每张图像均有对应的分割 Mask ，可用于评估病变检测和边界精度。ISIC 2017 数据集在此基础上进一步扩展，包含 2000 张训练图像、150 张验证图像和 600 张测试图像，所有图像均带有分割 Mask 。ISIC 2018 数据集的规模进一步扩大，包含 2594 张训练图像、100 张验证图像和 1000 张测试图像，确保了对分割模型的更稳健评估。

4.2 指标

在评估这些数据集上的皮肤病变分割性能时，广泛使用的指标包括：Dice相似性系数（Dice）、   、像素精度（ACC）、灵敏度（SE）和特异度（SP）。这些指标用于评估预测分割 Mask（SM）与真实分割 Mask（GM）之间的重叠程度和分类准确性。

 表示共享像索的数量，  和   分别表示每个 Mask 中的总像索数，TP、 TN、FP 和 FN 分别表示真阳性、真阴性、假阳性和假阴性。

4.3 实现细节

SegDT使用Adam优化器，学习率为1e-4，批大小为32，训练了100个epoch。在50个epoch后，学习率降低了10倍。训练过程中未使用任何显式的数据增强技术。实验在两台NVIDIA RTX 3090 GPU上进行。输入图像通过双线性缩放至目标尺寸256×256像素，然后归一化到[0, 1]范围。在模型中，这些尺寸对应于C=3、H=256、W=256。最终的分割 Mask 是通过使用固定阈值0.2对重建图像进行二值化获得的。该阈值是根据在保留验证集上的性能选择的。作者探索了从0.1到0.5的阈值，以0.05的增量进行选择，并选择了0.2，因为它在验证集的Dice分数上表现最佳。DiT模块的mlp_ratio和hidden_size参数分别设置为4和192。

4.4 结果与讨论

为验证SegDT的准确性和泛化能力，作者在ISIC 2016、2017和2018皮肤病变图像数据集上进行了广泛的实验。作者将SegDT的性能与几种最先进的方法进行了比较，并分析了不同方法在不同场景和挑战性设置下的性能。

分割性能表1展示了在三个ISIC数据集上的分割结果。SegDT在所有数据集和大多数指标上均达到了最先进或极具竞争力的性能。在ISIC 2016数据集上，SegDT获得了最高的Dice分数（94.76%）、IoU（IoU，91.40%）和准确率（97.08%）。尽管DU-Net+在灵敏度上略高，但SegDT表现出显著更高的特异性（99.44%），这表明其能够准确识别健康组织，这在临床环境中至关重要。

在ISIC 2017中，SegDT再次取得了具有竞争力的结果，其Dice分数最高（91.70%），准确率最高（95.49%）。尽管DU-Net+显示出略好的IoU和灵敏度，但SegDT保持了显著更高的特异性（98.74%）。

在最大的数据集ISIC 2018中，SegDT达到了最高的Dice分数（94.51%）和IoU（IoU，90.43%），并具有具有竞争力的准确性和灵敏度。特别是，SegDT达到了最高的特异性（97.43%），进一步突显了其在识别健康组织方面的优势。

表1中GFLOPs和参数数量的包含突出了SegDT的效率。与DU-Net+相比，SegDT在显著更少的GFLOPs（3.68 vs. 54.00）和参数（9.95M vs. 39.00M）下实现了显著更好的性能。这表明SegDT能够在更轻的计算负载下取得优异结果，使其更适合资源受限的环境。尽管MobileUNETR和SLP-Net具有更低的GFLOPs和参数数量，但SegDT在性能和效率之间提供了更好的平衡。与最高效的模型相比，SegDT的性能仅通过相对较小的计算成本增加实现，但在分割精度方面，尤其是在特异性方面，取得了显著提升。这种权衡在医学图像分析中通常是可取的，因为准确性至关重要。

与IDDPM相比，SegDT在推理速度上具有显著优势。SegDT仅需15步推理即可实现与IDDPM相当的分段质量，而IDDPM则需要35步。得益于SegDT采用的整流流技术，推理步骤大幅减少，推理时间几乎缩短了2倍，这使得SegDT显著更快，对于实际应用至关重要。

定性分析图3展示了ISIC数据集中的示例分割结果，包括分割良好的案例和具有挑战性的案例。左侧示例展示了SegDT在形状、大小和纹理变化的情况下准确界定病灶边界的能力。右侧的具有挑战性的案例突显了一些局限性。例如，在某些案例中，SegDT难以处理具有不规则边界或非常小的病灶。这可能是由于Transformer模块的感受野有限，或者难以捕捉极小病灶中的细粒度细节。需要进一步研究以解决这些局限性。

05 结论

本文提出了SegDT，一种用于高效医学图像分割的新型DiT模型，特别针对皮肤病变分割。SegDT利用紧凑型DiT架构，并结合整流流技术，以加速资源受限GPU上的推理过程。yinqing-SegDT_2507解决了现有扩散模型的局限性，这些模型通常存在计算成本高和推理时间长的问题，从而阻碍了它们在实际临床环境中的实际应用。

SegDT在ISIC数据集（2016、2017、2018）上展示了其有效性，通过实现具有竞争力的、最先进的分割性能和显著更快的推理速度。这种效率，得益于整流流的使用，对于现实世界应用中的快速临床分析至关重要。此外，SegDT紧凑的架构使其能够在低成本GPU上部署，拓宽了High-Level医学图像分割的可及性，并促进了尖端诊断工具的民主化。

未来的工作将集中于进一步优化SegDT的架构，并探索额外的技术以提升其性能和效率。作者计划研究该模型在其他医学图像分割任务和数据集上的泛化能力。此外，作者将探索整合额外的临床信息，例如患者元数据，以进一步提高分割精度和临床实用性。作者相信，SegDT代表了开发高效且精确的医学图像分割模型，用于实际临床应用的具有前景的一步。

公众号后台回复“数据集”获取100+深度学习各方向资源整理