Wave-GMS: Lightweight Multi-Scale Generative Model for Medical Image Segmentation
This addresses the problem of deploying high-performance AI tools in hospitals with limited GPU resources, though it appears incremental as it builds on existing lightweight and generative model approaches.
The paper tackles the need for efficient deep segmentation networks for equitable AI deployment in healthcare by proposing Wave-GMS, a lightweight multi-scale generative model that achieves state-of-the-art segmentation performance with only ~2.6M trainable parameters, as demonstrated on four public datasets.
For equitable deployment of AI tools in hospitals and healthcare facilities, we need Deep Segmentation Networks that offer high performance and can be trained on cost-effective GPUs with limited memory and large batch sizes. In this work, we propose Wave-GMS, a lightweight and efficient multi-scale generative model for medical image segmentation. Wave-GMS has a substantially smaller number of trainable parameters, does not require loading memory-intensive pretrained vision foundation models, and supports training with large batch sizes on GPUs with limited memory. We conducted extensive experiments on four publicly available datasets (BUS, BUSI, Kvasir-Instrument, and HAM10000), demonstrating that Wave-GMS achieves state-of-the-art segmentation performance with superior cross-domain generalizability, while requiring only ~2.6M trainable parameters. Code is available at https://github.com/ATPLab-LUMS/Wave-GMS.