cs.LGOct 7, 2026

Coverage-Aware Reasoning with Medical Tokens for Diagnosis Prediction

Authors: Kaisong Zhang, Haotian Fang, Junmeng Zhou, Hang Lv, Yulan Pan, Yanchao Tan

Organizations: Fuzhou University · Shanghai Jiao Tong University · Beihang University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Large language models (LLMs) offer promising potential for next-visit diagnosis prediction, owing to their ability to integrate longitudinal clinical evidence and reason over it in natural language. However, reinforcement learning for LLM reasoning commonly rewards each trajectory according to the correctness of its final answer. In next-visit diagnosis prediction, multiple diagnoses can be simultaneously valid, but independently rewarding one diagnosis per trajectory does not distinguish repeated hits from coverage of different diagnoses. The policy can therefore concentrate on a few correct diagnoses, leaving others uncovered. Meanwhile, LLM tokenizers can split ICD codes into several generic tokens with limited clinical meaning, requiring multiple decoding steps to predict each diagnosis and hindering reasoning over a large disease vocabulary. To address both challenges, we propose CARing, a framework that represents diagnoses with compositional Semantic IDs (SIDs) and optimizes reasoning trajectories for multi-label coverage. Concretely, we first encode ontology-enriched disease semantics into compact SIDs through residual quantization, and ground the resulting SID tokens in natural language and longitudinal EHR contexts through multi-task alignment and reasoning-enriched training to unlock transferable LLM reasoning. CARing further improves unordered multi-label prediction through a coverage reward for reinforcement learning and multi-positive supervision. At inference time, the model supports both efficient direct constrained decoding and multi-chain reasoning with rank fusion. On MIMIC-III and MIMIC-IV, CARing exceeds all EHR-trained baselines in weighted F1 and attains the highest top-k recall at every reported cutoff, including R@30 of 46.04% and 46.52% in reasoning mode. Our codes and logs are available at https://github.com/zmlxzyh/CARing-Codes-Logs.

Figures & tables

Figure 1: Motivation for multi-label reinforcement learning in next-visit diagnosis prediction. Standard RL can collapse onto one correct diagnosis, leaving others uncovered. Our coverage reward discounts repeated predictions and encourages coverage of distinct correct diagnoses.
Figure 2: Overview of CARing. The framework constructs compositional medical SIDs, aligns them with clinical language and reasoning, combines a coverage reward with multi-positive answer supervision, and supports direct and multi-chain reasoning inference.
ModelMIMIC-IIIMIMIC-IV
 w- F1F_{1}R@10R@20R@30R@40w- F1F_{1}R@10R@20R@30R@40
RETAIN21.2524.1334.3040.1245.9119.5425.4233.8238.9142.78
KAME18.5822.5531.4337.6742.5410.5917.9725.0129.5633.15
StageNet18.1321.3630.3736.6441.1011.4618.4125.2930.2033.82
CGL21.7925.5135.2841.5246.7520.0824.0733.9640.0444.31
TRANS19.1023.3232.2638.3943.1415.8420.5528.7134.3738.67
Table 1: Overall next-visit diagnosis prediction performance on MIMIC-III and MIMIC-IV. All values are percentages. Best and second-best results are shown in bold and underlined, respectively.
MetricLayerName onlyName + ICD path
R@30 (%) 38.9739.73
Codebook utilization (%)1100.00100.00
 298.83100.00
 397.27100.00
Token distribution entropy17.9437.951
 27.5427.919
Table 2: SID quality and MIMIC-III recall with disease names alone or names plus ICD ontology paths. Entropy is in bits; its maximum is 8 for a 256-entry codebook.
Figure 3: Effect of the four stages of reasoning-enriched training on MIMIC-III under non-thinking direct inference.
Figure 4: Training ablation on MIMIC-III in thinking mode.
Figure 5: MIMIC-III trajectories for S5 (EM reward) and S6 (coverage reward), both initialized from S4.
Figure 6: MIMIC-III trajectories for S6 and S7. The missing S7 value at step 42 is linearly interpolated between steps 41 and 43.
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

DatasetMIMIC-IIIMIMIC-IV
# of patients5,44979,393
# of visits14,141408,990
Avg. # of visits per patient2.605.15
Max. # of visits per patient29170
# of unique diagnoses3,87437,917
# of CCS codes285842
Appendix
Table 3: Dataset statistics after cohort preprocessing and before subsampling MIMIC-IV.
Figure 7: ICD ontology subgroup counts for four randomly sampled first SID tokens. Subgroups are the second nodes below the ontology root.
MetricName onlyName + ICD path
First-token groups pure by ICD / 25658143
First-token groups ≥80%\geq 80\% pure by ICD / 256108212
Top-level ICD weighted purity (%)69.3690.51
Same ICD group given shared token (%)57.6086.49
First-token groups pure by CCS / 2561344
First-token groups ≥80%\geq 80\% pure by CCS / 25636103
Appendix
Table 4: Semantic organization of first SID tokens over 4,491 distinct diseases. ICD groups are the first nodes below the ontology root.
ModelIFEvalGSM8K
Base Qwen3-1.7B0.68200.8832
S2 = S1 + Enriched Alignment0.14600.5238
S3 = S2 + General Reasoning0.47130.6300
Appendix
Table 5: General-ability evaluation across training stages. Higher scores are better.

