PaperlayerView exact revision on arXiv

Source-linked AI summary

Coverage-Aware Reasoning with Medical Tokens for Diagnosis Prediction

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

arXiv:2610.10641v1cs.LGcs.AI

TL;DR

Next-visit diagnosis prediction requires coverage of multiple valid diagnoses, but trajectory-level rewards can favor repeated answers and ICD codes can be difficult for LLMs to generate and reason over. CARing uses compositional medical SIDs and coverage-oriented training; on MIMIC-III and MIMIC-IV, it exceeds EHR-trained baselines in weighted F1 and leads recall at every reported cutoff.

  • Problem

    Next-visit diagnosis prediction has multiple simultaneously valid labels, yet rewarding each trajectory only for its own correctness does not distinguish repeated hits from coverage of different diagnoses.

    Source: §1 · p. 1, §5 · p. 9

  • Method

    CARing represents diagnoses with compositional medical SIDs and trains reasoning trajectories with a coverage reward and multi-positive supervision.

    Source: §1 · p. 2, §3.2 · p. 4, §3.4 · p. 5

  • Results

    On MIMIC-III and MIMIC-IV, CARing exceeds all EHR-trained baselines in weighted F1 and leads recall at every reported cutoff, with reasoning-mode R@30 of 46.04% and 46.52%, respectively.

    Source: §5 · p. 10

  • Takeaways & Limitations

    CARing combines natural-language reasoning with generation of multiple diagnoses from a large vocabulary using compositional medical SIDs.

    Source: §5 · p. 10

    The clinical accuracy and faithfulness of generated traces, including possible artifacts from teacher training with reference diagnoses, require closer study.

    Source: §D.2 · p. 28

Abstract

from arXiv · show

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.

1 INTRODUCTION

Next-visit diagnosis prediction is an unordered multi-label task, but trajectory-level correctness rewards can favor repeated diagnoses over broader coverage. CARing combines compositional medical SIDs with coverage-oriented training and reports stronger results than EHR-trained baselines on MIMIC-III and MIMIC-IV.

1 INTRODUCTION

  • A correct-answer reward for each trajectory can treat repeated hits like distinct diagnoses, allowing the policy to leave valid labels uncovered.The task is unordered and can have multiple correct diagnoses for one visit.Source: §1 · p. 1
  • CARing uses compositional medical SIDs to connect clinical language, patient histories, and diagnosis outputs in one autoregressive model.Disease semantics and ICD ontology paths are compressed through residual quantization into shared codebook tokens, then grounded through six alignment tasks and teacher-generated narratives.Source: §1 · p. 2
  • A coverage reward discounts each correct diagnosis by its frequency across sampled trajectories, while multi-positive supervision trains on correct SID answers.Together, these objectives encourage trajectories to recover complementary diagnoses.Source: §1 · p. 2
  • Direct inference uses constrained decoding, while reasoning inference samples multiple chains and fuses their constrained-decoding rankings.Both modes restrict predictions to valid medical SIDs.Source: §1 · p. 2
  • CARing exceeds all EHR-trained baselines in weighted F1 and achieves the highest top-k recall; reasoning-mode R@30 is 46.04% and 46.52%.The reported evaluation uses MIMIC-III and MIMIC-IV.Source: §1 · p. 2
  • Ablations show semantic enrichment improves direct prediction and that coverage reward and multi-positive supervision play complementary roles.The supplied passage does not specify the individual ablation values.Source: §1 · p. 3

2 RELATED WORK

Prior work uses EHR architectures to score fixed diagnosis vocabularies, while LLM reasoning and medical tokenization offer different routes to prediction. CARing brings compositional medical SIDs and coverage-aware training to multi-label diagnosis prediction.

2 RELATED WORK

  • EHR models encode visit sequences and may use ontologies or patient graphs, but generally score a fixed diagnosis vocabulary rather than generate diagnoses through explicit reasoning.Examples include recurrent or attention architectures (Choi et al., 2016a;b; Gao et al., 2020; Yang et al., 2023b) and ontology- or graph-based approaches (Ma et al., 2018; Lu et al., 2021; Chen et al., 2024; Tan et al., 2025).Source: §2 · p. 3
  • Medical tokenizers mainly provide discrete disease codes to downstream predictors, whereas MERA assigns one vocabulary token per disease.Semantic IDs instead compose items from shared codebook tokens, and recent systems let LLMs reason in text before generating identifiers (Liu et al., 2026; Team et al., 2026).Source: §2 · p. 3
  • Trajectory-level correctness rewards can concentrate reasoning policies on a few answers, while diagnosis recall depends on covering multiple correct diagnoses.CARing's group reward counts distinct correct SIDs, and its multi-positive loss supplies learning signals for correct diagnoses that no trajectory samples.Source: §2 · p. 3

