TCM-FTP: Fine-Tuning Large Language Models for Herbal Prescription Prediction ^††thanks: We gratefully acknowledge the support of NVIDIA Corporation with the GPU computational resources used for this research. ^††thanks: This work was supported by the Natural Science Foundation of Beijing (No. L232033), the National Natural Science Foundation of China (No. 82374302), and the Science and Technology Innovation Project of China Academy of Traditional Chinese Medicine (No. CI12021A00513). ^††thanks: Xingzhi Zhou and Xin Dong have contributed equally to this work. ^††thanks: This work was done while Xingzhi Zhou was an intern at NVIDIA.

We now formally define the prescription prediction problem. Given a prescription dataset $\mathcal{P}_{\text{train}}$, an element in $\mathcal{P}_{\text{train}}$ consists of a symptom description $s$ and its corresponding prescription $\{h_{i},w_{i}\}^{i\in[k]}$. Here, $h$ denotes the herb name, $w$ represents the herb dosage, and $[k]$ signifies a list ranging from 1 to $k$. Our objective is to train a model $\mathcal{M}$ such that $\mathcal{M}(s)$ reproduce $\{h_{i},w_{i}\}^{i\in[k]}$ accurately. For alignment with the language generation task, we concatenate $\{h_{i},w_{i}\}^{i\in[k]}$ using a comma separator to form a single sentence. Concrete examples are presented in Fig. 1.

Research on TCM prescription prediction mainly falls into three categories: topic model-based, graph model-based, and language model-based approaches. Topic model-based approaches treat relationships between symptoms and herbs as that of documents and topics [16, 17, 18]. These approaches rely on statistical relationships, lacking a semantic understanding of symptoms. Graph-based approaches construct a medical knowledge graph to model the relationships between symptoms and herbs [19, 14, 20]. However, these approaches also lack consecutive semantic information regarding symptoms. Language model-based approaches are a more promising way to model the complicated relationships between symptoms and herbs.

Language model-based prescription prediction models take patient symptom descriptions as input and generate herbal prescriptions sequentially. TCM Translator [21] uses transformer architectures to distill context vectors from symptoms and LSTM [22] as the decoder. AttentiveHerb [7] employs a seq2seq [6] model with dual attention mechanisms to distinguish primary from secondary symptoms and map herb-symptom interactions using clinical data. Herb-Know [8] utilizes herb descriptions to model associations with symptoms and evaluates whether herb effects align with symptom descriptions. TCMBERT [9] integrates transfer learning by initially training an ALBERT [23] on TCM-related documents and fine-tuning an LSTM-based seq2seq model for context vector extraction. RoKEPG [10] incorporates additional herb knowledge for fine-tuning prescription prediction models.

LLMs have made significant strides in various NLP tasks. Notable examples like ChatGPT and GPT-4.0 [24] have attracted considerable attention, although specifics about their architectures, parameters, and training strategies from OpenAI remain undisclosed. The emergence of open-sourced LLMs has sparked widespread interest among researchers. Prominent open-sourced models include LLaMA [25], Bloom [26], and ChatGLM [27], among others. While LLMs excel in standard NLP tasks, they often struggle with specialized domains requiring specific knowledge, such as medicine or finance. Supervised fine-tuning has become a standard approach to enhance LLMs for these domains by incorporating specialized knowledge.

In the context of TCM, several innovative models have been proposed. Bentsao [28] utilizes supervised fine-tuning on LLaMA, integrating structured and unstructured knowledge from CMeKG [29]. Huatuo [30] leverages both original and distilled data from ChatGPT, incorporating a Reinforcement Learning from Artificial Intelligence Feedback (RLAIF) mechanism and employing Proximal Policy Optimization (PPO) during fine-tuning. Zhongjing [31] applies continual pretraining to inject domain knowledge and uses supervised fine-tuning on a Chinese multi-turn medical dialogue dataset, complemented by Reinforcement Learning from Human Feedback (RLHF). ShenNong [32] builds on LLaMA with LoRA-based fine-tuning on an instructional dataset derived from ChatGPT and a traditional Chinese medicine knowledge graph, benefiting from a large-scale dataset of over 110,000 instructions.

Parameter Efficient Fine-Tuning (PEFT) utilizes a small amount of parameters to fine-tune a large language model effectively. Assuming there is a pretrained model $f_{\theta}(y|x)$, PEFT seeks to adjust a limited number of parameters, $\Delta\theta$, such that $|\theta|\gg|\Delta\theta|$. In contrast to conventional fine-tuning, which updates all parameters, denoted as $|\theta|=|\Delta\theta|$, and requires significant computational resources, PEFT selectively updates a limited number of learnable parameters to achieve results comparable to complete fine-tuning.

