A medical-language AI that explains predictions through concepts

Clinical notes contain clues about a patient’s condition, but language models often compress those notes into complex numerical representations that are hard to inspect. One proposed remedy is to make a model predict through named medical concepts—such as fractures or heart failure—so that its reasoning can be expressed in more familiar terms. The catch is that conventional concept-based systems often need people to annotate which concepts appear in each note, a costly step at clinical scale.

In this study, Jifan Gao and colleagues propose a model that aims to avoid those concept labels. Instead, it compares note tokens with descriptions of medical concepts, estimates which concepts relate to which parts of the text, and uses the resulting concept scores in a linear prediction step. The researchers used three-digit ICD-10 categories as the concept vocabulary. A pretrained language model represents both the note and the concept descriptions; a cross-attention module links them. The final prediction is calculated from concept relevance scores and learned concept-specific weights.

That design makes the calculation easier to trace: for an individual note, the model can show which concepts contributed to its prediction and which words were associated with those concepts. Across a dataset, researchers can rank concepts by their influence on predictions. The authors also tested whether the concepts seemed grounded in the text and whether removing concept-associated words affected the performance of separate black-box models.

Two prediction tasks

The team tested the approach on chest-trauma severity and in-hospital mortality. The trauma task used a private dataset from UW Health, with notes and radiology reports recorded in the first eight hours after admission. The mortality task used MIMIC-III, a publicly available critical-care dataset. For each task, the researchers compared the concept-based system with a black-box model using the same text-encoder backbone. They tested three backbones: SapBERT, ClinicalBERT and BGE. The notes supplied to the models were limited to 512 tokens.

Across these tests, the proposed model’s AUROC—a measure of how well a system distinguishes between outcome groups—was within 0.02 of its corresponding black-box model, according to the paper. This suggests that, in these evaluations, adding a concept-based prediction layer did not require a large loss in that measure. It does not show that the model is ready for clinical use, or that it would perform similarly in other hospitals or on other tasks.

The explanations were not equally consistent across settings. In the chest-trauma task, SapBERT and ClinicalBERT highlighted concepts such as rib fractures and pneumothorax, which fit the task. The concepts identified for mortality varied more across the three backbones. In a grounding check, researchers assessed whether selected concepts were explicitly documented in notes. For SapBERT, rib fractures were confirmed in 294 of 418 notes where that concept was a top contributor; hypotension was confirmed in 38 of 49 mortality notes. Other backbone-and-concept combinations had lower confirmation rates. These checks indicate that a concept can help drive a prediction even when the note does not clearly support the concept label.

The authors also masked words associated with highly contributing concepts and compared the result with masking the same number of randomly selected words. In their reported tests, removing concept-associated words reduced black-box performance more than the random-word control, suggesting that these words captured signals used by those models. This is evidence about model behavior, not proof that the concepts represent causes of an outcome or a clinician’s reasoning.

What remains uncertain

The work is an arXiv preprint, so its claims should be treated as provisional rather than as established clinical evidence. The evaluation covers two binary classification tasks, one using private data and the other MIMIC-III; it does not establish performance across hospitals, populations or other types of prediction. The private trauma data also limit independent replication from the paper’s data alone. The authors note that longer records may lose relevant information through the 512-token limit, and that decoder-only models and generation-focused tasks were not tested.

Most importantly, concepts are inferred from language-model representations rather than verified with labels for every note. The authors’ grounding results show variation across backbones and tasks, so an explanation that looks clinically plausible should not be taken as proof that the model understood the note correctly. The study offers a practical route toward more inspectable predictions, but further validation is needed before such explanations can be relied on in care.