Uni-Med: A Unified Medical Generalist Foundation Model For Multi-Task Learning Via Connector-MoE

Thanks for your insightful comments!

Q1: The ablation studies show that certain configurations (e.g., using a high number of projection experts) might lead to overfitting. This aspect could be discussed in more detail, including strategies to mitigate overfitting.

A1: Thank you for your suggestion.

• Choosing the optimal parameter configuration in multi-modal multi-task scenarios has always been a concern, which aims to achieve a balance between performance and computational efficiency [1-2]. It is a challenging study as the complexity of the scenario can end up overfitting to simpler modalities or tasks or underfitting complex ones.
• As shown in Table 2 (c), in the experiment exploring the key parameter of the number of experts, increasing the number of experts still brings performance gains on some datasets, but the average gain tends to stabilize across all tasks and datasets.
• To our knowledge, recent studies [3-5] have also conducted ablation experiments on these key parameters. However, there are different observations in different scenarios (number of modalities/tasks/data volumes). Therefore, we believe that selecting the optimal parameter settings using the development set is a simple and effective method to achieve the balance between performance and computational efficiency.

Thanks to your feedback, we will provide more detailed discussions in the revised version. We will continue to focus on this in our future research.

Q2: While the interpretation analysis only focuses on visual features on different tasks, the analysis of visual features on different medical image modalities should be considered.

A2: Thank you for your suggestion. We also use t-SNE method to visualize the distribution of visual features on medical image modalities and provide the results in Figure Re.1.

• Specifically, we first observe the visual feature distribution of different modalities under the same task in Figure Re.1 (a-c). We find that the feature distributions of CT and MRI modalities in the REG task have good discriminability after passing through the frozen visual encoder. After passing through the connector, the improvement in Silhouette score (from 0.3049 to 0.3335) is relatively limited.
• In addition, we select 100 samples from each of the 8 modalities and observe their visual feature distributions after passing through different visual encoders in Figure Re.1 (d-f). It can still be observed that the majority of modality distributions are ordered and tightly packed.

Based on the above observations, the distinction of medical image modalities is achieved effectively through the visual encoder, while task differentiation requires the well-designed connector. We will add the analysis of visual features on different medical image modalities in the revised version.

Q3: It could be better to add more work on evaluation of medical vision-language models in the section of related work to make sure that the relevant work is fully discussed.

A3: Thank you for your suggestion. We attach great importance to your suggestions and will add cutting-edge developments [6-7] on evaluation of medical vision-language models to the related work section in the revised version.

References

[1] Liu Q, Wu X, Zhao X, et al. When MOE Meets LLMs: Parameter Efficient Fine-tuning for Multi-task Medical Applications[C]//Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval. 2024: 1104-1114.

[2] Chen S, Jie Z, Ma L. Llava-mole: Sparse mixture of lora experts for mitigating data conflicts in instruction finetuning mllms[J]. arXiv preprint arXiv:2401.16160, 2024.

[3] Chen T, Zhang Z, JAISWAL A K, et al. Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers[C]//The Eleventh International Conference on Learning Representations.

[4] Dou S, Zhou E, Liu Y, et al. Loramoe: Revolutionizing mixture of experts for maintaining world knowledge in language model alignment[J]. arXiv preprint arXiv:2312.09979, 2023, 4(7).

[5] Gou Y, Liu Z, Chen K, et al. Mixture of cluster-conditional lora experts for vision-language instruction tuning[J]. arXiv preprint arXiv:2312.12379, 2023.

[6] Xia P, Chen Z, Tian J, et al. CARES: A Comprehensive Benchmark of Trustworthiness in Medical Vision Language Models[J]. arXiv preprint arXiv:2406.06007, 2024.

[7] Yan Q, He X, Yue X, et al. Worse than Random? An Embarrassingly Simple Probing Evaluation of Large Multimodal Models in Medical VQA[J]. arXiv preprint arXiv:2405.20421, 2024.

Thanks for your insightful comments!

Q1: The model currently supports only 2D images, whereas most commonly used medical imaging modalities, such as CT and MRI, are in 3D.

