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, compared to the local convolutional-based design. 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 and has linear complexity with respect to the input. 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 complexity. 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 significantly reducing parameters and FLOPs by over 71%, compared to the best method in the literature. Our code and models are available at: https://tinyurl.com/2p87x5xn .

中文翻译：

---

UNETR++：深入研究高效、准确的 3D 医学图像分割


由于 transformer 模型的成功，最近的工作研究了它们在 3D 医疗分割任务中的适用性。在 transformer 模型中，与基于局部卷积的设计相比，自注意力机制是努力捕获长距离依赖关系的主要构建块之一。然而，自注意力操作具有二次复杂度，这被证明是一个计算瓶颈，尤其是在体积医学成像中，其中输入是具有大量切片的 3D。在本文中，我们提出了一种名为 UNETR++ 的 3D 医学图像分割方法，该方法既能提供高质量的分割掩码，又能在参数、计算成本和推理速度方面提供效率。我们设计的核心是引入一种新的高效配对注意力 （EPA） 块，它使用一对基于空间和通道注意力的相互依赖的分支有效地学习空间和通道方面的判别特征。我们的空间注意力公式是有效的，并且相对于输入具有线性复杂性。为了实现以空间和通道为中心的分支之间的通信，我们共享查询和关键映射函数的权重，这些函数提供了互补的好处（配对关注），同时也降低了复杂性。我们对 Synapse、BTCV、ACDC、BraTS 和 Decathlon-Lung 这五个基准进行了广泛的评估，揭示了我们在效率和准确性方面的贡献的有效性。在 Synapse 上，我们的 UNETR++ 以 87.2% 的 Dice 分数创下了新的水平，同时与文献中的最佳方法相比，参数和 FLOP 显著降低了 71% 以上。我们的代码和模型可在以下网址获得： https://tinyurl.com/2p87x5xn .