Prediction-Augmented Generation for Automatic Diagnosis Tasks
Authors: Chan-Yang Ju, Dong-Ho Lee
Accepted at: acl25

Prediction-Augmented Generation for Automatic Diagnosis Tasks

Chan-Yang Ju*, and Dong-Ho Lee
Department of Applied Artificial Intelligence, Hanyang University, Korea
{karunogi,dhlee72}@hanyang.ac.kr

Abstract
Most Large language models (LLMs) adopt an autoregressive architecture, predicting the next word token based on the preceding context. While this approach is robust for language generation tasks such as writing and summarization, it has limitations for high-level reasoning tasks, such as prediction and decision-making. To overcome these limitations, we introduce a new method called Prediction-Augmented Generation (PAG). PAG can improve the generation quality and predictive accuracy of large language models in inference-driven tasks by integrating task-specific predictive models as external tools, enabling more structured and precise reasoning. Moreover, our method does not simply copy the inferences of a predictive model, but improves the inference results with knowledge from the large language model to create better predictions. We comprehensively evaluate our proposed method on diverse datasets for automatic diagnosis tasks requiring extensive domain knowledge and advanced reasoning.

1 Introduction
Large language models (LLMs) have gained attention for their ability to perform human’s tasks like writing, summarizing, and translation. These models are trained on vast text datasets encompassing books, articles, and web content (Brown, 2020; Radford et al., 2019), enabling them to understand context, syntax, and semantics. LLMs are becoming an important factor in both industry and academia due to its potential to assist and efficiently human work with large amounts of knowledge and language understanding ability.

However, LLMs use an autoregressive structure that predicts the next word token based on the previous tokens, which is less similar to human thinking and relies more on generating natural linguistic flow through memorization (Bender et al., 2021).

In addition, unlike prediction-based models that directly output the correct label in inference tasks, LLMs rely on indirectly predicting the word tokens that collectively represent the correct answer, which can degrade their predictive performance overall. This limitation often leads to “hallucinations” where the model generates unsupported text (Jiang et al., 2024b; McKenna et al., 2023). Moreover, even under identical conditions, small changes in the input words can lead to variations in the generated token probabilities and the resulting conclusions, highlighting the system’s inherent instability. This problem is present in most language models, making generative models less consistent and reliable in their answers.

Retrieval-augmented generation (RAG) tries to solve this issue by providing rationale and reducing the generation scope to minimize inaccuracies (Lewis et al., 2020; Nakano et al., 2021). However, the tasks expected of LLMs require a huge amount of experience and predictive ability, such as disease diagnosis, judgment, and strategic decisions, beyond knowledge-intensive tasks. LLMs are not expected to perform these advanced tasks adequately due to their limitations (Huang et al., 2023; Ullah et al., 2024), even with RAG. Figure 1 compares the performance of predictive models, generative models, and RAG in disease prediction tasks, showing that predictive models generally per-

![Performance comparison of Predictive models (light blue), single LLMs (light green), and RAG (light red) on the PolyMed dataset (Ju and Lee, 2023).](page_1016_384_384_256.png)

Figure 1: Performance comparison of Predictive models (light blue), single LLMs (light green), and RAG (light red) on the PolyMed dataset (Ju and Lee, 2023).
form better. Despite LLMs being trained on vast general data and medical knowledge, we found that this knowledge was not effectively utilized in the prediction task.

To overcome this limitation, we propose Prediction-Augmented Generation (PAG), illustrated in Figure 2, which can facilitate LLM’s prediction task beyond RAG. Inspired by RAG’s strategy of supplementing LLMs with external knowledge, we propose a novel method that leverages predictive models to make label predictions and then integrates these results with the knowledge provided by LLMs. Instead of merely copying the predicted labels, our method refines the results by leveraging the LLM’s extensive knowledge, leading to improved predictive accuracy. We also propose a novel method for evaluating generative models for traditional classification tasks, overcoming the difficulty of evaluating LLMs on label-centric prediction problems. Our method enables automated evaluation of LLMs even in conventional label prediction tasks by matching generated text with label text. We evaluate PAG in a comprehensive disease diagnosis setting, covering both single and multiple tasks to enable a thorough performance analysis from multiple perspectives.

Our contributions are summarized as follows:

• To our knowledge, our method is the first to introduce the concept of prediction-augmented generation, providing analysis and groundwork for future research.

• Our proposed PAG achieves state-of-the-art performance on the disease diagnosis task, surpassing existing predictive models and LLMs.

• We introduce a novel evaluation method, facilitating a more quantitative and objective evaluation of language models rather than relying solely on abstract assessments.

![Key Concept of PAG](page_186_134_579_246.png)

Figure 2: Key Concept of PAG

2 Related Work

External Tools for Enhanced LLMs Early research on LLMs focused on their proficiency in language understanding and generation tasks such as machine translation, summarization, and question and answer (Kenton and Toutanova, 2019; Brown, 2020). While these models demonstrate remarkable abilities in generating contextually relevant text, they also face challenges related to the cost of retraining for updated knowledge and the occurrence of hallucinations in their outputs. To overcome these limitations, previous work has explored the use of external tools to compensate for the weaknesses of LLMs, such as retrieval (Guu et al., 2020; Lewis et al., 2020), extensive external tool support via API calls (Schick et al., 2024), and computational assistance via code executors (Gao et al., 2023). These external tools can extend the context processing capabilities using external modules, thus moving away from the pure language modeling paradigm (Mialon et al., 2023).

Evaluation of LLMs Evaluating the performance of LLMs is an important topic in the natural language processing community. Early methods relied on word-count-based metrics (Papineni et al., 2002; Lin, 2004) that compare surface overlap with reference text, failing to capture deeper meaning. Vector-based similarity measures (Papineni et al., 2002) later improved semantic evaluation by embedding words in continuous spaces, though they remain dependent on the quality of the underlying embedding models.

To evaluate LLMs automatically, some researchers developed multiple-choice question and answering datasets (Clark et al., 2018; Hendrycks et al., 2020; Rein et al., 2023). While this approach can automatically evaluate LLMs, it requires constructing a dedicated dataset for evaluation. In addition, LLMs may still rely on statistical patterns without genuine comprehension or deep reasoning.

In recent years, there has been a shift toward using LLMs themselves as evaluation tools (Chiang and Lee, 2023; Shi et al., 2023). This approach involves using one LLM to evaluate the outputs of another, based on pre-defined criteria. While this method allows for more flexible and nuanced assessments, it introduces new challenges that LLMs tend to produce median scores or neutral judgments (Liu et al., 2023).
Automatic Diagnosis System
I’m a 34-year-old man and recently I’ve started experiencing pain on the right side of my lower abdomen. I also feel bloated and have lost my appetite since overeating a few days ago. What could be causing these symptoms?
These symptoms can sometimes be related to gastrointestinal issues such as indigestion or a possible appendicitis, which is a concern given the location of the pain.

Figure 3: Overall process of Prediction-Augmented Generation for automatic diagnosis system

3 Preliminaries

3.1 Challenges of ADS
Automatic diagnosis systems (ADS) have gained attention for predicting diseases using a patient’s self-report, such as age, gender, and symptoms. One challenge in ADS research is the difficulty of collecting diagnostic data, especially for statistically less informed diseases. Existing research has tried to solve this problem by combining independent knowledge inference result with prediction (Ju et al., 2023) or injecting knowledge into neural networks to help improve performance (Xu et al., 2019a; Lin et al., 2021; Zhang et al., 2023; Zhao et al., 2022; Tiwari et al., 2022). Motivated by this challenge, we expected that PAG could contribute to this work by combining LLM’s vast knowledge with predictive models to solve the problem. Our method is elaborated in the next section.

3.2 Problem Formulation
Figure 3 illustrates the overall process of PAG for ADS. In the input stage, unstructured patient utterances \( U \) are processed by the large language model \( \mathcal{F}^E \) for feature extraction to produce structured data \( E^I \). Next, during the feature alignment stage, the extracted data \( E^I \) is mapped to standardized features \( E^M \) using the medical term dictionary \( T \) and the mapping function \( \mathcal{F}^M \). In the alignment stage, the mapped features \( E^M \) are transformed into numerical representations \( X \) for the predictive model \( \mathcal{F}^P \). The predictive model then operates on these numerical representations in the prediction stage, predicting disease labels \( L^N \) as outputs, which also have corresponding textual labels \( L^T \). Following this, the knowledge generation stage uses large language models to enrich the predicted disease text \( L^T \) with relevant medical knowledge \( K \). In the diagnosis stage, the patient utterances \( U \), predicted disease text \( L^T \), and generated knowledge \( K \) are integrated to produce an initial diagnosis.

\[
P(D^T|U) = LLM(U, L^T, K)
\]
(1)

\[
P(D^L|D^T, U) = LLM(U, L^T, K, D^T)
\]
(2)

where \( D^L \) denotes the disease labels about textual diagnosis \( D^T \). To derive the final diagnosis, an aggregation mechanism combines the predictive model’s outputs and the LLM-generated diagnosis:

\[
D = \argmax_{d \in L^N, D^L} \sum_{s \in L^N, D^L} S_s(d)
\]
(3)

where \( S_s(d) \) represents the score assigned to disease \( d \) by source of \( s \).

4 Prediction-Augmented Generation

4.1 Feature Extraction and Alignment
The pipeline begins with Feature Extraction (Step 1), where relevant medical concepts are extracted from a patient utterance \( U \) according to a predefined schema matching the predictive model’s input. LLMs perform this extraction, filling in missing attributes and normalizing the concepts to match the schema. Since the extracted medical concepts are generated in natural language by LLMs, they require transformation into a predefined token suitable for the predictive model. To address this problem, a feature alignment (Step 2) function, denoted as \( F^M \), maps the extracted
medical concepts \( E^T \) to a predefined medical term dictionary \( T = \{t_1, t_2, ..., t_m\} \). Each of the extracted information is inserted into a query set \( Q = \{q_1, q_2, ..., q_n\} \) for word matching. The \( n \) is the number of extracted query terms, and \( m \) is the number of predefined medical terms in the dictionary. To match the query terms to medical terminology, a two-stage matching approach is employed. First, an exact match search is conducted. If a query term \( q_n \) is found in \( T \), it is directly assigned. If no exact match is found, the process moves to the second stage, which ranks term candidates based on lexical-based and semantic-based similarity scores. The lexical-based scoring method measures the overlap between query words and medical terms. The score for a given query-key pair is computed as:

\[
\text{Lexical}(q_n, t_m) = |W(q_n) \cap W(t_m)| \tag{4}
\]

where \( W(q_n) \) and \( W(t_m) \) represent the sets of words in a query term and a predefined medical term, respectively. For example, if the query is "pain on right side of the lower abdomen" and the key is "right side lower abdomen pain", the lexical score of this key would be 5. The semantic-based scoring method measures conceptual similarity using word embeddings. Each query \( q_n \) and key \( t_m \) is transformed into a word embedding representation, denoted as \( \mathbf{v}_{q_n} \) and \( \mathbf{v}_{t_m} \), respectively. The semantic similarity score is calculated as:

\[
\text{Semantic}(q_n, t_m) = \frac{\mathbf{v}_{q_n} \cdot \mathbf{v}_{t_m}}{\|\mathbf{v}_{q_n}\| \|\mathbf{v}_{t_m}\|} \tag{5}
\]

This method captures conceptual similarities, such as recognizing "headache" and "pain in the head" as semantically equivalent. The final ranking score for each query-key pair is computed by combining lexical and semantic similarity scores:

\[
\text{Score}(q_n, t_m) = \text{R}_L(q_n, t_m) + \text{R}_S(q_n, t_m) \tag{6}
\]

where \( \text{R}_L(q_n, t_m) \) and \( \text{R}_S(q_n, t_m) \) denote the ranking positions of \( t_m \) in the lexical and semantic similarity rankings, respectively. The top-ranked term receives a score of \( k \), the second-ranked term is assigned \( k - 1 \), and this pattern continues decrementally until the lowest-ranked term within the top \( k \) receives a score of 1. The top-\( k \) candidate terms, denoted as \( C = \{c_1, c_2, \ldots, c_k\} \), are provided to the LLM through a structured prompt, where the model selects the most contextually appropriate medical term as:

\[
t^* = \text{LLM}(q_n, C) \tag{7}
\]

where \( q_n \) represents the query term, and \( C \) is the set of candidate terms. The selected term \( t^* \) is then used for word matching, ensuring consistency with the predefined terminology schema. The matched medical information is subsequently utilized as input for the predictive model. To prepare the input \( X \) for the predictive model, the selected medical term can be transformed using one-hot encoding, word embeddings, or statistical representations.

4.2 Prediction and Knowledge Integration

