PRDAGE: a prescription recommendation framework for traditional Chinese medicine based on data augmentation and multi-graph embedding

Problem definition

The dataset comprises of classical TCM prescription records that contain essential information, such as patient symptoms, formulas and herbs. Let D denote the dataset, which consists of L cases, M symptoms, and K herbs. C = {c₁, c₂, …, c_L} represents the set of all cases in D, where L is the size of the case set. S = {s₁, s₂, …, s_M} is the set of all symptoms in the dataset, where M is the size of the symptom set. H = {h₁, h₂, …, h_K} denotes the set of all herbs in the dataset, where K is the size of the herb set. In the dataset D, a single case is represented by $c=\left(\left\{s_1,s_2,\dots ,s_i\right\},\left\{h_1,h_2,\dots ,h_k\right\}\right)$ and is composed of subsets of S and H. $sc=\left\{s_1,s_2,\dots ,s_i\right\}$ and $hc=\left\{h_1,h_2,\dots ,h_k\right\}$ denote the symptom set and herb set respectively for that case. Therefore, $c=\left\{sc,hc\right\}$ can also represent this case.

The diagnostic process in TCM involves extracting symptoms and physical signs from the patient’s body through various methods, such as observation, auscultation, inquiry, and palpation (the four diagnostic methods). The physician then uses a mental model to analyze the extracted information and prescribe a formula that addresses the condition. This task is similar to a prescription recommendation task. The model takes in symptom information as input and outputs a prescription. The learning function of the model is denoted as $y=f_h\left(C^{\prime }|\Theta \right)$, where C′ represents the set of training cases from D, Θ denotes the trainable parameters of the model, and the predicted outcome of y is a probability distribution of herbs(H) with a size of K.

Framework

The PRDAGE framework, shown in Fig. 1, comprises the following components: (1) Data preprocessing module. This module involves primarily extracting medical cases from TCM paper records, extracting vital information such as symptoms and herbs, and standardizing these elements. (2) The symptom and herb embedding module comprises the SBert embedding for symptoms, the representation of the four qi and five flavors, and meridian entry of herbs, as well as the fusion of semantic relation embeddings for symptoms and herbs based on TCM-GCN. (3) The data augmentation module focuses on enhancing symptom-centered case data, while (4) the prescription recommendation module uses the symptom embeddings obtained from (2) and trains with long short-term memory network (LSTM) and fully connected layers to generate the predictive model.

First, after data preprocessing, standardized symptoms and herb names are obtained. The symptoms are fed into the fine-tuned SBert to generate their embeddings, while the herbs are fed into the “Semantic knowledge Embedding” to obtain the embeddings of the four qi, five flavors, and meridian tropism. Subsequently, the embeddings of symptoms and herbs included in the medical case set are fed into the TCM-GCN module, resulting in a set of symptom and herb embeddings that capture the relationships between symptoms and herbs. Second, after the medical case data is fed into the data augmentation module, a medical case feature dataset is generated. Each medical case in this dataset contains several symptom features and corresponding herb features, which are all extracted from the set of symptom and herb embeddings obtained in the first part. Third, based on this medical case feature dataset, it is fed into LSTM and multiple fully connected layers for training to obtain a prediction model. During testing, when n test medical cases are fed in, n herb feature vectors of size m (where m is the size of the entire herb set) are generated. In the generated herb feature vectors, each value represents the predicted probability of the corresponding herb at that position. We select the top K herbs with the highest probabilities as the predicted results for analysis.

Data preprocessing

Data preprocessing is the process of extracting medical case information from TCM classic prescription books. It involves extracting key information such as symptoms, herbs, and formulas for training purposes. The process includes medical case text extraction, symptom and herb entity extraction, data cleaning, and standardization of symptoms and herbs, as shown in Fig. 1B. The primary data source for this study is TCM classic prescription books, including Experimental Records of Classic Prescriptions ( ), Historical Cases of Treatment on Cold Damage ( ), Selected Case Studies of Famous Physicians on Cold Damage Diseases ( ), Collection of Medical Cases on Cold Damage Diseases ( ), Selected Medical Cases on Cold Damage Diseases ( ), Commentary on One Hundred Families’ Medical Cases of the Golden Chamber Prescriptions ( ), Selected Case Studies of Famous Physicians on the Golden Chamber Prescriptions ( ), and Selected Case Studies of National Medical Masters on Classic Prescriptions ( ). We use a CZUR ET16 device to quickly scan the books into images and integrate them into a PDF format with images. Text is recognized using OCR methods, and a total of 3,092 classic medical case data entries are extracted.

