PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2310.18339v2cs.CLcs.AI

TL;DR

Medical LLM fine-tuning must accommodate diverse tasks and the high cost of updating large models. MOELoRA addresses this by combining low-rank experts with a task-motivated gate, and experiments on a multi-task Chinese medical dataset report superiority over existing parameter-efficient methods.

  • Problem

    Medical fine-tuning must address heterogeneous clinical tasks and the prohibitive time and computational cost of updating large LLMs.

  • Method

    MOELoRA combines LoRA-style low-rank trainable experts with a task-motivated gate to produce distinct parameters for different medical tasks.

  • Results

    MOELoRA consistently outperforms the other evaluated methods on average scores across all tasks.

  • Takeaways & Limitations

    MOELoRA provides a unified parameter-efficient framework for multi-task medical fine-tuning.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Medical LLM fine-tuning must handle diverse clinical tasks while avoiding the prohibitive cost of updating all parameters. MOELoRA combines mixture-of-experts with low-rank adaptation to learn task-specific parameters efficiently.

  • General-purpose LLMs often lack specialized medical knowledge, motivating fine-tuning for medical applications.
  • Medical fine-tuning faces task variety because clinical applications span heterogeneous tasks with different inputs and outputs.Examples include recommendation, diagnosis prediction, named entity recognition, and report generation.
  • Full-parameter fine-tuning of large LLMs incurs prohibitive time and computational expenses.
  • Mixture-of-experts uses separate experts and a gate to learn task-shared and task-specific knowledge, but existing frameworks are mainly designed for full fine-tuning.
  • MOELoRA combines LoRA-style parameter efficiency with multiple trainable experts and a task-motivated gate for distinct task parameters.The framework is presented as a unified multi-task parameter-efficient fine-tuning approach.

2 PRELIMINARY

Medical tasks are reformulated as linguistic input-output problems so LLMs can process them uniformly. The resulting datasets support joint conditional-language-modeling fine-tuning across multiple tasks.

  • LLMs require medical tasks to be reformulated because their inputs and outputs are typically linguistic.
  • Instruction templates are added to medical text to guide LLMs through task-specific processing.For medical NER, the template asks the model to recognize medical named entities in a sentence.
  • Task outputs are formatted as linguistic text rather than plain targets, enabling LLM generation for medical tasks.
  • Multi-task fine-tuning represents medical tasks as T = {T1, . . . , Tj, . . . , TM} with corresponding datasets Dj.
  • Because task data share a standardized linguistic format, all training instances use conditional language modeling in one joint optimization process.

3 METHOD

MOELoRA combines LoRA’s parameter-efficient updates with multiple experts and a task-motivated gate to learn task-specific parameters for multi-task medical applications. Its fine-tuning and inference process freezes the pretrained LLM, trains expert low-rank matrices, recovers task-specific parameters, and applies them for prediction.

  • MOELoRA: MOELoRA integrates multiple low-rank experts into LLM dense layers, extending LoRA to capture diverse knowledge across medical tasks.Each expert is constructed from two decomposed low-rank matrices, while pretrained LLM parameters remain frozen.
  • MOELoRA: MOELoRA retains LoRA’s parameter efficiency because its experts use low-rank trainable matrices rather than full dense-layer updates.The framework is designed to keep the trainable parameter size compact while providing separate expert components.
  • Task-Motivated Gate Function: The task-motivated gate takes task identity as input and assigns expert contribution weights to generate distinct updated parameters for each task.The gate uses task embeddings, a linear transformation, and softmax normalization; a sparse variant is also explored.
  • Task-Motivated Gate Function: Task-conditioned gating enables task-level parameter recovery, whereas input-conditioned gating would produce sample-specific parameters that are not recoverable per task.The authors associate recoverable task parameters with task customization and inference efficiency.
  • Fine-tuning and Inference: The training procedure configures MOELoRA layers and experts, freezes pretrained LLM parameters, and optimizes the combined model using forward computation and a loss function.During inference, contribution weights and task-specific parameters are recovered before the corresponding LLM parameters are applied for prediction.
  • Fine-tuning and Inference: The inference procedure calculates task-specific expert weights, recovers MOELoRA parameters for each task, and applies those parameters to prediction.These operations are summarized in Algorithm 1.

4 EXPERIMENT

The experiment evaluates MOELoRA on a multi-task Chinese medical benchmark using several baseline groups, task-specific metrics, and repeated runs. It examines performance, architecture and gate effects, expert settings, and fine-tuning or inference efficiency.

  • The evaluation addresses five questions covering comparative performance, architecture and gate effects, training strategies, expert and rank choices, efficiency, and expert specialization.
  • Dataset: PromptCBLUE contains 16 prompted Chinese medical tasks, but experiments randomly select 8 because of computational constraints.The competition test set is unreleased, so the development set serves as test data and part of the training set becomes validation data.
  • Baselines: Baselines span unfine-tuned LLMs, parameter-efficient fine-tuning, model editing, and cross-task generalization methods.The compared methods include ChatGPT, Huatuo, LoRA variants, Task-Arithmetic, LoRAHub, and MoLoRA.
  • Implementation: Experiments use ChatGLM-6B on Tesla V100 32G GPUs, with LoRA rank r=16, batch size 64, and up to 8,000 training steps.MOELoRA uses 8 experts, and the sparse gate searches K from 1 to 7, selecting K=2 as optimal.
  • Evaluation: Metrics are task-specific: Micro-F1 evaluates high-cardinality classification or entity tasks, Macro-F1 evaluates lower-cardinality tasks, and Rouge-L evaluates generation.Average performance across tasks is also reported, with scores averaged over three runs using random seeds {42, 43, 44}.