Once the input \( X \) is prepared, it is fed into the predictive model (Step 3), yielding \( k \) prediction results. The predicted numerical disease labels are represented as \( L^N \), while the corresponding textual labels are denoted as \( L^T \). To enhance the explainability and provide context, \( L^T \) is further processed by an LLM (Step 4), which generates detailed medical knowledge for each predicted disease. This generated knowledge for a disease is represented as \( K_d \).

Subsequently, the initial diagnosis is formed by integrating the patient's original utterances \( U \), the predictive results \( L^T \), and the generated medical knowledge \( K_d \). This integrated information is provided to the LLM to produce the refined diagnosis text \( D^T \) and the corresponding disease label \( D^L \) (Step 5).

The final step (Step 6) involves an aggregation mechanism to combine the predictive model's output and the LLM-generated diagnosis. The predicted disease labels \( L^N \) and the LLM-generated labels \( D^L \) are ranked based on their likelihood. For each source, the most probable disease is assigned the highest score \( k \), the second most probable receives \( k - 1 \), and this continues until the \( k \)-th disease, which receives a score of 1.

The score for each disease is computed as:

\[
\text{Vote}(d) = \sum_{s \in L^N, D^L} S_s(d) \tag{8}
\]

where \( S_s(d) \) represents the score assigned to disease \( d \) by source \( s \), reflecting the confidence or rank provided by each model. The final diagnosis \( D \) is the disease with the highest aggregated score:

\[
D = \argmax_{d \in L^N, D^L} \text{Vote}(d) \tag{9}
\]

This final diagnosis represents a comprehensive conclusion that integrates insights from both the predictive model and the LLM.
<table>
  <tr>
    <th>Models</th>
    <th>MDD</th>
    <th>MZ</th>
    <th>DXY</th>
    <th>Mean</th>
  </tr>
  <tr>
    <td>GPT-4o</td>
    <td>87.40</td>
    <td>71.69</td>
    <td><b>96.88</b></td>
    <td>84.58</td>
  </tr>
  <tr>
    <td>GPT-3.5</td>
    <td>81.38</td>
    <td>65.44</td>
    <td>92.66</td>
    <td>79.72</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub></td>
    <td>85.61</td>
    <td>69.19</td>
    <td>93.21</td>
    <td>82.31</td>
  </tr>
  <tr>
    <td>llama3.1<sub>8b</sub></td>
    <td>79.79</td>
    <td>67.96</td>
    <td>89.21</td>
    <td>78.12</td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub></td>
    <td>86.21</td>
    <td><b>75.14</b></td>
    <td>95.10</td>
    <td>82.19</td>
  </tr>
  <tr>
    <td>mistral<sub>8×22b</sub></td>
    <td>85.96</td>
    <td>68.72</td>
    <td>94.19</td>
    <td>83.11</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub></td>
    <td><b>90.12</b></td>
    <td>74.19</td>
    <td>95.06</td>
    <td><b>86.08</b></td>
  </tr>
</table>

Table 1: Feature alignment recall (FAR) scores of LMMs across MDD, MZ, and DXY datasets.

5 Experiments

5.1 Experimental Settings

