Prediction-Augmented Generation for Automatic Diagnosis Tasks (2025)

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.

## 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 @xcite , 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 efficientize 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 @xcite . 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 @xcite . 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. ## 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.