To guarantee the quality of the dataset, we cleaned the data collected in by removing duplicate values and discarding cases containing missing or anomalous data. We used named entity recognition methods to extract information such as symptoms from the medical case text. Standardization of symptoms and herbs is necessary to construct a symptom-herb network. The TCM Clinical Common Symptom Terminology Standards (Revised Edition) ( ), Syndrome Differentiation and Treatment Study ( ), and TCM Symptomology Research ( ) were used to standardise all symptoms, while the Chinese Traditional Medicine Thesaurus ( ) and Chinese Herbal Medicine Dictionary ( ) were used to standardise all herbs. Following standardisation, there were a total of 427 symptoms and 429 herbs. Figures 2A and 2B show the frequency distribution of each symptom and herb in the medical cases.

Embedding of symptom and herb

Establishing prescription recommendation models requires a crucial foundation of embedding symptoms and herbs. To represent the semantic features of symptom texts, we employed an initial embedding using SBert for symptoms and utilized the four qi, five flavors, meridian entry, for embedding herbs. To capture the complex relationships between symptoms and herbs, we utilized TCM-GCN for embedding both symptom and herb nodes. The embedding process for symptoms and herbs is depicted in Fig. 1D.

Data augmentation

Another significant portion of our dataset is the formula. A medical case typically corresponds to one formula (with a small number of samples having two formulas). To analyze the relationship between samples and formula labels, we conducted a comparative analysis of the number of medical cases and the number of formulas. The analysis revealed that most of the formula labels are less than 10 samples, and only nine formula labels have more than 40 samples. The details are shown in Table 1. Labels with fewer than five samples account for one-third of the total label count, including 20 labels with a single sample and 26 labels with two samples. The proportion of labels with low sample counts is relatively large, as showed in Fig. 3A.

During the prediction training, symptoms in a medical case are treated as a sentence and input into the LSTM layer for training. As the input for LSTM is a sequential time series, symptom sequences with different orderings are considered distinct sequences. Therefore, the order of symptoms can be adjusted within a medical case to achieve sample data augmentation. To maintain balance among the expanded samples, we use the median value of 30 in Fig. 3C as the anchor point and expand the samples for labels with fewer than 30 samples. The expansion factor is determined by the formula k = Round(30/n), where n represents the current label’s sample size, and Round() denotes rounding down. If the number of symptoms in a sample is less than or equal to 4, augmentation is performed based on the factorial of the number of symptoms in the medical record. For more than four symptoms, random reordering augmentation is applied according to the value of k, as detailed in Algorithm 1 . Following data augmentation, most medical case samples have over 30 samples (as shown in Fig. 3B), and the center of sample distribution shifts from 2 to 30 (as shown in Fig. 3C). The total number of samples after data augmentation is 9,829.

 
_______________________ 
Algorithm 1 Data Augmentation Algorithm____________________________________________ 
Require: Medical Case Dataset (C = {Ci}|C| 
    i=1):  A collection of classical pre- 
     scription medical cases. Ci = {Si,Fi,Hi} is a case in C.Si is a set of symp- 
     toms present in a medical case, denoted as (s1,s2,...,sm),where si denotes 
     the ith symptom;  Fi  is the fomula corresponding to this case;  Hi  is the 
     fomula corresponding to this case.  MAX_NUM, denoting the maximum 
     anchoring factor, taking values 30, 40, 50, 60. 
Ensure: The Dataset C′ after augmentation 
  1:  C′ = ∅ 
 2:  calculate the count of samples for Fi, the result is stored in f_count, which 
     is a dictionary type variable 
  3:  for each Ci in C do 
 4:     if f_count(Fi) <= MAX_NUM then 
 5:         K = Roundup(MAX_NUM/|Si|) 
  6:     end if 
 7:     if K > MAX_NUM//2 and |Si| <= 4 then 
 8:         Ci′ = itertools.permutations(C 
i) 
  9:         C′ = C′∪ Ci′ 
