MaskMed: Decoupled Mask and Class Prediction for Medical Image Segmentation

2025-11-19

AI Summary

In medical image segmentation, conventional point-wise convolutional heads rigidly bind output channels to specific classes, hindering feature sharing and semantic generalization. To address this, we propose a unified decoupled segmentation head that decomposes multi-class segmentation into class-agnostic mask generation and independent category classification, augmented by shared object queries to enable cross-class feature reuse. Furthermore, we design a full-scale-aware deformable Transformer that efficiently fuses full-resolution features using deformable attention guided by low-resolution features. Our method achieves state-of-the-art performance, improving Dice scores by 2.0% on AMOS 2022 and 6.9% on BTCV over nnUNet. The core contributions lie in the synergistic innovation of (1) a decoupled architectural design, (2) a shared object query mechanism for enhanced feature generalization, and (3) multi-scale deformable attention for robust hierarchical feature integration.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Feature Construction/Reformulation
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Systems and Infrastructure for Web, Mobile and WoT: Decentralized Web and Fediverse systems

Abstract

Medical image segmentation typically adopts a point-wise convolutional segmentation head to predict dense labels, where each output channel is heuristically tied to a specific class. This rigid design limits both feature sharing and semantic generalization. In this work, we propose a unified decoupled segmentation head that separates multi-class prediction into class-agnostic mask prediction and class label prediction using shared object queries. Furthermore, we introduce a Full-Scale Aware Deformable Transformer module that enables low-resolution encoder features to attend across full-resolution encoder features via deformable attention, achieving memory-efficient and spatially aligned full-scale fusion. Our proposed method, named MaskMed, achieves state-of-the-art performance, surpassing nnUNet by +2.0% Dice on AMOS 2022 and +6.9% Dice on BTCV.

Problem

• Decouples mask prediction from class prediction
• Enables feature sharing and semantic generalization
• Achieves memory-efficient full-scale feature fusion

Innovation

• Decoupled mask and class prediction using shared queries
• Full-scale deformable attention for feature fusion
• Memory-efficient transformer module for medical segmentation