[Paper Review] MISSFormer: An Effective Medical Image Segmentation Transformer

Xiaohong Huang, Zhifang Deng+2 more|arXiv (Cornell University)|Sep 15, 2021

Advanced Neural Network Applications40 references170 citations

TL;DR

MISSFormer introduces a hierarchical U-shaped transformer with an Enhanced Transformer Block and an Enhanced Transformer Context Bridge, trained from scratch, achieving state-of-the-art medical image segmentation on Synapse and ACDC datasets.

ABSTRACT

The CNN-based methods have achieved impressive results in medical image segmentation, but they failed to capture the long-range dependencies due to the inherent locality of the convolution operation. Transformer-based methods are recently popular in vision tasks because of their capacity for long-range dependencies and promising performance. However, it lacks in modeling local context. In this paper, taking medical image segmentation as an example, we present MISSFormer, an effective and powerful Medical Image Segmentation tranSFormer. MISSFormer is a hierarchical encoder-decoder network with two appealing designs: 1) A feed-forward network is redesigned with the proposed Enhanced Transformer Block, which enhances the long-range dependencies and supplements the local context, making the feature more discriminative. 2) We proposed Enhanced Transformer Context Bridge, different from previous methods of modeling only global information, the proposed context bridge with the enhanced transformer block extracts the long-range dependencies and local context of multi-scale features generated by our hierarchical transformer encoder. Driven by these two designs, the MISSFormer shows a solid capacity to capture more discriminative dependencies and context in medical image segmentation. The experiments on multi-organ and cardiac segmentation tasks demonstrate the superiority, effectiveness and robustness of our MISSFormer, the experimental results of MISSFormer trained from scratch even outperform state-of-the-art methods pre-trained on ImageNet. The core designs can be generalized to other visual segmentation tasks. The code has been released on Github: https://github.com/ZhifangDeng/MISSFormer

Motivation & Objective

• Address the limitation of CNNs in modeling long-range dependencies in medical image segmentation.
• Propose a position-free, hierarchical U-shaped transformer for accurate segmentation.
• Design Enhanced Transformer Block and Enhanced Transformer Context Bridge to capture local and global context.
• Evaluate MISSFormer on multi-organ CT and cardiac MRI datasets to demonstrate robustness and the ability to train from scratch.

Proposed method

• Redesigned feed-forward network called Enhanced Mix-FFN to improve feature discrimination and local/global context integration.
• Enhanced Transformer Block combining LayerNorm, Efficient Self-Attention, and Enhanced Mix-FFN to model long-range and local information with reduced complexity.
• Enhanced Transformer Context Bridge that fuses multi-scale features by concatenating flattened multi-level tokens and processing them through an Enhanced Transformer Block.
• Hierarchical encoder-decoder with overlapping 4x4 patches, patch merging/expansion, and skip connections forming a U-shaped architecture.
• Training from scratch on medical datasets with standard data augmentation, using SGD optimizer and polynomial learning rate policy.

Experimental results

Research questions

• RQ1 Can MISSFormer surpass state-of-the-art medical image segmentation methods trained from scratch on Synapse and ACDC datasets?
• RQ2 Does the Enhanced Transformer Block and the Context Bridge improve discrimination and context modeling over prior Transformer/MLP-based approaches?
• RQ3 How does multi-scale feature fusion via the Enhanced Transformer Context Bridge affect segmentation accuracy and edge delineation?
• RQ4 What is the impact of different skip connections and recursive steps in the Enhanced Mix-FFN on convergence and performance?

Key findings

• MISSFormer achieves state-of-the-art performance on Synapse and ACDC datasets, often surpassing methods pre-trained on ImageNet.
• The Simple Enhanced Mix-FFN and recursive skip connections improve training stability and segmentation accuracy compared to baseline SegFormer-based variants.
• In ablations, MISSFormer with Enhanced Transformer Context Bridge and multi-scale fusion shows notable Dice-Sørensen Coefficient (DSC) gains and edge delineation improvements.
• MISSFormer_S (without multi-scale bridge) underperforms compared to MISSFormer, highlighting the benefit of multi-scale information integration.
• MISSFormer demonstrates strong edge predictions and robust performance across hard examples, with competitive or superior per-organ results on Synapse and overall robustness on ACDC.