SA-MDKIF: A Scalable and Adaptable Medical Domain Knowledge Injection Framework for Large Language Models

TL;DR

This work addresses the gap in medical-domain reasoning by injecting specialized knowledge into general-purpose LLMs through a two-stage framework (skill training and skill adaptation). It defines 12 medical skills trained with AdaLoRA and dynamically fuses them with a frozen LLM via a differentiable skill router, enabling task-conditioned knowledge injection with minimal parameter updates. Empirical results across 9 medical tasks, including unseen tasks and few-shot settings, show consistent 10–30% improvements over baselines and state-of-the-art performance on several tasks, while using only a fraction of the parameters required by full fine-tuning. The approach offers practical impact by delivering scalable, adaptable medical LLMs and has potential to extend to other domains beyond medicine, aided by explicit MoE-like skill routing and efficient parameterization such as AdaLoRA.

Abstract

Recent advances in large language models (LLMs) have demonstrated exceptional performance in various natural language processing (NLP) tasks. However, their effective application in the medical domain is hampered by a lack of medical domain knowledge. In this study, we present SA-MDKIF, a scalable and adaptable framework that aims to inject medical knowledge into general-purpose LLMs through instruction tuning, thereby enabling adaptability for various downstream tasks. SA-MDKIF consists of two stages: skill training and skill adaptation. In the first stage, we define 12 basic medical skills and use AdaLoRA to train these skills based on uniformly formatted instructional datasets that we have constructed. In the next stage, we train the skill router using task-specific downstream data and use this router to integrate the acquired skills with LLMs during inference. Experimental results on 9 different medical tasks show that SA-MDKIF improves performance by 10-20% compared to the original LLMs. Notably, this improvement is particularly pronounced for unseen medical tasks, showing an improvement of up to 30%.

Figures (6)

• Figure 1: The goal of this work is to inject medical knowledge into the general-purpose LLM for its adaptation to different downstream medical tasks.
• Figure 2: The framework of SA-MDKIF. It consists of two stages: skill training and skill adaptation. In the skill router computation of Stage-II, we use gradient descent and CMA-ES to compute the router for normal settings and few-shot settings, respectively. The red and green lines represent the adaptation and inference processes, respectively.
• Figure 3: Examples of converting several categories of medical data to the unified instruction format.
• Figure 4: Details of our proposed fusion process.
• Figure 5: F1 scores of the few shot learning methods comparison on different number of labels in a multi-label classification task (ICD coding).