Universal Medical Image Representation Learning with Compositional Decoders

In the field of medical image analysis, three critical tasks are predominantly recognized: lesion detection, classification, and segmentation (Chen et al. 2022b). Detection entails identifying the location of a lesion within an image based on a textual query (Tiu et al. 2022), which is crucial for clinical finding and localization of abnormalities. Lesion classification involves assigning a class label to an image (Wang et al. 2022a) or a specified target region (Murtaza et al. 2020). Segmentation requires generating pixel-level labels for an entire image (Ronneberger, Fischer, and Brox 2015), aiding in the clinical demarcation of lesion boundaries. Multi-task models for the aforementioned tasks typically use a shared visual backbone to produce visual embeddings (Liu et al. 2023), followed by individual branches tailored for each specific task. While these task-specific learning frameworks are effective for particular datasets, they lack generality and necessitate designing from scratch for new tasks.

The emergence of large-scale models has significantly revolutionized the field of medical image analysis (Rajpurkar et al. 2022). Recent studies have been increasingly focused on developing general-purpose medical artificial intelligence (AI) models (Moor et al. 2023). A notable trend is the incorporation of the Segment Anything Model (SAM) (Wu et al. 2023; Ma and Wang 2023; Cheng et al. 2023; Zhang and Liu 2023), which amalgamates medical domain knowledge for medical image segmentation and their application across various segmentation tasks. Concurrently, the emergence of image-text models has garnered considerable attention (Liu et al. 2023; Ye et al. 2023; Zhao et al. 2023; Wang et al. 2022c). These models interpret image features through task-specific prompts, encompassing diverse modalities and domains, and represent a stride toward prompt-driven universal models. Most current methodologies predominantly concentrate on medical image segmentation (Butoi et al. 2023) without adequately acknowledging the interconnectedness and uniformity across various medical tasks. Our research aims to bridge this gap by proposing a unified approach capable of handling all three tasks simultaneously, enabling training with diverse annotations and tasks, and building the foundation for more versatile and universally applicable medical image analysis.

UniMed comprises a visual-language encoder and a dual-output decoder architecture (Fig. 2). Features are learned through visual and textual encoders upon receiving an input image. Subsequently, guided by the task labels, a decoder is employed to autoregressively predict the sequence.

Visual Encoder. To standardize the input space into discrete tokens, the encoder must be transformer-based. Therefore, the visual encoder $\texttt{Enc}_{v}$ utilizes the Swin-Transformer as its backbone, given its widely proven effectiveness. Given an input image $I$, this component extracts its layer features $V_{l}$ to derive the final multi-scale visual feature $V$ representation:

$$V=\texttt{Enc}_{v}(I)=[V_{1},V_{2},...,V_{L}]$$

(1)

where L is the number of layers.

Text Encoder. The text encoder is designed to capture annotation semantics and learn a broad spectrum of corpus knowledge. Since the specific annotations vary from dataset to dataset, the focus may be on lesions, instruments, or multiple annotation types, etc. For a piece of text $T$ generated from annotations, SentencePiece (Kudo and Richardson 2018) is first employed to divide the words and convert them into discrete token sequences. The text encoder, $\texttt{Enc}_{t}$, consists of multiple layers of Transformers that process the input text sequence. This process forms a text input queue $Q_{t}$, as follows:

$$Q_{t}=\texttt{Enc}_{t}(T)=[q_{t}^{1},q_{t}^{2},...,q_{t}^{n}]$$

(2)

where $n$ is the length of the query.

Decomposed Decoder. The input to the model includes visual features $V$, a text queue $Q_{t}$, and a general queue $Q_{g}=[q_{g}^{1},q_{g}^{2},...,q_{g}^{m}]$, while the output consists of a pixel $O_{p}$ and a semantic $O_{s}$. The flexible combination of these inputs and outputs is capable of supporting both general and referring tasks. The decomposed decoder is composed of stacked Transformers. It initially captures the global features of the image by computing the masked cross-attention $\texttt{Att}_{cross}$ among the three inputs (Cheng et al. 2022) and a self-attention $\texttt{Att}_{self}$ mechanism to generate the queue for the subsequent layer:

$$[\hat{Q}_{t}^{l},\hat{Q}_{g}^{l}]=\texttt{Att}_{cross}([Q_{t}^{l-1},Q_{g}^{l-1}],V)$$

(3)

$$[Q_{t}^{l},Q_{g}^{l}]=\texttt{Att}_{self}([\hat{Q}_{t}^{l-1},\hat{Q}_{g}^{l-1}])$$

(4)