3 METHOD

CARing builds ontology-enriched medical SIDs, grounds them in clinical language and patient histories, and trains reasoning trajectories for multi-label diagnosis coverage. It supports both direct constrained decoding and multi-chain reasoning at inference.

3.1 TASK FORMULATION

  • CARing evaluates the unordered next-visit diagnosis set, while case-level instances pair one history with one target for alignment and reasoning activation.Reinforcement learning and visit-level evaluation retain the complete target set.Source: §3.1 · p. 3

3.2 ONTOLOGY-ENRICHED COMPOSITIONAL MEDICAL TOKENIZATION

  • Ontology paths and disease titles are compressed through residual quantization into compact discrete SIDs that preserve disease semantics for autoregressive modeling.The experiments use three codebook levels.Source: §3.2 · pp. 3–4

3.3 MULTI-TASK SID–CLINICAL CONTEXT ALIGNMENT

  • Six next-visit prediction and concept-translation tasks jointly ground SID tokens in clinical language, histories, disease names, and SID outputs.The tasks cover both directions between disease names and SIDs, alongside textual or SID histories paired with either output form.Source: §3.3 · p. 4

3.4 REASONING-ENRICHED TRAINING

  • Teacher-generated disease narratives and case-level reasoning data enrich SID–language alignment, followed by supervised fine-tuning to generate a rationale before the target SID.Disease narratives interleave SIDs with clinical language and ontology metadata; case-level reasoning paths use the target as private teacher guidance.Source: §3.4 · pp. 4–5

3.5 REINFORCEMENT LEARNING WITH MULTI-POSITIVE SUPERVISION

  • Visit-level GRPO rewards distinct correct diagnoses rather than repeated hits, while multi-positive supervision supplies answer-token learning signals for positives absent from sampled rollouts.The multi-positive objective retains a rollout’s reasoning prefix, replaces its answer with a target SID, and weights selected SIDs equally in total.Source: §3.5 · pp. 5–6

3.6 TWO INFERENCE SETTINGS

  • Direct inference uses constrained beam search on SID histories, while multi-chain inference fuses diagnosis rankings from sampled reasoning chains with reciprocal rank fusion.Both settings rank individual diagnosis SIDs and return valid, deduplicated candidates.Source: §3.6 · p. 6

4 EXPERIMENTS

Experiments on MIMIC-III and MIMIC-IV compare CARing with EHR and LLM baselines and test ontology enrichment, training stages, and coverage-oriented objectives. CARing leads EHR-trained baselines in weighted F1 and has the highest recall at every reported cutoff, while ablations examine how its components affect prediction and diversity.

4 EXPERIMENTS

  • Experiments use MIMIC-III and 10,000 sampled MIMIC-IV patients, report weighted F1 and top-k recall, and compare neural EHR, PLM-based, LLM-assisted, and zero-shot LLM baselines.The backbone is Qwen3-1.7B, fine-tuned throughout all stages.Source: §4 · p. 7

4.1 MAIN RESULTS (RQ1)

  • CARing has the highest w-F1 among EHR-trained models on both datasets and the highest recall at every reported cutoff; DeepSeek-V4-Flash has the highest w-F1 overall.The thinking setting leads at seven of eight recall cutoffs, while direct inference is 0.09 points higher on MIMIC-III R@40.Source: §4.1 · pp. 7–8

4.2 ICD ONTOLOGY PATHS IN MEDICAL SIDS (RQ2)

  • Adding ICD ontology paths raises MIMIC-III R@30 from 38.97 to 39.73 after multi-task alignment and yields higher SID entropy in every layer.The name-plus-path variant fully uses all three codebooks.Source: §4.2 · p. 8, Table 2 · p. 8

4.3 ABLATION STUDY (RQ3)

  • For non-thinking direct inference, semantic enrichment adds 2.04 w-F1 and 1.21 R@30, while reasoning activation adds 4.45 w-F1 and 2.52 R@30 over S3.General corpora add 0.37 w-F1 and 0.39 R@30; the reported instruction-following and reasoning gains are evaluated separately on IFEval and GSM8K.Source: §4.3 · p. 8, Figure 3 · p. 8
  • In thinking mode, S6’s coverage reward improves over S5’s EM reward from 18.88 to 25.41 w-F1 and from 34.67 to 40.27 R@30; S7 reaches 27.58 and 46.04.S7 adds multi-positive loss to S6; Figure 4 presents the training ablation in thinking mode.Source: §4.3 · pp. 8–9, Figure 4 · p. 8

