[论文解读] UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation

Abdelrahman Shaker, Muhammad Maaz+4 more|arXiv (Cornell University)|Dec 8, 2022

Advanced Neural Network Applications被引用 52

一句话总结

UNETR++ 引入高效成对注意力（EPA）模块，在分层3D分割网络中联合建模空间与通道特征，以显著降低参数量与 FLOPs 的同时实现最先进的准确性。

ABSTRACT

Owing to the success of transformer models, recent works study their applicability in 3D medical segmentation tasks. Within the transformer models, the self-attention mechanism is one of the main building blocks that strives to capture long-range dependencies. However, the self-attention operation has quadratic complexity which proves to be a computational bottleneck, especially in volumetric medical imaging, where the inputs are 3D with numerous slices. In this paper, we propose a 3D medical image segmentation approach, named UNETR++, that offers both high-quality segmentation masks as well as efficiency in terms of parameters, compute cost, and inference speed. The core of our design is the introduction of a novel efficient paired attention (EPA) block that efficiently learns spatial and channel-wise discriminative features using a pair of inter-dependent branches based on spatial and channel attention. Our spatial attention formulation is efficient having linear complexity with respect to the input sequence length. To enable communication between spatial and channel-focused branches, we share the weights of query and key mapping functions that provide a complimentary benefit (paired attention), while also reducing the overall network parameters. Our extensive evaluations on five benchmarks, Synapse, BTCV, ACDC, BRaTs, and Decathlon-Lung, reveal the effectiveness of our contributions in terms of both efficiency and accuracy. On Synapse, our UNETR++ sets a new state-of-the-art with a Dice Score of 87.2%, while being significantly efficient with a reduction of over 71% in terms of both parameters and FLOPs, compared to the best method in the literature. Code: https://github.com/Amshaker/unetr_plus_plus.

研究动机与目标

• 在3D医学影像中阐明并解决分割准确性与模型效率之间的权衡。
• 提出一个建立在 UNETR 基础上的统一混合架构，具有参数和计算效率。
• 引入高效成对注意力（EPA）模块，以捕获丰富的空间和通道依赖关系。
• 在多个基准上评估 UNETR++，以展示准确性提升和效率改进。

提出的方法

• 提出一个基于 UNETR 的具有四个编码器/解码器阶段的分层编码-解码器架构。
• 开发高效成对注意力（EPA）模块，包含两个并行注意力模块（空间和通道），共享 Q/K 权重但具有独立的 V 路径。
• 使空间注意力在较低维空间中运作，以实现以输入令牌数量为线性复杂度的计算。
• 在空间和通道分支之间共享 Q/K 权重，以减少参数并实现互补特征学习。
• 在最终体素级预测之前，通过 1x1x1 和 3x3x3 卷积融合 EPA 输出。
• 使用结合软 Dice 损失和交叉熵损失的训练方式来优化分割质量。

实验结果

研究问题

• RQ1 高效成对注意力（EPA）模块是否在降低计算复杂度的同时保持或提升分割准确性？
• RQ2 在编码器和解码器中都引入 EPA 的分层 UNETR++ 架构是否能在多种基准上超越最先进的3D医学分割方法？
• RQ3 在多个数据集（Synapse、BTCV、ACDC、BRaTs、Decathlon-Lung）上，UNETR++ 在分割准确性（DSC）和效率（参数量、FLOPs）方面的表现如何？

主要发现

• 在 Synapse 数据集上，UNETR++ 达到 87.22% 的 Dice 分数，相较基线 UNETR，在参数量（42.96M）和 FLOPs（47.98G）上有显著降低。
• 仅在编码器中集成 EPA 可获得 85.17% DSC，进一步在解码器中添加 EPA 可将 DSC 提升至 87.22%，相比基线减少约54%参数量和约37% FLOPs。
• UNETR++ 在 Synapse 上优于 nnFormer，同时参数量和 FLOPs 下降超过70%，表明在准确性与效率之间取得有利平衡。
• 在 BTCV 上，UNETR++ 的平均 DSC 为 83.28%，FLOPs 为 31.0 GFLOPs，与 nnUNet（平均 DSC 83.16%，但 358 GFLOPs）相比具有竞争力。
• ACDC 结果显示 UNETR++ 的平均 DSC 为 82.83%（nnFormer 为 92.06%，UNETR 为 86.61%），表明在保持较高效率的同时具有较强的性能。
• BRaTs 和肺部数据集相比于近期的基于 Transformer 的方法，显示出有利的分割性能与效率权衡。