Multimodal Representation Distribution Learning for Medical Image Segmentation

1 Introduction

Medical image segmentation, as a pivotal component of auxiliary disease diagnosis, holds a crucial role in medical image applications. Enormous advances in medical image segmentation benefit from the more and more annotated datasets. However, it is still commonly impression that medical imaging datasets are too small to develop robust deep learning

models [Liu et al., 2023]. One reason for this is the expensive spending of high-quality pixel-level annotations. Expert annotators typically spend several hours to annotate a medical image. In addition, traditional methods only adopt vision modality to train the model, which limits the ability of encoded knowledge [Zhao et al., 2023; Wang et al., 2021c; Wang et al., 2021b]. In contrast, medical text annotation is naturally rich in semantics. These annotations are much easier to obtain. Hence, it is intuitive to integrate text annotation into medical image segmentation.

∗Corresponding authors.

Previous multi-modal methods adopt element addition to fuse text and vision features. As shown in Fig. 1 (a), the element addition method can be seen to fuse two distributions into one in joint feature space H. The fused distribution contains all the information of text and vision features. However, both task-related and irrelevant information is included in the fused distribution, which results in information redundancy in the fused features [Hosseini et al., 2024]. Thus, the performance improvement is limited. Recently, contrastive learning has been commonly used to align text and image features [Li et al., 2023a]. Specifically, it considers text caption as a linguistic view of the image under consideration as shown in Fig. 1 (b). It then pulls the pair of image and text feature distribution (p and q, respectively) close in semantic space. Nevertheless, contrastive learning only models the common information between text and image. It ignores the modality-specific information. For example, the edge information in the vision modal. Thus, the value of distinctive perspectives is neglected due to the suppression of modality-specific information [Zhang et al., 2024]. In addition, a large number of image-text pairs data are necessary to train a contrastive learning model. Thus, contrastive learning is possibly unsuitable for medical image segmentation due to high-quality image-text pairs lacking in many medical domains [Thawkar et al., 2023; Huang et al., 2022a; Huang et al., 2022b].

Furthermore, recent studies [Li et al., 2023a; Wang et al., 2023] show that text-image models perform poorly in spatial understanding tasks but excel in semantic understanding tasks. We also observe this phenomenon in experiments. Remarkable spatial perception is essential for accurate segmentation of object boundaries, which is important in medical image segmentation. The delineation of boundaries can signify the transitions between different human tissues or anatomical structures, thus providing essential information for accurately separating these instances [Wei et al., 2023; Huang et al., 2021b; Huang et al., 2021a].

To address these issues, we proposed a novel multi-modal feature distribution method for medical image segmentation. Firstly, medical text annotation is incorporated into our method to compensate for the lack of quality in image data. It also enhances the ability of the model to encode knowledge. Secondly, we propose a novel multi-modal feature fusion method, aiming to reduce redundant information. Modality-specific information is also retained in the proposed method. Specifically, as shown in Fig 1 (c), a set of learnable features is adopted to extract the modality-specific information of two modalities by maximizing the similarity of corresponding distributions. Then a distribution is adopted to model these learnable features. It can maximize discriminative information about tasks. Finally, the boundary regions in images demonstrate a strong correlation with features in the frequency domain. Consequently, we propose a frequency prompt encoder to improve the spatial perception of the model. It can fully leverage high-frequency information. These features serve as spatial information prompts, which along with the multi-modal features are decoded into segmentation results.

Our main contributions are summarized as follows:

• We propose a novel multi-modal method that leverages the text and image features for medical image segmentation. The medical text annotation serves as additional information, which can enhance the ability of model to encode knowledge.
• We propose a novel multi-modal fusion paradigm, aiming to reduce redundancy while retaining the modalityspecific information.
• A novel frequency prompt encoder is proposed to leverage high-frequency information for accurately segmenting boundaries.

2 RelatedWork

2.1 Medical Image Segmentation