PEFT encompasses three primary methodologies: adapter [33], p-tuning [34, 35, 36], and LoRA [15]. Adapter tuning [33] introduces PEFT by appending additional learnable modules, named adapters, to every layer of the pretrained model, exclusively fine-tuning these adapters while maintaining the original parameters. In contrast, p-tuning leveraged prompt engineering techniques, fine-tuning learnable vectors as new tokens and inserting these tokens into the multi-head self-attention layer (MSA) [34, 35, 36]. In this work, we employ LoRA [15] to efficiently fine-tune LLMs on the target dataset.

To better model the intricate relationships between symptoms and herbs, we propose TCM-FTP, which employs a pre-trained LLM coupled with an auto-regressive mechanism for TCM prescription prediction. For efficient fine-tuning of the LLM, we employ the LoRA technique[15], which optimizes the LLM with limited parameters and computational resources. To enhance the dataset, we permute the herbs in prescriptions to leverage the order-agnostic property inherent in TCM prescriptions.

Given a prescription dataset $\mathcal{P}_{train}$ with $\{s,\{h_{i},w_{i}\}^{i\in[k]}\}$, our primary goal is to use a language model to maximize the estimation of conditional probability:

where $s$ is the symptom description, $k$ is the number of herbs in the prescription, and $(h_{i},w_{i})$ represents the i-th (herb, dosage) pair with $i\in\{1,2,...,k\}$. $h_{<i}$ and $w_{<i}$ represent a set of herbs and a set of weights with index less than $i$, respectively. $\theta$ represents our model parameters.

To leverage the auto-regressive mechanism in LLMs, we concatenate the herbs and dosages into one string. We denote concatenated string $\textsc{cat}(\{h_{i},w_{i}\}^{i\in[k]})$ as $y$. Consequently, our objective function becomes:

where $\{y_{<j}\}$ denotes the set of the tokens from index $1$ to index $j-1$ in $y$, and $|y|$ represents the number of tokens in $y$.

To demonstrate the performance of our model, we compare it with the following baselines:

(1) Topic Models and Multi-label Classification Models:

• •

  LinkLDA [37] uses Link Latent Dirichlet Allocation [38] to model relationships between symptoms and corresponding herbs.

• •

  LinkPLSALDA [39] combines LDA [40] and PLSA [41] for topic modeling.

• •

  MLKNN [42] is an enhanced K-nearest neighbor (KNN) approach for multi-label classification (MLC).

(2) TCM Prescription Prediction Models:

• •

  PTM [5] excels in recommending herbs based on symptoms and predicting symptoms from herbs.

• •

  TCMPR [43] utilizes sub-networks for symptom description and information extraction.

• •

  KDHR [44] integrates herb properties and features using multi-level GCN layers.

• •

  PresRecST [12] follows a systematic approach with three stages for clinical predictions in TCM.

(3) Pre-trained Language Models:

• •

  Mengzi (T5-base) [45] is a T5-based [46] seq2seq model pre-trained on Chinese text and fine-tuned on the specific dataset as a competitive baseline.

• •

  GPT-3.5 and GPT-4.0 [24]: We also include the pre-trained LLMs without fine-tuning as baselines.

We implement the fine-tuning process with the Transformers package from Hugging Face²²2https://huggingface.co/docs/transformers. The number of training epochs is set to 10. A cosine scheduler is adopted, with the rank of LoRA being 16 and the alpha of LoRA at 32. We consider the number of herb permutations, K, as 20 and 50, respectively. The default learning rate is set at 1e-3. We leverage the gradient accumulation technique to increase the batch size, with 8 accumulation steps and a batch size of 16. We employ two foundation models Chinese-Alpaca-Plus-7B ³³3https://github.com/ymcui/Chinese-LLaMA-Alpaca/tree/main and ShenNong ⁴⁴4https://huggingface.co/michaelwzhu/ShenNong-TCM-LLM. Chinese-Alpaca-Plus-7B is a variant of LLaMA [25] with continual pre-training and supervised fine-tuning on Chinese corpus, denoted as LLaMA+ for simplicity. ShenNong is a further refinement of LLaMA+, fine-tuned with TCM instruction datasets. We use an 8-V100 GPU machine for fine-tuning and the running time for TCM-FTP ($K=50$) is 146 hours. During the inference stage, We employ top-k and top-p combinations for decoding, setting the top-k to 50 and the top-p to 0.7, with a default temperature of 0.95.

For the proposed TCM-FTP, we evaluate it from both quantitative and qualitative perspectives, including the following evaluation metrics.