4.4 COVERAGE AND DIVERSITY IN TRAINING (RQ4)

  • By step 100, coverage reward retains 2.89 unique correct SIDs, versus 1.52 with EM reward, and produces 7.03 unique SIDs per group versus 3.06.S6 also has higher mean inverse frequency among correct predictions, indicating less repetition.Source: §4.4 · p. 9, Figure 5 · p. 9
  • At step 60, S7 produces 12.40 unique SIDs per group versus 7.02 for S6; its batch-level SID entropy is 5.58 versus 4.58.Mean inverse frequency among correct predictions is 0.60 for S7 and 0.32 for S6.Source: §4.4 · p. 9, Figure 6 · p. 9

5 CONCLUSION

CARing addresses next-visit prediction’s need to recover multiple valid diagnoses by combining compositional medical SIDs with coverage-oriented reasoning. It reports stronger weighted F1 than EHR-trained baselines and the highest recall at every reported cutoff on MIMIC-III and MIMIC-IV.

5 CONCLUSION

  • R@30 reached 46.04% on MIMIC-III and 46.52% on MIMIC-IV in reasoning mode; CARing had the highest recall at every reported cutoff and exceeded all EHR-trained baselines in weighted F1.The conclusion also notes that the framework supports direct ranking and multiple reasoning chains.Source: §5 · p. 10
  • CARing uses compositional medical SIDs grounded in clinical language and patient histories, while coverage reward and multi-positive supervision encourage complementary diagnoses across trajectories.The coverage reward discounts repeated correct predictions; multi-positive supervision supplies signals for correct diagnoses that sampling may miss.Source: §5 · p. 10, §A.5 · p. 18

A.1 LONGITUDINAL DIAGNOSIS PREDICTION

  • Next-visit prediction spans distinct settings: selected-outcome prediction, future-concept generation, and recovery of the full diagnosis set impose different demands on models.The reviewed approaches include sequence models, ontology- and graph-based methods, pretrained record models, and generative predictors.Source: §A.1 · p. 16

A.2 MEDICAL CODE REPRESENTATIONS AND GENERATION

  • Compositional SIDs quantize disease descriptions and ontology paths into shared tokens that can serve as inputs and diagnosis outputs for one autoregressive model.Unlike a dedicated token per disease, the representation uses sequences of shared codebook tokens.Source: §A.2 · p. 16

A.3 REASONING FOR CLINICAL PREDICTION

  • Forecasting multiple diagnoses from earlier visits differs from coding conditions in a document, motivating evaluation of diagnosis coverage across sampled reasoning trajectories.The reviewed clinical-reasoning approaches use retrieval, reasoning supervision, or EHR-guided trajectories in other prediction settings.Source: §A.3 · p. 17

A.4 LARGE REASONING MODELS AND POLICY COLLAPSE

  • Prior findings distinguish individual-answer accuracy from coverage: RLVR can improve pass@1 while base models attain higher pass@k at large k, and reward conditions can concentrate outputs.These findings motivate measuring correctness and coverage when multiple outputs can be valid.Source: §A.4 · p. 17

A.5 CORRECTNESS AND DIVERSITY IN REINFORCEMENT LEARNING

  • CARing’s group reward gives each distinct correct diagnosis one unit, while an auxiliary multi-positive loss trains correct SIDs that sampled groups omit.This differs from methods that score set-level diversity or order-independent label sequences; here each trajectory ends in one diagnosis and the group is compared with the full next-visit set.Source: §A.5 · pp. 17–18

B ADDITIONAL METHOD DETAILS

This appendix section introduces representative input–output examples for the task formats used in training.

B ADDITIONAL METHOD DETAILS

B.1 MULTI-TASK SID–LANGUAGE ALIGNMENT

The six alignment tasks connect longitudinal histories and disease names with SID prediction in both directions. Visits remain chronological, while diagnoses within each visit are treated as unordered.

B.1 MULTI-TASK SID–LANGUAGE ALIGNMENT

  • Tasks T1–T4 cover SID-history-to-SID, text-history-to-SID, SID-history-to-disease-name, and text-history-to-disease-name mappings.The examples show the model returning one possible next-visit diagnosis or its SID.Source: §B.1 · pp. 18–19
  • Tasks T5–T6 establish bidirectional mappings between disease names and compositional SIDs.The examples pair a disease name with its SID and then map that SID back to the disease name.Source: §B.1 · pp. 18–19
  • All six tasks are jointly optimized with the standard next-token prediction objective.Diagnoses within visits are unordered, but visits are listed chronologically.Source: §B.1 · p. 18