4.2 Overall Performance (RQ1)

MOELoRA(D) achieves the strongest average performance across tasks among the evaluated methods. The comparisons also indicate that shared knowledge, task-motivated gating, and dense expert usage are important in this multi-task setting.

  • MOELoRA(D) consistently outperforms all competing methods on average across tasks.The comparison is based on the comprehensive results in Table 2.
  • Baseline comparisons: Unfine-tuned LLMs significantly lag behind fine-tuned methods, highlighting the importance of incorporating task-specific medical knowledge.
  • Parameter-efficient fine-tuning: LoRA-based methods outperform P-Tuning, while LoRA (Full) exceeds LoRA (Single), supporting shared knowledge across tasks.LoRA (Full+TP) slightly underperforms LoRA (Full), possibly because task prompts lengthen inputs and truncate informative words.
  • Cross-task generalization: Cross-task generalization methods underperform relatively because they require vast task data that conflicts with the multi-task setting.The experiments evaluate only 8 tasks, which may contribute to their relative underperformance.
  • Gate design: MOELoRA outperforms MoLoRA, while dense gating surpasses sparse gating on most tasks.Sparse gating performs best on two tasks, whereas dense gating uses all experts and benefits shared medical knowledge.

4.3 Ablation Study (RQ2)

The ablation study evaluates MOELoRA’s mixture-of-experts architecture, gate function, gate configuration, and training strategies. Removing or altering these components generally reduces performance or efficiency relative to the full design.

  • Component ablations: Removing MOE, thereby reverting to LoRA (Full), produces inferior performance, underscoring the importance of the expert architecture.
  • Component ablations: Removing the gate and using uniform expert weights also reduces performance, supporting the effectiveness of task-motivated expert weighting.
  • Component ablations: Using a separate gate for each MOELoRA layer is comparable on several tasks but is slightly worse overall because of over-parameterization.
  • Training strategies: Batching samples by task or randomly selecting a task per batch both degrade performance relative to the training strategy used by MOELoRA.
  • Robustness analysis: With attention-only LoRA layers, MOELoRA(D)-QKV outperforms LoRA (Full)-QKV on most tasks, while applying more MOELoRA layers further improves performance.

4.4 Hyper-parameter Analysis (RQ3)

Hyper-parameter analysis shows that more experts improve performance up to a point, while increasing LoRA rank improves performance at the cost of more trainable parameters. The authors therefore select a practical rank setting that balances efficiency and performance.

  • Expert number: With overall LoRA rank r fixed at 16, performance improves as expert number N rises from 0 to 8 but declines marginally at N=16.The decline is attributed to each expert receiving a smaller LoRA rank, weakening its low-rank matrices.
  • LoRA rank: Increasing LoRA rank r consistently improves performance but proportionately increases the number of trainable parameters.
  • LoRA rank: The analysis identifies r=16 as a practical choice for balancing efficiency and performance.

4.5 Efficiency Analysis (RQ4)

MOELoRA achieves training and inference efficiency comparable to LoRA (Full) while retaining task-motivated gating. Its gate design avoids the additional inference cost incurred by sample-dependent expert weighting.

  • Efficiency comparison: MOELoRA trains no more than 0.48% of the LLM parameters while achieving training and inference efficiency comparable to LoRA (Full).
  • Efficiency comparison: MoLoRA and MOELoRA(M) require more trainable parameters because they add gates to each trainable low-rank layer.
  • Inference efficiency: All models except MoLoRA have the same inference latency because MoLoRA cannot recover fine-tuned parameters when expert weights vary across samples.
  • Inference efficiency: The task-motivated gate avoids efficiency degradation by preventing the additional forward computation required by MoLoRA.

4.6 Case Study (RQ5)

Expert contributions vary substantially across medical tasks, indicating task-specific specialization alongside shared knowledge. Similar expert-weight patterns for CHIP-CDN and KUAKE-QIC suggest that related tasks can use shared knowledge.

  • Expert specialization: Expert weights differ substantially across tasks, indicating that different experts specialize in distinct facets of medical knowledge.
  • Shared task knowledge: CHIP-CDN and KUAKE-QIC have largely similar expert-weight patterns, except for experts 3 and 4.
  • Shared task knowledge: The similar weights for CHIP-CDN and KUAKE-QIC suggest that MOELoRA harnesses shared knowledge for related tasks.

5 RELATED WORKS

Prior work applies LLMs to medical domains but often focuses on dialogue or individual tasks, while multi-task coverage and efficient adaptation remain limited. PEFT methods reduce tuning costs, yet LoRA performs poorly for multi-task medical applications.

  • Medical LLM research includes benchmarks, prompting strategies, medication recommendation, and model editing, but commonly centers on dialogue or single tasks.
  • Multi-task medical applications remain underexplored because prior methods often neglect several important tasks simultaneously and require substantial fine-tuning costs.
  • PEFT methods reduce trainable parameters and computational complexity through adapters, virtual tokens, or low-rank matrices.
  • LoRA can match full fine-tuning without additional inference computation, but performs inferiorly for multi-task medical applications.

6 CONCLUSION

The paper presents MOELoRA as a multi-task parameter-efficient framework for LLM-driven medical applications. It uses low-rank experts and a task-motivated gate, with experiments verifying its effectiveness on a multi-task Chinese medical dataset.

  • MOELoRA combines low-rank experts with a task-motivated gate to learn task-related knowledge efficiently and produce distinct fine-tuned parameters.
  • The framework is evaluated on a multi-task Chinese medical dataset, where comprehensive experiments verify its effectiveness.
  • Future work will investigate combining explicit medical knowledge, such as knowledge graphs, with LLM fine-tuning.
Loading 2310.18339v2…