Medical image segmentation refers to the identification of organ or lesion pixels from medical images [Yao et al., 2023]. U-Net [Ronneberger et al., 2015] is commonly considered as a typical segmentation model for medical images. TransUNet [Chen et al., 2021] and Swin-UNet [Cao et al., 2022] combine Transformer and UNet for achieving better performance. However, the performance of traditional segmentation methods depend on the high-quality annotated data. Additionally, these methods only are trained with visual modality data, which potentially constrains model to encoded knowledge. Several approaches [Cherti et al., 2023; Sun et al., 2023] explore leveraging medical text annotation to reduce reliance on high-quality annotations while maintaining performance. Inspired by these methods, we propose a novel method that adopts medical text annotation to strengthen the ability of model to encoded knowledge. In this way, it can achieve solid performance on the dataset with limited high-quality annotations.

2.2 Multi-Modal Learning

Recent studies have shown that text information is useful for medical image analysis. Thus, it is essential to design an effective multimodal fusion paradigm to fuse text and vision features. LAVT [Yang et al., 2022] and VLT [Ding et al., 2022] integrate text feature into vision feature by a language encoder. LViT [Li et al., 2023b] directly adopts element-wise addition to fuse text and vision features. The problem with these methods is that containing redundant information in the fused features. Several methods such as MedCLIP [Wang et al., 2022b], TPRO [Zhang et al., 2023], and others [Wu et al., 2023; Cao et al., 2023] base CLIP [Radford et al., 2021] adopt contrastive learning to align text and image features. However, these methods only model the commonalities between different modalities. It ignores the wealth of model-specific information. In addition, the contrastive learning model requires a large amount of image-text pair data for training. This paper proposes a novel multi-modal fusion method. It can effectively reduce redundancy while retaining modality-specific information.

2.3 Spectral Representation

Deep neural networks are biased towards learning lowfrequency representations [Mildenhall et al., 2021]. To better leverage high-frequency information, several methods [Mao et al., 2023] investigate the Fast Fourier Transform (FFT)- based frequency representation. FFT is a powerful tool, and it can leverage the strengths of both spectral and spatial representations [Tang et al., 2021]. Recent research [Zhu et al., 2023] has focused on the parallel extraction of spatial and frequency information within the encoder block. However, the spatial and computational complexity is increased due to this parallel extraction for the high-dimension tensor. In this paper, we propose a novel frequency prompt encoder to fully aggregate high-frequency information from the frequency domain. The proposed method avoids parallel processing in the vision encoder, which requires fewer parameters and lower computational cost.

Frequency Representation

Figure 2: Overview of the proposed method. First, given the medical image and corresponding medical text annotation, Text encoder and vision encoder to generate text and vision features. Then a set of learnable features is adopted to reduce redundancy. A distribution is adopted to model these learnable features. The Fused feature is sampled from this distribution and a cross-modality processer is adopted to further process it. Finally, the segmentation output is obtained by decoding the post-processed features and the frequency prompt. Additionally, 'snow' means freezing the parameters of the text encoder.

3 Methodology

3.1 Overview

As shown in Fig. 2, the proposed method consists of four components. The first part integrates a vision encoder EncI and a text encoder EncT to extract the vision and text features. EncT is a pre-trained CLIP text encoder. The second part is a multi-modal feature fusion module. It leverages a distribution to reduce redundant information while retaining the modality-specific information. The third part is a frequency prompt encoder EncP . Specifically, considering the frequency information is significant for boundary segmentation, we designed EncP to fully aggregate high-frequency information from the frequency domain. Detailed explanations of these two parts are provided in the following sections. The last part is a segmentation decoder Dec. It leverages the fused features from a cross-modality processor PCM and frequency prompt fromEncP to generate segmentation results.

3.2 Multi-Modal Feature Fusion

Previous multi-modal fusion methods usually ignore the redundancy in the fused features. Consequently, it is hard to capture the discriminative information. To address these issues, we propose a novel multi-modal feature fusion method with a learnable distribution. Specifically, given the medical text annotation T and medical image I, EncT and Enc are adopted to extract the text and vision features FT and FV :

Moreover, a convolution block ConvB is used to map visual and text features to the joint feature space:

where ConvB consists of a convolution layer and a batch normalization layer. ReLU is activation function.

