Mobile U-ViT: Revisiting large kernel and U-shaped ViT for efficient medical image segmentation

TL;DR

Mobile U-ViT is a lightweight, efficient vision transformer architecture designed specifically for medical image segmentation on resource-constrained mobile devices, achieving state-of-the-art results across multiple datasets.

Contribution

The paper introduces Mobile U-ViT, a novel mobile-friendly architecture with hierarchical patch embedding and a large-kernel local-global block tailored for medical imaging tasks.

Findings

Achieves state-of-the-art performance on 8 public datasets.

Operates efficiently on resource-constrained devices.

Generalizes well to unseen datasets.

Abstract

In clinical practice, medical image analysis often requires efficient execution on resource-constrained mobile devices. However, existing mobile models-primarily optimized for natural images-tend to perform poorly on medical tasks due to the significant information density gap between natural and medical domains. Combining computational efficiency with medical imaging-specific architectural advantages remains a challenge when developing lightweight, universal, and high-performing networks. To address this, we propose a mobile model called Mobile U-shaped Vision Transformer (Mobile U-ViT) tailored for medical image segmentation. Specifically, we employ the newly purposed ConvUtr as a hierarchical patch embedding, featuring a parameter-efficient large-kernel CNN with inverted bottleneck fusion. This design exhibits transformer-like representation learning capacity while being lighter and faster. To enable efficient local-global information exchange, we introduce a novel Large-kernel Local-Global-Local (LGL) block that effectively balances the low information density and high-level semantic discrepancy of medical images. Finally, we incorporate a shallow and lightweight transformer bottleneck for long-range modeling and employ a cascaded decoder with downsample skip connections for dense prediction. Despite its reduced computational demands, our medical-optimized architecture achieves state-of-the-art performance across eight public 2D and 3D datasets covering diverse imaging modalities, including zero-shot testing on four unseen datasets. These results establish it as an efficient yet powerful and generalization solution for mobile medical image analysis. Code is available at https://github.com/FengheTan9/Mobile-U-ViT.