Dataset We evaluate our method on two diagnostic tasks. For single disease prediction, we employ three publicly available datasets—MDD\footnote{https://competitions.codalab.org/competitions/29706}, Muzhi (MZ) (Wei et al., 2018), and Dxy (Xu et al., 2019b)—using standard train/test splits that combine both implicit and explicit symptom information. Additionally, the PolyMed dataset (Ju and Lee, 2023) challenges the model under three scenarios: Single test (predicting diseases seen during training), Unseen test (predicting diseases absent from training), and Multi test (predicting multiple diseases simultaneously). Further details are provided in Appendix A.2.

Metrics We evaluate our model using Recall@K, Precision@K, and NDCG (Normalized Discounted Cumulative Gain)@K—which measures ranking quality by assigning lower importance to correct predictions at lower ranks. To aggregate the three PolyMed tests, we use the Weighted arithmetic mean (WAM) score, detailed in Appendix A.3.

Baselines We employ Linear Discriminant Analysis (LDA), TabNet (Arik and Pfister, 2021), XBNNet (Sarkar, 2022), M2DIA (Ju et al., 2023), and MAEA (Ju and Lee, 2024) as predictive models. For LLMs, we use closed-source models gpt-4o-2024-08-06(GPT-4o) (OpenAI, 2024) and gpt-3.5-turbo-0125(GPT-3.5-turbo) (OpenAI, 2023), alongside quantized open-source models from the llama3.1 family (Dubey et al., 2024) (llama3.1:8b, llama3.1:70b), gemma2:27b (Team et al., 2024), mistral:8x22b (Jiang et al., 2024a), and mistral-large (Mistral AI Team, 2024). Further details on each model and RAG are in Appendix A.1 and Appendix A.4.

<table>
  <tr>
    <th>Base LLM</th>
    <th>Models</th>
    <th>MDD</th>
    <th>MZ</th>
    <th>DXY</th>
  </tr>
  <tr>
    <td rowspan="4">GPT-4o</td>
    <td>LDA</td>
    <td>67.78</td>
    <td>54.92</td>
    <td>70.19</td>
  </tr>
  <tr>
    <td>PAG(LDA)</td>
    <td><b>79.07</b></td>
    <td><b>58.45</b></td>
    <td><b>76.92</b></td>
  </tr>
  <tr>
    <td>XBNet</td>
    <td>69.87</td>
    <td>59.15</td>
    <td>67.30</td>
  </tr>
  <tr>
    <td>PAG(XBNet)</td>
    <td><b>74.89</b></td>
    <td><b>64.08</b></td>
    <td><b>72.11</b></td>
  </tr>
  <tr>
    <td rowspan="4">mistral 8×22b</td>
    <td>TabNet</td>
    <td>78.66</td>
    <td>50.00</td>
    <td>75.96</td>
  </tr>
  <tr>
    <td>PAG(TabNet)</td>
    <td><b>82.42</b></td>
    <td><b>62.67</b></td>
    <td><b>76.92</b></td>
  </tr>
  <tr>
    <td>LDA</td>
    <td>67.78</td>
    <td>54.92</td>
    <td>73.07</td>
  </tr>
  <tr>
    <td>PAG(LDA)</td>
    <td><b>78.66</b></td>
    <td><b>67.60</b></td>
    <td><b>75.96</b></td>
  </tr>
  <tr>
    <td rowspan="4">mistral-large 123b</td>
    <td>XBNet</td>
    <td>69.03</td>
    <td>58.45</td>
    <td>69.23</td>
  </tr>
  <tr>
    <td>PAG(XBNet)</td>
    <td><b>73.22</b></td>
    <td><b>64.78</b></td>
    <td><b>72.11</b></td>
  </tr>
  <tr>
    <td>TabNet</td>
    <td>76.98</td>
    <td>51.40</td>
    <td>76.92</td>
  </tr>
  <tr>
    <td>PAG(TabNet)</td>
    <td><b>81.59</b></td>
    <td><b>65.49</b></td>
    <td><b>76.92</b></td>
  </tr>
  <tr>
    <td rowspan="4">mistral-large 123b</td>
    <td>LDA</td>
    <td>60.25</td>
    <td>52.81</td>
    <td>65.38</td>
  </tr>
  <tr>
    <td>PAG(LDA)</td>
    <td><b>66.94</b></td>
    <td><b>64.78</b></td>
    <td><b>70.19</b></td>
  </tr>
  <tr>
    <td>XBNet</td>
    <td>61.92</td>
    <td>57.04</td>
    <td>64.42</td>
  </tr>
  <tr>
    <td>PAG(XBNet)</td>
    <td><b>67.78</b></td>
    <td><b>63.38</b></td>
    <td><b>75.96</b></td>
  </tr>
  <tr>
    <td rowspan="4">mistral-large 123b</td>
    <td>TabNet</td>
    <td>69.45</td>
    <td>57.74</td>
    <td>68.26</td>
  </tr>
  <tr>
    <td>PAG(TabNet)</td>
    <td><b>71.96</b></td>
    <td><b>66.90</b></td>
    <td><b>72.11</b></td>
  </tr>
</table>

Table 2: Recall@1 score of PAG and competitors on restored MDD, MZ, and DXY datasets, aligned as in Table 1. The Base LLM is used for data restoration and the PAG process.

5.2 Data Restoration Test

We conducted a data restoration experiment to assess the extraction of structured data from unstructured patient text and its mapping to predefined terms. Virtual patient utterances were generated from structured symptom data in the MDD, MZ, and DXY datasets, and then processed through our alignment system to verify accurate reconstruction, detailed in Appendix A.1. Table 1 shows the Feature alignment recall (FAR) scores for each dataset and model, as defined in Equation 10. The results indicate that the mistral large model performs best, followed by GPT-4o. The performance of each dataset is influenced by the quality of the predefined terminology dictionary. Redundant terms (e.g., "poor appetite" versus "loss of appetite") sometimes caused contextual confusion, hindering accurate reconstruction. Moreover, smaller models tended to exhibit text generation errors, such as typos and inaccurate token outputs, which lowered restoration performance. Detailed generated case are provided in Table 12 and Table 13 in Appendix. Table 2 presents the results of the comparison between the performance of a single predictive model and PAG on the restored data with each large lan-
<table>
  <tr>
    <th rowspan="2">Models</th>
    <th colspan="4">MDD</th>
    <th colspan="4">MZ</th>
    <th colspan="4">DXY</th>
  </tr>
  <tr>
    <th>R@1</th><th>R@3</th><th>ND@3</th><th>R@1</th><th>R@3</th><th>ND@3</th><th>R@1</th><th>R@3</th><th>ND@3</th>
  </tr>
  <tr>
    <td>LDA</td><td>75.73</td><td>84.93</td><td>81.10</td><td>67.60</td><td>98.59</td><td>86.78</td><td>79.80</td><td>97.11</td><td>90.34</td>
  </tr>
  <tr>
    <td><b>gemma2<sub>27b</sub> + PAG(LDA)</b></td><td><b>82.00</b></td><td><b>89.95</b></td><td><b>86.75</b></td><td><b>71.83</b></td><td><b>98.59</b></td><td><b>88.43</b></td><td><b>81.73</b></td><td><b>98.07</b></td><td><b>91.41</b></td>
  </tr>
  <tr>
    <td>PM</td>
    <td>XBNNet</td><td>77.82</td><td>90.37</td><td>85.19</td><td>69.71</td><td>97.18</td><td>86.67</td><td>71.15</td><td>98.07</td><td>87.51</td>
  </tr>
  <tr>
    <td></td><td><b>gemma2<sub>27b</sub> + PAG(XBNNet)</b></td><td><b>80.75</b></td><td><b>92.05</b></td><td><b>87.71</b></td><td><b>71.12</b></td><td><b>99.29</b></td><td><b>88.43</b></td><td><b>76.92</b></td><td><b>99.03</b></td><td><b>90.75</b></td>
  </tr>
  <tr>
    <td></td><td>TabNet</td><td><b>86.61</b></td><td>92.46</td><td>90.03</td><td>72.53</td><td><b>100</b></td><td>89.67</td><td>81.73</td><td>95.19</td><td>89.72</td>
  </tr>
  <tr>
    <td></td><td><b>gemma2<sub>27b</sub> + PAG(TabNet)</b></td><td><b>86.61</b></td><td><b>94.97</b></td><td><b>91.39</b></td><td><b>73.94</b></td><td><b>100</b></td><td><b>89.83</b></td><td><b>82.69</b></td><td><b>97.11</b></td><td><b>91.28</b></td>
  </tr>
  <tr>
    <td>GPT-4o</td><td>39.33</td><td>83.26</td><td>65.57</td><td>38.02</td><td>97.18</td><td>73.41</td><td>31.73</td><td>86.53</td><td>63.54</td>
  </tr>
  <tr>
    <td>GPT-4o + RAG</td><td>65.27</td><td>89.95</td><td>79.53</td><td>66.19</td><td>95.07</td><td>83.58</td><td>66.34</td><td><b>98.07</b></td><td>84.60</td>
  </tr>
  <tr>
    <td><b>GPT-4o + PAG(TabNet)</b></td><td><b>86.19</b></td><td><b>95.39</b></td><td><b>91.28</b></td><td><b>69.71</b></td><td><b>99.29</b></td><td><b>87.91</b></td><td><b>79.80</b></td><td><b>97.11</b></td><td><b>90.22</b></td>
  </tr>
  <tr>
    <td>GPT-3.5</td><td>37.65</td><td>81.17</td><td>63.35</td><td>34.50</td><td>95.77</td><td>71.41</td><td>31.73</td><td>87.50</td><td>64.90</td>
  </tr>
  <tr>
    <td>GPT-3.5 + RAG</td><td>60.25</td><td>83.68</td><td>74.15</td><td><b>66.19</b></td><td>92.95</td><td>82.80</td><td><b>66.34</b></td><td>95.19</td><td><b>83.28</b></td>
  </tr>
  <tr>
    <td><b>GPT-3.5 + PAG(TabNet)</b></td><td><b>75.73</b></td><td><b>94.14</b></td><td><b>86.96</b></td><td>61.26</td><td><b>98.59</b></td><td><b>84.17</b></td><td><b>62.50</b></td><td><b>96.15</b></td><td><b>82.60</b></td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub></td><td>33.89</td><td>71.12</td><td>55.79</td><td>23.94</td><td>79.57</td><td>55.72</td><td>21.15</td><td>55.76</td><td>40.47</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub> + RAG</td><td>58.15</td><td>82.42</td><td>72.21</td><td>60.56</td><td>95.07</td><td>81.59</td><td>42.30</td><td>90.38</td><td>69.74</td>
  </tr>
  <tr>
    <td><b>gemma2<sub>27b</sub> + PAG(TabNet)</b></td><td><b>86.61</b></td><td><b>94.97</b></td><td><b>91.39</b></td><td><b>73.94</b></td><td><b>100</b></td><td><b>89.83</b></td><td><b>82.69</b></td><td><b>97.11</b></td><td><b>91.28</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>18b</sub></td><td>29.70</td><td>74.89</td><td>56.24</td><td>40.84</td><td>92.95</td><td>72.43</td><td>31.73</td><td>81.73</td><td>60.12</td>
  </tr>
  <tr>
    <td>llama3.1<sub>18b</sub> + RAG</td><td>58.99</td><td>84.51</td><td>74.00</td><td>58.45</td><td><b>95.07</b></td><td>80.17</td><td>65.38</td><td>93.26</td><td>81.71</td>
  </tr>
  <tr>
    <td><b>llama3.1<sub>18b</sub> + PAG(TabNet)</b></td><td><b>84.51</b></td><td><b>93.72</b></td><td><b>89.94</b></td><td><b>68.30</b></td><td><b>92.95</b></td><td><b>83.03</b></td><td><b>71.15</b></td><td><b>95.19</b></td><td><b>85.18</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub></td><td>41.84</td><td>82.42</td><td>65.85</td><td>41.54</td><td>95.77</td><td>74.00</td><td>43.26</td><td>85.57</td><td>67.94</td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub> + RAG</td><td>63.59</td><td>89.53</td><td>78.86</td><td>61.26</td><td>95.07</td><td>82.04</td><td>64.42</td><td>98.07</td><td>83.64</td>
  </tr>
  <tr>
    <td><b>llama3.1<sub>70b</sub> + PAG(TabNet)</b></td><td><b>83.26</b></td><td><b>94.14</b></td><td><b>89.68</b></td><td><b>70.42</b></td><td><b>98.59</b></td><td><b>88.10</b></td><td><b>71.15</b></td><td><b>99.03</b></td><td><b>87.99</b></td>
  </tr>
  <tr>
    <td>mistral<sub>8×22b</sub></td><td>36.40</td><td>71.96</td><td>57.58</td><td>39.43</td><td>90.84</td><td>70.02</td><td>17.30</td><td>71.15</td><td>48.76</td>
  </tr>
  <tr>
    <td>mistral<sub>8×22b</sub> + RAG</td><td>63.59</td><td>86.61</td><td>77.29</td><td>64.08</td><td>92.95</td><td>81.93</td><td>68.26</td><td>95.19</td><td>83.87</td>
  </tr>
  <tr>
    <td><b>mistral<sub>8×22b</sub> + PAG(TabNet)</b></td><td><b>84.51</b></td><td><b>94.56</b></td><td><b>90.58</b></td><td><b>72.53</b></td><td><b>99.29</b></td><td><b>89.05</b></td><td><b>76.92</b></td><td><b>98.07</b></td><td><b>89.51</b></td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub></td><td>18.82</td><td>59.41</td><td>42.84</td><td>28.16</td><td>83.80</td><td>61.88</td><td>25.00</td><td>68.27</td><td>49.90</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub> + RAG</td><td>62.76</td><td>87.86</td><td>77.50</td><td>63.38</td><td>94.36</td><td>82.19</td><td>71.15</td><td>97.11</td><td>86.65</td>
  </tr>
  <tr>
    <td><b>mistral-large<sub>123b</sub> + PAG(TabNet)</b></td><td><b>78.66</b></td><td><b>94.56</b></td><td><b>88.25</b></td><td><b>71.12</b></td><td><b>99.29</b></td><td><b>88.62</b></td><td><b>75.00</b></td><td><b>98.07</b></td><td><b>88.93</b></td>
  </tr>
</table>

Table 3: Experimental results of Predictive Models (PM), LLMs, and PAG on MDD, Muzhi, and DXY datasets, using Recall(R), Precision(P), and NDCG(ND) metrics. Extended results are presented in Table 10 of the Appendix.

guage model used in Table 1, using the Recall@1 score. We have conducted experiments on the three large language models that performed the highest in Table 1 and found that PAG significantly improves the diagnosis performance. The PAG with GPT-4o and LDA combination recorded the largest performance improvement, and the GPT-4o-based PAG-applied model performed the best overall. The Mistral-large model had the highest data restoration performance in Table 1, but it performed relatively poorly compared to GPT-4o and Mistral in Table 2. This suggests that data restoration performance does not ensure final prediction performance. Since PAG involves a variety of tasks, including word matching, knowledge generation, and prediction refinement, a model that maintains balanced performance at each stage is likely to be more suitable for our method.

5.3 Single Diagnosis Task

We evaluated the performance of disease prediction using original data from MDD, MZ, and DXY to validate the performance of pure PAG, excluding the impact of data restoration quality. We used the prediction performance of Recall@1 and 3, and the NDCG@3 metrics. The experimental results in Table 3 show that PAG significantly improves prediction performance across all datasets. Although some older and smaller models exhibit performance degradation, PAG still yields an overall enhancement in predictive accuracy.

We evaluated the LLM’s performance by aligning its generated diagnostic text with the correct disease labels using our feature alignment process. In the single LLM performance, models with poor prediction performance, such as Mistral-large, showed significant performance improvement after applying PAG. This suggests that PAG can be utilized
<table>
  <tr>
    <th>Models</th>
    <th>R@1</th>
    <th>R@3</th>
    <th>R@5</th>
    <th>P@1</th>
    <th>P@3</th>
    <th>P@5</th>
    <th>ND@1</th>
    <th>ND@3</th>
    <th>ND@5</th>
  </tr>
  <tr>
    <td>XBNNet</td>
    <td>18.96</td>
    <td>38.47</td>
    <td>49.59</td>
    <td><b>27.36</b></td>
    <td>17.33</td>
    <td>13.07</td>
    <td>21.31</td>
    <td>33.71</td>
    <td>38.75</td>
  </tr>
  <tr>
    <td>GPT-4o + PAG(XBNNet)</td>
    <td><b>20.00</b></td>
    <td><b>45.76</b></td>
    <td><b>57.30</b></td>
    <td>26.77</td>
    <td><b>19.92</b></td>
    <td><b>14.84</b></td>
    <td><b>21.87</b></td>
    <td><b>38.02</b></td>
    <td><b>43.19</b></td>
  </tr>
  <tr>
    <th rowspan="2">PM</th>
    <td>MAEA</td>
    <td><b>25.78</b></td>
    <td>47.56</td>
    <td>58.82</td>
    <td><b>34.65</b></td>
    <td>20.75</td>
    <td>15.22</td>
    <td><b>28.21</b></td>
    <td>41.90</td>
    <td>46.96</td>
  </tr>
  <tr>
    <td>GPT-4o + PAG(MAEA)</td>
    <td>24.65</td>
    <td><b>50.13</b></td>
    <td><b>63.72</b></td>
    <td>32.07</td>
    <td><b>21.56</b></td>
    <td><b>16.34</b></td>
    <td>26.66</td>
    <td><b>42.64</b></td>
    <td><b>48.73</b></td>
  </tr>
  <tr>
    <th rowspan="2">M2DIA</th>
    <td>M2DIA</td>
    <td><b>24.73</b></td>
    <td>45.75</td>
    <td>56.11</td>
    <td><b>33.19</b></td>
    <td>19.95</td>
    <td>14.55</td>
    <td><b>27.06</b></td>
    <td>40.24</td>
    <td>44.89</td>
  </tr>
  <tr>
    <td>GPT-4o + PAG(M2DIA)</td>
    <td>24.09</td>
    <td><b>49.18</b></td>
    <td><b>61.38</b></td>
    <td>31.55</td>
    <td><b>21.17</b></td>
    <td><b>15.73</b></td>
    <td>26.12</td>
    <td><b>41.93</b></td>
    <td><b>47.42</b></td>
  </tr>
  <tr>
    <th rowspan="4">GPT-4o</th>
    <td>GPT-4o</td>
    <td>12.50</td>
    <td>33.96</td>
    <td>43.28</td>
    <td>16.01</td>
    <td>14.45</td>
    <td>11.04</td>
    <td>13.45</td>
    <td>26.77</td>
    <td>30.95</td>
  </tr>
  <tr>
    <td>GPT-4o + RAG</td>
    <td>14.68</td>
    <td>31.39</td>
    <td>42.13</td>
    <td>18.55</td>
    <td>13.27</td>
    <td>10.70</td>
    <td>15.66</td>
    <td>25.95</td>
    <td>30.82</td>
  </tr>
  <tr>
    <td>GPT-4o + PAG(M2DIA)</td>
    <td>24.09</td>
    <td><b>49.18</b></td>
    <td><b>61.38</b></td>
    <td>31.55</td>
    <td><b>21.17</b></td>
    <td><b>15.73</b></td>
    <td>26.12</td>
    <td>41.93</td>
    <td><b>47.42</b></td>
  </tr>
  <tr>
    <td>GPT-4o + PAG(M2DIA) w/o \( K_d \)</td>
    <td><b>25.20</b></td>
    <td>48.20</td>
    <td>59.03</td>
    <td><b>33.55</b></td>
    <td>20.86</td>
    <td>15.22</td>
    <td><b>27.50</b></td>
    <td><b>41.99</b></td>
    <td>46.87</td>
  </tr>
  <tr>
    <th rowspan="4">GPT-3.5</th>
    <td>GPT-3.5</td>
    <td>12.46</td>
    <td>28.50</td>
    <td>36.21</td>
    <td>15.92</td>
    <td>12.14</td>
    <td>9.24</td>
    <td>13.39</td>
    <td>23.33</td>
    <td>26.78</td>
  </tr>
  <tr>
    <td>GPT-3.5 + RAG</td>
    <td>15.12</td>
    <td>29.68</td>
    <td>37.49</td>
    <td>19.16</td>
    <td>12.57</td>
    <td>9.53</td>
    <td>16.17</td>
    <td>25.28</td>
    <td>28.78</td>
  </tr>
  <tr>
    <td>GPT-3.5 + PAG(M2DIA)</td>
    <td><b>22.53</b></td>
    <td><b>46.84</b></td>
    <td><b>57.71</b></td>
    <td>29.40</td>
    <td><b>20.26</b></td>
    <td><b>14.88</b></td>
    <td>24.40</td>
    <td><b>39.69</b></td>
    <td><b>44.58</b></td>
  </tr>
  <tr>
    <td>GPT-3.5 + PAG(M2DIA) w/o \( K_d \)</td>
    <td>22.51</td>
    <td>44.76</td>
    <td>56.10</td>
    <td><b>29.94</b></td>
    <td>19.45</td>
    <td><b>14.51</b></td>
    <td><b>24.57</b></td>
    <td>38.53</td>
    <td>43.62</td>
  </tr>
  <tr>
    <th rowspan="4">gemma27b</th>
    <td>gemma27b</td>
    <td>14.13</td>
    <td>35.17</td>
    <td>44.22</td>
    <td>18.06</td>
    <td>14.96</td>
    <td>11.28</td>
    <td>15.17</td>
    <td>28.14</td>
    <td>32.20</td>
  </tr>
  <tr>
    <td>gemma27b + RAG</td>
    <td>13.65</td>
    <td>30.37</td>
    <td>40.39</td>
    <td>17.40</td>
    <td>12.84</td>
    <td>10.26</td>
    <td>14.59</td>
    <td>25.10</td>
    <td>29.64</td>
  </tr>
  <tr>
    <td>gemma27b + PAG(M2DIA)</td>
    <td>22.93</td>
    <td><b>49.78</b></td>
    <td><b>62.22</b></td>
    <td>30.18</td>
    <td><b>21.36</b></td>
    <td><b>15.97</b></td>
    <td>24.90</td>
    <td><b>41.77</b></td>
    <td><b>47.34</b></td>
  </tr>
  <tr>
    <td>gemma27b + PAG(M2DIA) w/o \( K_d \)</td>
    <td><b>23.57</b></td>
    <td>45.89</td>
    <td>56.18</td>
    <td><b>31.19</b></td>
    <td>19.98</td>
    <td>14.56</td>
    <td><b>25.73</b></td>
    <td>39.80</td>
    <td>44.43</td>
  </tr>
  <tr>
    <th rowspan="4">LLM</th>
    <td>llama3.1s</td>
    <td>9.10</td>
    <td>21.74</td>
    <td>30.69</td>
    <td>11.58</td>
    <td>9.24</td>
    <td>7.81</td>
    <td>9.77</td>
    <td>17.64</td>
    <td>21.69</td>
  </tr>
  <tr>
    <td>llama3.1s + RAG</td>
    <td>12.87</td>
    <td>29.63</td>
    <td>38.68</td>
    <td>16.20</td>
    <td>12.56</td>
    <td>9.84</td>
    <td>13.71</td>
    <td>24.11</td>
    <td>28.21</td>
  </tr>
  <tr>
    <td>llama3.1s + PAG(M2DIA)</td>
    <td><b>21.53</b></td>
    <td><b>44.47</b></td>
    <td><b>57.86</b></td>
    <td>28.40</td>
    <td><b>19.27</b></td>
    <td><b>14.91</b></td>
    <td>23.40</td>
    <td><b>37.81</b></td>
    <td><b>43.81</b></td>
  </tr>
  <tr>
    <td>llama3.1s + PAG(M2DIA) w/o \( K_d \)</td>
    <td><b>21.98</b></td>
    <td>43.51</td>
    <td>55.17</td>
    <td><b>29.23</b></td>
    <td>18.96</td>
    <td>14.30</td>
    <td><b>23.99</b></td>
    <td>37.50</td>
    <td>42.73</td>
  </tr>
  <tr>
    <th rowspan="4">llama3.1t0b</th>
    <td>llama3.1t0b</td>
    <td>11.91</td>
    <td>27.91</td>
    <td>36.91</td>
    <td>15.26</td>
    <td>11.87</td>
    <td>9.41</td>
    <td>12.80</td>
    <td>22.53</td>
    <td>26.58</td>
  </tr>
  <tr>
    <td>llama3.1t0b + RAG</td>
    <td>14.09</td>
    <td>30.82</td>
    <td>40.53</td>
    <td>17.74</td>
    <td>13.03</td>
    <td>10.30</td>
    <td>15.02</td>
    <td>25.44</td>
    <td>29.81</td>
  </tr>
  <tr>
    <td>llama3.1t0b + PAG(M2DIA)</td>
    <td><b>21.67</b></td>
    <td><b>47.35</b></td>
    <td><b>60.16</b></td>
    <td><b>28.59</b></td>
    <td><b>20.44</b></td>
    <td><b>15.48</b></td>
    <td><b>23.55</b></td>
    <td><b>39.66</b></td>
    <td><b>45.40</b></td>
  </tr>
  <tr>
    <td>llama3.1t0b + PAG(M2DIA) w/o \( K_d \)</td>
    <td>21.19</td>
    <td>44.83</td>
    <td>57.42</td>
    <td>28.25</td>
    <td>19.47</td>
    <td>14.84</td>
    <td>23.12</td>
    <td>37.86</td>
    <td>43.55</td>
  </tr>
  <tr>
    <th rowspan="4">mixtral8x22b</th>
    <td>mixtral8x22b</td>
    <td>14.02</td>
    <td>32.13</td>
    <td>40.61</td>
    <td>17.90</td>
    <td>13.65</td>
    <td>10.37</td>
    <td>15.05</td>
    <td>26.09</td>
    <td>29.88</td>
  </tr>
  <tr>
    <td>mixtral8x22b + RAG</td>
    <td>10.83</td>
    <td>24.93</td>
    <td>34.00</td>
    <td>13.84</td>
    <td>10.59</td>
    <td>8.66</td>
    <td>11.58</td>
    <td>20.39</td>
    <td>24.56</td>
  </tr>
  <tr>
    <td>mixtral8x22b + PAG(M2DIA)</td>
    <td>22.63</td>
    <td><b>47.37</b></td>
    <td><b>59.21</b></td>
    <td>29.46</td>
    <td>20.35</td>
    <td><b>15.26</b></td>
    <td>24.47</td>
    <td>39.99</td>
    <td>45.34</td>
  </tr>
  <tr>
    <td>mixtral8x22b + PAG(M2DIA) w/o \( K_d \)</td>
    <td><b>23.46</b></td>
    <td>46.97</td>
    <td>57.74</td>
    <td><b>31.50</b></td>
    <td>20.42</td>
    <td>14.90</td>
    <td><b>25.69</b></td>
    <td><b>40.50</b></td>
    <td><b>45.36</b></td>
  </tr>
  <tr>
    <th rowspan="4">mistral-large123b</th>
    <td>mistral-large123b</td>
    <td>13.58</td>
    <td>31.85</td>
    <td>40.89</td>
    <td>17.43</td>
    <td>13.55</td>
    <td>10.42</td>
    <td>14.60</td>
    <td>25.77</td>
    <td>29.85</td>
  </tr>
  <tr>
    <td>mistral-large123b + RAG</td>
    <td>10.83</td>
    <td>27.61</td>
    <td>36.06</td>
    <td>13.83</td>
    <td>11.69</td>
    <td>9.17</td>
    <td>11.58</td>
    <td>21.96</td>
    <td>25.81</td>
  </tr>
  <tr>
    <td>mistral-large123b + PAG(M2DIA)</td>
    <td>22.68</td>
    <td><b>48.59</b></td>
    <td><b>60.12</b></td>
    <td>29.35</td>
    <td><b>20.92</b></td>
    <td><b>15.45</b></td>
    <td>24.48</td>
    <td>40.78</td>
    <td>45.97</td>
  </tr>
  <tr>
    <td>mistral-large123b + PAG(M2DIA) w/o \( K_d \)</td>
    <td><b>24.89</b></td>
    <td>47.96</td>
    <td>58.11</td>
    <td><b>33.13</b></td>
    <td>20.76</td>
    <td>15.00</td>
    <td><b>27.14</b></td>
    <td><b>41.69</b></td>
    <td><b>46.29</b></td>
  </tr>
</table>

Table 4: Overall WAM score for Predictive Models (PM), LLMs, and our PAG with three tests of PolyMed. Metrics "R", "P", and "ND" denote Recall, Precision, and NDCG respectively. "w/o \( K_d \)" denotes PAG that omits the knowledge generation module. Extended results are presented in Table 11 of the Appendix.

as a robust external tool to complement the performance of LLMs in prediction tasks.

We compared PAG with a Retrieval-augmented generation (RAG) approach and found that PAG outperformed RAG overall. RAG retrieves the top five disease documents by calculating similarity between the user’s symptoms embedding and disease information from medical organizations and uses these documents to generate diagnosis text. In the MZ and DXY datasets having less than five disease labels, PAG still outperformed RAG in prediction, even though we provided information for all disease labels to exclude the impact of retrieval performance. This suggests that PAG can improve the accuracy of generation over RAG in the prediction task.

5.4 Multiple Diagnosis Task

Table 4 shows experimental results in multiple diagnosis task of PolyMed. We measured WAM score (Equation 13) to aggregate the single, unseen, and multi test results. The experimental results show that the model with PAG improves the prediction performance in all experiments. Especially in Top-5 predictions, PAG significantly improves scores, demonstrating superior alignment of multiple disease candidates compared to predictive models. Predictive models typically trained for a single label output, so they often fail to produce multiple candidates. By leveraging LLM knowledge, PAG refines predictions and enhances the candidate pool, which is especially beneficial for automatic diagnosis systems offering multiple suspected diseases.
<table>
  <tr>
    <th>Models</th>
    <th>A@1</th>
    <th>R@5</th>
    <th>P@5</th>
    <th>ND@5</th>
  </tr>
  <tr>
    <td>M2DIA</td>
    <td>12.63</td>
    <td>31.04</td>
    <td>6.20</td>
    <td>22.03</td>
  </tr>
  <tr>
    <td>GPT-4o</td>
    <td>11.15</td>
    <td><b>46.09</b></td>
    <td><b>9.21</b></td>
    <td>29.79</td>
  </tr>
  <tr>
    <td>gemma2:27b</td>
    <td><b>15.61</b></td>
    <td>45.91</td>
    <td>9.18</td>
    <td><b>31.94</b></td>
  </tr>
  <tr>
    <td>mistral-large123b</td>
    <td>13.94</td>
    <td>43.68</td>
    <td>8.73</td>
    <td>29.56</td>
  </tr>
  <tr>
    <td><b>PAG(GPT-4o,M2DIA)</b></td>
    <td><b>16.17</b></td>
    <td><b>43.86</b></td>
    <td><b>8.77</b></td>
    <td><b>30.66</b></td>
  </tr>
  <tr>
    <td>PAG(GPT-4o,M2DIA) w/o \( K_d \)</td>
    <td>14.12</td>
    <td>35.68</td>
    <td>7.13</td>
    <td>25.62</td>
  </tr>
  <tr>
    <td><b>PAG(gem2,M2DIA)</b></td>
    <td><b>14.86</b></td>
    <td><b>44.79</b></td>
    <td><b>8.95</b></td>
    <td><b>30.77</b></td>
  </tr>
  <tr>
    <td>PAG(gem2,M2DIA) w/o \( K_d \)</td>
    <td>12.63</td>
    <td>31.22</td>
    <td>6.24</td>
    <td>22.57</td>
  </tr>
  <tr>
    <td><b>PAG(mistral,M2DIA)</b></td>
    <td><b>17.28</b></td>
    <td><b>40.70</b></td>
    <td><b>8.14</b></td>
    <td><b>29.74</b></td>
  </tr>
  <tr>
    <td>PAG(mistral,M2DIA) w/o \( K_d \)</td>
    <td>14.31</td>
    <td>34.38</td>
    <td>6.87</td>
    <td>25.12</td>
  </tr>
</table>

Table 5: Experiments result of the unseen test dataset. Metrics “R”, “P”, and “ND” denote Recall, Precision, and NDCG respectively. A@1 denotes the top-1 score for all metrics. In PAG(a,b), a and b are LLMs and predictive models respectively.

The RAG approach uses association rules to rank diseases by counting the overlap between PolyMed knowledge data and patient symptoms, then retrieves the top 20 disease information to generate diagnoses. However, the RAG-based diagnosis generation method showed overall lower performance than PAG. This suggests that in disease diagnosis, inference-based tools are more crucial for performance than knowledge retrieval.

5.5 LLM Knowledge Impact

Table 4 also evaluates the effect of LLM knowledge by comparing the full PAG model to a variant without its knowledge generation component (w/o \( K_d \)). The full model outperforms in Top-3 and Top-5 metrics, while the variant shows an advantage in Top-1, suggesting that LLM knowledge enhances the overall quality of the prediction pool. Table 5 focuses on the PolyMed unseen test, where diseases absent from the training data must be inferred using medical knowledge. In this test, the PAG without the knowledge generation, which is more reliant on the predictive model, performs worse. By integrating the predictive model with LLM-derived knowledge, the full PAG model demonstrates enhanced performance, even surpassing the standalone performance of the LLM on some metrics. Moreover, LLM knowledge generation improves both predictive accuracy and explainability. This is especially valuable in applications like disease prediction, where understanding the rationale behind a prediction is critical. Examples of LLM-generated knowledge and prediction cases are provided in Table 14 in the Appendix.

![Graph showing PAG performance as the number of disease candidates varies. The dotted lines represent individual open-source LLM (with M2DIA) performance, while the black line shows the average.](page_1012_120_410_312.png)

Figure 4: Graph showing PAG performance as the number of disease candidates varies. The dotted lines represent individual open-source LLM (with M2DIA) performance, while the black line shows the average.

![Ablation study presenting average Top-1 and Top-5 metrics for three model groups (Mistral123b, GPT-4o, Mistral8x22b). Bars compare the full PAG model with variants excluding LLM knowledge (\( K_d \)), the predictive model (PM), or the large language model (LLM).](page_1012_670_410_312.png)

Figure 5: Ablation study presenting average Top-1 and Top-5 metrics for three model groups (Mistral123b, GPT-4o, Mistral8x22b). Bars compare the full PAG model with variants excluding LLM knowledge (\( K_d \)), the predictive model (PM), or the large language model (LLM).

5.6 Impact of Information Quantity

We evaluated performance by varying the number of disease predictions fed to the open-source LLM with PAG(M2DIA) using the mean of all metrics’ WAM scores on the PolyMed dataset (see Figure 4). We excluded the LLM’s knowledge generation of PAG (w/o \( K_d \)) to solely assess its ability to process increasing amounts of disease candidates. While performance initially improves, gains plateau or even decline beyond a certain point. Optimal performance is achieved with 25 to 35 disease predictions. This behavior likely reflects the LLM’s limited capacity to process excessive information and inherent diagnostic constraints, indicating that the number of predictions in PAG should be treated as a hyperparameter to optimize.
5.7 Ablation Study

We conducted an ablation study to evaluate the contributions of components in the PAG framework. We compare the full PAG model with three ablated variants: (1) removing the knowledge generation module (w/o \( K_d \)), (2) using only the LLM without the predictive model and knowledge(w/o \( K_d \), PM), and (3) using only the predictive model without LLM and knowledge (w/o \( K_d \), LLM). Performance was measured using the average of Recall, Precision, and NDCG at Top-1 and Top-5 ranks.

As shown in Figure 5, the full PAG model achieves better Top-1 performance than the LLM-only variant but is below the predictive model alone and the version without knowledge. This is expected, as predictive models are trained to optimize single-label accuracy, while LLMs are better suited for broader reasoning. In Top-5 metrics, however, full PAG consistently outperforms all other variants. This demonstrates that integrating LLM-generated knowledge enhances the overall quality of candidate predictions.

5.8 Case Study

To illustrate how PAG refines disease prediction through step-wise reasoning, we present a case in Table 6. The input includes structured symptom data indicating a burning sensation behind the breastbone and stomach ache, with the absence of nausea and acid reflux. Although the predictive model identifies several potential conditions, including coronary heart disease and esophagitis, the final diagnosis is refined through knowledge and reasoning powered by large language models.

Using retrieved knowledge about symptom profiles, the LLM identifies esophagitis as the most plausible diagnosis, based on the specificity of the chest-burning symptom and lack of cardiac-related indicators. This demonstrates the ability of PAG to incorporate both statistical predictions and contextual medical reasoning. The refined prediction matches the ground truth, highlighting the value of LLM knowledge in narrowing down ambiguous or overlapping diagnostic candidates. Additionally, our PAG enhances explainability by providing the rationale behind its reasoning, based on the knowledge it has generated about each condition, thereby clarifying the cognitive process that leads to the final diagnosis. These strengths demonstrate the potential of PAG for real-world deployment in medical tasks where trust and reliability are critical.

<table>
  <tr>
    <th>Component</th>
    <th>Content</th>
  </tr>
  <tr>
    <td>Meta Data</td>
    <td>( 'symptoms': { "Burning sensation behind the breastbone": True, "Nausea": False, "Acid reflux": False, "Stomach ache": True }, 'disease': 'Esophageal' )</td>
  </tr>
  <tr>
    <td>Target Disease</td>
    <td>Esophagitis</td>
  </tr>
  <tr>
    <td>PM Prediction</td>
    <td>[Coronary heart disease, Esophagitis, Rhinitis, Pneumonia, Enteritis, ...] Coronary heart disease: [Chest pain, Shortness of breath, Fatigue, Nausea, Light-headedness, Pain in arms or shoulders] Esophagitis: [Difficulty swallowing, Chest pain, Painful swallowing, Heartburn, Acid regurgitation, Sore throat, Hoarseness, Cough, Nausea, Vomiting, Food getting stuck in the esophagus]</td>
  </tr>
  <tr>
    <td>LLM Reasoning</td>
    <td>Burning sensation behind the breastbone and a stomach ache, in the absence of nausea and acid reflux, align most closely with esophagitis. While nausea and acid reflux are common in esophagitis, the central symptom of chest pain and a burning sensation (heartburn and stomach ache are highly indicative of inflammation in the esophagus) are more specific to the diagnosis of esophagitis. Other related conditions, such as coronary heart disease, are less likely due to the absence of other critical symptoms like shortness of breath, pain in arms or shoulders, and light-headedness. This diagnosis fits the noted symptoms and rules out the presence of symptoms more indicative of other diseases on the list. : 1 : Esophagitis</td>
  </tr>
  <tr>
    <td>Refined Result</td>
    <td>Esophagitis</td>
  </tr>
</table>

Table 6: Case study illustrating step-wise reasoning from prediction to LLM-based refinement in PAG.

6 Conclusion

We proposed the Prediction-Augmented Generation (PAG) technique, which utilizes predictive models as an external tool to support large language models in inference-driven tasks. We evaluated our method on a wide range of disease diagnosis tasks and found that it can effectively improve not only the diagnosis generation of large language models, but also the prediction performance of predictive models. In particular, we found that the intervention of large language models effectively improves the quality of the disease prediction pool. As a result, PAGs can be utilized as a key technology to effectively combine prediction models and LLMs to improve the reliability of medical diagnosis, and can contribute as a powerful tool for practical medical applications. Our approach is based on recent research trends suggesting that generation through multi-step reasoning can enhance the predictive performance of large language models, and we demonstrated that integrating the output of a predictive model into this process enables more accurate predictions. In future work, we will refine the integration method and evaluate PAG’s effectiveness across various prediction tasks.

Limitations

Knowledge Integration. Our PAG effectively combines predictive models with LLM-derived knowledge to enhance predictive capability and overall performance. However, further research is needed to optimize the fusion of prediction results with LLM knowledge. Our experiments show that while PAG excels at generating multiple disease candidates, its ability to pinpoint the single
most probable diagnosis is comparatively weaker. This trade-off indicates a need for improved fusion strategies, and future work should focus on developing knowledge integration methods that can simultaneously optimize both Top-1 and Top-5 performance.

Controllability. Despite the robust language generation capabilities of LLMs, they often struggle with controllability. Our case studies in Table 12 show that LLM-generated text sometimes fails to follow the rules and structure provided in the prompts, and smaller, quantized models can make errors like typos. Addressing these issues could lead to more robust and reliable performance.

Explainability. PAG offers meaningful explainability by incorporating explicit knowledge generation and rationale into its diagnostic output. However, this explainability is derived not from a direct analysis of the model’s internal computations, but from conventional post-hoc methods that generate coherent textual explanations. Thus, while PAG’s explainability enhances trust and transparency in its predictions, it does not fully elucidate all internal processes. To overcome this limitation, further integration of inherently explainable AI techniques is required.

Acknowledgments

This work was partly supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No.RS-2022-00155885, Artificial Intelligence Convergence Innovation Human Resources Development (Hanyang University ERICA)) and the MSIT(Ministry of Science and ICT), Korea, under the Convergence security core talent training business support program(IITP-2024-RS-2024-00423071) supervised by the IITP(Institute of Information & Communications Technology Planning & Evaluation)