10:     else 
11:         for i in range(K) do 
12:            S′i = random.shuffle((s1,s2,...,sm)) 
13:            C′ = C′∪{ 
 Si′,F 
i,Hi} 
14:         end for 
15:     end if 
16:  end for 
17:  return C′_______________________________________________    

Prescription recommendations

We input multiple symptoms of a medical case as a sentence into the LSTM layer, and then pass the output to two fully connected layers for training. The output of the fully connected layers, after being processed by the Softmax function, yields the probabilities of the herbs. We select the top K herbs with the highest probabilities as the predicted prescription.

Prescription recommendation is the process of predicting a set of herbs given an input set of symptoms. This task can be understood as calculating the probability of each herb based on a given set of symptoms. It is essentially a multi-label classification task, and therefore we use the binary cross-entropy loss function for model training, as shown in Eq. (9). (9)${\displaystyle \begin{array}{c}\hfill L\left(y,\hat{y}\right)=-\frac{1}{N}\sum _{i=1}^N\left[y_{i}log\left({\hat{y}}_i\right)+\left(1-y_i\right)log\left(1-{\hat{y}}_i\right)\right]\hfill \\ \hfill \hat{y}=f_h\left(c|\Theta \right)\hfill \end{array}}$

where L represents the loss function, N is the number of medical case samples in the training set, y_i is the set of true herbs corresponding to the ith case, ${\hat{y}}_i$ is the predicted probability of herbs for the ith case, c is the dataset of traditional herbal formulas, and Θ represents the trainable parameters of the model. The predicted result is a set of herb prediction probabilities of size K, denoted as H.

Data sets

Through the methods described in Data Preprocessing for data collection and preprocessing, we obtained 3,052 medical case data, which is named Dataset-A. After data augmentation with a 30-fold increase, a total of 9,829 data entries were obtained, named Dataset-A-30. The dataset contains 427 different symptoms and 429 different herbs, and is divided into two subsets for training and testing purposes. The training and testing subsets have proportions of 0.7 and 0.3, respectively. To test the model’s ability to generalize to real medical cases, we selected medical case data from outside the dataset. We extracted 278 traditional herbal formula medical cases from Selected Medical Cases of Liu Duzhou ( ) and Essential Collection of Hu Xishu’s Medical Cases ( ), and named this dataset Dataset-B. We then conducted secondary testing on Dataset-B to further validate the model’s generalization ability. Table 2 presents the statistical information of the experimental datasets.

Experimental setup

The experiments were conducted using TensorFlow-GPU 2.9.0 on an Intel(R) Xeon(R) Silver 4314 CPU @ 2.40 GHz 2.39 GHz (2 processors) with 32 GB of memory. GPU-accelerated training was performed using NVIDIA Tesla T4. The symptom embeddings were set to 256 dimensions, the LSTM layer dimensions were set to 256, and the fully connected layer dimensions were set to 256 and 512 respectively. The output layer dimensions were consistent with the number of herbs, which was 429. The initial learning rate for the model was set to 1e−4, and the Adam optimizer was used to update the parameters with momentum parameters of β1 = 0.9 and β2 = 0.99. The model was trained for 100 epochs with a batch_size of 8. If the loss and hit rate on the test set remained unchanged after 10 epochs, training was terminated early, and the optimal model would be saved.

Baseline model

In this study, we employ the following baseline models to validate the proposed method:

KGETM (Wang et al., 2019): The Knowledge Graph Embedding Enhanced Topic Model is a model used for herb recommendation. It leverages TransE to learn semantic knowledge of symptoms and herbs in TCM knowledge graphs, constructing a graph embedding method for TCM knowledge.

SMGCN (Jin et al., 2020): The Symptom-Multi-Graph Convolutional Network is a model that constructs multiple graphs between symptoms and herbs to obtain embeddings for symptoms and herbs. It also derives representations of syndromes and trains alongside symptoms to recommend herbs.

KDHR (Yang et al., 2022): A graph convolutional network-based herb recommendation model that integrates auxiliary properties of herbs on the basis of a TCM knowledge graph. It obtains symptom feature and herb feature representations through a ‘symptom-herb’ graph.