A set of learnable features FL is defined in the joint feature space H. It is adopted to capture the discriminative information about task-related in FT and FV . Specifically, texts with similar semantics have high cosine similarity, which means they are close to each other in the feature space [Nlong Zhao et al., 2023]. Previous research [Lu et al., 2022; Huang et al., 2020] further demonstrates that text describing identical categories cluster together in the feature space. Inspired by this phenomenon, a prompt TP that is similar in semantics to T is defined. TP is encoded by EncT and ConvB, which generates a prompt feature FP in H, defined as:

FP is close to TP in H. Then FP is repeated K times to obtain FP = {F1P ,F2P , · · · ,FKP }. However, these features are identical, directly initializing them as learnable features would cause them to converge to the same vector. To avoid this trouble, Gaussian noise N(0, I) is added to FiP ,

This operation can introduce variation to the N features while ensuring that every FiL is still close to FT in H. Consequently, FL = {F1L,F2L, · · · ,FKL } is initialized as learnable features. To model the text and vision feature distributions (p, q, respectively) in H, we assume that p, q follow Gaussian distributions p ∼ N(µp, σ2 p ) and p ∼ N(µq,σ2 q ). With N learnable features FL, we can estimate the µp and σp of p as:

Similarity, µq and σq of q can be estimated as:

As a result, FL can extract the modality-specific information in FT and FV .

Then p and q are adopted to reduce the redundant information by modeling the discriminative information Specifically, we leverage Kullback-Leibler divergence LKL to measure the similarity between p and q. FL can extract modality-common information by minimizing LKL. LKL is defined as :

In the proposed method, LKL can be represented as:

To further maximize discriminative information and reduce redundant information, we use a Gaussian distribution N(µr,σr), where µr = µ(FL) and σr = σ(FL) are the mean and covariance of r, respectively. The fused feature F is sampled from r. F effectively fuses modality-common and modality-specific information. Nevertheless, sampling F from r is not differentiable for optimization. Thus, the reparameterization trick is adopted, which is similar to VAE [Kingma and Welling, 2013]. Formally, the process of sampling F is rewrite as:

where ω ∼ N(0, I) has same dimension as µr and σr.

3.3 Frequency Prompt Encoder

To improve the spatial perception ability of the text-image model, we design a frequency prompt encoder EncP . It is designed to fully utilize high-frequency information, which aims to prompt the boundaries information for the segmentation decoder. In contrast to previous methods, the proposed EncP is not integrated into basic blocks. Thus it can avoid parallel processing in the vision encoder, which requires fewer parameters and lower computational cost.

The Discrete Fourier Transform (DFT) is a powerful tool for converting images to the frequency domain. In practice, the FFT is used to compute the DFT efficiently, defined as:

where fh,w dnotes the spectrum represention of I.

Furthermore, the amplitude and phase spectrum of Fu,v can be computed as |Fu,v| and arg(Fu,v). The phase spectrum contains information about the edges and overall structure [Ghiglia and Pritt, 1998]. Consequently, the phase spectrum outperforms the amplitude spectrum in terms of representation ability in ablation studies. Therefore, we adopt the phase spectrum Fpha as the default in the proposed method. As shown in Fig. 2, EncP consists of two components. The first component EncP1 is consists of sequences of convolution layers, layer normalization, and activations. It can be represented as:

where Z iP is the output of layer i and i = {1, 2, · · · , n}. Convi , LN i , and GeLU i are the i-th convolution layer, GeLU activation function, and layer normalization, respectively. Similarity, the second component EncP2 can be defined as:

where Fpha is the frequency prompt of the I, ConvT i is the i-th transposed convolution layer. In experiments, we set N = 2.

3.4 Segmentation Decoder

Like LViT [Li et al., 2023b], the EncI and Dec are U-shape architecture. Moreover, the cross-modality processor PCM is also a U-shape architecture. More details about the network architectures of EncI , Dec, and PCM are provided in the supplementary materials. Last decoder bolck outputs the segmentation result O, formulated as:

where FMV is the output of the last layer of EncI.

3.5 Loss Function

The loss function used in this work are Dice loss LDice, crossentropy loss LCE, and LKL, formulated as:

where S is the number of pixels in the image, C is the number of classes, oij and dij are the ground truth and predicted segmentation output, respectively.

The final loss function L is defined as:

Base on experiments, we set λ1 = 0.5, λ2 = 0.5, and λ3 = 1.

4 Experiments

4.1 Datasets