B.2 REASONING-ENRICHED TRAINING

The appendix details how CARing grounds SIDs and trains reasoning before SID prediction, then specifies the coverage reward and catalog-constrained inference procedures. It also describes limits on what the reward guarantees.

B.2 REASONING-ENRICHED TRAINING

  • Reasoning activation uses history-only prompts and supervises a teacher-generated reasoning path followed by the target SID; the prompt tokens are masked from the loss.Disease-level and case-level enriched corpora are also trained as plain-text next-token sequences.Source: §B.2 · p. 19, §B.2.1 · p. 21

B.3 PROMPTS FOR REASONING-ENRICHED DATA

  • Two enrichment routes create supervision: disease-level narratives ground a coded concept in clinical knowledge, while case-level data pair history-supported reasoning with observed diagnoses.The case-level reasoning path is reused as reasoning supervision for activation.Source: §B.2 · pp. 20–21, §B.3 · p. 21

B.4 COVERAGE REWARD PROPERTY

  • The raw group reward equals the number of distinct correct predictions: repeating a correct SID or adding an incorrect prediction does not increase it.A correct SID appearing c times contributes c(1/c) = 1 to the group total.Source: §B.4 · p. 24
  • Finite sampling and within-group advantage normalization do not guarantee a globally uniform policy over all ground-truth diagnoses.When all rollouts have the same reward, their normalized advantages are zero.Source: §B.4 · p. 24

B.5 GRPO OPTIMIZATION DETAILS

  • The optimization accumulates GRPO and KL gradients from original rollouts alongside multi-positive gradients from answer-replaced branches before one optimizer step.Source: §B.5 · p. 24

B.6 CATALOG-CONSTRAINED BEAM SEARCH

  • Catalog-constrained beam search masks tokens that cannot extend a valid SID prefix, then checks completed paths against the catalog and deduplicates them.Direct inference disables explicit thinking; reasoning inference samples chains and fuses their constrained rankings.Source: §B.6 · pp. 24–25

C.1 DATASET PREPROCESSING

The experiments use patient-level longitudinal diagnosis prediction splits on MIMIC-III and a sampled MIMIC-IV cohort. CARing is compared with conventional and language-assisted EHR predictors and a zero-shot LLM under a specified prediction protocol.

C.1 DATASET PREPROCESSING

  • Patients with at least two visits were split chronologically by patient at 7:1:2, and the MIMIC-IV experiment sampled 10,000 patients while preserving that ratio.The task predicts the complete diagnosis set at the next visit from diagnoses in prior visits.Source: §C.1 · p. 25

C.2 BASELINES

  • The baseline set spans conventional EHR models RETAIN, KAME, StageNet, CGL, and TRANS, language-assisted BoxLM and HieraMed, and zero-shot DeepSeek-V4-Flash.EHR-trained baselines use 64-dimensional disease embeddings and aggregated patient representations.Source: §C.2 · pp. 25–26
  • DeepSeek-V4-Flash receives chronological prior diagnosis codes and is asked for 40 unique ranked diagnoses, without diagnosis names, target diagnoses, examples, or task-specific training.Thinking and non-thinking modes use the same prompt; predictions are evaluated by top-k recall after normalization, catalog matching, deduplication, and invalid-code removal.Source: §C.2 · p. 26

C.3 IMPLEMENTATION DETAILS

Implementation uses ontology-enriched disease metadata to construct SIDs, trains a Qwen3 backbone through alignment and reasoning activation, and then applies GRPO with coverage-related supervision.

C.3 IMPLEMENTATION DETAILS

  • Qwen3-Embedding-8B encodes ontology-enriched disease metadata for SID construction, while Qwen3-1.7B is the fully optimized backbone.The alignment stage jointly trains SID and visit-boundary token embeddings; DeepSeek-V4-Flash generates disease-level and longitudinal corpora.Source: §C.3 · p. 26
  • Reasoning activation starts from the alignment checkpoint and retains its training setup while reducing the learning rate to 1 × 10^-5.The alignment stage uses AdamW, effective batch size 1024, learning rate 3 × 10^-4, and a maximum of five epochs with early stopping.Source: §C.3 · p. 26
  • GRPO samples 16 rollouts per input and uses a 1 × 10^-6 learning rate with multi-positive loss coefficient λMP = 5 × 10^-3.Training uses batch size 256, at most ten epochs, and four NVIDIA A800 GPUs.Source: §C.3 · p. 26