A1: Thank you for your advice. Same as most medical MLLMs, we input 2D slices and corresponding questions for 3D images such as CT and MRI. We acknowledge that Uni-Med has certain limitations in handling genuine 3D medical image inputs. The primary challenge lies in the need for different visual encoders to process 3D images effectively [1]. Replacing the visual encoder to handle 3D images would compromise our ability to process 2D image datasets effectively. This remains an area for future research.

Q2: For the report generation task, the evaluation should include metrics like RadGraph Score and RadCliQ, as BLEU and ROUGE cannot fully assess the semantic accuracy.

A2: Thank you for your suggestion. Firstly, we learne the concepts of the two metrics mentioned above:

• RadGraph-based metrics. The RadGraph model [2] parses radiology reports into graphs containing clinical entities and relations between them. The RadGraph F1 metric computes the overlap in entities and relations separately, then reports their average.

• RadCliQ. Radiology Report Clinical Quality (RadCliQ) is a composite metric that integrates RadGraph F1 and BLEU score in a linear regression model to predict the total number of errors in a report [3].

Secondly, we calculate and report RadGraph entity F1, RadGraph relation F1, RadCliQ-v0 and RadCliQ-v1 on MIMIC-CXR dataset in Table Re.1, using the code released by Yu et al. [3]. The improvement of RadGraph-based metrics and the decrease of RadCliQ both indicate that Uni-Med achieves better semantic accuracy in the report generation task.

Q3: Is the number of projection experts correlated with the number of tasks or the number of image modalities?

A3: Thank you for your question. We hold the opinion that both are important, and the answer to this question needs to be based on the actual situation. We analyze the following three scenarios:

• Single modal, multi-task. A same image may need to complete different tasks, and the number of experts should be related to the number of tasks.

• Multi-modal, single task. The setting of the number of experts should consider the number of modalities.

• Multi-modal, multi-task. In this scenario, further analysis of the data is required. Taking Uni-Med as an example, when we try to visualize visual features separately by task and modality, we find that the distribution of features by task is more chaotic, while the distribution of features by modality is more orderly (Detailed in Q4 & A4). Therefore, CMoE needs to consider task information more, and the number of experts is more related to the task.

Q4: Figure 5 shows that visual features of the same task are more tightly distributed. How would t-SNE behave for different modalities? Why are visual features related to tasks? For a single image, it can be used for both classification and report generation.

A4: Thank you for your question.

• In Figure 5, we visualize the distribution of visual features by tasks. It can be clearly observed that the distribution of features by task is chaotic in Figure 5 (a), which means that there is no obvious discrimination between different tasks after passing through the frozen visual encoder. Visual features of the same task are more tightly distributed after CMoE in Figure 5 (c) than MLP in Figure 5 (b).

• We use t-SNE to visualize the distribution of visual features by modalities and provide the results in Figure Re.1. Specifically, we first observe the visual feature distribution of different modalities under the same task in Figure Re.1 (a-c). We find that the feature distributions of CT and MRI modalities in the REG task have good discriminability after passing through the frozen visual encoder. After passing through the connector, the improvement in Silhouette score (from 0.3049 to 0.3335) is relatively limited. In addition, we select 100 samples from each of the 8 modalities and observe their visual feature distributions after passing through different visual encoders in Figure Re.1 (d-f). It can still be observed that the majority of modality distributions are ordered and tightly packed.

• The above findings also provide a new perspective for explaining the explicit task conditioned projection in CMoE. When aligning visual and language embedding spaces through the connector in Uni-Med scenario, task information is more difficult to distinguish than modality information.

• As described in the question, for a single image, it can be used for different tasks. If the visual features are not related to the task, then the tokens of the same image input into LLM are exactly the same (through linear, MLP, and token-level CMOE in Table 2). In this case, achieving multitasking relies entirely on the capabilities of LLM. On the contrary, we assume that different tasks require attention to different image features. After passing through the connector, the features of the same image for different tasks have adaptively changed, which can alleviate the negative impact of the tug of war problem in multi-task learning on LLM. The significant improvement in experimental results confirms the latter assumption.