References

Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. 2019. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD ’19, page 2623–2631, New York, NY, USA. Association for Computing Machinery.

Sercan Ö Arik and Tomas Pfister. 2021. Tabnet: Attentive interpretable tabular learning. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 6679–6687.

Emily M Bender, Timnit Gebru, Angelina McMillan-Major, and Shmargaret Shmitchell. 2021. On the dangers of stochastic parrots: Can language models be too big?. In Proceedings of the 2021 ACM conference on fairness, accountability, and transparency, pages 610–623.

Tom B Brown. 2020. Language models are few-shot learners. arXiv preprint arXiv:2005.14165.

Cheng-Han Chiang and Hung-Yi Lee. 2023. Can large language models be an alternative to human evaluations? In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15607–15631.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457.

Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.

Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. 2023. Pal: Program-aided language models. In International Conference on Machine Learning, pages 10764–10799. PMLR.

Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. 2020. Retrieval augmented language model pre-training. In International conference on machine learning, pages 3929–3938. PMLR.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.

Lei Huang, Weijiang Yu, Weitao Ma, Weihong Zhong, Zhangyin Feng, Haotian Wang, Qianglong Chen, Weihua Peng, Xiaocheng Feng, Bing Qin, et al. 2023. A survey on hallucination in large language models: Principles, taxonomy, challenges, and open questions. arXiv preprint arXiv:2311.05232.

Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. 2024a. Mixtral of experts. arXiv preprint arXiv:2401.04088.