where $l$ represents the $l\texttt{-th}$ layer. For general tasks, the last general query $Q_{g}$ is utilized as the global image representation. For referring tasks, the text query $Q_{t}$ serves as the referring feature, while the general query $Q_{g}$ is combined with it to produce the final representation.

For pixel-level output, the decoder utilizes global image features from the general queue to produce output $O_{p}=[O_{p}^{1},O_{p}^{2},...,O_{p}^{m}]$, facilitating a nuanced understanding of the image at a fine-grained level. Moreover, for semantic-level output, the decoder relies on both the general and text queues $O_{s}=[O_{s}^{1},O_{s}^{2},...,O_{s}^{m+n}]$ to facilitate higher-level semantic understanding and generation.

Overall operation. UniMeds’s encoder encompasses the visual $\texttt{Enc}_{v}$ and text $\texttt{Enc}_{t}$ feature extraction branch, whereas the decomposed decoder Dec utilizes visual features $V$, textual queues $Q_{t}$, and general $Q_{g}$ queues to forecast both pixel-level $O_{p}$ and semantic-level $O_{s}$ outputs (Cheng et al. 2022). The overall operation can be expressed as:

$$[O_{p},O_{s}]=\texttt{Dec}(V,(Q_{t},Q_{g}))$$

(5)

Since the decomposed decoder represents the output in terms of semantic and pixel outputs, the composed decoder unifies the labels of different tasks into a format that can be expressed through these two outputs.

Unified annotations. For classification, the model outputs the image category, with the corresponding annotation being the category text. We match the text with the semantic output path by tokenizing it using SentencePiece. For detection, the model identifies the location of the target area, with the annotation being the diagonal coordinates of the bounding box. To encode this sparse structure, we encode the sparse structure by expanding the vocabulary with 1000 special tokens (Chen et al. 2021b). The bounding box is then represented by four tokens, two indicating the upper-left corner and the other two representing the lower-right corner, together with the category, serve as semantic outputs. For segmentation, the model generates a result for each pixel, with the label being the pixel-level mask. The category annotation is processed using the method described earlier. The label image is encoded into discrete tokens, enabling the simultaneous generation of both pixel-level and semantic-level predictions.

General Classification/Detection (Fig. 3 a)). General classification and detection rely on the input image for prediction, only leveraging visual features and general queues as input to the decomposed decoder. Through composed decoders, all annotations are tokenized, and the classification or detection task directly outputs the prediction results through the semantic path. Hence, expressed as:

$$[O_{s}]=\texttt{Dec}(V,(Q_{g}))$$

(6)

General Segmentation (Fig. 3 b)). The input of its decomposed decoder is consistent with the Eq.2. The image is encoded and make predictions by simultaneously generating pixel-level and semantic-level outcomes. The operation is as follows:

$$[O_{p},O_{s}]=\texttt{Dec}(V,(Q_{g}))$$

(7)

Referring Classification/Detection (Fig. 3 c)). The referring task requires a combination of visual features, text, and general queues (Eq. 1) to derive corresponding segmentation results. This enables clinical practice to flexibly obtain precise localization and diagnostic predictions of specified lesions by giving additional text prompts.

$$[O_{s}]=\texttt{Dec}(V,(Q_{t},Q_{g}))$$

(8)

Referring Segmentation (Fig. 3 d)). It requires latent query and text query as input, so the formula is the same as Eq.5.

$$[O_{p},O_{s}]=\texttt{Dec}(V,(Q_{t},Q_{g}))$$

(9)

Compared with Eq. 7, the referring segmentation can be regarded general task with the language conditions.

Through the combined arrangement of queues and outputs, UniMed can support a variety of medical imaging tasks (Fig. 3). This paper advocates for achieving unity through functional cohesion rather than interface specifications, thereby maximizing the shared utilization of common components across diverse tasks while preserving independence for each task.

For medical image pre-training, relying solely on labeled data is far from enough, especially compared with the millions of data available in natural images. Hence, we delve into strategies for leveraging unlabeled data for training, aiming to bridge the disparity. Guided by this principle and drawing inspiration from the self-supervised contrastive learning paradigm, this work devises a joint training methodology, enabling the learning of labeled and unlabeled data simultaneously.

Pipeline. As for labeled data, the corresponding labels undergo standard transformations as outlined in Sec. Composed Decoders. These labels are categorized into semantic and pixel types, and training is conducted utilizing both semantic loss and pixel loss. For unlabeled data, follow the dense contrastive learning strategy (Wang et al. 2021). For each image, we generate two sets of random views via data augmentation and feed them into the encoder to obtain two sets of features. These features are separately passed through downstream dense projection heads, and the same encoder is trained by computing the contrastive learning loss between the two sets of features. During training, we adopt an exponential moving average to update the parameters and retain the encoder part after training is completed, discarding the dense header. During inference, task predictions are executed through the adaptable combination of visual, text, and latent queue input terminals, alongside semantic and pixel output terminals. The total training loss can be expressed as a combination of these:

