When MOE Meets LLMs: Parameter Efficient Fine-tuning for Multi-task Medical Applications

Qidong Liu,Xian Wu,Xiangyu Zhao,Yuanshao Zhu,Derong Xu,Feng Tian,Yefeng Zheng

2024-05-31

Abstract:The recent surge in Large Language Models (LLMs) has garnered significant attention across numerous fields. Fine-tuning is often required to fit general LLMs for a specific domain, like the web-based healthcare system. However, two problems arise during fine-tuning LLMs for medical applications. One is the task variety problem, which involves distinct tasks in real-world medical scenarios. The variety often leads to sub-optimal fine-tuning for data imbalance and seesaw problems. Besides, the large amount of parameters in LLMs leads to huge time and computation consumption by fine-tuning. To address these two problems, we propose a novel parameter efficient fine-tuning framework for multi-task medical applications, dubbed as MOELoRA. The designed framework aims to absorb both the benefits of mixture-of-expert (MOE) for multi-task learning and low-rank adaptation (LoRA) for parameter efficient fine-tuning. For unifying MOE and LoRA, we devise multiple experts as the trainable parameters, where each expert consists of a pair of low-rank matrices to retain the small size of trainable parameters. Then, a task-motivated gate function for all MOELoRA layers is proposed, which can control the contributions of each expert and produce distinct parameters for various tasks. We conduct experiments on a multi-task medical dataset, indicating MOELoRA outperforms the existing parameter efficient fine-tuning methods. The code is available online.

Computation and Language,Artificial Intelligence

What problem does this paper attempt to address?

This paper attempts to solve two main problems encountered when multi - task fine - tuning large language models (LLMs) in medical applications:

1. **Task diversity problem**: Medical scenarios in the real world involve a variety of different tasks, such as doctor recommendation, diagnosis prediction, drug recommendation, medical named - entity recognition, clinical report generation, etc. The inputs and outputs of these tasks vary greatly, making it difficult to use a single model to optimize and fine - tune all tasks. In addition, due to data imbalance and the seesaw effect, although it is feasible to directly train models for each task separately, it requires a great deal of professional knowledge and labor, and lacks the learning of task - shared knowledge.

2. **High fine - tuning cost**: The number of parameters in large language models is huge, which leads to extremely high time and computational costs for full - parameter fine - tuning. Therefore, for large language models with more than 10 billion parameters, they are usually closed - source, and the cost of fine - tuning is extremely high. This has prompted researchers to look for more efficient parameter fine - tuning methods.

To solve the above problems, the paper proposes a new parameter - efficient fine - tuning framework - MOELoRA, which combines the advantages of Mixture of Experts (MOE) and Low - Rank Adaptation (LoRA). Specifically, MOELoRA learns task - specific knowledge by designing multiple experts and controls the contribution weight of each expert through a task - driven gate function, thereby generating independent parameter sets for different tasks. This method can not only effectively deal with the task diversity problem, but also significantly reduces the cost of fine - tuning.