References

[1] Bai F, Du Y, Huang T, et al. M3d: Advancing 3d medical image analysis with multi-modal large language models[J]. arxiv preprint arxiv:2404.00578, 2024.

[2] S Jain, A Agrawal, A Saporta, et al. RadGraph: Extracting clinical entities and relations from radiology reports. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks, volume 1, December 2021.

[3] Yu F, Endo M, Krishnan R, et al. Evaluating progress in automatic chest x-ray radiology report generation[J]. Patterns, 2023, 4(9).

Thank you for the response. The authors have adequately addressed my primary concerns, and I have no further questions. I will maintain my previous rating.

Thanks for your insightful comments!

Q1: I disagree with the claim that there is limited research on connecting modalities in multi-modal models.

A1: There is no doubt about that researches about connecting modalities in multi-modal models is popular and extensive. As mentioned in lines 4-6, 50-52, and 104-105, our viewpoint is that research focus on the field of connectors in MLLMs. We will highlight this scope of our research in the revised version to reduce your misunderstanding.

Q2: I don't think that it is a unique medical issue, and the benchmarking of a novel architecture would be better served using more common datasets to the ML community. As the authors note, the medical datasets and models involved were hard for them to control for data leakage. Why not compare the routing technique against non-medical models on non-medical datasets?

A2: Thank you for your suggestion.

• First, we agree that this is not a unique medical issue. However, when we construct medical MLLMs based on multi-modal and multi-task scenarios, we observe the tug of war problem at the connector level within standard MLLM architectures. This is the background of our proposal of Uni-Med.

• Second, it is of great significance to research and develop the medical generalist foundation model. Providing a superior solution to the tug-of-war problem which is particularly serious due to the diversity of image modalities and tasks in the medical field becomes our basic motivation. We believe that "solving problem in prominent fields" and "exploring generalization in general fields" are equally important contributions to the machine learning community.

• Third, not only in the medical field, but also in general field, model evaluation is plagued by data leakage issue. In our work, the data leakage issue is only observed on the MPx-Single using the model provided by RadFM. We ensure that all experiments of Uni-Med are free of data leakage issue and the results are reliable.

• Fourth, we conduct preliminary exploration of generalization of our method in general field. We fully follow the training strategy of LLaVA-1.5 and report metrics on 9 benchmarks with/without CMoE in Table Re.2. The results show that the introduction of CMoE brings significant improvements on all benchmarks.

Q3: The results are hard to follow in two very large Table 2 and 3. The explanation of these comparisons is quite brief, and not well described in the text despite being the primary comparison for the paper.

A3: Thank you for your suggestion. Due to the page limit of the paper, we acknowledge the inadequacies in our writing regarding the explanation of experimental results. We will add these details in the revised version. We hope to make the readers feel that the presentation is clear and easy to understand.

Q4: If I understand correctly, Uni-Med is being tested on held-out samples from datasets it was trained on, while the other LLMs are being tested on a mixture of data that in some cases was even in their training datasets?

A4: Your understanding is accurate. For other medical MLLMs, we use readily available model checkpoints for testing. We clearly know that a completely fair comparison would be to utilize the same datasets across model architectures. But we have already achieved relative fairness:

• We use the official test set split for all datasets, except for Slake-VQA, as we utilize it to build data for other tasks. In this case, there is an unfair comparison between Slake-VQA, but it's actually unfair to Uni-Med because part of the test data is used for training in other models.
• As for the fact that we report RadFM's data leakage in the MPx-Single dataset (RadFM opens source this dataset and provides split), we conduct testing strictly according to the split, which only indicates that RadFM's model checkpoint may not have been trained according to this split.
• None of the ablation experiments has data leakage issues, and the effectiveness of CMoE in any configuration is reliable.

Q5: Line 97-105 clearly disambiguates the usage of the term MoE, and it might be helpful to do this sooner in the introduction to help clarify this work for readers.