$$L_{total}=\begin{matrix}\underbrace{L_{s}+L_{p}}\\ L_{label}\end{matrix}+\lambda\begin{matrix}\underbrace{L_{c}+L_{dc}}\\ L_{unlabel}\end{matrix}$$

(10)

Where $\lambda$ acts as a weight to balance the two terms. The semantic output $L_{s}$ is the cross entropy loss, the pixel output $L_{e}$ includes the binary cross entropy loss and the dice loss, and the unlabeled learning loss includes the contrast $L_{c}$ and the dense contrast loss $L_{dc}$.

This study takes the endoscopic modality of medical images as an example and conducts a comprehensive investigation by collecting datasets from various research groups worldwide. The database established contains 12 datasets and covers all 3 tasks. The unlabeled datasets are endoscopic videos that are difficult to label, Colonoscopic (Mesejo et al. 2016), Hyper-Kvasir (Borgli et al. 2020), Kvasir-Capsule (Smedsrud et al. 2021), LDPolypVideo (Ma et al. 2021), and ColonVideo (private, from Jiangsu Provincial People’s Hospital). For detection, evaluation is performed on the SUN (Misawa et al. 2021) colonoscopy public dataset, the largest benchmark for polyp detection. The classification task utilizes the ColonCG (Private, from Jiangsu Provincial People’s Hospital), which is the most comprehensive dataset for colon disease classification, including five categories: normal, polyp, adenoma, cancer, and ulcerative colitis. Segmentation tasks are evaluated on the common public polyp segmentation datasets: CVC-ClinicDB (Bernal et al. 2015), CVC-ColonDB (Tajbakhsh, Gurudu, and Liang 2015), Kvasir-SEG (Jha et al. 2020), ETIS-LaribPolypDB (Silva et al. 2014), and EndoScene (Vázquez et al. 2017). Evaluation metrics include mean average precision (mAP), mean accuracy (mAcc), and Dice corresponding to the three tasks respectively.

We employ Focal-T (Yang et al. 2022) as the backbone of the visual encoder, utilizing a transformer text encoder with causal masking (Radford et al. 2021) as the language encoder. For training, the AdamW (Loshchilov and Hutter 2017) optimizer with a base learning rate set to 1e-4, a weight decay of 0.05, and a linear decay learning rate scheduler are applied. The training procedure spans 50 epochs with a batch size of 8. A total of two data loaders are used: one for labeled data and another for unlabeled data, maintaining a sampling ratio of 1:1. The final loss function comprises both supervised and unsupervised components, with a ratio of 10:1 to balance their contributions effectively. During the fine-tuning process, the model’s input and output are controlled through configuration files, allowing for customized settings tailored to specific task requirements.

UniMed undergoes comparison with existing specialized and universal methods across various tasks. Fine-tuning is performed on 8 datasets for 3 common medical image analysis downstream tasks, with performance reported in Table 1, 2, 3. The analysis yields the following observations.

UniMed v.s. Specialized Detection Methods. In a comprehensive evaluation, UniMed is compared with various state-of-the-art methods (Table 1), including polyp detection (Wu et al. 2021), two-stage (He et al. 2017), single-stage (Jocher et al. 2020), and transformer-based detection methods (Zhou et al. 2022; Carion et al. 2020; Dai et al. 2021a). UniMed surpasses all specialized detection methods (Fig. 4), achieving the highest mean Average Precision (mAP) score of 56.8%, outperforming the next-best result by a significant margin of 3.2%.

UniMed v.s. Specialized Classification Methods. A total of six methods are compared (Table 2), encompassing CNN-based (He et al. 2016; Tan and Le 2019) , Transformer-based (Zhai et al. 2022; Liu et al. 2022), and CNN-Transformer hybrid approaches (Dai et al. 2021b; Wortsman et al. 2022). UniMed emerges as the top performer, achieving the highest classification results with an average accuracy of 90.8% across five categories. Notably, among these methods, Model Soups achieves commendable results, trailing behind the proposed UniMed model by on 1.7 points. This observation underscores the effectiveness of averaging multiple weights in Model soups, contributing to its competitive performance despite being sub-optimal.

