Boosting Convolution with Efficient MLP-Permutation for Volumetric Medical Image Segmentation

TL;DR

PHNet tackles the challenge of efficient volumetric medical image segmentation by integrating CNNs for local feature extraction with an MLPP module that captures global context through axial decomposition and token segmentation. The 2.5D encoder addresses anisotropy, while the lightweight decoder ensures high-resolution reconstructions. Key contributions include IP-MLP, AA-MLP, and TP-MLP, plus a token-segmentation strategy that preserves positional information and reduces computation. Empirical results on COVID-19-20, Synapse, LiTS, and MSD BraTS show state-of-the-art performance with favorable computational costs, supported by comprehensive ablations.

Abstract

Recently, the advent of vision Transformer (ViT) has brought substantial advancements in 3D dataset benchmarks, particularly in 3D volumetric medical image segmentation (Vol-MedSeg). Concurrently, multi-layer perceptron (MLP) network has regained popularity among researchers due to their comparable results to ViT, albeit with the exclusion of the resource-intensive self-attention module. In this work, we propose a novel permutable hybrid network for Vol-MedSeg, named PHNet, which capitalizes on the strengths of both convolution neural networks (CNNs) and MLP. PHNet addresses the intrinsic isotropy problem of 3D volumetric data by employing a combination of 2D and 3D CNNs to extract local features. Besides, we propose an efficient multi-layer permute perceptron (MLPP) module that captures long-range dependence while preserving positional information. This is achieved through an axis decomposition operation that permutes the input tensor along different axes, thereby enabling the separate encoding of the positional information. Furthermore, MLPP tackles the resolution sensitivity issue of MLP in Vol-MedSeg with a token segmentation operation, which divides the feature into smaller tokens and processes them individually. Extensive experimental results validate that PHNet outperforms the state-of-the-art methods with lower computational costs on the widely-used yet challenging COVID-19-20 and Synapse benchmarks. The ablation study also demonstrates the effectiveness of PHNet in harnessing the strengths of both CNNs and MLP. The code is available on Github: \href{https://github.com/xiaofang007/PHNet}{https://github.com/xiaofang007/PHNet}.

Figures (11)

• Figure 1: Left: Performance v.s. throughput of PHNet and other SOTA methods. Right: performance v.s. FLOPs of CNN, Transformers, and MLP, in different model capacity.
• Figure 2: Overview of the proposed PHNet.
• Figure 3: Detailed network architecture of PHNet.
• Figure 4: Illustration of the anisotropic problem. (a) Anatomical planes of Lung including (b) transverse plane (c) sagittal plane (d) coronal plane. Better visualization with zooming in.
• Figure 5: Illustration of multi-layer permute perceptron (MLPP) module.