A5: Thank you for your suggestion. As mentioned in line 50-53, current research to mitigate the tug-of-war problem mainly tailors the MoE approach to the language model components, overlooking the potential benefits of exploring and enhancing the connector. We will provide a clearer presentation of our motivation and the usage of MoE in the introduction section.

Q6: LLaVA learns to implicitly condition the projected visual tokens on task and is the appropriate benchmark for this explicit routing framework.

A6: In fact, all experiments with linear and MLP connectors in Table 2 have model architectures consistent with LLaVA and LLaVA-1.5, respectively. In addition, CMoE with token-level router strategy is also an implicitly condition projection architecture. We hope these explanations are helpful for you to understand our benchmark settings.

Q7: The paper would benefit from framing it as an investigation into novel connectors with which to build any foundation model. Framing it as a novel medical foundation model focuses on the wrong thing, and makes me wonder why it isn't tested on a broader range of medical tasks and datasets.

A7: We elaborate on the background and motivation for why we chose the medical field for research in the first and second point in A2. With current computing resources, we have trained and tested on as wide a range of medical tasks as possible. Compared to the existing medical MLLMs, Uni-Med have added more diverse tasks and datasets. We do indeed look forward to having more data and a wider variety of tasks to validate our method. We will continue to focus on this in our future research.

Thanks for your insightful comments!

Q1: While the reviewer appreciates the acknowledgment of several limitations of this work, it would be better to mention them in the main text, especially limitation no. 5 in lines 662-663. If the space is not sufficient, at least they should be briefly mentioned in the main text, and a reference to detailed limitations should be provided.

A1: Thank you for your suggestion. We have realized that the limitations of our work and related references should be mentioned in the main text. We will add a limitation section in the main text of the revised version.

Q2: Confusion about the training/fine-tuning details: for models presented in Table 3, are they individually fine-tuned for each dataset based on the split introduced in the appendix? How is the fine-tuning implemented? Is it end-to-end or LoRA fine-tuning? For the Uni-Med, the reviewer understands that it is trained on all datasets appearing in Table 3, but there is no individual fine-tuning. Is this understanding correct?

A2: Thank you for your questions. We will elaborate on the implementation details of the models in Table 3 to reduce your confusion.

• The understanding of Uni-Med is basically accurate. Uni-Med achieves joint training on 6 six distinct medical tasks and 12 datasets, requiring only one-stage training on a single A800 GPU and no task/dataset fine-tuning. It strictly follows the dataset split introduced in the appendix.

• For model comparison, we use readily available model checkpoints for testing. The details are as follows: (1) About training data and model type. Except for Med-Flamingo, the raw training data of the comparison models all contain some of the datasets in Table 3. For example, LLaVA-Med uses full parameter fine-tuning on the Slake-VQA and Path-VQA, respectively, which means it offers different dataset-specific model checkpoints. XrayGPT is task-specific model, and it traines on MIMIC-CXR. RadFM is a generalist foundation model and its training data includes Slake-VQA, MIMIC-CXR and MPx-Single. (2) About datasets split. We use the official test set split for all datasets, except for Slake-VQA, as we utilize it to build data for other tasks. In this case, there is an unfair comparison between Slake-VQA, but it's actually unfair to Uni-Med because part of the test data is used for training in other models.

• A completely fair comparison across different model architectures is to use the same dataset split for training and testing. The medical MLLMs used for comparison all followe the standard architecture consists of a vision encoder, a connector (e.g. XrayGPT: linear layer; LLaVA-Med: MLP), and an LMM. From this perspective, the experiments of the connector using linear layer or MLP in Table 2, to some extent, represent the results of XrayGPT and LLaVA-Med frameworks that are fully consistent with the Uni-Med training and fine-tuning strategy, respectively.

Thank you for recognizing our extensive experiments and analyses. To reduce the confusion of readers, we will add more implementation details in the revised version.

Q3: Line 143 is missing an introduction about the soft router.

A3: Thank you for your meticulous review and reminder. The soft router receives input tokens and calculates the routing weights for each expert. We will add the missing introduction about the soft router in the revised version.