TCMPR (Dong et al., 2021): A TCM prescription recommendation method based on a sub-network symptom word mapping approach. It extracts sub-network structures between symptoms from the knowledge network, effectively representing the embedding features of symptom terminology.

PRDAGE: Our proposed model is a prescription recommendation based on data augmentation and multi-graph embedding for TCM (PRDAGE). It is designed to capture the textual semantic information of symptoms and represent the complex relationships between symptoms and herbs.

Evaluation metrics

To quantitatively assess the effectiveness of the PRDAGE method, we employ three performance evaluation metrics: precision, recall, and F1-score, to measure the model’s performance. For medical case c in the test set, the evaluation metrics are defined as follows: (10)${\displaystyle Precision@K=\frac{|Top\left(Se{t}_{pre},K\right)\cap Se{t}_{label}|}{K}}$ (11)${\displaystyle Recall@K=\frac{|Top\left(Se{t}_{pre},K\right)\cap Se{t}_{label}|}{\left|Se{t}_{label}\right|}}$ (12)${\displaystyle F1-score@K=\frac{2\ast Precision@K\ast Recall@K}{Precision@K+Recall@K}.}$

In predicting herbs, Top(Set_pre, K) represents the K herbs with the highest prediction scores. Set_label denotes the actual herbs present in hc. Precision at k (Precision@K) indicates the proportion of correctly predicted herbs among the top K predicted herbs. Recall at k (Recall@K) represents the proportion of the top K predicted herbs that are actually present in hc. The F1-score is the weighted average of the first two metrics, providing a more objective representation of the model’s performance. The value of K ranges from 1 to 20.

Experimental results

The experiments were conducted with the model trained on Dataset A-30. To avoid overfitting caused by augmented data, model testing was performed on the unaugmented dataset (Dataset-B). We compared the PRDAGE model with four other models: KGETM, SMGCN, KDHR, and TCMPR. The test results shown in Table 3, where bold and underlined values depict the best and second-best results, respectively. The results indicate that the PRDAGE model demonstrates the best overall performance in Top@5, Top@10, and Top@20. Compared to KGETM, PRDAGE achieved a 9.26% increase in precision, a 13.53% increase in recall, and an 11.50% increase in F1 score at Top@5. At Top@10, the accuracy improved by 12.43%, recall by 12.93%, and F1 score by 12.60%. At Top@20, the accuracy, recall, and F1 score increased by 13.94%, 14.84%, and 14.15%, respectively. These improvements are statistically significant. Although DAGEPR’s performance enhancement is not significant compared to the second-best model, there is still some improvement. Specifically, at Top@5, the accuracy and recall increased by 0.77% and 2.36%, respectively. At Top@10, the accuracy and recall improved by 1.69% and 3.81%, respectively.

Finally, we analyzed GPU memory consumption, computational time and model complexity of DAGEPR and baseline methods. GPU memory consumption and computational time is measured by the GPU memory used by the models during training. Model complexity is evaluated by the number of parameters.

Table 4 shows the comparison between DAGEPR and the four baseline methods in terms of computational time, number of parameters, and GPU memory usage: (1) GPU memory usage: the GPU memory usage of DAGEPR was 852 MB, which was much lower than that of TCMPR, and KDHR. (2) Number of parameters: the number of parameters of DAGEPR was 0.9427 M, which was less than that of TCMPR. (3) Computational Time: the computational time of DAGEPR was 197.27 s, much less than KDHR and TCMPR. In summary, DAGEPR was better than most of the baselines in terms of computational efficiency and resource usage, and the number of parameters was moderate.

Ablation experiments

PRDAGE consists of two crucial components: the embedding module for symptoms and herbs, and the data augmentation module. To validate the effectiveness of these two components in model training, we conducted tests under the following four scenarios (as shown in Table 5): (1) PRDAGE without symptom/herb embedding and data augmentation (PRDAGE-NoE-NoA), (2) PRDAGE without symptom/herb embedding but with data augmentation (PRDAGE-NoE), (3) PRDAGE with symptom/herb embedding but without data augmentation (PRDAGE-NoA), and (4) PRDAGE with both symptom/herb embedding and data augmentation. The experimental comparison results are presented in Fig. 4. It can be observed from the figure that the overall performance is the worst when neither augmentation nor embedding is applied (PRDAGE-NoE-NoA), and the best when both augmentation and embedding are utilized (PRDAGE), indicating that data augmentation and embeddings significantly contribute to the improvement of model performance. The performance with only augmentation (PRDAGE-NoE) and only embedding (PRDAGE-NoA) is relatively close, but the performance with only embedding is slightly better than that with only augmentation, which suggests that the embedding module plays a more important role in enhancing model performance.