Che Jiang, Biqing Qi, Xiangyu Hong, Dayuan Fu, Yang Cheng, Fandong Meng, Mo Yu, Bowen Zhou, and Jie Zhou. 2024b. On large language models’ hallucination with regard to known facts. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics:
Human Language Technologies (Volume 1: Long Papers), pages 1041–1053.

Chan-Yang Ju and Dong-Ho Lee. 2024. Model-agnostic confidence measurement for aggregating multimodal ensemble models in automatic diagnostic systems. Methods, 231:103–114.

Chan-Yang Ju, Ji-Sung Park, and Dong-Ho Lee. 2023. M2-dia: Enhancing diagnostic capabilities in imbalanced disease data using multimodal diagnostic ensemble framework. In 2023 IEEE International Conference on Bioinformatics and Biomedicine (BIBM), pages 1194–1201. IEEE.

ChanYang Ju and Dongho Lee. 2023. PolyMed: A Medical Dataset Addressing Disease Imbalance for Robust Automatic Diagnosis Systems.

Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. 2019. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of naacl-HLT, volume 1, page 2. Minneapolis, Minnesota.

Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. 2020. Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33:9459–9474.

Chin-Yew Lin. 2004. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pages 74–81.

Shuai Lin, Pan Zhou, Xiaodan Liang, Jianheng Tang, Ruihui Zhao, Ziliang Chen, and Liang Lin. 2021. Graph-evolving meta-learning for low-resource medical dialogue generation. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 13362–13370.

Yang Liu, Dan Iter, Yichong Xu, Shuohang Wang, Ruochen Xu, and Chenguang Zhu. 2023. G-eval: Nlg evaluation using gpt-4 with better human alignment. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 2511–2522.

Nick McKenna, Tianyi Li, Liang Cheng, Mohammad Javad Hosseini, Mark Johnson, and Mark Steedman. 2023. Sources of hallucination by large language models on inference tasks. arXiv preprint arXiv:2305.14552.

Grégoire Mialon, Roberto Dessì, Maria Lomeli, Christoforos Nalmpantis, Ram Pasunuru, Roberta Raileanu, Baptiste Rozière, Timo Schick, Jane Dwivedi-Yu, Asli Celikyilmaz, et al. 2023. Augmented language models: a survey. arXiv preprint arXiv:2302.07842.

Mistral AI Team. 2024. Large enough: Mistral large 2 announcement. Accessed: 19-Sept-2024.

Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christopher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, et al. 2021. Webgpt: Browser-assisted question-answering with human feedback. arXiv preprint arXiv:2112.09332.

OpenAI. 2023. Gpt-3.5 turbo. Documentation for GPT-3.5-turbo, optimized for performance and cost-efficiency.

OpenAI. 2024. Hello gpt-4o. Introducing GPT-4o, a variant of GPT-4o optimized for efficiency and cost.

Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 311–318.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9.

David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R Bowman. 2023. Gpqa: A graduate-level google-proof q&a benchmark. arXiv preprint arXiv:2311.12022.

Tushar Sarkar. 2022. Xbnet: An extremely boosted neural network. Intelligent Systems with Applications, 15:200097.

Timo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettle-moyer, Nicola Cancedda, and Thomas Scialom. 2024. Toolformer: Language models can teach themselves to use tools. Advances in Neural Information Processing Systems, 36.

Xiaoming Shi, Jie Xu, Jinru Ding, Jiali Pang, Sichen Liu, Shuqing Luo, Xingwei Peng, Lu Lu, Haihong Yang, Mingtao Hu, et al. 2023. Llm-mini-cex: Automatic evaluation of large language model for diagnostic conversation. arXiv preprint arXiv:2308.07635.

Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. 2024. Gemma 2: Improving open language models at a practical size. arXiv preprint arXiv:2408.00118.

Abhisek Tiwari, Sriparna Saha, and Pushpak Bhat-tacharyya. 2022. A knowledge infused context driven dialogue agent for disease diagnosis using hierarchical reinforcement learning. Knowledge-Based Systems, 242:108292.

Ehsan Ullah, Anil Parwani, Mirza Mansoor Baig, and Rajendra Singh. 2024. Challenges and barriers of using large language models (llm) such as chat-gpt for diagnostic medicine with a focus on digital pathology—a recent scoping review. Diagnostic pathology, 19(1):43.
Zhongyu Wei, Qianlong Liu, Baolin Peng, Huaxiao Tou, Ting Chen, Xuan-Jing Huang, Kam-Fai Wong, and Xiang Dai. 2018. Task-oriented dialogue system for automatic diagnosis. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 201–207.

Lin Xu, Qixian Zhou, Ke Gong, Xiaodan Liang, Jian-heng Tang, and Liang Lin. 2019a. End-to-end knowledge-routed relational dialogue system for automatic diagnosis. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 7346–7353.

Lin Xu, Qixian Zhou, Ke Gong, Xiaodan Liang, Jian-heng Tang, and Liang Lin. 2019b. End-to-end knowledge-routed relational dialogue system for automatic diagnosis. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 7346–7353.

Jing Zhang, Xiaokang Zhang, Daniel Zhang-Li, Jifan Yu, Zijun Yao, Zeyao Ma, Yiqi Xu, Haohua Wang, Xiaohan Zhang, Nianyi Lin, et al. 2023. Glim-dialog: Noise-tolerant pre-training for knowledge-grounded dialogue generation. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 5564–5575.

