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

Summary

The paper introduces Uni-Med, a medical generalist foundation model designed for multi-task learning across six different medical tasks. The proposed CMoE module leverages a mixture of projection experts to align visual and language embedding spaces effectively. The model demonstrates significant performance improvements across diverse medical tasks, validated through extensive experiments and ablation studies.

Strengths

- The introduction of the CMoE module to address the tug-of-war problem at the connector level is novel and well-executed. - The paper provides thorough ablation studies to validate the effectiveness of the proposed CMoE module under various configurations. - Uni-Med achieves impressive performance with minimal training computational overhead, highlighting its efficiency in handling large-scale multi-modal medical data.

Weaknesses

- 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. - While the interpretation analysis only focus on visual features on different tasks, the analysis of visual features on different medical image modalities should be considered. - 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, such as [1,2]. [1] 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. [2] 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.

Questions

- More detailed ablation studies. - The complete interpretation analysis. - More complete reference work.

Rating

6

Confidence

5

Soundness

3

Presentation

4

Contribution

3

Limitations

N/A

Summary

The paper presents Uni-Med, a medical generalist foundation model designed to perform multiple medical tasks efficiently through multi-task learning. This model introduces a Connector-Mixture-of-Experts (CMoE) module to mitigate the tug-of-war problem in multi-modal, multi-task optimization, which is a common issue in current models. Uni-Med achieves competitive or superior performance across six medical tasks without requiring task-specific fine-tuning.

Strengths

1. Multi-modal multi-task optimization is a complex and important question for large multimodal models, the introduction of the Connector-Mixture-of-Experts (CMoE) module, employs a mixture of projection experts to align visual and language embedding spaces, shows superior performance on multiple tasks. 2. The paper conducts a comprehensive interpretation analysis of the problem from the perspective of gradient optimization and parameter statistics. 3. Extensive experiments demonstrate Uni-Med's effectiveness across multiple tasks and datasets.

Weaknesses

1. The model currently supports only 2D images, whereas most commonly used medical imaging modalities, such as CT and MRI, are in 3D. 2. 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.

Questions

Is the number of projection experts correlated with the number of tasks or the number of image modalities? 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 genertaion.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Summary

The authors propose to build a medical generalist multi-modal foundation model using a novel "connector mixture of experts" module to solve the problem of "multi-task" learning. Their connector-MOE technique introduces a projection and routing module from the visual encoder into the LLM that is explicitly conditioned on the underlying task. The model is tested on a set of medical tasks, and benchmarked against several other multimodal medical models demonstrating improved performance.

Strengths

The explicit task conditioned projection is novel, and integrates an older concept of MoEs into a SOTA MLLM framework. The authors do an excellent job sourcing and assembling a large, multi-task, multi-modal set of medical benchmarks and performing an extensive set of ablations. The paper is benchmarked broadly across multiple tasks - fitting the definition of a foundation model.

Weaknesses

I disagree with the claim that there is limited research on connecting modalities in multi-modal models. This is an area of immense interest and extensive research broadly within the field of machine learning. 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. This is particularly important because, 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 where this issue won't be the case? The results are hard to follow in two very large Tables Table 2 and 3, particularly Table 3 which is the comparison to existing standards. The explanation of these comparisons is quite brief, and not well described in the text despite being the primary comparison for the paper. The paper cites data leakage as being an issue, but isn't dataset shift an issue too? 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? A fairer comparison would be to utilize the same datasets across model architectures.

Questions

In the discussion of MOEs and the paper motivation I think that it could be clarified substantially. 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. Generalist foundation model seems to be redundant? Isn't a "Foundation model" by definition "generalist"? A simple linear projection and purely autoregressive design with visual instruction fine-tuning like LLaVA learns to implicitly condition the projected visual tokens on task and is the appropriate benchmark for this explicit routing framework. I feel like the paper would, in general, benefit from framing it as an investigation into novel connectors with which to build any foundation model - medical or otherwise. 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, scaled up and down, and so forth.

Rating

3

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

No limitations are mentioned, and I think that this is a missed opportunity. There are clearly limitations with regards to the comparisons to other models, datasets involved with these comparisons, and overall size of the involved models utilized in Uni-Med.

Summary

This paper introduces Uni-Med, which applies mixture of experts at the connector level for efficient training toward a unified medical multi-modal foundation model. The contributions of this work include 1.) curation of indexes to quantify tug-of-war problem in multi-modal multi-task modal; 2.) a novel perspective of applying Connector MoE for multi-modal multi-task model, which enables efficient training; 3.) comprehensive ablative studies to evaluate various configurations for different modules; 4) commitment to providing open-source code and weights of the proposed method.

Strengths

1. It is a technically-solid work. The mitigation of the tug-of-war problem is justified from multiple perspectives, including developed indexes, parameter statistics scores, routing weights, and tSNE feature visualization. 2. The proposed framework is evaluated thoroughly in ablation studies. 3. Consistent improvements over existing open-source medical foundation models are observed. 4. The presentation is clear and easy to follow.

Weaknesses

1. 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. 2. 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?

Questions

1. Line 143 is missing an introduction about the soft router.

Rating

8

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Most of the limitations are discussed in the Appendix, while some of them are not addressable.

Official Comment by Reviewer wK4x

Thank you for the rebuttal, which addressed some of my concerns. I have increased my score and am looking forward to reading your revision in future venue.

Response to the rebuttal

I thank the author for their time and effort preparing for the rebuttal. After reading the rebuttal, I have a follow-up question about the comparison with open-source MLLM and Table 3 (correspond to Q2 and A2 above). Other models, such as LLaVA-Med and Med-Flamingo, are never trained on some of the other datasets. For instance, the authors used the LLaVA-Med checkpoints that are not trained on MIMIC-CXR and also did not fine-tune it on the MIMIC-CXR dataset. Rather, that checkpoint was directly applied to the test set of MIMIC-CXR, and the results were reported in Table 3. Is this understanding correct?