Patcher: Patch Transformers with Mixture of Experts for Precise Medical Image Segmentation

TL;DR

Patcher is a novel encoder-decoder Vision Transformer architecture that uses patch segmentation and a mixture-of-experts decoder to improve medical image segmentation accuracy, combining local detail and global context.

Contribution

It introduces Patcher blocks with overlapping large patches and a MoE-based decoder, enhancing spatial modeling and feature specialization in medical image segmentation.

Findings

Outperforms state-of-the-art methods on stroke lesion segmentation

Achieves superior results on polyp segmentation

Demonstrates effective feature extraction from local to global levels

Abstract

We present a new encoder-decoder Vision Transformer architecture, Patcher, for medical image segmentation. Unlike standard Vision Transformers, it employs Patcher blocks that segment an image into large patches, each of which is further divided into small patches. Transformers are applied to the small patches within a large patch, which constrains the receptive field of each pixel. We intentionally make the large patches overlap to enhance intra-patch communication. The encoder employs a cascade of Patcher blocks with increasing receptive fields to extract features from local to global levels. This design allows Patcher to benefit from both the coarse-to-fine feature extraction common in CNNs and the superior spatial relationship modeling of Transformers. We also propose a new mixture-of-experts (MoE) based decoder, which treats the feature maps from the encoder as experts and selects a suitable set of expert features to predict the label for each pixel. The use of MoE enables better specializations of the expert features and reduces interference between them during inference. Extensive experiments demonstrate that Patcher outperforms state-of-the-art Transformer- and CNN-based approaches significantly on stroke lesion segmentation and polyp segmentation. Code for Patcher is released with publication to facilitate future research.