[Paper Review] CohortGPT: An Enhanced GPT for Participant Recruitment in Clinical Study

Zihan Guan, Zihao Wu+6 more|arXiv (Cornell University)|Jul 21, 2023

Artificial Intelligence in Healthcare and Education15 citations

TL;DR

CohortGPT uses a knowledge graph and reinforcement-learning–driven dynamic Chain-of-Thought sampling to improve LLM-based medical text classification for clinical study participant recruitment in few-shot settings.

ABSTRACT

Participant recruitment based on unstructured medical texts such as clinical notes and radiology reports has been a challenging yet important task for the cohort establishment in clinical research. Recently, Large Language Models (LLMs) such as ChatGPT have achieved tremendous success in various downstream tasks thanks to their promising performance in language understanding, inference, and generation. It is then natural to test their feasibility in solving the cohort recruitment task, which involves the classification of a given paragraph of medical text into disease label(s). However, when applied to knowledge-intensive problem settings such as medical text classification, where the LLMs are expected to understand the decision made by human experts and accurately identify the implied disease labels, the LLMs show a mediocre performance. A possible explanation is that, by only using the medical text, the LLMs neglect to use the rich context of additional information that languages afford. To this end, we propose to use a knowledge graph as auxiliary information to guide the LLMs in making predictions. Moreover, to further boost the LLMs adapt to the problem setting, we apply a chain-of-thought (CoT) sample selection strategy enhanced by reinforcement learning, which selects a set of CoT samples given each individual medical report. Experimental results and various ablation studies show that our few-shot learning method achieves satisfactory performance compared with fine-tuning strategies and gains superb advantages when the available data is limited. The code and sample dataset of the proposed CohortGPT model is available at: https://anonymous.4open.science/r/CohortGPT-4872/

Motivation & Objective

• Address the challenge of identifying eligible participants from unstructured clinical notes and radiology reports.
• Leverage a domain knowledge graph to guide LLM reasoning for disease label prediction.
• Improve few-shot learning performance with a dynamic CoT sample selection strategy optimized by reinforcement learning.
• Demonstrate competitiveness against fine-tuning approaches with limited labeled data and provide reproducible code/dataset.
• Show extensibility of the framework to other medical NLP tasks.

Proposed method

• Embed medical knowledge via knowledge graph prompts (KG-as-Tree, KG-as-Relation, KG-as-Rules) into LLM inputs.
• Adopt a policy-gradient–based dynamic CoT sample selection to choose CoT prompts for each medical report.
• Use a BioGPT-based status encoder to compute similarities between the report and CoT candidates to inform the policy.
• Optimize a reward function that reflects correct/incorrect label predictions to train the CoT selector.
• Default to KG-as-Rule for KG-to-prompt transformation based on ablation results.
• Evaluate on IU-RR and MIMIC-CXR datasets against fine-tuning baselines (BioBERT, BioGPT) and few-shot LLM baselines (Alpaca, BloomZ, ChatGPT, GPT-4).

Figure 1: A policy model will be trained on a small number of training samples to dynamically select CoT samples from a CoT candidate pool. A knowledge graph containing the hierarchical information of the disease labels will be transformed into a series of executable rules. Then the dynamic CoT samp

Experimental results

Research questions

• RQ1 RQ1: How does the proposed method perform in medical report classification under few-shot settings compared to fine-tuning approaches?
• RQ2 RQ2: How do different KG-to-prompt strategies and CoT sample selection methods affect performance?
• RQ3 RQ3: Can the approach provide explainable, stepwise reasoning through selected CoT samples as evidenced by case studies?

Key findings

• In low-data regimes, the method with ChatGPT or GPT-4 outperforms traditional fine-tuning (IU-RR: 0.69 F1 with 5-shot ChatGPT; 0.81 with 5-shot GPT-4 vs BioBERT 0.44 and BioGPT 0.25).
• On IU-RR and MIMIC-CXR, few-shot LLMs show advantages over fine-tuning when data are scarce, with similar trends on larger splits.
• KG-as-Rule yields the best overall prompt performance (Exact Match 0.56, Precision 0.73, Recall 0.72, F1 0.69, HL 0.04).
• Dynamic CoT sample selection consistently outperforms Random, Manual, and Most-Similar CoT strategies in multiple metrics.
• Performance improves with more training samples and a larger CoT candidate pool; optimal k-shot values observed around 5 or 8; too many CoT samples can harm performance.
• Case studies indicate selected CoT samples enable ChatGPT to deduce answers through multi-step reasoning as designed.

Figure 2: A knowledge graph was created by [ 29 ] to represent relationships between diseases, organs, or tissues. In this graph, disease labels are represented by nodes in solid boxes, corresponding organs or tissues are represented by nodes in dotted boxes, and the edges linking the nodes represen