[Paper Review] 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 AI11 citations

TL;DR

AgileFormer introduces deformable patch embedding, spatially dynamic self-attention, and multi-scale deformable positional encoding to a ViT-UNet, achieving state-of-the-art medical image segmentation across three datasets with strong scalability.

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}.

Motivation & Objective

• Address heterogeneity in medical image targets with varying shapes and sizes in segmentation tasks.
• Develop a spatially dynamic ViT-UNet that improves localization and global context handling.
• Propose deformable patch embedding and deformable positional encoding to enhance representation learning.

Proposed method

• Replace rigid patch embedding with deformable patch embedding based on deformable convolution.
• Use a spatially dynamic self-attention block that alternates deformable multi-head attention and neighborhood multi-head attention.
• Introduce multi-scale deformable positional encoding to model irregular grids in self-attention.
• Adopt a U-Net style encoder-decoder with skip connections and deep supervision.
• Compare AgileFormer variants AgileFormer-T and AgileFormer-B with and without deep supervision.
• Evaluate on Synapse (2D multi-organ), ACDC (2D cardiac), and Decathlon brain tumor (3D) datasets.

Experimental results

Research questions

• RQ1 Can spatially dynamic components enable ViT-UNets to better handle heterogeneous organ shapes and sizes in medical images?
• RQ2 Do deformable patch embedding and multi-scale deformable positional encoding improve segmentation accuracy and scalability over fixed-window ViT-UNets?
• RQ3 How does AgileFormer scale in performance with model size compared to other SOTA ViT-UNets across multiple datasets?
• RQ4 What is the impact of deep supervision on AgileFormer's segmentation performance?
• RQ5 How do the proposed components perform across 2D and 3D medical segmentation tasks?

Key findings

• AgileFormer outperformed recent SOTA methods on all three datasets (Synapse, ACDC, Decathlon).
• Model scaling from AgileFormer-T to AgileFormer-B yields notable DSC gains with moderate parameter and FLOPs increases.
• Ablation studies show deformable patch embedding and MS-DePE contribute to performance gains.
• Spatially dynamic attention (DMSA/NMSA) and deformable components provide improved handling of varying object shapes/sizes.
• Deep supervision further enhances performance with notable DSC improvements.
• 3D experiments show AgileFormer-T+DS achieving competitive results against 3D ViT-UNets.