Quantitative Evaluation. For evaluation metrics, we use precision, recall, and F1-score as herb prediction metrics and designed NMSE for herb dosage evaluation.

• •

  Precision: The proportion of correctly predicted herbs out of all predicted herbs, reflecting the accuracy of positive predictions.

• •

  Recall: The proportion of correctly predicted herbs out of all herbs in the ground truth, indicating the model’s ability to identify all relevant items.

• •

  F1-score: The harmonic mean of precision and recall, giving a balanced measure of overall accuracy.

• •

  Normalized mean square error (NMSE): To tackle the problem above, we design this metric to evaluate the accuracy of predicted dosage by normalizing the squared differences using the original weights. For a given prescription $p$ and a predicted prescription and $\hat{p}$, we suppose they are composed by a set of pairs of a herb and a dosage, $p=\{(h,w)\}$, where $h,w$ refer to herb and dosage respectively, and $\hat{p}=\{(h^{\prime},w^{\prime})\}$ indicates the generated result from model. The calculation of NMSE is as follows,

  $$NMSE=\frac{1}{Z}\sum_{\begin{subarray}{c}(h,w)\in p,\\ \left(h^{\prime},w^{\prime}\right)\in\hat{p}\end{subarray}}\mathbf{1}\left[h=h^{\prime}\right]\left(\frac{w^{\prime}-w}{w}\right)^{2},$$

  (9)

  where $Z$ is the number of correctly predicted herbs, and $\mathbf{1}[\cdot]$ is the indicator function. For evaluation, the average dosage of each herb from the training data is used as a baseline for dosage predictions. For herbs unseen in the training data, the dosage is predicted as the average dosage of all known herbs. This approach is referred to as NMSE_base in the NMSE calculations.

Qualitative Evaluation. Existing quantitative evaluation metrics assess model quality based solely on sample labels, neglecting the compatibility and symptom-specific effectiveness of prescriptions. To comprehensively evaluate our model, we engaged five TCM experts to conduct an expert qualitative evaluation (EQE) of selected prescriptions generated by our model. Each doctor independently assessed the prescriptions for herbal effectiveness in treating symptoms (SHE) and herbal compatibility (HC), assigning scores on a scale of 0 to 5. Higher scores indicate greater effectiveness or compatibility.

The overall comparison results between TCM-FTP and the baselines are presented in Table II. The proposed TCM-FTP outperforms all baseline models in the herb prediction task on DigestDS, achieving an F1-score of 0.8016 using LLaMA+ as the foundation model and 0.8031 using ShenNong. This highlights the superior capability of TCM-FTP in modeling the intricate relationships between symptoms and herbs in prescriptions. Unlike previous approaches, TCM-FTP also includes herb dosage prediction, which is crucial in TCM due to the significant impact of dosage combinations. As shown in Table II, TCM-FTP achieves a much lower NMSE compared to the baseline using average statistics (0.0604 for TCM-FTP versus 0.0754 for Mengzi (T5-base) [45]). This enhances the practicality of TCM-FTP in clinical TCM prescription generation.

Our proposed TCM-FTP shows significant advantages over various baseline models across all aspects.

• •

  Compared to the pre-trained language model Mengzi (T5-base) [45], TCM-FTP significantly enhances prediction performance with improved precision (from 0.7332 to 0.7951), recall (from 0.7474 to 0.8113), and F1-score (from 0.7403 to 0.8031). Performances on GPT-3.5 and GPT-4.0 [24] were notably poor, indicating the limitations of general llms in tasks requiring specialized knowledge.

• •

  In comparison to other TCM prescription prediction models, TCM-FTP outperforms PTM [5], which leads among these baselines, followed by TCMPR [43] in accuracy, with KDHR [44] and PresRecST [12] showing poorer performance. PTM, a topic model-based approach, is computationally complex, while TCMPR, KDHR, and PresRecST are graph-based models sensitive to graph-related factors. TCM-FTP excels these approaches in generating accurate herb recommendations due to advanced language modeling capability.

• •

  Compared to multi-label classification (MLC) and topic model approaches, our model demonstrates superior performance. MLKNN [42] and two topic model methods perform similarly, whereas MLC and topic model-based approaches achieve lower performance due to computational intensity and limited consideration of inter-herb relationships. This underscores the proposed TCM-FTP proves more suitable for real-world herb prediction tasks, highlighting its enhanced performance over various baseline methods.