Yu Zhao, Yunxin Li, Yuxiang Wu, Baotian Hu, Qingcai Chen, Xiaolong Wang, Yuxin Ding, and Min Zhang. 2022. Medical dialogue response generation with pivotal information recalling. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 4763–4771.

A Experimental Details

We run all quantized open-source LLMs using the Ollama framework, which provides streamlined deployment and execution of large language models with quantized weights. Experiments are conducted on a workstation equipped with three NVIDIA A6000 GPUs (48 GB each). All LLMs are executed using Ollama’s default inference parameters, including decoding strategy, temperature, and token limits.

A.1 Implementations Details

Data Restoration Test. We constructed a patient symptom table by integrating implicit and explicit symptom attributes from the MDD, MZ, and DXY datasets. We generated synthetic patient utterance data using the GPT-4o model based on this consolidated information. To ensure data quality, we employed a dual-validation process. Initially, the GPT-4o model generated patient utterances based on the integrated symptom data. Subsequently, we prompted GPT-4o to review its own output by comparing them with the original data, verifying that no symptom information was omitted. Only the utterances that passed this validation process were used in our experiments. The prompt used for data generation is provided in Table 15, and sample utterances are shown in Table 13. Additionally, OpenAI’s text-embedding-ada-002 was employed for obtaining word embeddings during the Feature Alignment process.

Structured Data Evaluation. To evaluate PAG on the original data of MDD, MZ, DXY, and PolyMed, we replaced the patient utterance information with structured patient data. Using these original data, a predictive model generated 30 disease predictions, which were then used by the LLM to produce related knowledge. Finally, we generate the diagnosis of PAG based on structured patient data, inferences from the predictive model, and the disease knowledge generated by LLM.

LLM Prediction. For the label prediction task of all evaluation datasets, we utilized our proposed feature alignment module. Patient information was provided as a prompt to the LLM, which then generated five candidate disease texts in a structured format. The generated disease texts were compared with a predefined disease label dictionary to retrieve the 15 most similar terms. Consequently, the retrieved candidate terms and the target disease text are combined in the prompt, enabling the LLM to select the most appropriate term. OpenAI’s text-embedding-ada-002 was also used for word embedding.

RAG Prediction. We employed two Retrieval-augmented generation (RAG) methods in this study. The first method is based on a document embedding search and is applied to the MDD, MZ, and DXY datasets, which do not provide explicit disease knowledge. In this approach, we first integrate the implicit and explicit symptom attributes from the datasets to construct a comprehensive representation of patient symptoms, which is then embedded. Concurrently, we collect documents containing disease overviews and symptom information from medical institutions (See Table 8) and embed these as well. By comparing the patient symptom embedding with the document embeddings, we retrieve the five most similar documents and provide them as input to the LLM for diagnosis. The embedding model used is OpenAI’s
text-embedding-ada-002. The second method employs association rule-based retrieval on the PolyMed dataset, which consists of predefined medical terms for each disease. This approach constructs all combinations of patient symptoms and disease symptoms, calculates the intersection count between these sets to rank the diseases, and then supplies the top 20 disease knowledge entries to the LLM prompt.

A.2 Dataset

<table>
  <tr>
    <th>Dataset</th>
    <th>#Train</th>
    <th>#Test</th>
    <th>#Diseases</th>
    <th>#Symptoms</th>
  </tr>
  <tr>
    <td>MDD</td>
    <td>1912</td>
    <td>223</td>
    <td>12</td>
    <td>118</td>
  </tr>
  <tr>
    <td>MZ</td>
    <td>568</td>
    <td>142</td>
    <td>4</td>
    <td>66</td>
  </tr>
  <tr>
    <td>DXY</td>
    <td>423</td>
    <td>104</td>
    <td>5</td>
    <td>44</td>
  </tr>
</table>

Table 7: Statistic information of MDD, MZ, DXY. "#Train" and "#Test" are the total number of train and test data, and "#Diseases" and "#Symptoms" indicate the number of unique diseases and symptoms.

Source
https://www.hopkinsmedicine.org/health/conditions-and-diseases
https://my.clevelandclinic.org/health/diseases
https://www.mayoclinic.org/diseases-conditions
https://my.clevelandclinic.org/health/articles
https://www.msdmanuals.com/professional/ear-nose-and-throat-disorders
https://www.ninds.nih.gov/health-information/disorders
https://www.stanfordchildrens.org/en/topic
https://kidshealth.org/en/parents

Table 8: Knowledge source of MDD, MZ, DXY for RAG.

<table>
  <tr>
    <th>Dataset</th>
    <th>#Data</th>
    <th>#Diseases</th>
    <th>#Symptoms</th>
  </tr>
  <tr>
    <td>Train</td>
    <td>3636</td>
    <td>57</td>
    <td>352</td>
  </tr>
  <tr>
    <td>Single test</td>
    <td>909</td>
    <td>57</td>
    <td>248</td>
  </tr>
  <tr>
    <td>Unseen test</td>
    <td>538</td>
    <td>74</td>
    <td>290</td>
  </tr>
  <tr>
    <td>Multi test</td>
    <td>867</td>
    <td>53</td>
    <td>179</td>
  </tr>
  <tr>
    <th>Total</th>
    <th>5950</th>
    <th>131</th>
    <th>433</th>
  </tr>
</table>

Table 9: Statistic information of PolyMed-case. "#Data" is the total number of data, and "#Diseases" and "#Symptoms" indicate the number of unique diseases and symptoms.

MDD, MZ, DXY We evaluated disease diagnosis capability in a single-task setting using the public datasets MDD, MZ, and DXY. These datasets were originally designed to assess conversational automatic diagnosis systems, including tasks that extract implicit symptoms from explicit symptoms. However, to measure PAG’s diagnostic performance, we employ a symptom set that integrates both explicit and implicit symptoms. Detailed information about each dataset is provided in Table 7.

PolyMed The PolyMed dataset is used to evaluate automatic diagnosis systems from three perspectives.

The first is a Single test to evaluate the disease prediction performance of the model on diseases learned from the Train dataset. This evaluation is equivalent to evaluating the Train-test split data, as is the evaluation of a typical learning-based model.

The second is the Unseen test, which requires the model to make predictions on diseases that it has not learned from the Train data. This task is solved by utilizing disease knowledge data contained in PolyMed data, which is derived from the challenge of supplementing the lack of disease data with knowledge information to overcome the limitation that training data for many common ADS diseases are not available.

The last task is Multi test, which is a test for cases where a patient is diagnosed with more than one suspected disease, and requires prediction for multiple related diseases.

These tasks can evaluate ADS from different angles and allow for a comprehensive assessment. Table 9 provides statistical information about the PolyMed data.

A.3 Metrics

Recall@K Recall@K measures the fraction of correct diseases retrieved in the top K predictions relative to the total number of correct diseases. It is defined as:

\[
\text{Recall@K} = \frac{\sum_{i=1}^{K} \mathbf{1}\{\text{disease}_i \text{ is correct}\}}{\text{Total number of correct diseases}}
\]

where \( \mathbf{1}\{\cdot\} \) is an indicator function. In the data restoration test, "diseases" can be replaced with "symptoms".

Precision@K Precision@K quantifies the proportion of correct diseases among the top K predictions:

\[
\text{Precision@K} = \frac{\sum_{i=1}^{K} \mathbf{1}\{\text{disease}_i \text{ is correct}\}}{K}
\]

NDCG@K The Normalized Discounted Cumulative Gain at K (NDCG@K) evaluates the ranking quality of predicted diseases. In our setting, every
correct disease is assigned the same relevance score (i.e., \( rel_i = 1 \)). Consequently, the Discounted Cumulative Gain at \( K \) (DCG@K) simplifies to:

\[
\text{DCG}@K = \sum_{i=1}^{K} \frac{1\{\text{disease}_i \text{ is correct}\}}{\log_2(i + 1)},
\]

The Ideal DCG at \( K \) (IDCG@K) is computed by placing all correct diseases in the top \( n \) positions:

\[
\text{IDCG}@K = \sum_{i=1}^{n} \frac{1}{\log_2(i + 1)},
\]

where \( n = \min(K, \text{Total number of correct diseases}) \). NDCG@K is then given by:

\[
\text{NDCG}@K = \frac{\text{DCG}@K}{\text{IDCG}@K}
\] (12)

Weighted Arithmetic Mean To collectively evaluate the three assessments of PolyMed data, we used the Weighted Arithmetic Mean method to calculate a composite score. WAM score can calculated by:

\[
R = \{Single, Unseen, Multi\},
\]
\[
w_i = 1 - \frac{R_i}{\sum_{i \in R} R_i},
\]
\[
WAM = \frac{\sum_{i \in R} w_i R_i}{\sum_{i \in R} w_i}
\] (13)

where \( R_i \) is the score of each \( Recall@k \), \( Precision@k \), and \( NDCG@k \) results for the three test datasets. This weighting system penalizes imbalanced improvements, leading to a decrease in the \( WAM \) score if there is an unequal enhancement in the scores. This approach prevents an exclusive focus on specific task optimization to improve model performance.

A.4 Details of baselines

A.4.1 Predictive Models

LDA, XBNet, TabNet For LDA, we used a Singular Value Decomposition solver to fit the data. XBNet (Sarkar, 2022) was trained with the Adam optimizer at a learning rate of 0.01 over 500 epochs using a two-linear-layer architecture, and the best-performing model was selected. TabNet (Arik and Pfister, 2021) was optimized using Optuna (Akiba et al., 2019) with a maximum of 3000 epochs and a batch size of 256; from 1000 training runs, the model yielding the highest performance was chosen.

M2DIA M2DIA (Ju et al., 2023) is a predictive model for the PolyMed task, utilizing a multimodal ensemble approach that combines predictive modeling with knowledge inference techniques to achieve superior diagnostic performance. In our implementation, we adopted the original model’s ensemble structure, which integrates six prediction models with their corresponding weights.

MAEA MAEA (Ju and Lee, 2024) is a state-of-the-art model for the PolyMed task that refines the aggregation approach of M2-DIA. It introduces a model-agnostic confidence measurement to quantify the reliability of each modality’s contribution and dynamically weight ensemble predictions.

A.4.2 Large Language Models

Llama 3.1 Meta’s Llama 3.1 was used in two configurations: an 8B model (ollama’s llama3.1:8b-instruct-fp16, FB16 quantization, 16GB memory) and a 70B model (ollama’s llama3.1:70b-instruct-q8_0, Q8 quantization, 75GB memory).

Gemma 2 Gemma 2 is a 27B model from Google, deployed as ollama’s 27b-instruct-fp16 with FB16 quantization (54GB memory).

Mixtral Mixtral comprises eight 22B models, implemented as ollama’s mixtral:8x22b with Q4 quantization (80GB memory).

Mistral-large Mistral-large is a 123B model, provided as ollama’s mistral-large:123b with Q4 quantization (69GB memory).
<table>
  <tr>
    <th rowspan="2">Models</th>
    <th colspan="3">MDD</th>
    <th colspan="3">MZ</th>
    <th colspan="3">DXY</th>
  </tr>
  <tr>
    <th>R@1</th><th>R@3</th><th>ND@3</th>
    <th>R@1</th><th>R@3</th><th>ND@3</th>
    <th>R@1</th><th>R@3</th><th>ND@3</th>
  </tr>
  <tr>
    <td>LDA</td><td>75.73</td><td>84.93</td><td>81.10</td><td>67.60</td><td>98.59</td><td>86.78</td><td>79.80</td><td>97.11</td><td>90.34</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub> + PAG(LDA)</td><td><b>82.00</b></td><td>89.95</td><td><b>86.75</b></td><td><b>71.83</b></td><td>98.59</td><td><b>88.43</b></td><td><b>81.73</b></td><td>98.07</td><td><b>91.41</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub> + PAG(LDA)</td><td>70.71</td><td><b>90.37</b></td><td>82.46</td><td>68.30</td><td><b>100.00</b></td><td>87.56</td><td>74.03</td><td>98.07</td><td>88.57</td>
  </tr>
  <tr>
    <td>mixtral<sub>8×22b</sub> + PAG(LDA)</td><td>81.59</td><td><b>90.37</b></td><td><b>86.75</b></td><td>67.60</td><td>97.18</td><td>85.99</td><td>78.84</td><td>97.11</td><td>90.12</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub> + PAG(LDA)</td><td>66.94</td><td>89.95</td><td>80.75</td><td>64.08</td><td>97.88</td><td>84.76</td><td>75.00</td><td><b>99.03</b></td><td>89.53</td>
  </tr>
  <tr>
    <td>XBNet</td><td>77.82</td><td>90.37</td><td>85.19</td><td>69.71</td><td>97.18</td><td>86.67</td><td>71.15</td><td>98.07</td><td>87.51</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub> + PAG(XBNet)</td><td><b>80.75</b></td><td>92.05</td><td><b>87.71</b></td><td><b>71.12</b></td><td><b>99.29</b></td><td><b>88.43</b></td><td>76.92</td><td><b>99.03</b></td><td><b>90.75</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub> + PAG(XBNet)</td><td>76.56</td><td><b>93.72</b></td><td>86.68</td><td>68.30</td><td>95.07</td><td>84.73</td><td>66.34</td><td><b>99.03</b></td><td>86.09</td>
  </tr>
  <tr>
    <td>mixtral<sub>8×22b</sub> + PAG(XBNet)</td><td>76.56</td><td>92.88</td><td>86.26</td><td>69.01</td><td>97.18</td><td>86.23</td><td>76.92</td><td><b>99.03</b></td><td>90.12</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub> + PAG(XBNet)</td><td>66.10</td><td>92.46</td><td>82.02</td><td>66.90</td><td>97.18</td><td>85.73</td><td><b>77.88</b></td><td>98.07</td><td>90.12</td>
  </tr>
