D-LKA Attention | D-LKA注意力同时拥有SA注意力和大卷积核的能力

点击下方卡片，关注「集智书童」公众号

医学图像分割在Transformer模型的应用下取得了显著的改进，这些模型在捕捉广泛的上下文和全局背景信息方面表现出色。然而，这些模型的计算需求不断增加，与Token数量的平方成正比，限制了它们的深度和分辨率能力。大多数当前的方法逐片处理D volumetric图像数据（称为伪3D），错过了重要的切片间信息，从而降低了模型的整体性能。

为了解决这些挑战，作者引入了Deformable Large Kernel Attention（D-LKA Attention）的概念，这是一种简化的注意力机制，采用大卷积核来充分理解 volumetric 上下文。该机制在与自注意力类似的感受野内操作，同时避免了计算开销。此外，作者提出的注意力机制还受益于可变形卷积，可以灵活地扭曲采样网格，使模型能够适应不同的数据模式。作者设计了D-LKA注意力的2D和3D适应版本，其中后者在跨深度数据理解方面表现出色。这些组件共同构成了作者的新型分层Vision Transformer架构，即D-LKA Net。

作者在流行的医学分割数据集（Synapse、NIH胰腺和皮肤病变）上对作者的模型进行评估，结果表明它具有卓越的性能。

代码：https://github.com/mindflow-institue/deformableLKA

1、简介

医学图像分割在计算机辅助诊断中起着至关重要的作用，帮助医疗专业人员分析复杂的医学图像。这个过程不仅减轻了手动任务的繁重工作和对医疗专业知识的依赖，还能实现更快速和更准确的诊断。分割的自动化为更快速和更准确的诊断结果提供了潜力，有助于制定适当的治疗策略，并支持执行图像引导的手术程序。因此，创建快速和精确的分割算法的必要性成为推动这项研究的动力。

自2010年代中期以来，卷积神经网络（CNNs）已成为许多计算机视觉应用的首选技术。它们能够从原始数据中自动提取复杂的特征表示，无需手动进行特征工程，这引起了医学图像分析社区的极大兴趣。许多成功的CNN架构，如U-Net、全卷积网络、DeepLab或SegCaps（分割胶囊），已经被开发出来。这些架构在语义分割任务中取得了巨大成功，先前的最新方法已经被超越。

在计算机视觉研究中，不同尺度下的目标识别是一个关键问题。在CNN中，可检测目标的大小与相应网络层的感受野尺寸密切相关。如果一个目标扩展到超出这个感受野的边界，这可能会导致欠分割结果。相反，与目标实际大小相比使用过大的感受野可能会限制识别，因为背景信息可能会对预测产生不必要的影响。

解决这个问题的一个有希望的方法涉及在并行使用具有不同尺寸的多个Kernel，类似于Inception块的机制。然而，由于参数和计算要求的指数增长，将Kernel大小增加以容纳更大的目标在实践中受到限制。因此，出现了各种策略，包括金字塔池化技术和不同尺度的扩张卷积，以捕获多尺度的上下文信息。

另一个直观的概念涉及将多尺度图像金字塔或它们的相关特征表示直接纳入网络架构。然而，这种方法存在挑战，特别是在管理训练和推理时间方面的可行性方面存在挑战。在这个背景下，使用编码器-解码器网络，如U-Net，已被证明是有利的。这样的网络在较浅的层中编码外观和位置，而在更深的层中，通过神经元的更广泛的感受野捕获更高的语义信息和上下文信息。

一些方法将来自不同层的特征组合在一起，或者预测来自不同尺寸的层的特征以使用多尺度的信息。此外，出现了从不同尺度的层中预测特征的方法，有效地实现了跨多个尺度的见解整合。然而，大多数编码器-解码器结构面临一个挑战：它们经常无法在不同尺度之间保持一致的特征，并主要使用最后一个解码器层生成分割结果。

语义分割是一项任务，涉及根据预定义的标签集为图像中的每个像素预测语义类别。这项任务要求提取高级特征同时保留初始的空间分辨率。CNNs非常适合捕获局部细节和低级信息，尽管以忽略全局上下文为代价。视觉Transformer（ViT）架构已经成为解决处理全局信息的视觉任务的关键，包括语义分割，取得了显著的成功。

