SegMate: Asymmetric Attention-Based Lightweight Architecture for Efficient Multi-Organ Segmentation

TL;DR

SegMate is a lightweight, efficient multi-organ segmentation framework that maintains high accuracy while significantly reducing computational resources, making it suitable for resource-limited clinical environments.

Contribution

We introduce SegMate, a novel asymmetric attention-based architecture that achieves state-of-the-art accuracy with lower computational costs across multiple datasets.

Findings

Reduces GFLOPs by up to 2.5x compared to vanilla models.

Decreases VRAM usage by up to 2.1x.

Achieves a Dice score of 93.51% on TotalSegmentator.

Abstract

State-of-the-art models for medical image segmentation achieve excellent accuracy but require substantial computational resources, limiting deployment in resource-constrained clinical settings. We present SegMate, an efficient 2.5D framework that achieves state-of-the-art accuracy, while considerably reducing computational requirements. Our efficient design is the result of meticulously integrating asymmetric architectures, attention mechanisms, multi-scale feature fusion, slice-based positional conditioning, and multi-task optimization. We demonstrate the efficiency-accuracy trade-off of our framework across three modern backbones (EfficientNetV2-M, MambaOut-Tiny, FastViT-T12). We perform experiments on three datasets: TotalSegmentator, SegTHOR and AMOS22. Compared with the vanilla models, SegMate reduces computation (GFLOPs) by up to 2.5x and memory footprint (VRAM) by up to 2.1x, while generally registering performance gains of around 1%. On TotalSegmentator, we achieve a Dice score of 93.51% with only 295MB peak GPU memory. Zero-shot cross-dataset evaluations on SegTHOR and AMOS22 demonstrate strong generalization, with Dice scores of up to 86.85% and 89.35%, respectively. We release our open-source code at https://github.com/andreibunea99/SegMate.