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

📄 arXiv: 2409.17508v2 📥 PDF

作者: Xun Zhu, Ying Hu, Fanbin Mo, Miao Li, Ji Wu

分类: cs.CV, cs.AI, cs.LG

发布日期: 2024-09-26 (更新: 2024-11-01)

🔗 代码/项目: GITHUB

💡 一句话要点

Uni-Med：通过Connector-MoE实现多任务学习的统一医学通用基础模型

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 医学多模态学习 多任务学习 混合专家模型 连接器模块 医学影像 大型语言模型 Uni-Med

📋 核心要点

  1. 现有医学多模态大语言模型在多任务学习中存在“拔河”问题，即不同任务之间相互干扰，导致整体性能难以提升。
  2. Uni-Med提出了一种新颖的连接器混合专家(CMoE)模块，通过在连接器处使用混合专家路由，有效缓解了多任务干扰。
  3. 实验表明，Uni-Med在六种医学任务上取得了显著的性能提升，平均提升高达8%，并在多种任务上超越了现有最佳模型。

📝 摘要（中文）

多模态大型语言模型(MLLM)在各种视觉和语言任务中表现出令人印象深刻的通用能力。然而，在医学领域构建用于多任务学习的统一MLLM仍然是一个棘手的挑战。为了缓解MLLM中多模态多任务优化中的“拔河”问题，最近的进展主要集中在改进LLM组件，而忽略了连接模态之间差距的连接器。本文介绍Uni-Med，一种新型医学通用基础模型，它由通用视觉特征提取模块、连接器混合专家(CMoE)模块和LLM组成。受益于所提出的CMoE，它利用精心设计的路由器和连接器处的混合投影专家，Uni-Med实现了对“拔河”问题的有效解决方案，并且可以执行六种不同的医学任务，包括问答、视觉问答、报告生成、指代表达式理解、指代表达式生成和图像分类。据我们所知，Uni-Med是首次尝试在MLLM的连接器处解决多任务干扰。大量的消融实验验证了在任何配置下引入CMoE的有效性，平均性能提升高达8%。我们进一步从梯度优化和参数统计的角度提供了对“拔河”问题的解释分析。与以往最先进的医学MLLM相比，Uni-Med在不同的任务上实现了有竞争力的或优越的评估指标。

🔬 方法详解

问题定义：现有的医学多模态大语言模型(MLLM)在处理多任务学习时，由于不同任务之间存在梯度冲突和参数竞争，导致模型在优化过程中出现“拔河”现象，即提升一个任务的性能可能会牺牲另一个任务的性能。这种现象严重阻碍了医学MLLM在多任务场景下的应用。

核心思路：Uni-Med的核心思路是在连接视觉特征提取模块和语言模型(LLM)的连接器(Connector)处引入混合专家(Mixture-of-Experts, MoE)机制。通过MoE，不同的任务可以选择不同的专家进行处理，从而避免了参数竞争，缓解了“拔河”问题。这种设计允许模型为每个任务学习更专业的表示，从而提高整体性能。

技术框架：Uni-Med的整体架构包括三个主要模块：1) 通用视觉特征提取模块：用于提取医学图像的视觉特征。2) 连接器混合专家(CMoE)模块：这是Uni-Med的核心创新，包含一个路由器(Router)和多个投影专家(Projection Experts)。路由器根据输入选择合适的专家组合，专家负责将视觉特征投影到LLM的输入空间。3) LLM：用于处理投影后的视觉特征和文本信息，完成各种医学任务。

关键创新：Uni-Med最重要的技术创新点在于CMoE模块。传统的MLLM通常使用单一的线性层或MLP作为连接器，所有任务共享相同的参数。而CMoE模块为每个任务提供了不同的专家组合，从而实现了任务特定的特征转换。这种设计能够有效缓解多任务干扰，提高模型的泛化能力。

关键设计：CMoE模块的关键设计包括：1) 路由器：使用一个可学习的门控网络，根据输入特征计算每个专家的权重。2) 投影专家：每个专家都是一个线性层或MLP，负责将视觉特征投影到LLM的输入空间。3) 损失函数：除了任务相关的损失函数外，还可以引入辅助损失函数来平衡不同专家的使用频率，避免某些专家被过度使用。

📊 实验亮点

Uni-Med在六种不同的医学任务上进行了评估，包括问答、视觉问答、报告生成、指代表达式理解、指代表达式生成和图像分类。实验结果表明，Uni-Med在多个任务上取得了显著的性能提升，平均性能提升高达8%。与之前的state-of-the-art医学MLLM相比，Uni-Med在不同的任务上实现了有竞争力的或优越的评估指标，验证了CMoE的有效性。

🎯 应用场景

Uni-Med具有广泛的应用前景，可以应用于医学影像诊断、医学报告生成、医学知识问答等多个领域。通过构建一个统一的医学通用基础模型，可以降低开发成本，提高模型的可复用性，并促进医学人工智能的发展。未来，Uni-Med有望成为医生辅助诊断的重要工具，提高医疗效率和质量。

📄 摘要（原文）

Multi-modal large language models (MLLMs) have shown impressive capabilities as a general-purpose interface for various visual and linguistic tasks. However, building a unified MLLM for multi-task learning in the medical field remains a thorny challenge. To mitigate the tug-of-war problem of multi-modal multi-task optimization in MLLMs, recent advances primarily focus on improving the LLM components, while neglecting the connector that bridges the gap between modalities. In this paper, we introduce Uni-Med, a novel medical generalist foundation model which consists of a universal visual feature extraction module, a connector mixture-of-experts (CMoE) module, and an LLM. Benefiting from the proposed CMoE that leverages a well-designed router with a mixture of projection experts at the connector, Uni-Med achieves efficient solution to the tug-of-war problem and can perform six different medical tasks including question answering, visual question answering, report generation, referring expression comprehension, referring expression generation and image classification. To the best of our knowledge, Uni-Med is the first effort to tackle multi-task interference at the connector in MLLMs. Extensive ablation experiments validate the effectiveness of introducing CMoE under any configuration, with up to an average 8% performance gains. We further provide interpretation analysis of the tug-of-war problem from the perspective of gradient optimization and parameter statistics. Compared to previous state-of-the-art medical MLLMs, Uni-Med achieves competitive or superior evaluation metrics on diverse tasks. Code and resources are available at https://github.com/tsinghua-msiip/Uni-Med.