ViT的基础是注意力机制，它有助于在整个输入序列上聚合信息。这种能力使网络能够合并远程的上下文提示，超越了CNN的有限感受野尺寸。然而，这种策略通常会限制ViT有效建模局部信息的能力。这种局限可能会妨碍它们检测局部纹理的能力，这对于各种诊断和预测任务至关重要。这种缺乏局部表示可以归因于ViT模型处理图像的特定方式。

ViT模型将图像分成一系列Patch，并使用自注意力机制来模拟它们之间的依赖关系。这种方法可能不如CNN模型中的卷积操作对感受野内提取局部特征有效。ViT和CNN模型之间的这种图像处理方法的差异可能解释了CNN模型在局部特征提取方面表现出色的原因。

近年来，已经开发出创新性方法来解决Transformer模型内部局部纹理不足的问题。其中一种方法是通过互补方法将CNN和ViT特征结合起来，以结合它们的优势并减轻局部表示的不足。TransUNet是这种方法的早期示例，它在CNN的瓶颈中集成了Transformer层，以模拟局部和全局依赖关系。HiFormer提出了一种解决方案，将Swin Transformer模块和基于CNN的编码器结合起来，生成两个多尺度特征表示，通过Double-Level Fusion模块集成。UNETR使用基于Transformer的编码器和CNN解码器进行3D医学图像分割。CoTr和TransBTS通过Transformer在低分辨率阶段增强分割性能，将CNN编码器和解码器连接在一起。

增强局部特征表示的另一种策略是重新设计纯Transformer模型内部的自注意力机制。在这方面，Swin-Unet在U形结构中集成了一个具有线性计算复杂性的Swin Transformer块作为多尺度 Backbone 。MISSFormer采用高效Transformer来解决视觉Transformer中的参数问题，通过在输入块上进行不可逆的降采样操作。D-Former引入了一个纯Transformer的管道，具有双重注意模块，以分段的方式捕获细粒度的局部注意和与多元单元的交互。然而，仍然存在一些特定的限制，包括计算效率低下，如TransUNet模型所示，对CNN Backbone 的严重依赖，如HiFormer所观察到的，以及对多尺度信息的忽略。

此外，目前的分割架构通常采用逐层处理3D输入 volumetric 的方法，无意中忽视了相邻切片之间的潜在相关性。这一疏忽限制了对 volumetric 信息的全面利用，因此损害了定位精度和上下文集成。此外，必须认识到，医学领域的病变通常在形状上发生变形。因此，用于医学图像分析的任何学习算法都必须具备捕捉和理解这些变形的能力。与此同时，该算法应保持计算效率，以便处理3D volumetric数据。

作者的贡献：

为了解决上述提到的挑战，作者提出了一个解决方案，即可变形大卷积核注意力模块（Deformable LKA module），它是作者网络设计的基本构建模块。这个模块明确设计成在有效处理上下文信息的同时保留局部描述符。作者的架构在这两个方面的平衡增强了实现精确语义分割的能力。

值得注意的是，作者的模型引入了一种基于数据的感受野的动态适应，不同于传统卷积操作中的固定滤波器Mask。这种自适应方法使作者能够克服与静态方法相关的固有限制。这种创新方法还扩展到了D-LKA Net架构的2D和3D版本的开发。

在3D模型的情况下，D-LKA机制被量身定制以适应3D环境，从而实现在不同 volumetric 切片之间无缝信息交互。最后，作者的贡献通过其计算效率得到进一步强调。作者通过仅依靠D-LKA概念的设计来实现这一点，在各种分割基准上取得了显著的性能，确立了作者的方法作为一种新的SOTA方法。

2. 方法

在本节中，作者首先概述方法论。首先，作者回顾了由Guo等人引入的大卷积核注意力（Large Kernel Attention，LKA）的概念。然后，作者介绍了作者对可变形LKA模块的创新探索。在此基础上，作者介绍了用于分割任务的2D和3D网络架构。

2.1. 大卷积核注意力

大卷积核提供了与自注意力机制类似的感受野。可以通过使用深度卷积、深度可扩展卷积和卷积来构建大卷积核，从而减少了参数和计算量。构建输入维度为和通道数的卷积核的深度卷积和深度可扩展卷积的卷积核大小的方程如下：

具有卷积核大小和膨胀率。参数数量和浮点运算（FLOPs）的计算如下：