Performance comparison with different data augmentation factors

To seek the optimal data augmentation coefficient, we tested datasets augmented with different coefficients. We augmented the datasets with various augmentation coefficients, K_Aug = 10, 20, 30, 40, 50, 60. In order to enhance the training speed, the learning rate was set to 0.0001, and the epoch number was set to 100. The final results are shown in Fig. 5A. It can be observed from the figure that as the augmentation coefficient increases, the model’s accuracy also improves. When K_Aug = 30, the model achieves optimal performance in terms of precision, recall, and F1 score. However, when the augmentation coefficient is further increased, the overall performance of the model declines. Therefore, we select 30 as the anchor value for the medical case data augmentation coefficient.

In order to assess the model’s performance on augmented datasets, the data was augmented with different K_Aug coefficients, and the dataset was split into training and testing sets at a ratio of 0.7:0.3. The model’s performance on the augmented test set is shown in Fig. 5B, while the results on the unaugmented test set are presented in Fig. 5A. It can be observed from the figures that models trained on augmented datasets achieve over 70% Top@1 precision and Top@20 recall. When the K_Aug coefficient is 50 and 60, the model’s Top@1 precision even exceeds 90%. However, when tested on unaugmented datasets, the accuracy significantly decreases when the K_Aug coefficient exceeds 30. This phenomenon can be attributed to the augmented dataset’s increased proportion of similar data, which may result in overfitting during model training.

Data augmentation is an effective method to address data scarcity, but it can also lead to overfitting issues. In this study, several aspects were considered to reduce the risk of overfitting. Firstly, experiments were conducted with different augmentation factors (K_Aug =10, 20, 30, 40, 50, 60), and it was found that the model achieved the best performance in terms of accuracy, recall, and F1-score when K_Aug =30. Beyond this value, the model’s performance began to decline, indicating the onset of overfitting due to the introduction of too many similar samples. Secondly, an early stopping strategy was implemented during the training process to prevent overfitting. Lastly, to ensure that our model did not overfit the augmented data, its performance was evaluated on an independent, unaugmented test dataset (Dataset-B). The results showed that the model maintained good performance on this dataset, indicating that it had learned robust features rather than just fitting the augmented data.

Analysis of the overall model performance

The results in Fig. 5A indicate that the selection of an appropriate data augmentation coefficient is crucial. Only models trained with a suitable data augmentation coefficient exhibit good generalization capabilities. Otherwise, there is a risk of the model performing exceptionally well during training but poorly on other datasets.

From Fig. 5, it can be observed that when we consider the number of herbs in the prescription for augmentation (i.e., as Top@K increases), the precision shows a declining trend, while the recall rate increases significantly, and the F1-score also improves. The F1 score increases in the Top@K range of [1–5], remains relatively stable in the range of [5–10], and then shows a marked decline after 10. This trend is likely related to the number of herbs in the dataset samples. In this dataset, the majority of medical case samples have a number of herbs within the range of [5–10], hence the F1 score also reaches its optimum and remains stable in this interval. In the reference [6], the mean number of herbs per prescription in their dataset is 11.3, and their model’s F1 score begins to stabilize from Top@12, which supports our observations.

Although the model achieved optimal results compared to the baseline models, its overall performance remains relatively average, with a noticeable gap from the requirements of clinical practice. Generally, the overall performance of prescription recommendation models is relatively low, which may be attributed to the fact that prescriptions consist of combinations of multiple herbs. Some herbs have similar effects and can substitute for each other, a factor that is not considered in the medical case data and evaluation. Additionally, another reason may be that the amount of medical case samples used for training is not large enough to fully reflect the complex associations between symptoms and herbs. Therefore, in future research, it is necessary to continue collecting medical case data and expand the sample size of medical cases.