Explore similar work

CardsList
Sep 21, 2026cs.AI

Fathom-Vaidya: Advancing Medical Reasoning with Rubric-Based Rewards

Deploying Large Language Models (LLMs) in healthcare requires robust performance across two complementary dimensions - diagnostic reasoning: the convergent, evidence-driven task of inferring a patient's condition from clinical data to produce a diagnosis, and clinical healthcare reasoning: the broader, navigational judgment required to communicate, plan, and adapt across multi-turn clinical interactions where a single correct answer may not exist. Recent benchmarks such as HealthBench and MedXpertQA reveal persistent weaknesses in both areas, exposing failures in complex diagnostic scenarios and limitations in contextual, patient-centered dialogue. We introduce a sequential training framework that targets these facets using synthetic data and rubric-based reinforcement learning. First, we improve diagnostic reasoning using MedBullets-derived questions with rule- and rubric-guided Reinforcement Learning (RL). We then shift to clinical reasoning by generating 5.3k synthetic multi-turn scenarios, each paired with multi-dimensional rubrics to comprehensively assess the response. This approach yields over 10% improvement on MedXpertQA, and our 30B model achieves 50.1% accuracy on HealthBench-Hard, surpassing proprietary baselines including GPT-5 (thinking). Our results show that targeted synthetic datasets and rubric-based training can systematically improve both diagnostic and interactive clinical reasoning in medical LLMs.

Kalash Shah, Kunal Singh, Snehan J +1

Fractal AI Research

RL for Language Model ReasoningRubric-Based RLClinical ReasoningMedical ReasoningMedical LLMs
AskPDF
Apr 27, 2026cs.CL

MultiDx: A Multi-Source Knowledge Integration Framework towards Diagnostic Reasoning

Diagnostic prediction and clinical reasoning are critical tasks in healthcare applications. While Large Language Models (LLMs) have shown strong capabilities in commonsense reasoning, they still struggle with diagnostic reasoning due to limited domain knowledge. Existing approaches often rely on internal model knowledge or static knowledge bases, resulting in knowledge insufficiency and limited adaptability, which hinder their capacity to perform diagnostic reasoning. Moreover, these methods focus solely on the accuracy of final predictions, overlooking alignment with standard clinical reasoning trajectories. To this end, we propose MultiDx, a two-stage diagnostic reasoning framework that performs differential diagnosis by analyzing evidence collected from multiple knowledge sources. Specifically, it first generates suspected diagnoses and reasoning paths by leveraging knowledge from web search, SOAP-formatted case, and clinical case database. Then it integrates multi-perspective evidence through matching, voting, and differential diagnosis to generate the final prediction.~Extensive experiments on two public benchmarks demonstrate the effectiveness of our approach.

Yimin Deng, Zhenxi Lin, Yejing Wang +9

Xi’an Jiaotong University · City University of Hong Kong · Tencent Jarvis Lab +1

Knowledge Augmentation for Language ModelsMedical DiagnosisClinical ReasoningClinical PredictionDifferential DiagnosisMedical LLMs
AskPDF
May 8, 2026cs.CL

MedAction: Towards Active Multi-turn Clinical Diagnostic LLMs

Most existing LLM diagnoses are evaluated on static, single-turn settings where complete patient information is provided upfront, an oversimplification of real clinical practice. We study active diagnosis: the real-life clinical process of starting from initial observation, ordering tests, interpreting results, and updating a differential diagnosis across multiple turns. Through systematic analysis, we identify three recurring failure modes in current LLMs: ungrounded test ordering, unreliable diagnostic update, and degraded multi-turn coherence. Together, these failures reveal a core deficit: existing medical training data teaches models to reason from complete information but not to act under evolving, partial evidence. To address this gap, we introduce MedAction, a tree-structured distillation pipeline that synthesizes diverse and high-quality multi-turn diagnostic trajectories via LLM-environment interaction. We propose two knowledge-graph-grounded metrics to filter trajectory quality: Disease Trajectory Consistency (DTC), which tracks whether the model's hypothesis converges toward the correct diagnosis, and Reasoning-Action Consistency (RAC), which verifies that belief updates are driven by gathered evidence. Using this pipeline, we construct MedAction-32K, a dataset of 32,681 trajectories from 2,896 PMC cases. Fine-tuning an 8B model on MedAction-32K achieves state-of-the-art performance among open-source models on both MedR-Bench and our curated MedAction-300-Hard benchmark, pushing the edge for open-source medical LLMs.

Hsin-Ling Hsu, Zizheng Wang, Donghua Zhang +9

1National Chengchi University · 2Georgetown University · University of Michigan +3

HealthcareMedical DiagnosisClinical ReasoningLLM Agent TrainingDifferential DiagnosisMedical LLMs
AskPDF