[论文解读] AgileFormer: Spatially Agile Transformer UNet for Medical Image Segmentation

Peijie Qiu, Yang Jin+3 more|arXiv (Cornell University)|Mar 29, 2024

COVID-19 diagnosis using AI被引用 11

一句话总结

AgileFormer 引入可变形补丁嵌入、空间动态自注意力，以及多尺度可变形位置编码到 ViT-UNet，在三个数据集上实现先进的医学影像分割，具有良好的可扩展性。

ABSTRACT

In the past decades, deep neural networks, particularly convolutional neural networks, have achieved state-of-the-art performance in a variety of medical image segmentation tasks. Recently, the introduction of the vision transformer (ViT) has significantly altered the landscape of deep segmentation models. There has been a growing focus on ViTs, driven by their excellent performance and scalability. However, we argue that the current design of the vision transformer-based UNet (ViT-UNet) segmentation models may not effectively handle the heterogeneous appearance (e.g., varying shapes and sizes) of objects of interest in medical image segmentation tasks. To tackle this challenge, we present a structured approach to introduce spatially dynamic components to the ViT-UNet. This adaptation enables the model to effectively capture features of target objects with diverse appearances. This is achieved by three main components: extbf{(i)} deformable patch embedding; extbf{(ii)} spatially dynamic multi-head attention; extbf{(iii)} deformable positional encoding. These components were integrated into a novel architecture, termed AgileFormer. AgileFormer is a spatially agile ViT-UNet designed for medical image segmentation. Experiments in three segmentation tasks using publicly available datasets demonstrated the effectiveness of the proposed method. The code is available at \href{https://github.com/sotiraslab/AgileFormer}{https://github.com/sotiraslab/AgileFormer}.

研究动机与目标

• 通过分割任务中形状和大小各异的目标，解决医学影像异质性。
• 开发一个空间动态 ViT-UNet，提升定位和全局上下文处理。
• 提出可变形补丁嵌入和可变形位置编码，以增强表征学习。

提出的方法

• 用基于可变形卷积的可变形补丁嵌入替代刚性补丁嵌入。
• 使用一个空间动态自注意力模块，交替使用可变形多头注意力和邻域多头注意力。
• 引入多尺度可变形位置编码，以建模自注意力中的不规则网格。
• 采用 U-Net 風格的编码器-解码器，带跳跃连接和深度监督。
• 对 AgileFormer 的变体 AgileFormer-T 与 AgileFormer-B，在有无深度监督的情况下进行对比。
• 在 Synapse（2D 多器官）、ACDC（2D 心脏）以及 Decathlon 脑肿瘤（3D）数据集上进行评估。

实验结果

研究问题

• RQ1 空间动态组件是否能使 ViT-UNets 更好地处理医学影像中异质的器官形状和大小？
• RQ2 可变形补丁嵌入和多尺度可变形位置编码是否在分割精度和可扩展性方面优于固定窗口的 ViT-UNets？
• RQ3 与其他 SOTA ViT-UNets 的比较中， AgileFormer 在模型规模增长时的性能如何？
• RQ4 深度监督对 AgileFormer 的分割性能有何影响？
• RQ5 所提出的组件在 2D 与 3D 医学分割任务中的表现如何？

主要发现

• AgileFormer 在三个数据集（Synapse、ACDC、Decathlon）上均超越了最近的 SOTA 方法。
• 将 AgileFormer-T 扩展到 AgileFormer-B，带来显著的 DSC 增益，参数和 FLOPs 增长适中。
• 消融研究显示可变形补丁嵌入和 MS-DePE 有助于性能提升。
• 空间动态注意力（DMSA/NMSA）及可变形组件在处理不同形状/大小对象方面提供改进。
• 深度监督进一步提升性能，DSC 显著改善。
• 3D 实验显示 AgileFormer-T+DS 在对比 3D ViT-UNets 时达到具有竞争力的结果。