UniSeg: A Prompt-driven Universal Segmentation Model as well as A Strong Representation Learner (2304.03493v1)

Abstract: The universal model emerges as a promising trend for medical image segmentation, paving up the way to build medical imaging large model (MILM). One popular strategy to build universal models is to encode each task as a one-hot vector and generate dynamic convolutional layers at the end of the decoder to extract the interested target. Although successful, it ignores the correlations among tasks and meanwhile is too late to make the model 'aware' of the ongoing task. To address both issues, we propose a prompt-driven Universal Segmentation model (UniSeg) for multi-task medical image segmentation using diverse modalities and domains. We first devise a learnable universal prompt to describe the correlations among all tasks and then convert this prompt and image features into a task-specific prompt, which is fed to the decoder as a part of its input. Thus, we make the model 'aware' of the ongoing task early and boost the task-specific training of the whole decoder. Our results indicate that the proposed UniSeg outperforms other universal models and single-task models on 11 upstream tasks. Moreover, UniSeg also beats other pre-trained models on two downstream datasets, providing the community with a high-quality pre-trained model for 3D medical image segmentation. Code and model are available at https://github.com/yeerwen/UniSeg.

• The paper introduces a prompt-driven universal model that leverages early task awareness to improve multi-task segmentation across diverse imaging modalities.
• The model employs a novel FUSE module and a modified nnUNet backbone to dynamically generate task-specific prompts from a universal prompt.
• The paper demonstrates superior performance with higher Dice coefficients and robust transfer learning on various medical imaging datasets.

UniSeg: A Universal Segmentation Model for Multi-Task Medical Imaging

This paper presents UniSeg, a universal model for multi-task medical image segmentation that also serves as a strong representation learner. Unlike traditional approaches that either tackle different segmentation tasks independently or treat them as multi-class problems, UniSeg addresses the core issues of task correlations and early task awareness to enhance segmentation outcomes across diverse modalities and domains.

Model Design and Methodology

The challenge in medical image segmentation lies in both handling multiple tasks with limited data and recognizing the inherent correlations between those tasks. UniSeg improves upon previous universal segmentation models by incorporating a novel prompt-driven strategy. The architecture of UniSeg is based on a vision encoder, a specially designed fusion and selection (FUSE) module, and a prompt-driven decoder.

1. Universal Prompt and Task Awareness: The model introduces a learnable universal prompt that captures correlations among different tasks. This universal prompt interacts with the feature maps extracted by the encoder to generate task-specific prompts. By feeding this task-specific prompt into the decoder as part of its input, the model becomes 'aware' of the ongoing task early in the process, enhancing the overall task-specific training of the decoder.
2. Dynamic Task Prompts: The FUSE module is responsible for converting the universal prompt and encoded features into task-specific prompts. This mechanism allows the model to accommodate multiple tasks with a unified decoder and segmentation head, thus mitigating the redundancy seen in multi-head networks with one decoder per task.
3. Encoder-Decoder Backbone: UniSeg leverages a modified version of nnUNet for its core architecture. This backbone effectively captures image features while supporting multi-modal input, conducive to segmentation of 3D medical images with different characteristics.

Performance Evaluation

UniSeg's efficacy was validated on a broad set of medical datasets, encompassing tasks such as organ and tumor segmentation from CT, MR, and PET modalities. The model outperforms both universal models like DoDNet and single-task models across 11 upstream tasks, exhibiting a notable performance improvement. For instance, UniSeg achieved compelling results with Dice coefficients surpassing others on datasets like Liver, Kidney, and Pancreas, demonstrating strong generalization capabilities.

Furthermore, UniSeg's representation learning capability was tested using transfer learning on two downstream datasets. The results confirmed its robust performance compared to both unsupervised and supervised pre-trained models. UniSeg not only reached but often exceeded the performance benchmarks set by existing pre-trained segmentation models.

Implications and Speculations for Future Developments

The introduction of a universal prompt in UniSeg suggests significant potential for future developments in AI-driven medical imaging. By utilizing task correlations earlier in the processing pipeline, the model sets a precedent for more efficient and accurate multi-task learning frameworks. As medical datasets continue to grow in variety and complexity, strategies akin to UniSeg's prompt-driven approach could become critical in developing scalable, resource-efficient models.

In addition to improving medical image segmentation, this research could influence the design of universal models in other fields requiring multi-task learning. The methodology could be adapted or extended to accommodate more diverse data types and tasks outside the current domain. Future work might explore extensions to accommodate real-time adaptations or increased task diversity without loss of segmentation quality.

In conclusion, UniSeg represents a significant step forward in universal segmentation model design, showing promising performance and adaptability across various medical imaging tasks. The research provides a strong foundation for future advancements in multi-task learning, reinforcing the value of recognizing and integrating task correlations early in model architectures.