Human-AI Co-design for Clinical Prediction Models

TL;DR

HACHI presents a human-in-the-loop framework that leverages AI agents to convert unstructured clinical notes into a compact set of interpretable yes/no concepts, forming a $k$-concept Clinical Prediction Model (CPM). The approach alternates between rapid AI-guided concept discovery and expert feedback, enabling transparent review, refinement, and alignment with clinical goals. In two real-world tasks—Traumatic Brain Injury (TBI) and Acute Kidney Injury (AKI)—HACHI outperformed baselines, revealed novel clinically relevant predictors, and demonstrated improved generalizability across sites and time periods, while also highlighting the crucial role of the clinical AI team in guiding prompts, addressing data leakage, and mitigating bias. The work emphasizes that simple, interpretable CPMs learned through iterative human-AI co-design can be deployed locally with minimal resources and open-source tooling, though prospective validation and fairness considerations remain essential before clinical deployment.$k$-concept CPMs, $AUC$, and site/time-period generalizability are central to the framework's reported gains and practical impact.

Abstract

Developing safe, effective, and practically useful clinical prediction models (CPMs) traditionally requires iterative collaboration between clinical experts, data scientists, and informaticists. This process refines the often small but critical details of the model building process, such as which features/patients to include and how clinical categories should be defined. However, this traditional collaboration process is extremely time- and resource-intensive, resulting in only a small fraction of CPMs reaching clinical practice. This challenge intensifies when teams attempt to incorporate unstructured clinical notes, which can contain an enormous number of concepts. To address this challenge, we introduce HACHI, an iterative human-in-the-loop framework that uses AI agents to accelerate the development of fully interpretable CPMs by enabling the exploration of concepts in clinical notes. HACHI alternates between (i) an AI agent rapidly exploring and evaluating candidate concepts in clinical notes and (ii) clinical and domain experts providing feedback to improve the CPM learning process. HACHI defines concepts as simple yes-no questions that are used in linear models, allowing the clinical AI team to transparently review, refine, and validate the CPM learned in each round. In two real-world prediction tasks (acute kidney injury and traumatic brain injury), HACHI outperforms existing approaches, surfaces new clinically relevant concepts not included in commonly-used CPMs, and improves model generalizability across clinical sites and time periods. Furthermore, HACHI reveals the critical role of the clinical AI team, such as directing the AI agent to explore concepts that it had not previously considered, adjusting the granularity of concepts it considers, changing the objective function to better align with the clinical objectives, and identifying issues of data bias and leakage.

Figures (2)

• Figure 1: The HACHI framework uses LLMs with humans-in-the-loop to build an effective clinical prediction model (CPM). (a) HACHI is composed of an outer loop where the clinical AI team provides guidance and feedback to the AI agent on how to learn a CPM, and an inner loop where the AI agent follows instructions from the clinical AI team to find $k$ concepts (formally defined as yes/no questions) that maximize the CPM's predictive accuracy. The inner AI-guided CPM learning procedure is broken down into three steps: 1. Initialize the CPM by brainstorming clinical concepts from keyphrases extracted out of clinical notes; 2. Propose candidate concepts by analyzing which keyphrases are most associated with the outcome of interest; 3. Evaluate the candidate concepts by annotating each concept and selecting the best-performing concept(s). Steps 2 and 3 are repeated until convergence. Each round, the clinical AI team analyzes the results from the AI-guided CPM learning procedure and provides feedback on how the procedure can be improved, such as by modifying the prompts and clarifying which concepts are and are not of interest. HACHI improves over baselines and across rounds for two real-world clinical prediction tasks: (b) diagnosis of traumatic brain injury (TBI), where performance is evaluated in terms of AUC with respect to the overall population (shown in plot) and stratified across two sites (shown in table). (c) development of acute kidney injury (AKI), where performance is evaluated in an internal validation set (Period 1, shown in plot and table) and a later, temporally disjoint test dataset (Period 2, shown in table). Error bars show standard errors.
• Figure 2: PHI-compliant web interface for auditing the AI-agent CPM learning procedure. It shows (i) the learned factors, (ii) the LLM concept annotations at a patient/note level, (iii) which patients received an incorrect prediction, and (iv) the performance of the CPM on a held-out test set.