The proposed method is evaluated on three medical image segmentation datasets: MoNuSeg [Kumar et al., 2017], MosMedData+ [Li et al., 2023b], and GlaS [Sirinukunwattana et al., 2017]. The first two datasets are the same benchmark datasets used in [Li et al., 2023b]. MoNuSeg contains 30 digital microscopic tissue images of several patients, while MosMedData+ [Morozov et al., 2020; Hofmanninger et al., 2020] contains 2729 CT scan slices of lung infections. The ratio of training, validation, and test sets are the same as in [Li et al., 2023b]. GlaS has 85 images for training and 80 for testing.

4.2 EvaluationMetrics

Dice score and IoU [Li et al., 2023b] are used to evaluate the performance of the proposed method and other SOTA methods. Moreover, we have reported the time complexity and space complexity of the proposed method. Specifically, the time complexity is evaluated in Floating Point Operations per second (FLOPs), and the space complexity is measured by the number of parameters.

4.3 Implementation Details

The proposed method is optimized by AdamW [Loshchilov and Hutter, 2017] and cosine annealing learning rate scheduler is adopted. The initial learning rate is set to 1e-3 for all datasets. Image input sizes are 224 × 224 both for MoNuSeg, GlaS, and MosMedData+. An early stop mechanism is adopted until the performance of the model does not increase for 50 epochs. The batch size is 2 for MoNuSeg and GlaS and 24 for MosMedData+. The default number of learnable features K is set to 32. All experiments are conducted on a single NVIDIA RTX 3090 GPU with 24GB memory.

4.4 Baseline

Baselines are divided into two sets: single-modal approaches and multi-modal approaches. Single-modal approaches include U-Net [Ronneberger et al., 2015], DoubleUnet [Ates et al., 2023], AttUNet [Wang et al., 2021a], Swin-UNet [Cao et al., 2022], TransUNet [Chen et al., 2021], UCTransNet [Wang et al., 2022a] MedSAM [Han et al., ] and SAM [Han et al., ]. Multi-modal approaches include ViLT [Kim et al., 2021], CLIP [Radford et al., 2021], TGANet [Tomar et al., 2022], LAVT [Yang et al., 2022], and LViT [Li et al., 2023b].

4.5 Comparison with State-of-the-Art Methods

The quantitative comparison results are shown in Tab. 1. From this table, we can clearly observe that the proposed method achieves the best performance on all datasets. In particular, the proposed method achieves an 80.24% Dice score and 67.15% IoU on MoNuSeg, which is outperformance in the single modal set. The proposed method maintains a similar performance to DCA on the GlaS dataset with a smaller number of parameters and lower computational complexity.

The proposed method achieves an 80.96% Dice score and 68.12% IoU on MoNuSeg, which respectively increases by 0.81% and 0.65% compared to the suboptimal method LViT-T. Similarly, the proposed method obtains 91.08% Dice score and 84.00% MIoU on GlaS. It outperforms the second best method by 1.06% and 1.32%, respectively. Surprisingly, we observe notable improvements in the MosMedData+ dataset. In particular, the performance respectively increases by 1.45% and 1.70% in Dice score and MIoU. This performance improvement benefits from reducing redundant information and maximizing the discriminative information.

It is worth noting that the proposed method with text annotation achieves better performance than the proposed method without text annotation. In detail, the proposed method with text annotation respectively increased by 0.72% and 0.97% in Dice score and MIoU on MoNuSeg. Similarly, the performance improvements are 0.92% and 1.36% on GlaS, and 0.45% and 0.70% on MosMedData+. It demonstrates that medical text annotation can enhance the robustness of encoded semantics. Moreover, as shown in Fig. 3, the proposed method achieves more higher-quality segmentation results.

4.6 Ablation Study

Effectiveness of Supervised Components

Tab. 2 reports the ablation experimental results of the effectiveness of supervised components. These results indicate that all of these components are effective for the proposed method. It is worth noting that incorporating text annotation significantly improves segmentation performance.