FLOPs的数量与输入图像的大小成线性增长。参数的数量随通道数和卷积核大小的增加而呈二次增长。然而，由于它们通常都很小，因此它们不是限制因素。

为了最小化对于固定卷积核大小K的参数数量，可以将方程3对于膨胀率的导数设定为零：

例如，当卷积核大小为时，结果是。将这些公式扩展到3D情况是直接的。对于大小为和通道数C的输入，3D情况下参数数量和FLOPs 的方程如下：

具有卷积核大小和膨胀。

2.2. 可变形大卷积核注意力

利用大卷积核进行医学图像分割的概念通过引入可变形卷积得以扩展。可变形卷积可以通过整数偏移自由调整采样网格以进行自由变形。额外的卷积层从特征图中学习出变形，从而创建一个偏移场。基于特征本身学习变形会导致自适应卷积核。这种灵活的卷积核形状可以提高病变或器官变形的表示，从而增强了目标边界的定义。

负责计算偏移的卷积层遵循其相应卷积层的卷积核大小和膨胀。双线性插值用于计算不在图像网格上的偏移的像素值。如图2所示，D-LKA模块可以表示为：

其中输入特征由表示，。表示为注意力图，其中每个值表示相应特征的相对重要性。运算符 表示逐元素乘法运算。值得注意的是，LKA不同于传统的注意力方法，它不需要额外的规范化函数，如或。这些规范化函数往往忽视高频信息，从而降低了基于自注意力的方法的性能。

在该方法的2D版本中，卷积层被可变形卷积所替代，因为可变形卷积能够改善对具有不规则形状和大小的目标的捕捉能力。这些目标在医学图像数据中常常出现，因此这种增强尤为重要。

然而，将可变形LKA的概念扩展到3D领域会带来一定的挑战。主要的约束来自于需要用于生成偏移的额外卷积层。与2D情况不同，由于输入和输出通道的性质，这一层无法以深度可分的方式执行。在3D环境中，输入通道对应于特征，而输出通道扩展到，其中是卷积核的大小。大卷积核的复杂性导致沿第3D的通道数扩展，导致参数和FLOPs大幅增加。因此，针对3D情况采用了另一种替代方法。在现有的LKA框架中，深度卷积之后引入了一个单独的可变形卷积层。这种战略性的设计调整旨在减轻扩展到3D领域所带来的挑战。

2.3. 2D D-LKA网络

2D网络的架构如图1所示。第一变种使用MaxViT作为编码器组件，用于高效特征提取，而第二变种则结合可变形LKA层进行更精细、卓越的分割。

在更正式的描述中，编码器生成4个分层输出表示。首先，卷积干扰将输入图像的维度减小到。随后，通过4个MaxViT块的4个阶段进行特征提取，每个阶段后跟随降采样层。随着过程进展到解码器，实施了4个阶段的D-LKA层，每个阶段包含2个D-LKA块。然后，应用Patch扩展层以实现分辨率上采样，同时减小通道维度。最后，线性层负责生成最终的输出。

2D D-LKA块的结构包括LayerNorm、可变形LKA和多层感知器（MLP）。积分残差连接确保了有效的特征传播，即使在更深层也是如此。这个安排可以用数学方式表示为：

其中输入特征，层归一化LN，可变形LKA注意力，深度卷积，线性层和GeLU激活函数。

2.4. 3D D-LKA网络

3D网络架构如图1所示，采用编码器-解码器设计进行分层结构化。首先，一个Patch嵌入层将输入图像的维度从（）减小到（）。在编码器中，采用了3个D-LKA阶段的序列，每个阶段包含3个D-LKA块。在每个阶段之后，通过降采样步骤将空间分辨率减半，同时将通道维度加倍。中央瓶颈包括另一组2个D-LKA块。解码器结构与编码器相对称。

为了将特征分辨率加倍，同时减少通道数，使用转置卷积。每个解码器阶段都使用3个D-LKA块来促进远距离特征依赖性。最终的分割输出由一个卷积层产生，后面跟随一个卷积层以匹配特定类别的通道要求。

为了建立输入图像和分割输出之间的直接连接，使用卷积形成了一个跳跃连接。额外的跳跃连接根据简单的加法对来自其他阶段的特征进行融合。最终的分割图是通过和卷积层的组合产生的。

