点击下方卡片，关注 「AI视界引擎」 公众号
「AI视界引擎交流群」
在过去的几十年中，深度神经网络，尤其是卷积神经网络，已经在各种医学图像分割任务中取得了最先进的表现。近年来，视觉Transformer（ViT）的引入显著改变了深度分割模型的格局。由于它们出色的性能和可扩展性，对ViTs的关注日益增加。然而，我们认为目前基于视觉Transformer的UNet（ViT-UNet）分割模型的设计可能无法有效地处理医学图像分割任务中感兴趣对象的异质外观（例如，不同的形状和大小）。
为了应对这一挑战，我们提出了一种结构化的方法，将空间动态组件引入到ViT-UNet中。这种适应使得模型能够有效地捕捉到具有不同外观的目标对象特征。这是通过以下三个主要组件实现的：
可变形的Patch嵌入；
空间动态的多头注意力；
可变形的位置编码。
这些组件被集成到一个称为AgileFormer的新颖架构中。AgileFormer是一个为医学图像分割设计的空间灵活的ViT-UNet。在三个分割任务中使用公开可用数据集进行的实验证明了所提出方法的有效性。
医学图像分割任务在现代医学中至关重要，因为它是许多基于图像的诊断和分析的第一步[1, 2]。基于深度学习的自动化分割方法由于其高效率和最先进的表现而在这一领域占据主导地位。其中，卷积神经网络（CNNs）[3, 4, 5, 6, 7, 8]自从UNet[3]提出以来已成为最流行的选择。这是因为CNN在处理图像驱动任务时具有固有的优势，例如它们捕捉局部性和平移不变性的能力。然而，它们主要由于受限的感受野而在捕捉全局语义方面存在困难。
相比之下，最近提出的视觉Transformer（ViT）[9]通过自注意力机制缓解了这个问题，该机制捕捉图像块之间的依赖关系，而不管它们的空间距离如何。首个基于ViT的UNet（ViT-UNet）医学图像分割模型结合了ViT编码器和CNN解码器，因此被称为TransUNet[1]。然而，TransUNet由于参数量巨大（约1亿）和计算复杂度较高而受到限制。
这是因为它采用了标准自注意力，其时间和内存复杂度与输入标记大小成二次关系。利用窗口注意力[10]，它在所有图像块中对一个小窗口内的自注意力进行并行处理，SwinUNet[2]缓解了这一挑战，并且是第一个以自注意力作为主要特征提取器的纯ViT-UNet。然而，SwinUNet使用固定大小的窗口，这可能限制了其捕捉不同大小和形状目标对象的精确表示的能力。这可能限制了其在多类分割任务中的适应性和泛化能力。
此外，SwinUNet并没有表现出良好的扩展行为。换句话说，当增加模型大小时（如图1(a)所示），它们的性能并没有有效地扩展。这与使用自然图像的实证研究提供的证据形成对比，这些研究表明ViT表现出卓越的扩展行为。我们假设这一现象也可能归因于针对分割的目标对象在大小和形状上的异质性，这无法通过固定大小的窗口处理。
为了解决这些挑战，提高ViT-UNet的性能和可扩展性需要空间动态组件。尽管最近的工作已经探索将动态（例如，多尺度[11]，可变形[12; 13]）组件引入到医学成像的ViT-UNet中，但它们主要关注采用现有的可变形注意力。然而，它们没有研究将动态组件引入到ViT-UNet的其他部分（即，图像块嵌入和位置编码）。我们认为，增强这些其他部分的动态元素同样重要。关于AgileFormer如何从根本上与上述三篇工作不同，我们将在附录A中为读者提供深入讨论。
因此，我们开发了一种新颖的空间敏捷纯ViT-UNet，以捕捉医学图像分割任务中多样的目标对象。首先，我们用一种新颖的可变形图像块嵌入替换了ViT-UNet中的标准刚性正方形图像块嵌入。其次，我们采用空间动态自注意力[14]作为捕捉空间变化特征的建筑块。第三，我们提出了一种新颖的多尺度可变形位置编码，以在自注意力中建模不规则采样的网格。我们将这些动态组件集成到一个名为AgileFormer的新颖ViT-UNet架构中。广泛的实验证明了所提出方法在三个医学图像分割任务中的有效性。AgileFormer优于最近的医学图像分割领域的最先进的UNet模型，并展示了卓越的可扩展性。
在本节中，我们提供了一个从标准SwinUNet到所提AgileFormer的路线图（见图1(b)）。为了使本文内容自包含，我们首先简要介绍ViT-UNets中的基本要素。一个ViT-UNet（例如，SwinUNet）是一个带有跳跃连接的U形编码器-解码器网络，其中编码器和解码器的主要特征提取是通过自注意力机制实现的。标准ViT-UNet的构建块由三个基本组成部分组成：Patch嵌入、自注意力和位置编码。Patch嵌入将图像Patch投影到特征嵌入中。更近期的ViTs甚至将下采样（即在SwinUNet中的Patch合并）视为Patch嵌入的一部分。在本文中，我们将遵循相同的惯例。
自注意力，用于捕获图像Patch之间的依赖关系，用于主要特征提取。为了计算的可行性和局部性需求，最近的ViT-UNets使用基于窗口的自注意力机制。与卷积不同，自注意力丢弃了图像Patch之间的空间相关性，这阻碍了纯ViT在分割任务中的定位能力。位置编码用于解决这一限制。我们想指出，也有混合模型，其中编码器/解码器中的主要特征提取是通过CNN和ViT实现的。然而，本文关注于开发一个纯ViT-UNet。
本节的其余部分组织如下，介绍图2中提出的AgileFormer的三个主要组成部分：可变形Patch嵌入（第2.1节）、空间动态多头注意力（第2.2节）和可变形位置编码（第2.3节）。
可变形Patch嵌入是一种最近提出的技术，用于改进深度学习模型中视觉数据的表示。它特别适用于需要高分辨率输入的任务，例如语义分割和目标检测。
ViT-UNet首先将图像块转换为 tokens。这个过程通常包括将图像分割成一系列不重叠的 ()（例如，在SwinUNet中为 ()）固定大小的块。随后，这些块中的每一个都会被投影成一个1D特征向量。执行刚性（正方形）块嵌入的主要原因是其简单性，因为它可以很容易地通过标准卷积（卷积核大小 ；步长 ；膨胀系数 ）来实现。
然而，我们认为这种刚性块嵌入对于分割任务并不是最优的选择，主要有两个原因。首先，分割需要更精确的像素级定位。然而，刚性块嵌入只能提供块级定位。其次，在大多数医学图像分割任务中（例如，多器官分割任务），目标对象的形状和大小差异显著。
为了解决刚性Patch嵌入的限制，我们提出了一个可变形Patch嵌入（见图2），通过利用可变形卷积。可变形卷积定义为
其中 是一个具有 个位置的统一网格的 维特征图 ( 为2D； 为3D)。 是在网格 上操作的卷积核，定义了 的 -最近（即核大小）邻近位置。 是从不规则网格采样的偏移量。这些偏移量通过一个卷积层 学习得到。 是一个采样函数，它通过双线性/三线性插值来采样 中的位置 ，因为偏移量 通常为分数。
首次Patch嵌入。 我们用两个连续的可变形卷积层（；；）替换了标准SwinUNet中的单层刚性卷积Patch嵌入。这样做的原因是，两个连续的、重叠的可变形Patch嵌入可以提取更好的局部表示，这弥补了自注意力中的局部性不足。
下采样Patch嵌入。 我们还用单层卷积下采样（；；）替换了标准SwinUNet中的Patch合并，用于下采样。我们使用重叠核以更好地保留局部模式[22, 15, 20]，这也与重叠的可变形Patch嵌入相一致。
自我关注是ViT UNet的基本构建块。与卷积不同，自我关注不强制任何空间感应偏差，而是纯粹通过依赖标记之间的依赖关系（主要是相似性）来做决策，同时缺乏捕捉适应多类分割的空间特征的能力。因此，我们提出使用空间动态自我关注模块作为ViT-UNet的构建块。这个模块受到[15]的启发，它包括可变形多头注意力（DMSA）[15]和邻域多头注意力（NMSA）[23]。Transformer块是通过交替这两种注意力机制构建的（见图2）。我们还将在第三编码器块中分配更多的计算，其阶段比例为，而不是SwinUNet中的。这是因为编码器的第三层通常比其他层捕捉到更好的特征表示[24]。
第个头部的可变形多头（即个头部）注意力[15]被表述为
其中：
这里，我们重用了方程（1）中的表示法，其中是均匀的点网格，是为第个头部生成的偏移量，是一个插值函数。和是可训练的参数，和分别是DMSA中键和值的线性投影的隐藏维度。DMSA中的偏移量也是通过将查询通过一个卷积层生成的。方程（3）中表示的由此产生的非规则采样特征图记为。与可变形卷积类似，非规则采样的特征图随后通过使用非规则采样的键和值应用于自注意力（见方程（2））。
与标准自注意力机制相比，后者计算特征图 上给定位置 的每个元素与其它所有元素的相似性，而邻域注意力构造[23]仅利用位置 周围的 个最近邻的信息。我们在方程(2)中重新使用符号 ，因为NMSA的操作类似于卷积。因此，邻域注意力将标准自注意力的二次计算复杂度降低到大约与 的空间维度成线性关系，因为 通常较小（例如， ()）。
此外，这重新将局部操作引入自注意力机制，允许平移等方差，从而增强更好地保留局部信息的能力。按照方程(2)中的记法，位置 的邻域多头注意力计算如下：
其中：
这里， 表示给定位置 的第 个邻域位置。值得注意的是，所得到的注意力的维度是 ，其中 ，而不是标准自注意力或窗口注意力中的 。
在之前的ViT UNets中，位置编码（PE）的设计几乎没有被探索。大多数ViT-UNets要么忽略位置编码[1]，要么继承自其祖先模型[25; 2; 16]。具体来说，[25]使用了绝对位置编码（APE），为每个标记分配一个绝对值。而其他研究[2; 16; 20]使用了相对位置编码（RPE）[26]来编码标记之间的相对位置。然而，这些设计都是针对1D信号而忽略空间相关性。因此，它们并不很好地适应于具有空间相关性的2D/3D信号的建模。最近，条件位置编码（CPE）[27]被设计用于视觉任务，同时具有APE和RPE的能力。
更重要的是，标准的APE、RPE和CPE并不直接适用于不规则采样的网格，因为它们都是在刚性网格中公式化的。与本文提出的非规则采样DMSA相一致，我们提出了一个多尺度可变形位置编码（MS-DePE），旨在跨多个尺度编码非规则采样的位置信息（见图3）。
所提出的MS-DePE以条件位置编码的形式如下公式化：
其中 是输入特征图， 表示由 参数化的可训练位置编码层。 实现为具有不同核大小（即 和 ()）的多尺度可变形深度卷积层。为此，我们在将特征图 传递给 之前，首先恢复其空间分辨率。应用 之后，我们将特征图重新展平回其原始形状（详细信息见图3）。
类似于[1; 2; 11]中的方法，我们通过改变嵌入维度（）和注意力头数（）来开发AgileFormer的两个变体，但保持了主要结构不变（即编码器/解码器中的Transformer块数量）：AgileFormer-T（iny，小型）（；）和AgileFormer-B（ase，基础）（；）。我们还将在[7; 20]中描述的深度监督（DS）纳入到我们提出的方法中。
数据集。 与先前的工作[1; 2; 20]一致，我们在三个公开可用的医学图像分割数据集上验证了所提出方法的有效性：Synapse数据集[28]，自动化心脏诊断挑战（ACDC）数据集[29]，以及来自Decathlon挑战的大脑肿瘤分割数据集[30]。Synapse多器官数据集包括30个3D腹部CT扫描以及8个器官（即主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏、胃）的相应分割掩膜。ACDC数据集由100个3D心脏MRI扫描组成，每个扫描都有一个分割掩膜，标注了右心室（RV）、心肌（Myo）和左心室（LV）。Decathlon大脑肿瘤数据集包含了484个3D多模态大脑肿瘤MRI扫描以及分割掩膜，描绘了增强肿瘤、非增强肿瘤和水肿。为了与之前的工作[18; 20]保持一致，我们报告了整个肿瘤（WT）、增强肿瘤（ET）和肿瘤核心（TC）的结果。
实验设置与评估指标。 对于Synapse和ACDC数据集，我们遵循了[1; 2]中概述的实验协议，包括训练/测试集、输入图像大小（即）、数据增强、模型选择和评估。为了与先前文献进行公平比较，Synapse和ACDC数据集的实验在2D条件下进行。对于3D体积大脑肿瘤分割，我们遵循[18]中指定的实验协议，以便与先前文献进行公平比较。输入图像设置为。我们使用骰子相似系数（DSC）来评估分割性能。通过配对t检验估计了最佳模型与次佳模型平均DSC之间的统计显著性。我们强调，在Synapse、ACDC和Decathlon数据集上应用的所有竞争方法都遵循了相同的实验协议。
实施细节。 所有模型都使用了dice损失和交叉熵损失[2, 20]的组合进行训练，采用AdamW[36]优化器，并配合余弦学习率衰减和20个周期的热身。我们使用ImageNet预训练权重初始化模型参数。在2D实验中，使用了24的批处理大小，初始学习率设置为，模型训练了400个周期。对于3D实验，批处理大小和学习率分别设置为4和，模型训练了200个周期。我们想指出的是，我们使用nnFormer框架重新实现了所有3D基线模型（具体3D实验的细节请见附录B）。所有实验都使用PyTorch实现，并在配备32GB内存的Nvidia V100 GPU上执行。
我们提出的AgileFormer在所有三个数据集上超越了最近的现有技术水平（SOTA）方法（表1、2和3）。
具体来说，对于Synapse多器官、ACDC心脏和脑肿瘤分割任务，所提方法比第二名模型的平均DSC（%）分别高出0.84、0.23和0.80。这种改进被证明是统计学上显著的。具体而言，在Synapse数据集中，所提方法实现了85.74的平均DSC，超过了其他基线方法，能够分割出8个器官中的6个（表1）。对于剩下的两个器官，所提方法是表现第二佳的方法。在ACDC数据集中，所提方法实现了92.55的平均DSC，超过了所有基线方法，在分割RV、Myo和LV方面表现最佳（见表2）。
在3D体积脑肿瘤分割任务中，所提方法实现了85.7的平均DSC，超越了最近的SOTA ViT-UNet方法（表3）。此外，所提方法在这三个公共数据集上展示了最佳的定性分割质量（见图4）。
我们对模型设计变体进行了消融研究，包括Patch嵌入以及不同的空间动态注意力和位置编码选择（见表4）。首先，我们展示了将Patch合并替换为卷积下采样可以使平均DSC获得0.34%的小幅性能提升。然后，添加可变形嵌入可以使DSC性能提升1.9%。其次，我们通过交替窗口/可变形/邻域注意力消融了不同的空间动态自注意力选择。我们移除了位置编码以消除其目前的影响。将窗口注意力（WMSA）与DMSA和NSMA交替使用，分别带来了0.7%（DMSA）和0.4%（NSMA）的性能提升。
交替NMSA和DMSA则带来了1.4%的性能提升。第三，我们恢复了位置编码，展示了并非所有的位置编码都能带来性能提升，例如，添加绝对位置编码（APE）和循环位置编码（CPE）甚至导致了性能下降。这是因为这两种位置编码并未针对DMSA引入的不规则采样网格设计。相比之下，我们提出的MS-DePE提升了0.8%的性能。值得注意的是，与SwinUNet相比，我们提出的AgileFormer并没有带来巨大的计算负担，参数数量仅增加了1.1%，浮点运算量提升了15%。
我们比较了所提出方法与其他ViT-UNets的模型缩放行为。如图1(a)所示，在Synapse多器官分割任务中，随着模型尺寸的增加，所提出的方法表现出性能提升，该任务中不同的器官在形状和大小上表现出很大的变异性。具体来说，从AgileFormer-T到AgileFormer-B，性能在DSC(%)上提高了2.15，参数从大约28.85M增加到123.47M，增长了约四倍。我们还观察到，具有动态模块的方法比没有动态模块的方法缩放效果更好，除了CoTr和SDAUT。我们假设这可能是因为CoTr的主要骨干网络仍然是CNN。此外，CoTr和SDAUT中的可变形模块仅放置在瓶颈部分。
在本论文中，我们提出了AgileFormer，它系统地引入了空间动态组件（即可变形Patch嵌入、空间动态自注意力以及多尺度可变形位置编码）到UNet中，以捕获医学图像分割中不同目标对象的空间动态信息。广泛的实验证明了所提出方法在多种医学图像分割任务中的有效性。我们期望将空间动态组件系统地引入到ViT-UNet中的这一理念，能够指导未来在处理具有多目标的医学图像分割中如何提取空间动态表示的设计。
附录
附录A 与其他动态ViT-UNet模型的比较
CoTr。 与其祖先模型[37]类似，CoTr是一个混合模型，其主要骨架（编码器/解码器）仍然是卷积神经网络（CNN）。在[37]中概述的多尺度可变形注意力作为瓶颈。因此，它只能引入有限的可变形性来捕获空间上动态的特征表示，因为主要特征提取器没有空间上动态的组件。
此外，由于其主要骨架仍然是CNN，CoTr没有充分利用ViT来捕获长距离依赖。我们假设这些因素可能限制其可扩展性（见附录C中的证据）。相反，我们提出的方法使用ViT作为主要骨架（编码器/解码器），其中邻域注意力和可变形注意力作为主要构建块，这既保留了ViT的优势，同时具有捕获空间局部化和动态特征的额外优势。
SDAUT。 尽管SDAUT也使用与提出方法相同的可变形注意力机制[15]，但它主要是为图像恢复任务而提出的，而不是分割任务。SDAUT的最优架构结合了窗口注意力和可变形注意力，其中可变形注意力仅应用于第三个编码器和解码器。相反，提出的方法在每个层次（编码器/解码器）上都应用了可变形注意力，这有助于提取空间上动态的特征。此外，我们进一步引入了可变形Patch嵌入，从根本上改变了刚性的Patch嵌入。我们还引入了多尺度可变形位置编码来处理由可变形注意力引入的不规则采样网格。我们发现，这两个新颖组件与空间上动态注意力一样重要，以增强分割性能。
MERIT。 MERIT主要利用多分辨率信息捕获空间上动态（主要是大小变化）的特征，通过结合从不同大小输入提取的特征。相反，提出的方法处理大小和形状的变化，通过空间上动态注意力。而我们还使用多尺度可变形位置编码来进一步捕获多尺度信息。
附录B
我们发现nnFormer [20]中的比较是不公平的。nnFormer和SwinUNETR使用尺寸为的图像作为输入；而我们实现的UNETR的结果是在大小的输入下取得的。因此，所有方法的输入图像尺寸被设置为。我们还发现，nnFormer与其他方法（即TransUNet、TransBTS、UNETR、SwinUNETR和CoTr）之间存在巨大的性能差距，这与其他分割任务上的表现不一致。我们推测这种性能差距可能是由不同的预处理、训练和评估程序造成的。
因此，在我们的实验中，我们使用nnFormer的代码库重新训练了所有方法，以便标准化比较。对于提出的方法，批处理大小设置为4，初始学习率设置为。对于基线方法，我们遵循了它们原始论文中的默认超参数设置。我们发现，即使是较小的输入尺寸，我们的复现也使得TransUNet、TransBTS、UNETR、SwinUNETR和CoTr的性能比nnFormer中报告的结果要高。
点击上方卡片，关注 「AI视界引擎」 公众号