C.4.1 RANDOMLY SAMPLED FIRST-TOKEN CASES

Sampled first SID tokens group diseases into recognizable clinical themes, and full-vocabulary measurements show stronger ontology organization for SIDs built from disease names plus ICD paths.

C.4.1 RANDOMLY SAMPLED FIRST-TOKEN CASES

  • Figure 7 counts distinct diseases by second-level ICD subgroup for four uniformly sampled first SID tokens drawn from the 256-token name-plus-path variant.The sample used a fixed seed of 2024.Source: §C.4.1 · p. 26, Figure 7 · p. 27
  • Among sampled prefixes, 17 of 18 ⟨a86⟩ diseases concern malignancy history, 20 of 23 ⟨a184⟩ are nervous-system diagnoses, and all ⟨a23⟩ and ⟨a192⟩ diseases are spinal injuries and chronic ulcers, respectively.The ⟨a184⟩ prefix also groups three musculoskeletal diagnoses across ontology branches.Source: §C.4.1 · p. 27

C.4.2 SEMANTIC ORGANIZATION ACROSS THE DISEASE VOCABULARY

  • The full-vocabulary analysis measures first-token grouping across 4,491 diseases using top-level ICD groups and single-level CCS categories.The vocabulary spans 52 top-level groups and 264 CCS categories; the analysis includes majority purity, weighted purity, and pair agreement.Source: §C.4.2 · p. 27
  • Names-plus-path SIDs raise dominant top-level-group share from 69.36% to 90.51%, increase pure first-token groups from 58 to 143, and raise CCS pair agreement from 32.74% to 57.39%.The comparison is between SIDs learned from names plus ICD paths and names alone across 4,491 diseases.Source: Table 4 · p. 27, §C.4.2 · p. 28

C.5 GENERAL ABILITY EVALUATION

Enriched alignment reduces general-ability benchmark performance, while adding general reasoning data partially recovers it. The broader-impact discussion describes possible clinical-review uses and extensions to other coding systems and EHR prediction tasks.

C.5 GENERAL ABILITY EVALUATION

  • Enriched alignment (S2) reduces general-ability performance on both benchmarks, while adding general reasoning data (S3) partially recovers it.Table 5 reports scores for the base Qwen3-1.7B model and two intermediate training stages; higher scores are better.Source: §C.5 · p. 28, Table 5 · p. 28

D.1 BROADER IMPACTS

  • Ranked diagnoses and readable reasoning traces could help physicians revisit earlier visits, but clinical decisions would still rely on examination, testing, and professional judgment.The authors propose prospective studies to examine effects on case review and diagnostic decisions.Source: §D.1 · p. 28
  • The SID pipeline could be adapted to other coding systems, fuller records with medication, procedure, and laboratory events, and other EHR prediction tasks.Source: §D.1 · p. 28

D.2 LIMITATIONS

Generated traces may organize evidence toward complementary diagnoses, but their clinical meaning and accuracy remain unresolved. Future work proposes controlled tests of trace effects, clinician assessment of traces, and extensions to broader EHR events.

D.2 LIMITATIONS

  • Reasoning chains may organize temporal or comorbidity evidence toward complementary diagnoses, but Table 1 gains are only consistent with this proposed mechanism.The passage frames this as a possible explanation, not a demonstrated causal effect.Source: §D.2 · p. 28
  • Plausible traces may omit prediction-relevant factors, and faithfulness varies across models and tasks; clinical accuracy remains to be studied.Teacher traces receive a reference diagnosis during construction, whereas the student sees only H at inference.Source: §D.2 · p. 28

D.3 FUTURE WORK

  • A controlled comparison of pθ(y | H, τ) and pθ(y | H), before rank fusion and with matched sampling and computation, is proposed to separate trace conditioning, repeated sampling, and rank fusion.The analysis would track SID log scores, ranks, and additional target diagnoses entering top-k lists.Source: §D.3 · p. 28
  • Clinicians could assess teacher and student traces for factual accuracy, relevance to recorded history, and usefulness in reviewing candidate diagnoses.Changing selected evidence in H or τ could test whether predictions respond to clinical factors expressed in the trace.Source: §D.3 · p. 29
  • Future work will examine SID vocabularies across coding systems and add medication, procedure, and laboratory events, with possible extensions to recommendation tasks.The passage identifies medication recommendation and broader EHR event recommendation as potential applications.Source: §D.3 · p. 29

Related papers

Loading 2610.10641v1…