Specifically, the Dice score and MIoU are increased by 1.46% and 1.29% on MoNuSeg, respectively. The performance improvements are 0.52% in Dice score and 0.72% in MIoU on GlaS. The frequency prompt encoder is also effective for the proposed method. It respectively increases by 1.08% and 1.13% in Dice score and MIoU on MoNuSeg. Similarly, performance improvements are 0.25% and 0.97% on GlaS in Dice score and MIoU, respectively. The proposed method with both text annotation and frequency prompt encoder achieves the best performance.

Impact of Different Multi-modal Fusion Methods

Tab. 3 reports the ablation experimental results of the different multi-modal fusion methods. These results indicate that the distribution fusion method is more effective than the addition fusion method. In details, the Dice score and MIoU are increased by 0.47% and 0.65% on MoNuSeg, respectively. The similar trends are observed on GlaS dataset, which respectively increased by 0.87% and 0.97% in Dice score and MIoU. This significant improvement benefits from minimizing redundant information while retaining the modalityspecific information. Consequently, multi-modal features can be effectively fused and more valuable in improving segmentation results.

Impact of Amplitude and Phase Spectrum

Tab. 4 indicates the effectiveness of the frequency adapter, and it can be observed that phase spectrum information is more helpful for spectrum representation compared to phase information. In other words, the proposed frequency adapter can extract more valuable edge and overall information from the phase spectrum. Thus our proposed frequency adapters fully take advantage of amplitude information, which is more related to segmentation boundaries.

Impact of Different Balance Factors

We conduct the ablation study of different λ3 values. In particular, λ3 is respectively set as 0.1, 0.5, 1.0, and 1.5. The experimental results are shown in Tab. 5. It can be observed that the proposed method achieves the best performance when λ3 is set to 1.0 on MoNuSeg. And the best performance is achieved on GlaS when λ3 is 0.1. This is because of the size of the dataset.

On the one hand, the model would remove more information that it considers to be redundant if the size of the dataset is too small in the large lambda. Smaller datasets mean less information. In other words, excessive information compression within modalities may discard useful knowledge while λ3 are further increased in small datasets. On the one hand, the fused feature would contain more redundant information while λ3 are small in GlaS. It results in multi-modal features that could be effective to adopt and have a negative impact

(a)Image

(b) GT

(c) TGANet

(d) LViT

(e) UCTransNet

(f) Ours

Figure 3: Qualitative comparison of the proposed method and other SOTA methods on MoNuSeg, GlaS, and MosMedData+ datasets. The

green boxes highlight regions where the proposed method performs better than the other methods.

Feature Learnable Text Vision

Vision

Text Learnable

(a) MoNuSeg Feature Learnable Text Vision Vision Learnable Text

(b) GlaS

(c) MosMedData+

Figure 4: The visualization results of text, vision, and learnable features on three datasets. Specifically, red ellipses, yellow ellipses, and blue ellipses represent the distributions of text features, learnable features, and visual features, respectively. Different colored points represent the samples drawn from the corresponding distributions.

for decoder. Summarized, this variation could be attributed to differences in multimodal redundant levels in the disparity size of datasets, which conversely affects the optimal weighting to minimize redundant information in the training stage.

Interpretability of Proposed Fusion Method

To validate that the distribution of learnable features have modeled the discriminative information within multimodal data, we conduct random sampling from each distribution on three datasets. Subsequently, we reduce the obtained highdimensional vectors to a two-dimensional plane using t-SNE. As illustrated in Fig. 4, the distribution of learnable features lies between the text and visual feature distribution, indicating the effectiveness of the proposed approach. Specifically, the distribution of learnable features can reduce the redundancies (modality-specific task-irrelated infirmation) multimodal data by only modeling the most discriminative information. Meanwhile, the discriminative information extracted from text and vision feature contains both modality-common and modality-specific task-related information.

5 Conclusion

In this paper, a novel multi-modal method that leverages the text and image features for medical image segmentation. Specifically, we propose a novel multi-modal fusion paradigm, aiming to reduce the redundant information while retaining the modality-specific information. And the discriminative information can be extracted from the fused feature, which significantly improves the performance. Moreover, a novel frequency prompt encoder is designed to leverage highfrequency information for accurately segmenting boundaries. Experimental results demonstrate the effectiveness of the proposed method.

Acknowledgments

This work was supported by National Natural Science Foundation of China (62301621, 62301323) and Shenzhen Science and Technology Program (20231121172359002).