LG AI Research & KAIST at EHRSQL 2024: Self-Training Large Language Models with Pseudo-Labeled Unanswerable Questions for a Reliable Text-to- SQL System on EHR s (2024)

Abstract: Text-to-SQL models are pivotal for making Electronic Health Records (EHRs) accessible to healthcare professionals without SQL knowledge. With the advancements in large language models, these systems have become more adept at translating complex questions into SQL queries. Nonetheless, the critical need for reliability in healthcare necessitates these models to accurately identify unanswerable questions or uncertain predictions, preventing misinformation. To address this problem, we present a self-training strategy using pseudo-labeled unanswerable questions to enhance the reliability of text-to-SQL models for EHRs. This approach includes a two-stage training process followed by a filtering method based on the token entropy and query execution. Our methodology’s effectiveness is validated by our top performance in the EHRSQL 2024 shared task, showcasing the potential to improve healthcare decision-making through more reliable text-to-SQL systems.

## introduction electronic health records (ehrs) are relational databases storing patients' medical histories within hospitals, covering details from admission to discharge. common challenges with ehrs include difficulties in documenting and tracking health information, supporting team coordination, and sharing data @xcite . although ensuring the accurate capture of relevant information is crucial for addressing these challenges, accessing and querying these records often requires knowledge of sql, making it challenging for healthcare provides in practical settings without technical expertise. a solution to this problem is developing a textto-sql model that can translate natural language questions into sql queries to retrieve information from ehrs. recent advancements in large language models (llms) have expanded their utility beyond natural language processing to include code generation, enabling them to interpret text for table manipulation and translate descriptions into code effectively @xcite . these capabilities showcased by code-generating llms @xcite @xcite demonstrate their potential in text-to-sql applications. these developments suggest a promising horizon for leveraging llms to make ehr data more accessible to healthcare professionals, eliminating the prerequisite of sql knowledge and significantly simplifying information retrieval @xcite @xcite . however, in the healthcare domain, the reliability of text-to-sql models is crucial compared to other areas of nlp application. these models must not only generate precise sql queries from natural language but also identify unanswerable questions-queries that cannot be solved with the available database-to avoid potentially harmful outcomes. the risk of providing answers to such questions underscores the need for these models to err on the side of caution, by preferring not to provide an answer rather than risking the provision of incorrect information @xcite . this approach underlines the unique challenges faced in healthcare nlp, emphasizing the need for accuracy and the ability to recognize when it cannot provide a reliable answer. our work introduces pluq, an approach leveraging the self-training paradigm to improve the reliability of text-to-sql models for ehrs through training with pseudo-labeled unanswerable questions. self-training, a semisupervised learning technique, involves re-training a model using its own predictions on unlabeled data to boost its performance. our method adopts a twostage version of self-training process, where we initially fine-tune a seed model using a given training dataset. we then augment the training dataset by incorporating unanswerable questions that the finetuned model identifies from an unlabeled dataset. subsequently, we fine-tune the model once more using this augmented training data to produce the final model. self-training is commonly used in scenarios where unlabeled data is abundant but obtaining labeled data is costly. text-to-sql for ehrs exemplifies such a scenario. in this context, a real-world service can collect users' natural language queries without much effort, but determining the correct sql statement or verifying its answerability with the given database is time-consuming. we adopt the self-training approach to effectively address the issue of class imbalance between answerable and unanswerable questions, thus enhancing the robustness and reliability of the model's performance. after the two-stage self-training process, we apply a filtering strategy to eliminate uncertain predictions. this strategy employs two types of filtering: one based on the maximum entropy of tokens and the other on the execution results of queries. specifically, we assess the entropy in each token generated by the language model, designating the prediction's entropy as that of the token with the highest entropy. if a prediction's entropy exceeds a certain threshold, we consider it an unanswerable question, reflecting the model's lack of confidence in providing a correct answer. additionally, we remove sql queries that either produce errors or fail to retrieve valid values from the mimic-iv 1 dataset. this approach was validated by our performance in the ehrsql 2024 shared task @xcite , where we achieved the top ranking, demonstrating the effectiveness of our method in improving the reliability of text-to-sql systems in healthcare. in summary, our study contributes a method that enhances the reliability of text-to-sql systems for ehrs, addressing the shared task of handling unanswerable questions. this work supports better access to and utilization of ehrs, aiding in informed healthcare decision-making. the main contributions of our paper are: in the field of natural language processing (nlp), recent research has focused on text-to-sql and applying large language models (llms) to electronic health records (ehrs). these studies have advanced the handling of complex queries and the processing of healthcare data, setting the stage for our research on test-time data sample labeling and augmentation in ehrs. text-to-sql in the evolving field of natural language processing, the development of text-to-sql technologies represents a significant advancement. pioneering efforts in this area, @xcite and @xcite , harnessed the power of bert for column classification to tackle the wiki sql @xcite dataset, which is characterized by its simplistic select/where queries. for more complex scenarios, the spider dataset @xcite , comprising multi- table questions, necessitated a understanding of relationship between different tables. wang et al. (2020a) employed graph-based methods to integrating information, while lin et al. (2020) introduced schema linking as an input. moreover, fine-tuning pre-trained language models such as t5 (raffel et al., 2019) has yielded substantial performance improvements in this field. large language models in text-to-sql the emergence of llms has inspired novel approaches to text-to-sql tasks. dong et al. (2023) introduced efficient zero-shot framework, which capitalize on the robust understanding capabilities of 636 training set test set generate sql query null augmented training set ! : question " : sql query #: model llms, with a particular emphasis on prompt-based techniques that demonstrates remarkable efficiency. tai et al. (2023); nan et al. (2023); gao et al. (2024) explore optimal demonstrations based on methodologies like text dense similarity or query similarity selection. pourreza and rafiei (2023) enhances the integrity of generated sql through decomposition of queries and self-correction strategies. shi et al. (2024) proposed llms as an agent for generating code and executing it, leverage the few-shot learning capabilities of llms for solving the multitabular health record datasets. enhancement of llms through data augmentation and self-training amini et al. (2022) presents an extensive review of self-training methods, including consistency-based approaches and transductive learning. post the rise of llms, the field of self-training methods has garnered considerable attention. to enhance the capabilities of llms, some studies have focused on autonomous data generation. wang et al. (2023) stands out by generating synthetic data from a pool consisting not only of seed data but also data generated through an iterative process. seo et al. (2024) employs a few-shot learning approach, drawing samples from external sources to align the seed data in lowresource settings. the line of autonomously augmenting data broadens and enhances the model's capabilities. yuan et al. (2024) introduces using their own outputs to continuously improve both their instruction-following and reward-modeling abilities, demonstrating significant performance enhancements over traditional training methods. nlp in ehrs the application of nlp techniques in ehrs has been extensively explored, utilizing texts and structured knowledge. pampari et al. (2018) proposed a question-answering sys-tem based on unstructured clinical notes. more recently, many works have been developed in the development of generation task based on structured ehrs. wang et al. (2020b) construct the tablebased qa datasets using mimic-iii(johnson et al., 2016) . park et al. (2021) introduced a graph-based ehr qa system that leverages sparql queries from the mimic-sql dataset(wang et al., 2020b). raghavan et al. (2021) focuses on qa tasks using the structured patient records in the mimic-iii. lee et al. (2023) datasets containing multi-table queries and those involving null values, reflecting real-world scenarios in healthcare domain. in our research, we utilize a trained llm on the training dataset for test-time data sample labeling and subsequent augmentation. this approach is particularly focused on addressing the imbalance in the 'null' class. ## method we train a seed model using the original training dataset and then use this model for pseudo labeling on the test set. from this, we select only the samples labeled as unanswerable and augment them to the original training dataset to create the final dataset for self-training. our self-trained model, pluq generates sql queries. we apply postprocessing and two stages of filtering to these queries to ensure their reliability and produce the final answers. ## experiments the experiments are conducted on the development and test sets provided by the ehrsql 2024 shared tasks. all results presented are derived from runs on the official platform. section 4.1 details the models, datasets, and metrics used for training and inference, while section 4.2 discusses the experimental results. finally, in section 4.3, we conduct ablation studies on various components of pluq to examine their individual contributions and impacts. ## conclusion in our work, we develop a self-training strategy designed to enhance the reliability of text-to-sql models for electronic health records (ehrs) through the inclusion of pseudo-labeled unanswerable questions. this approach is particularly valuable in scenarios where there is an abundance of unlabeled data and labeling is costly, thus providing substantial clinical utility in real-world applications. our approach employs a two-stage training process alongside a filtering mechanism based on token entropy and query execution outcomes to improve the model's precision and its ability to identify unanswerable questions. the performance is validated by our leading performance in the ehrsql 2024 shared task. our method contributes towards rendering ehrs more accessible to healthcare professionals without sql knowledge, addressing a critical need for reliable information retrieval in healthcare. future research could explore how large language models facilitate the integration of unstructured medical texts into specific schemas, enhancing interoperability in varied healthcare settings.