[论文解读] Beyond Self-Attention: Deformable Large Kernel Attention for Medical Image Segmentation

Reza Azad, Leon Niggemeier+6 more|arXiv (Cornell University)|Aug 31, 2023

Advanced Neural Network Applications被引用 7

一句话总结

本文提出可变形大感受感受野注意力（D-LKA），一种轻量化、高效的注意力机制，用于在医学图像分割任务中替代视觉变换器（Vision Transformers）中的自注意力机制。通过使用具有可变形采样网格的大感受野卷积核，D-LKA 以线性复杂度捕捉全局上下文信息，从而实现优越的 3D 体积特征学习。所提出的 D-LKA Net 在 Synapse、NIH 胰腺和皮肤病变数据集上实现了最先进性能，且参数量和 FLOPs 更少。

ABSTRACT

Medical image segmentation has seen significant improvements with transformer models, which excel in grasping far-reaching contexts and global contextual information. However, the increasing computational demands of these models, proportional to the squared token count, limit their depth and resolution capabilities. Most current methods process D volumetric image data slice-by-slice (called pseudo 3D), missing crucial inter-slice information and thus reducing the model's overall performance. To address these challenges, we introduce the concept of extbf{Deformable Large Kernel Attention (D-LKA Attention)}, a streamlined attention mechanism employing large convolution kernels to fully appreciate volumetric context. This mechanism operates within a receptive field akin to self-attention while sidestepping the computational overhead. Additionally, our proposed attention mechanism benefits from deformable convolutions to flexibly warp the sampling grid, enabling the model to adapt appropriately to diverse data patterns. We designed both 2D and 3D adaptations of the D-LKA Attention, with the latter excelling in cross-depth data understanding. Together, these components shape our novel hierarchical Vision Transformer architecture, the extit{D-LKA Net}. Evaluations of our model against leading methods on popular medical segmentation datasets (Synapse, NIH Pancreas, and Skin lesion) demonstrate its superior performance. Our code implementation is publicly available at the: https://github.com/mindflow-institue/deformableLKA

研究动机与目标

• 解决视觉变换器在医学图像分割任务中自注意力机制带来的高计算成本问题。
• 通过实现跨切片特征融合，克服 3D 医学图像分析中 2D 切片处理方式的局限性。
• 在保留局部细节的同时提升全局上下文建模能力，突破卷积神经网络（CNN）的局部感受野限制。
• 开发一种轻量化、参数高效的注意力机制，在医学分割基准测试中保持高性能。

提出的方法

• 提出可变形大感受野注意力（D-LKA），利用大感受野卷积核以线性复杂度模拟自注意力的长程建模能力。
• 引入可学习、可微分的变形网格，动态扭曲采样位置，以适应复杂的解剖结构。
• 设计 D-LKA 的 3D 变体，用于捕捉体积医学图像中的切片间依赖关系，增强 3D 上下文理解能力。
• 构建一种分层视觉变换器架构（D-LKA Net），结合 D-LKA 模块与跳跃连接以实现特征精细化。
• 采用 CNN 与 Transformer 组件相结合的混合编码器-解码器结构，平衡局部细节与全局上下文信息。
• 通过跳跃连接实现多尺度特征聚合策略，以保留空间分辨率并提升分割精度。

Figure 1 : Proposed network architecture of the 3D D-LKA model on the left and the 2D D-LKA model on the right.

实验结果

研究问题

• RQ1 大感受野卷积机制是否能有效替代视觉变换器中的自注意力机制，同时降低计算成本？
• RQ2 在 3D 医学图像中，大感受野注意力的可变形采样如何提升对不规则形状器官的特征表示能力？
• RQ3 与 2D 伪 3D 方法相比，3D D-LKA 在跨切片特征学习方面提升了多少？
• RQ4 跳跃连接与可变形卷积在 D-LKA Net 架构中的最终分割性能中分别起到何种贡献？
• RQ5 轻量化注意力机制是否能在医学分割基准测试中，同时实现更高的准确率与模型效率，超越现有最先进模型？

主要发现

• 在 Synapse 数据集上，D-LKA Net 的 Dice 系数（DSC）达到 87.49%，优于 UNETR++（80.59%）和 UNETR（77.42%），且参数量减少 37.73M。
• 在 NIH 胰腺数据集上，D-LKA Net 的 DSC 达到 81.22%，较 UNETR++ 提升 0.63%，HD95 指标降低 1.04 个百分点。
• 采用可变形卷积的 3D LKA 变体相比非可变形版本，DSC 提升 0.63%，证明了自适应采样机制的优势。
• 消融实验表明，跳跃连接显著提升性能，仅使用最高层跳跃连接即可使 DSC 提升 0.42%。
• 在胰腺数据集上，该模型参数量（62.07M）为最先进方法中最低，体现出极高的效率。
• 定性结果表明，与 UNETR 和 UNETR++ 相比，D-LKA Net 更好地捕捉了胰腺的不规则形状，显著减少了分割伪影。

Figure 2 : Architecture of the deformable LKA module.