3D D-LKA块包括层归一化，后跟D-LKA注意力，应用了残差连接的部分。随后的部分采用了一个卷积层，后面跟随一个卷积层，两者都伴随着残差连接。这个整个过程可以总结如下：

带有输入特征 、层归一化 、可变形 LKA 、卷积层 和输出特征 的公式。是指一个前馈网络，包括2个卷积层和激活函数。

2.5. D-LKA的计算复杂度

表7显示了普通卷积和构建卷积的参数数量比较。尽管标准卷积的参数数量在通道数较多时急剧增加，但分解卷积的参数总体较低，并且增长速度不那么快。

与分解卷积相比，可变形分解卷积增加了大量参数，但仍然明显小于标准卷积。可变形卷积的主要参数是由偏移网络创建的。在这里，作者假设可变形深度卷积的Kernel大小为（5,5），可变形深度空洞卷积的Kernel大小为（7,7）。这导致了21×21大小的大Kernel的最佳参数数量。更高效地生成偏移量的方法将大大减少参数数量。

值得注意的是，引入可变形LKA确实会增加模型的参数数量和每秒的浮点运算次数（FLOPS）。然而，重要的是强调，这增加的计算负载不会影响作者模型的整体推理速度。

相反，对于Batch-size > 1，作者甚至观察到推理时间的减少，如图7所示。例如，基于作者的广泛实验，作者观察到对于Batch-size为16，具有可变形卷积和没有可变形卷积的推理时间分别为8.01毫秒和17.38毫秒。作者认为这是由于在2D中对可变形卷积的高效实现所致。为了测量时间，使用了大小为（）的随机输入。在GPU热身周期50次迭代之后，网络被推断了1000次。测量是在NVIDIA RTX 3090 GPU上进行的。

2.6. 性能与效率

为了充分利用性能与参数之间的权衡关系，作者在图8中可视化了在Synapse 2D数据集上报告的DSC和HD性能以及基于参数数量的内存消耗。D-LKA Net引入了相当多的参数，约为101M。这比性能第二好的方法ScaleFormer使用的111.6M参数要少。

与更轻量级的DAEFormer模型相比，作者实现了更好的性能，这证明了参数增加的合理性。大多数参数来自于MaxViT编码器；因此，将编码器替换为更高效的编码器可以减少模型参数。值得注意的是，在此可视化中，作者最初将HD和内存值都归一化到[0, 100]范围内。随后，作者将它们从100缩小，以增强更高值的表示。

3. 实验

3.1. 2D结果

Synapse数据集：

在表1中，作者对其他SOTA技术取得的领先性能与作者提出的方法进行了全面比较。

从Dice相似性系数（DSC）方面的结果显示，D-LKA Net表现优于先前建立的SOTA方法。具体来说，与ScaleFormer相比，作者的方法提高了1.41%，与DAEFormer相比，提高了1.64%，与其他方法相比，提高幅度更大。

值得注意的是，在特定解剖区域的分割方面取得了显著的改进，如右肾脏、左肾脏、胃和胰腺。特别值得注意的是，胰腺的分割结果得到了显著的改进，相对于第2的方法，提高了2.04%。考虑到对于小腹部器官（如胆囊或胰腺）的分割一直是现有SOTA方法的挑战，这一显著的性能改进代表了在实现更准确的分割结果方面迈出的重要一步。图3展示了不同方法的定性比较。

与DAEFormer相比，作者的方法对于胃的误分类较少。尽管Unet和Swin-UNet有时会将远处的组织分类为肝脏、胆囊或胃的一部分，但作者的方法减少了误分类，更好地表示了器官的形状。

皮肤病变分割结果：

与领先方法相比，作者的D-LKA Net在皮肤病变分割基准（包括ISIC 2017、ISIC 2018和PH2）上的比较结果详见表2。值得注意的是，无论在不同的数据集上，作者的D-LKA Net始终在各种评估指标上优于竞争对手。这一持续的优势跨越不同数据集，突显了D-LKA Net的强大泛化能力。

图4呈现了结果的定性比较。与基线方法相比，D-LKA Net更好地跟随病变的复杂轮廓。与Swin-UNet和HiFormer-B相比，后者倾向于过度或不足地分割某些区域，作者的方法实现了更精确的分割。

3.2. 3D结果

Synapse数据集：