UniMed v.s. Specialized Segmentation Methods. When comparing UniMed with state-of-the-art methods in polyp segmentation (Falk et al. 2019; Fan et al. 2020; Wei et al. 2021, 2022) and semantic segmentation (Table 3) (Isensee et al. 2021; Chen et al. 2021a; Cheng et al. 2022; Zhang et al. 2023), UniMed consistently achieves the best segmentation results (Fig. 4) across datasets with varying levels of segmentation difficulty. Specifically, UniMed outperforms suboptimal methods by 0.6, 1.0, 0.9, 6.6, and 2.1 points on the CVC-ClinicDB, CVC-ColonDB, Kvasir-SEG, ETIS-LaribPolypDB, EndoScene datasets, respectively. These results underscore the robustness of UniMed’s performance across diverse datasets, demonstrating its effectiveness in tackling segmentation challenges across different medical imaging scenarios.

UniMed v.s. Generalist Methods. Most notably, the proposed end-to-end architecture outperforms other general models across various medical tasks (Table 1, 2, 3). General methods (Chen et al. 2022a; Lu et al. 2022; Zhang et al. 2022; Li et al. 2023; Zou et al. 2023) tend to exhibit higher overall performance compared to specific methods, suggesting that the interaction of data and tasks fosters enhanced model learning. Additionally, models striving for high understanding performance often demonstrate lower localization performance (e.g., Uni-Perceiver v2), as it is not trivial to merge semantic and visual understanding into a single model. Similarly (Fig. 4), visual comparison results with other methods clearly show that our method has higher accuracy in boundary segmentation and localization detection.

UniMed is pre-trained, requiring only zero or a small number of parameters before its application to various downstream tasks. Thus, we assessed the model’s transferability to other tasks in both zero-shot and 100-shot settings.

Zero-shot Transfer. Experimental results demonstrate compelling evidence of UniMed’s substantial zero-shot capability in the medical field compared to other general methods (Table 4). This suggests that UniMed can be readily applied to various tasks without further adjustments. Moreover, UniMed surpassed sub-optimal results by 6.7% and 6.9% in detection and classification tasks, respectively. Additionally, the performance improvement in segmentation tasks even outperforms other methods by up to 7.2%.

100-shot Transfer. UniMed also demonstrates the superior overall performance of strong 100-shot on medical image analysis tasks (Table 4). In some instances, it even rivals fully supervised models trained with full-scale data, exemplified by the SUN dataset (100-shot AP 49.8%, compared to Mask RCNN’s 49.2%). In particular, comparing it with the X-decoder model, it can be seen that its performance is exceeded in both zero-shot and 100-shot cases. This underscores the key role of large amounts of unsupervised data in feature generalization.

As mentioned earlier, UniMed boasts a unique advantage of task interaction, enabling both single and joint-task reasoning. This distinctive capability enhances the model’s practicality in real clinical scenarios, particularly in customizing referrals and tasks. In Fig. 5, visualization results of single and joint task inference without architectural changes are showcased. For instance, when provided with a set of referrals such as [”polyp, adenoma, cancer”], UniMed seamlessly delivers both pixel-level and semantic-level predictions.

Ablation studies are performed to analyze the architecture of UniMed, and all experiments are tested on the segmentation task.

Backbone. Increasing the size of the backbone network indeed leads to performance improvements, as evidenced in Table 5. As the depth and embedding dimensions expand the visual encoder, performance improves across the board. This suggests that more powerful feature extractors facilitate prediction for downstream tasks.

Task Collaboration. To explore the relationship between task collaboration and interference, by eliminating individual tasks based on all tasks (Table 6), the following finding was drawn: 1) Learning multiple tasks together is more effective than focusing on a single task alone. 2) Detection tasks positively influence segmentation tasks, whereas classification tasks exhibit suboptimal performance on certain datasets. This discrepancy may be attributed to the interference caused by dividing images into multiple categories, particularly when dealing with challenging data that is difficult to classify accurately.

Balance between Labeled and Unlabeled Losses. Table 7 demonstrates that the best performance is achieved when the unsupervised loss is set to 0.1, suggesting that the supervised component holds greater significance in the overall training process. When the ratio of unsupervised to supervised losses is 1:1, there is a drop in performance. These observations indicate that 1) it is necessary to introduce a joint expression learning strategy, where unlabeled data helps the model learn more general features. 2) The supervised component effectively drives model learning, while the unsupervised component serves as a ”regularization” mechanism, guiding the model to acquire more robust knowledge.

Balance between Labeled and Unlabeled Data. Table 8 shows that during experimental loading, sampling ratios between labeled and unlabeled data below 1 perform better. That is, the sampling ratio of unlabeled data should be higher. This is because, in medical data, the amount of labeled data is very small, while unlabeled data can significantly increase the diversity of samples. Based on the above findings, this paper adopts a 1:1 ratio setting.