In summary, our approach exhibits significant advantages over various baseline methods and is better suited for real-world herb recommendation tasks. In addition, TCM knowledge embedded in foundation models enhances the fine-tuning process. Epoch-wise performance, shown in Fig. 3(a) for both loss and F1-score, indicates that although TCM-FTP (ShenNong, K=50) shows similar loss changes to TCM-FTP (LLaMA+, K=50), the F1-score line for TCM-FTP (ShenNong, K=50) consistency surpasses that of TCM-FTP (LLaMA+, K=50). This suggests that TCM knowledge significantly aids the fine-tuning process with the ShenNong model. Additionally, with a lower number of permutations K=20, TCM-FTP (ShenNong) achieves an F1-score of 0.8008, compared to 0.7652 for TCM-FTP (LLaMA+), highlighting the advantage of having TCM knowledge embedded in the ShenNong model for modeling symptom-herb relationships.

The number of permutation The herb permutation introduces the order-agnostic property to enhance the prediction performance. We present the results of TCM-FTP with varying numbers of herb permutations $K$ in Table III. As $K$ increases from 0 to 50, the f1-score improves gradually from 0.4885 to 0.8031, and the NMSE decreases from 0.1683 to 0.0604. This suggests the importance of the order-agnostic property as an effective inductive bias for prescription prediction.

Learning rates We validate the impact of learning rates by conducting experiments with K=10 and varying learning rates, as shown in Table IV. Higher learning rates correlate with improved precision, recall, and F1-score. This suggests that lower learning rates lead to underfitting, likely due to significant disparities between the pre-trained model parameters and the target parameters. Our objective is to fine-tune a task-specific generator without retaining the broad language capabilities of large language models (LLMs), necessitating substantial adjustments to achieve superior performance. However, excessively large learning rates cause model corruption, such as a learning rate of 0.0025.

Decoding paramters We investigate the robustness of decoding parameters by varying top-p, top-k, and temperature. Initially, in Fig. 3(b), we set top-p to 0.7 and chart the f1-score across various top-k and temperature settings. The results show that the model’s performance is stable across different temperatures and top-k values, with only a slight decline when the temperature exceeds 0.95. Next, in Fig. 3(c), with top-k fixed at 50, we graph the f1-score for different top-p and temperature values, observing that lower temperatures are more stable, leading to higher probabilities for the top predicted tokens, thus enhancing prediction accuracy. Lastly, in Fig. 3(d), with the temperature set at 0.95, we map the f1-score for various top-p and top-k combinations, noting that lower top-k values offer more stability in the face of top-p variations.

Since conducting expert qualitative evaluation (EQE) is time-consuming and labor-intensive for doctors, this study randomly selected 20 data points from the test set for evaluation. As shown in Table V, the TCM-FTP model outperformed other baseline models across all metrics, achieving an average SHE score of 4.08 (standard deviation 0.5628) and an average HC score of 4.03 (standard deviation 0.5404). This indicates that TCM-FTP can generate effective prescriptions that adhere to TCM compatibility principles. In contrast, GPT-3.5 and GPT-4.0 [24] scored below 3 for both SHE and HC, with total scores of 5.75 and 5.89, respectively, significantly lower than TCM-FTP and TCMPR [43]. This demonstrates that GPT-3.5 and GPT-4.0 perform poorly in generating prescriptions that adhere to TCM principles. TCMPR performed relatively well, with SHE and HC scores of 3.56 (standard deviation 0.6715) and 3.54 (standard deviation 0.6578), respectively, but still fell short of TCM-FTP. Overall, TCM-FTP excelled in prescription generation, receiving higher approval from doctors and showing less score variability, indicating better stability and robustness.

In order to visually showcase the predictive performance and capabilities of our TCM-FTP, we obtain the predicted herb results formed by each model on the test set and had them evaluated by TCM experts. Fig. 4 illustrates the results for two cases, including the input (chief complaints, present medical history), output (actual prescriptions provided by the doctor), and the predicted prescription results from TCM-FTP (ShenNong) and Mengzi (T5 base) [45] models, as well as GPT-3.5 and GPT-4.0 [24] predictions. Herb names/dosage weights marked in red indicate that the model’s prediction match the actual prescription provided by the doctor. Additionally, professional evaluations of each model are presented, ranking them using the “>” symbol (with models ranked higher indicating better performance).

The results show that the TCM-FTP (ShenNong) models achieve strong prediction performance. The herb label predictions closely match the doctor’s prescriptions, indicating effective data fitting after fine-tuning. Additionally, the dosage predictions align well with the actual dosages provided by the doctor, matching clinically relevant herb dosages. In contrast, GPT-3.5 and GPT-4.0 were also tested on the same cases but produced comparatively poor predictions with significant deviations from the actual results. This likely stems from their lack of specialized TCM training, hindering their performance.

In summary, the proposed TCM-FTP (ShenNong) outperforms general models. This highlights the importance of fine-tuning large language models on high-quality, domain-specific data to fully leverage their capabilities in specialized fields.