作者在Synapse数据集上将作者的3D方法与先前的SOTA方法进行了比较。结果见表3。作者在DSC方面比以前的SOTA方法UNETR++提高了0.27%。与nnFormer相比，作者实现了0.92%的改进。

对于HD95指标，D-LKA Net达到了第二佳结果。与UNETR++相比，脾脏、左肾和主动脉的性能略有提高。而右肾、胆囊和胰腺的分割结果显著提高。这些小器官的分割性能提高尤为重要。

在参数方面，作者的参数数量最少，仅为42.35M，同时仍然实现了出色的分割性能。FLOPs的数量为66.96 G，仅次于最少的UNETR++。与Swin-UNETR和nnFormer等SOTA方法相比，作者只需约17%和31%的计算量，同时实现了更好的性能。

图5显示了Synapse数据集上的定性结果。与nnFormer相比，作者更好地捕捉主动脉的整体情况，不会混淆其他组织作为器官。与UNETR++相比，作者更好地分割了胰腺，而UNETR++倾向于欠分割。此外，作者的方法比UNETR++更准确地分割了肝脏和胃，后者倾向于过分割这些器官。

胰腺数据集：

NIH胰腺数据集的结果见表4。作者的方法在所有4个指标上都取得了最佳性能。与最接近的竞争对手UNETR++相比，DSC提高了0.63%，Jaccard提高了0.82%，HD95减少了1.04，ASD减少了0.26。D-LKA Net的参数数量也最少，为62.07M。

图6显示了不同方法的定性比较。UNETR未能将胰腺分割为一个单独的目标。UNETR++的分割结果中有较小的伪影。作者的方法更好地跟随器官的高度不规则形状，优于其他方法。

3.3. 消融实验

稳健性

为增强作者的评估稳健性并分析统计显著性，作者在Synapse 2D版本上为每种方法进行了5次训练运行。这种做法不仅确保了更全面的评估，还使作者能够可视化性能的变化。

在作者的评估中，作者观察到主动脉、胆囊、左肾、右肾、肝脏、胰腺和胃的性能稳步提高，其中中位性能高于其他SOTA方法。只有脾脏的分割性能稍差。此外，还在胆囊、胰腺和胃方面取得了显著的性能改善。

可变形LKA的影响

作者继续进行消融研究，以确定D-LKA的有效性。为此，作者构建了一个使用3D LKA而没有可变形层的网络版本，另一个版本使用3D LKA，并在可变形层之外添加了一个额外的卷积层。该分析的结果见表5。引入额外的3D卷积层导致性能显著提高，与3D LKA基线相比，DSC提高了0.99%。

然而，这种修改还增加了网络内的参数数量。将3D卷积层替换为可变形卷积层导致性能进一步提升，DSC增加了0.63%。与之前的修改类似，这也引入了更多的参数和FLOPs到网络中。由于网络的规模仍然相对较小，这些指标的增加是可以接受的。

跳跃连接

最后，作者评估了跳跃连接对分割过程的影响。结果见表6。作者删除了所有的跳跃连接，并逐渐将它们添加到网络中，从最高级别的跳跃连接开始。结果表明，跳跃连接对于获得最佳的分割性能至关重要。此外，作者强调最高级别的跳跃连接对于实现最佳的分割结果至关重要，可以将DSC性能提高0.42%。

4. 参考

[1].Beyond Self-Attention: Deformable Large Kernel Attention for Medical Image Segmentation.

5、推荐阅读

小目标检测的福音 | 一文全览3年来Transformer是怎么在小目标领域大杀四方的？

中山大学提出ASAG | 简单有效的Query加权成就性能更强的目标检测

ExMobileViT | 优化轻量化ViT的不二选择，源于MobileViT又高于MobileViT！

扫码加入👉「集智书童」交流群

（备注：方向+学校/公司+昵称）

想要了解更多：

前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉（分割、检测）、多模态、目标跟踪、NerF」

行业技术方案👉「AI安防、AI医疗、AI自动驾驶」AI模型部署落地实战👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」

欢迎扫描上方二维码，加入「集智书童-知识星球」，日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑，期待交流！

免责声明凡本公众号注明“来源：XXX（非集智书童）”的作品，均转载自其它媒体，版权归原作者所有，如有侵权请联系我们删除，谢谢。
点击下方“阅读原文”，了解更多AI学习路上的「武功秘籍」