</table>

Table 10: Experimental results of Predictive Models (PM) and PAG on MDD, Muzhi, and DXY datasets, using Recall(R), Precision(P), and NDCG(ND) metrics.

<table>
  <tr>
    <th rowspan="2">Models</th>
    <th colspan="2">R@1</th>
    <th colspan="2">R@3</th>
    <th colspan="2">R@5</th>
    <th colspan="2">P@1</th>
    <th colspan="2">P@3</th>
    <th colspan="2">P@5</th>
    <th colspan="2">ND@1</th>
    <th colspan="2">ND@3</th>
    <th colspan="2">ND@5</th>
  </tr>
  <tr>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
    <th>R</th><th>P</th>
  </tr>
  <tr>
    <td>LDA</td><td>20.26</td><td>39.35</td><td>46.62</td><td>27.30</td><td>17.34</td><td>12.21</td><td>22.19</td><td>34.27</td><td>37.62</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub> + PAG(LDA)</td><td><b>22.36</b></td><td><b>42.78</b></td><td><b>51.87</b></td><td><b>29.63</b></td><td><b>18.59</b></td><td><b>13.42</b></td><td><b>24.33</b></td><td><b>37.22</b></td><td><b>41.33</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>8b</sub> + PAG(LDA)</td><td>18.88</td><td>37.88</td><td>47.95</td><td>25.17</td><td>16.62</td><td>12.52</td><td>20.59</td><td>32.58</td><td>37.16</td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub> + PAG(LDA)</td><td>19.43</td><td>41.19</td><td>50.97</td><td>25.88</td><td>18.01</td><td>13.26</td><td>21.18</td><td>34.85</td><td>39.32</td>
  </tr>
  <tr>
    <td>mixtral<sub>8×22b</sub> + PAG(LDA)</td><td>22.04</td><td>41.61</td><td>50.15</td><td>29.11</td><td>18.16</td><td>13.06</td><td>23.95</td><td>36.26</td><td>40.13</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub> + PAG(LDA)</td><td>21.48</td><td>41.25</td><td>50.45</td><td>28.43</td><td>17.98</td><td>13.13</td><td>23.35</td><td>35.85</td><td>40.06</td>
  </tr>
  <tr>
    <td>XBNet</td><td>18.96</td><td>38.47</td><td>49.59</td><td>27.36</td><td>17.33</td><td>13.07</td><td>21.31</td><td>33.71</td><td>38.75</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub> + PAG(XBNet)</td><td><b>20.57</b></td><td><b>46.67</b></td><td><b>58.40</b></td><td><b>27.77</b></td><td><b>20.29</b></td><td><b>15.09</b></td><td><b>22.53</b></td><td><b>38.79</b></td><td><b>44.11</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>8b</sub> + PAG(XBNet)</td><td>17.91</td><td>40.57</td><td>54.08</td><td>24.77</td><td>17.93</td><td>14.06</td><td>19.81</td><td>33.95</td><td>40.07</td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub> + PAG(XBNet)</td><td>18.75</td><td>44.00</td><td>56.80</td><td>25.53</td><td>19.20</td><td>14.70</td><td>20.63</td><td>36.24</td><td>42.02</td>
  </tr>
  <tr>
    <td>mixtral<sub>8×22b</sub> + PAG(XBNet)</td><td>20.51</td><td>44.76</td><td>56.10</td><td>27.26</td><td>19.50</td><td>14.53</td><td>22.37</td><td>37.49</td><td>42.62</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub> + PAG(XBNet)</td><td>19.43</td><td>45.87</td><td>57.25</td><td>25.81</td><td>19.98</td><td>14.77</td><td>21.17</td><td>37.72</td><td>42.85</td>
  </tr>
  <tr>
    <td>TabNet</td><td>17.05</td><td>31.77</td><td>42.04</td><td>25.30</td><td>14.92</td><td>11.35</td><td>19.35</td><td>28.68</td><td>33.32</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub> + PAG(TabNet)</td><td><b>20.03</b></td><td><b>42.74</b></td><td><b>53.78</b></td><td><b>27.84</b></td><td><b>18.82</b></td><td><b>13.98</b></td><td><b>22.18</b></td><td><b>36.35</b></td><td><b>41.36</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>8b</sub> + PAG(TabNet)</td><td>15.72</td><td>35.71</td><td>46.33</td><td>22.86</td><td>16.07</td><td>12.28</td><td>17.68</td><td>30.20</td><td>35.04</td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub> + PAG(TabNet)</td><td>16.49</td><td>37.90</td><td>49.79</td><td>23.55</td><td>17.05</td><td>13.13</td><td>18.45</td><td>31.92</td><td>37.30</td>
  </tr>
  <tr>
    <td>mixtral<sub>8×22b</sub> + PAG(TabNet)</td><td>18.33</td><td>39.63</td><td>50.16</td><td>25.34</td><td>17.66</td><td>13.17</td><td>20.26</td><td>33.74</td><td>38.52</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub> + PAG(TabNet)</td><td>18.45</td><td>40.13</td><td>50.12</td><td>25.43</td><td>17.84</td><td>13.17</td><td>20.37</td><td>34.14</td><td>38.65</td>
  </tr>
  <tr>
    <td>MAEA</td><td><b>25.78</b></td><td>47.56</td><td>58.82</td><td><b>34.65</b></td><td>20.75</td><td>15.22</td><td><b>28.21</b></td><td>41.90</td><td>46.96</td>
  </tr>
  <tr>
    <td>gemma2<sub>27b</sub> + PAG(MAEA)</td><td>25.43</td><td><b>52.29</b></td><td><b>63.79</b></td><td>32.91</td><td><b>22.42</b></td><td><b>16.36</b></td><td>27.46</td><td><b>44.21</b></td><td><b>49.38</b></td>
  </tr>
  <tr>
    <td>llama3.1<sub>8b</sub> + PAG(MAEA)</td><td>21.92</td><td>47.71</td><td>60.66</td><td>28.89</td><td>20.56</td><td>15.60</td><td>23.81</td><td>39.78</td><td>45.61</td>
  </tr>
  <tr>
    <td>llama3.1<sub>70b</sub> + PAG(MAEA)</td><td>23.05</td><td>48.45</td><td>60.95</td><td>30.21</td><td>20.90</td><td>15.67</td><td>25.00</td><td>40.90</td><td>46.50</td>
  </tr>
  <tr>
    <td>mixtral<sub>8×22b</sub> + PAG(MAEA)</td><td>23.73</td><td>49.12</td><td>62.18</td><td>30.68</td><td>21.14</td><td>15.97</td><td>25.59</td><td>41.75</td><td>47.63</td>
  </tr>
  <tr>
    <td>mistral-large<sub>123b</sub> + PAG(MAEA)</td><td>23.06</td><td>49.18</td><td>61.97</td><td>29.86</td><td>21.21</td><td>15.94</td><td>24.88</td><td>41.30</td><td>47.08</td>
  </tr>
</table>

Table 11: WAM score of Predictive Models (PM) and PAG with three tests of PolyMed. Metrics "R", "P", and "ND" denote Recall, Precision, and NDCG respectively.
Utterance I’ve been feeling really tired and without much energy lately. I’ve also been throwing up and having really watery stools. I don’t have much of an appetite, either.
Extraction {"symptoms": {"Fatigue": True, "Nausea": False, "Vomiting": True, "Diarrhea": True, "Loss of appetite": True}}
Validation {"symptoms": {"Fatigue": True, "Nausea": True, "Vomiting": True, "Diarrhea": True, "Loss of appetite": True}}

Utterance I’ve been having a lot of trouble lately with persistent coughing and I’m also coughing up quite a bit of phlegm. On top of that, I’m experiencing some shortness of breath, and there’s this uncomfortable feeling in my throat that just won’t go away. It’s really starting to worry me.
Extraction {"symptoms": {"Persistent coughing": True, "Coughing up phlegm": True, "Shortness of breath": True, "Uncomfortable feeling in the throat": True}}
Validation <NPB>

Utterance I’ve been having a really tough time with this persistent stuffiness lately. My nose feels completely congested, and there’s noticeable swelling inside. It’s like my nasal passages are constantly blocked and swollen, making it uncomfortable to breathe. However, I haven’t had any fever or experienced any double vision.
Extraction Here is the structured summary of the patients reported symptoms:{“symptoms”: {“Stuffy nose”: True, “Nasal congestion”: True, “Swelling in nasal passages”: True, “Difficulty breathing”: True, “Fever”: False, “Double vision”: False }}
Validation The following information from the text data is missing in the structured data: * "persistent stuffiness"* "noticeable swelling inside"* "nasal passages are constantly blocked and swollen"Corrected Structured Data:{“symptoms”: {“Stuffy nose”: True,”Nasal congestion”: True,”Swelling in nasal passages”: True,”Difficulty breathing”: True,”Fever”: False,”Double vision”: False,”Persistent stuffiness”: True,”Noticeable swelling inside”: True,”Constant blockage and swelling of nasal passages”: True}}

Utterance "I’ve been dealing with a fever and these red spots that have turned into blisters. My mouth has been sore, and I really haven’t felt like eating anything. I’m quite worried because it seems to be progressively getting worse. Can you help me understand what’s going on?"
Extraction It seems you haven’t provided any text related to a patient’s dialogue with their doctor. Could you please provide the necessary information so I can assist you in extracting a structured summary of the patient’s reported symptoms, medical history, and specific concerns?
Validation <NPB>

Table 12: Feature extraction examples. Utterance is a patient symptom utterance generated by Table 15 prompt. Extraction and Validation is corresponding extracted feature and validation using Table 16 prompt. The red text shows examples of incorrect information extraction.
Original {"Restlessness": true, "Fever": true, "Rash": true, "Blisters": true, "Lethargy": true, "Coughing": false, "Vomiting": false, "Anorexia (Loss of appetite)": false}
Draft I’ve been feeling really restless and I’ve had this fever that just won’t go away. There’s also this rash that’s appeared and some blisters have started forming. I’m also feeling more tired and lethargic than usual. It’s been pretty uncomfortable and I’m not sure what’s going on.
Validation Lately, I’ve been feeling really restless and can’t shake off this persistent fever. I’ve noticed a rash and some blisters developing, and overall, I’ve been feeling much more tired and lethargic than usual. Despite these symptoms, I haven’t been experiencing any coughing or vomiting, and my appetite remains normal without any signs of anorexia.

Original {"Chest tightness": true, "Pain behind the breastbone": true, "Diarrhea": true, "Acid reflux": true}
Draft I’ve been dealing with some really uncomfortable chest tightness and pain right behind my breastbone. On top of that, I’ve been experiencing diarrhea and a lot of acid reflux. It’s starting to get pretty concerning, and I’m not sure what’s causing all of this.
Validation <NPB>

