Research map / Large language models / Language, safety and interpretability

Memorize and Rank: Elevating Large Language Models for Clinical Diagnosis Prediction

AAAI 2025 · Large language models › Language, safety and interpretability · approach: LLM unlearning and benchmarks · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Mera fine-tunes generative LLMs to memorize ICD code definitions and then trains them with hierarchical contrastive learning and a dynamic EOV-threshold loss to rank diagnoses for next-visit prediction on MIMIC-III/IV.

Problem

Clinical diagnosis prediction from patient visit histories is hard due to scarce patient data and a huge ICD candidate space (13k+ codes), and generative LLMs have not been effectively used for it because they do not know medical codes and token-level training mismatches the ranking task.

Approach

Teach the LM to memorize bidirectional code-definition mappings and ICD hierarchy, then optimize the token-probability distribution over all code tokens with a level-wise InfoNCE loss over ICD disease groups plus a dynamic cross-entropy loss that places a special EOV token as the confidence threshold, with teacher-forced intra-visit training.

Key results

• On MIMIC-III ICD-9 diagnosis prediction, Mera (BioMistral-7B) reaches 33.24 weighted F1 and 49.01 recall@20, versus 27.35 and 41.29 for the prior best KGxDP.
• On MIMIC-IV, Mera (BioMistral-7B) reaches 36.16 weighted F1 and 49.09 recall@20, above KGxDP's 30.38 and 43.47.
• Zero-shot GPT-4 reaches only 45.16% code accuracy and 48.48% definition accuracy on ICD-9 memorization, while fine-tuned BioMistral-7B reaches 99.61% code accuracy.

Measured effects reported in the paper

Surprising findings

• Fine-tuning a generative LLaMA2-7B directly on diagnosis prediction yields recall@20 about 20 points below the GNN-based best model, showing vanilla generative LMs are poor at this task.
• Zero-shot GPT-4 is still far behind fine-tuned models on both code memorization and diagnosis prediction.

Limitations

• Relies on decoder-only 7B LLMs; smaller models such as GPT-2 medium perform poorly on diagnosis prediction.
• Requires ICD ontology hierarchy and code-definition mappings to build training data.

Experimental setup

Datasets: MIMIC-III MIMIC-IV ICD-9 ontology ICD-10 ontology

Benchmarks: MIMIC-III diagnosis prediction MIMIC-IV diagnosis prediction MIMIC-III heart failure prediction MIMIC-IV heart failure prediction

Metrics: weighted F1 recall@10 recall@20 AUC F1 code accuracy definition accuracy

Baselines: RETAIN Dipole Timeline HiTANet Deepr GRAM G-BERT CGL Chet MCDP KGxDP RoBERTa

Builds on: InfoNCE LLaMA2 BioMistral MedBERT

Closest papers by meaning

• Large Language Models Are Clinical Reasoners: Reasoning-Aware Diagnosis Framework with Prompt-Generated Rationales (AAAI 2024 · similarity 0.89)
• Reasoning-Enhanced Healthcare Predictions with Knowledge Graph Community Retrieval (ICLR 2025 · similarity 0.87)
• BoxLM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare (ICML 2025 · similarity 0.86)
• Hierarchical Divide-and-Conquer for Fine-Grained Alignment in LLM-Based Medical Evaluation (AAAI 2025 · similarity 0.86)
• Contrastive Learning for Clinical Outcome Prediction with Partial Data Sources (ICML 2024 · similarity 0.85)
• RAD: Towards Trustworthy Retrieval-Augmented Multi-modal Clinical Diagnosis (NeurIPS 2025 · similarity 0.85)