Table 13: Utterance generation examples. The Original is the symptom data of the ADS dataset. Draft and Validation are generated by the prompt of Table 15
User Prompt Disease information:
Mastitis - (51.25%) Breast swelling, breast pain, skin redness, warmth to touch, lump or thickening of tissue, fever, chills. Coronary heart disease - (44.31%) Chest pain (angina), Shortness of breath, Fatigue, Irregular heartbeat. Enteritis - (2.81%) Abdominal pain, Diarrhea, Nausea, Vomiting, Fever, Bloody stools, Loss of appetite, Weight loss, Fatigue. Pneumonia - (1.37%) cough, fever, chills, difficulty breathing/shortness of breath, chest pain when breathing or coughing, fatigue, nausea or vomiting, diarrhea, confusion (especially in older adults), loss of appetite. Asthma - (0.2%) Coughing (especially at night), Wheezing, Shortness of breath, Chest tightness or pain, Difficulty breathing, Fatigue. Esophagitis - (0.03%) Difficulty swallowing, chest pain, heartburn, regurgitation of food or liquids, sensation of a lump in the throat, chronic cough, hoarse voice, frequent hiccups. Traumatic brain injury - (0.02%) Headache, dizziness, confusion, nausea or vomiting, fatigue, speech difficulties, sleeping problems, memory loss, mood changes, sensitivity to light and sound, blurred vision, ringing in the ears, seizures. Thyroiditis - (0.01%) Neck swelling, Pain in the neck or throat, Fatigue, Weight gain, Cold intolerance, Constipation, Dry skin, Muscle weakness, Joint pain, Irregular menstrual periods, Depression, Thinning hair. Conjunctivitis - (0.0%) [Redness of the eye, Itching or burning sensation, Excessive tearing, Thick yellow discharge that crusts over the eyelashes during sleep, Blurred vision, Increased sensitivity to light, Swollen lymph nodes, Gritty feeling in the eyes] Rhinitis - (0.0%) ["Sneezing", "Runny nose", "Stuffy nose", "Coughing", "Itchy nose", "Postnasal drip", "Watery eyes", "Fatigue", "Headache"] External otitis - (0.0%) Ear pain, itching, discharge from ear, difficulty hearing, redness or swelling of outer ear. Dermatitis - (0.0%) [itchy skin, redness, swelling, rash, burning sensation, dry and scaly skin, blisters, crusting, oozing and weeping sores, skin color changes]
Patient symptom:
Lately, I’ve been feeling really fatigued and have a fever that just won’t go away. There are times when I feel dizzy and have even started vomiting. I’ve also noticed some discomfort in my throat and an unusual sensitivity to cold. However, I haven’t been experiencing a cough, incontinence, body aches, expectoration, or nausea. Can you help me figure out what’s going on? Diagnosis:
\( D^T \) Generation Based on the symptoms provided - fatigue, fever that won’t go away, dizziness, vomiting, throat discomfort, and unusual sensitivity to cold - the most likely diagnosis is Mastitis with a probability of 51.25%. The key symptoms that support this conclusion are the persistent fever, fatigue, and throat discomfort. These align well with the typical presentation of mastitis, which often includes breast swelling, pain, skin redness, warmth to touch, and fever or chills.
It is important to note that while these symptoms do not perfectly match those of Mastitis (there’s no mention of specific breast-related symptoms), they are the closest fit given the disease information provided. However, a comprehensive medical evaluation by a healthcare professional would be necessary for an accurate diagnosis.
\( D^L \) Generation Mastitis Traumatic brain injury Thyroiditis Rhinitis

Table 14: PAG Diagnostic Example. User Prompt supplies the LLM with diagnostic details, including the patient’s symptom utterance and “Disease information”—which aggregates the predictive model’s output with LLM-generated knowledge for each disease. The \( D^T \) Generation is the corresponding text-based rationale, while the \( D^L \) Generation represents the final disease label inferred from the user prompt and \( D^T \). The LLM Diagnosis was produced using system messages provided in Table 18.
<table>
  <tr>
    <th>Task</th>
    <th>Prompt</th>
  </tr>
  <tr>
    <td>Utterance Generation</td>
    <td>You are tasked with generating a patient’s utterance that conveys their symptoms to a doctor in a natural and convincing way. You will be provided with a diagnosis, a list of symptoms the patient has, and a list of symptoms the patient does not have. Your job is to create an utterance that reflects the patient’s condition by describing only the symptoms they experience, while avoiding the use of any disease names or direct references to medical diagnoses.<br>Make the response sound like a genuine patient seeking medical attention. Be concise but descriptive enough to convey the severity or concerns.<br>DO NOT INCLUDE INFORMATION THAT IS NOT IN THE PATIENT SYMPTOMS<br>DO NOT INCLUDE DOCTOR’S TALK. ONLY NEED PATIENT’S UTTERANCE<br>DO NOT INCLUDE DISEASE NAME IN THE UTTERANCE, ONLY INCLUDE SYMPTOMS<br>YOU MUST INCLUDE ALL THE PAITENT’S SYMTPOMS<br>Here is the patient’s information:</td>
  </tr>
  <tr>
    <td>Utterance Validation</td>
    <td>You are tasked with verifying whether a patient’s utterance fully captures both the symptoms they are experiencing and those they are not experiencing. You will be provided with two lists: one containing symptoms the patient has and one containing symptoms the patient does not have.<br>- If there is no missing symptoms, generate <NPB> token.<br>- If any symptoms (experienced or not experienced) are missing, rewrite the utterance to include all relevant information. The rewritten utterance should sound natural and reflect how a patient would describe both their symptoms and the absence of specific symptoms to the doctor.</td>
  </tr>
</table>

Table 15: system messages for utterance generation.
<table>
  <tr>
    <th>Task</th>
    <th>Prompt</th>
  </tr>
  <tr>
    <td>Feature Extraction</td>
    <td>Analyze this patient’s dialogue with their doctor and extract a structured summary of the patient’s reported symptoms, medical history, and any specific concerns they mention. Organize the information in a clear and concise format.<br>{#Feature schema example}<br>IF THERE IS NO CORRESPONDING INFORMATION, MARK IT AS NULL.<br>IF THERE ARE MULTIPLE ELEMENTS SEPARATE THEM WITH A COMMA.<br>ENTER SYMPTOMS AS A DICTIONARY TYPE AND VALUES AS True<br>IF THE PATIENT HAS THE SYMPTOM OR False IF THE PATIENT DOES NOT.<br>STRUCTURED DATA MUST BE DICTIONARY STRUCTURE.<br>EXTRACT INFORMATION WITHOUT EXPLANATIONS AND GREETINGS.<br>COLLECT INFORMATION ONLY ABOUT THE USER’S SYMTPOMS AND DO NOT COLLECT ANY OTHER INFORMATION SUCH AS MEDICAL HISTORY, SPECIFIC CONCERNS.<br>Here is the patient’s information:</td>
  </tr>
  <tr>
    <td>Feature Validation</td>
    <td>Given the following text data and corresponding structured data, check if all the information from the text data is accurately and completely included in the structured data. If any information is missing in structured data, specify which details are not present.<br>{#Feature schema example}<br>- If there is no missing information, generate <NPB> token.<br>- If there is missing information, add missing data to the structured data.<br>- Ensure there are no syntax errors. If any errors exist, correct them to make the output a valid Python dictionary.<br>STRUCTURED DATA MUST BE DICTIONARY STRUCTURE<br>EXTRACT STRUCTURED INFORMATION WITHOUT EXPLANATIONS AND GREETINGS AND TITLES</td>
  </tr>
</table>

Table 16: system messages for feature extraction.
<table>
  <tr>
    <th>Task</th>
    <th>Prompt</th>
  </tr>
  <tr>
    <td>Term match</td>
    <td>You are an expert in medical terminology. Given a query term and a list of retrieved terms, match the retrieved term that has the same meaning as the query term. Provide only the term that matches. Match the term with the same meaning as the query term. If there is no appropriate match, return <NW>.
Example Query Term: Hypertension
Retrieved Terms: High blood pressure Hypotension Tachycardia Bradycardia Arrhythmia Hyperglycemia Hypoglycemia Hypertrophy Hyperlipidemia Hypoxia
Correct Match: High blood pressure</td>
  </tr>
  <tr>
    <td>LLM Diagnosis</td>
    <td>You are an expert in medical diagnostics. Given patient information provided by a patient, predict up to 5 suspected diseases that the patient may have. Provide the list of predicted diseases without any numbering or explanation or hyphen.
Example Patient Information: Sex: Female Age: 45 Family History: Heart disease, high cholesterol Background: Non-smoker, regular exercise, no alcohol Underlying Disease: Hypertension Symptoms: Chest pain, shortness of breath, dizziness, fatigue, irregular heartbeat
Predicted Diseases: Common Cold Influenza Pneumonia Allergic Rhinitis Bronchitis</td>
  </tr>
  <tr>
    <td>RAG</td>
    <td>You are an expert in medical diagnostics. Given patient information provided by a patient, predict up to 5 suspected diseases that the patient may have by referring to the Disease Knowledge. Provide the list of predicted diseases without any numbering or explanation.
Example Patient Information: Sex: Female Age: 45 Family History: Heart disease, high cholesterol Background: Non-smoker, regular exercise, no alcohol Underlying Disease: Hypertension Symptoms: Chest pain, shortness of breath, dizziness, fatigue, irregular heartbeat
Disease Knowledge: Upper respiratory infections: Knowledge of Upper respiratory infections Influenza: Knowledge of Influenza Pneumonia: Knowledge of Pneumonia Allergic rhinitis: Knowledge of Allergic rhinitis Bronchitis: Knowledge of Bronchitis Sinusitis: Knowledge of Sinusitis Tuberculosis: Knowledge of Tuberculosis Asthma: Knowledge of Asthma COVID-19: Knowledge of COVID-19 Lung Cancer: Knowledge of Lung Cancer
Predicted Diseases: Upper respiratory infections Influenza Pneumonia Allergic Rhinitis Bronchitis</td>
  </tr>
</table>

Table 17: system messages for LLM prediction and RAG with Term match.
<table>
  <tr>
    <th>Task</th>
    <th>Prompt</th>
  </tr>
  <tr>
    <td><i>K_d</i> Generation</td>
    <td>What symptoms will appear in the {} disease? List only symptom, without any description, following form [symptom1, symptom2, ...].</td>
  </tr>
  <tr>
    <td><i>D^T</i> Generation</td>
    <td>Given the patient’s symptoms and relevant disease information, provide a concise diagnosis. Briefly explain the reasoning behind this diagnosis, identifying key symptoms, relevant medical history, or any notable patterns that support the conclusion. Do not include any disclaimers, warnings, or mentions of AI generation</td>
  </tr>
  <tr>
    <td><i>D^L</i> Generation</td>
    <td>You are an expert in medical diagnostics. Given patient information (pi) provided by a patient, predict suspected diseases that the patient may have by referring to the retrieved Disease information (di) and diagnosis (ds). Provide the list of predicted diseases without any numbering or explanation.<br>Example:<br>Patient information: pi<br>Disease information: di<br>Diagnosis: ds<br>Predicted Diseases: Upper respiratory infections Influenza Pneumonia Allergic Rhinitis Bronchitis</td>
  </tr>
  <tr>
    <td>PAG w/o <i>K_d</i></td>
    <td>You are an expert in medical diagnostics. Given a patient’s information, a query disease diagnosed by a predictive model, and a list of predicted diseases, remove any inaccurate or irrelevant diseases. If no modifications are necessary, return <NW>. Provide only the filtered list of diseases without any numbering or explanation.<br>Patient Information:<br>Sex: Female<br>Age: 35<br>Family History: History of respiratory diseases<br>Background: Non-smoker, no known allergies<br>Underlying Disease: None<br>Symptoms: Fever, cough, sore throat, fatigue, body aches<br>Predicted Diseases:<br>Common Cold Influenza Pneumonia Allergic Rhinitis Bronchitis Sinusitis Tuberculosis Asthma COVID-19 Lung Cancer<br>Filtered List:<br>Common Cold Influenza Pneumonia Allergic Rhinitis Bronchitis Sinusitis COVID-19</td>
  </tr>
</table>

Table 18: system messages for PAG.

Reviewer paper summaries: This paper proposes a novel approach called Prediction-Augmented Generation (PAG). PAG is designed to overcome the limitations of LLMs, particularly their lower performance in high-level reasoning and decision-making tasks. By leveraging an external predictive model, PAG enables more structured and accurate inference. PAG utilizes a task-specific predictive model as a tool to refine the outputs generated by LLMs, thereby improving predictive accuracy. This method demonstrates superior performance compared to the existing RAG approach. Experiments were conducted using various disease diagnosis datasets, evaluating performance through tasks such as data restoration, single disease diagnosis, and multiple disease diagnosis. The PAG methodology outperforms traditional predictive models and RAG, significantly enhancing prediction accuracy and reliability.
Summary of strengths: - Introduced the concept of Prediction-Augmented Generation (PAG) for the first time and provided a detailed analysis, laying the foundation for future research.
- Achieved state-of-the-art performance in disease diagnosis prediction, surpassing existing models.
Summary of weaknesses: - Since PAG enhances predictions by having the LLM generate additional explanations, the exact internal reasoning process may not be fully transparent. In particular, there is a risk that the explanations generated by the LLM could be overly general or exaggerated without proper medical grounding.
- PAG requires the simultaneous use of both a predictive model and an LLM, leading to significantly higher computational costs compared to using a standalone predictive model.
Comments, suggestions and typos: - It would be beneficial to include a qualitative evaluation of PAG results.
Overall assessment scores: 3
Metareview score: 3.5
Metareview: This paper proposes a approach called Prediction-Augmented Generation (PAG), which is designed to overcome the limitations of LLMs, particularly their lower performance in high-level reasoning and decision-making tasks. PAG achieved state-of-the-art performance in disease diagnosis prediction, surpassing existing models. However, PAG requires the simultaneous use of both a predictive model and an LLM, leading to significantly higher computational costs compared to using a standalone predictive model.
Summary of reasons to publish: PAG achieved state-of-the-art performance in disease diagnosis prediction, surpassing existing models.