Efficient and Effective Internal Memory Retrieval for LLM-Based Healthcare Prediction

Mingchen Li; Jiatan Huang; Zonghai Yao; Hong yu

cs.CLApr 8, 2026

Efficient and Effective Internal Memory Retrieval for LLM-Based Healthcare Prediction

Authors: Mingchen Li, Jiatan Huang, Zonghai Yao, Hong yu

Organizations: University of Massachusetts, Amherst · University of Connecticut · University of Massachusetts, Lowell · UMass Chan Medical School

Abstract

Large language models (LLMs) hold significant promise for healthcare, yet their reliability in high-stakes clinical settings is often compromised by hallucinations and a lack of granular medical context. While Retrieval Augmented Generation (RAG) can mitigate these issues, standard supervised pipelines require computationally intensive searches over massive external knowledge bases, leading to high latency that is impractical for time-sensitive care. To address this, we introduce Keys to Knowledge (K2K), a novel framework that replaces external retrieval with internal, key-based knowledge access. By encoding essential clinical information directly into the model's parameter space, K2K enables rapid retrieval from internal key-value memory without inference-time overhead. We further enhance retrieval quality through activation-guided probe construction and cross-attention reranking. Experimental results demonstrate that K2K achieves state-of-the-art performance across four benchmark healthcare outcome prediction datasets.

Figures & tables

Figure 1: Overview of K2K. The input X consists of longitudinal EHR diagnostic codes, and the output y represents healthcare prediction results. K2K first constructs an internal memory within the LLM by injecting external medical knowledge through the Internal Memory Construction module. Given the input, K2K then employs Activation-Guided Probe Construction to build a probe query that retrieves relevant information from the internal memory. Finally, the Cross-Attentive Reranking module dynamically integrates diverse retrieved knowledge. The aggregated memory knowledge, together with the learned input embeddings, is used for the final healthcare prediction.

Figure 2: Overview of the K2K framework, consisting of three steps: (1) Retrieval Memory Construction builds $K_{\text{doc}}=[K_{1},K_{2},\cdots,K_{dm}]$ and $K_{\text{graph}}=[K_{1}^{*},K_{2}^{*},\cdots,K_{dm}^{*}]$ ; (2) Activation-Guided Probe Query Construction (Blue box function f(.)) enhances the query representation for key retrieval from $K_{\text{doc}}$ and $K_{\text{graph}}$ ; (3) Cross-Attentive Reranking retrieves relevant document knowledge $H_{\text{doc}}^{w}$ and graph knowledge $H_{\text{graph}}^{w}$ for the enhanced query $Q_{w}^{+}=f(H_{w}^{+})$ , and integrates them with the original representation $H$ for final prediction. Here, $t\in{1,2,3}$ .

	III-Mort	III-Read	IV-Mort	IV-Read
Train	7,777	7,777	100,125	10,0125
Test	953	953	12,667	12,667
Dev	978	978	12,547	12,547

Table 1: Datasets Statistics, Mort refers to the Mortality. III refers to the MIMIC-III. Read refers to readmission.

		Mortality-MIMIC-III					Readmission-MIMIC-III
Type	Model	F1	Jac.	AUPRC	AUROC	Avg	F1	Jac.	AUPRC	AUROC	Avg
ML	GRU Chung et al. (2014)	13.87	7.45	8.03	53.50	20.71	68.28	51.84	52.94	50.00	55.77
	RETAIN Choi et al. (2016)	13.73	7.37	9.57	54.86	21.38	45.88	23.48	54.11	51.29	43.69
	Deepr Nguyen et al. (2016)	13.87	7.45	7.58	51.66	20.14	68.28	51.84	51.68	49.70	55.38
	AdaCare Ma et al. (2020)	12.90	6.89	7.80	50.69	19.57	63.49	46.51	52.83	52.27	53.77
	StageNet Gao et al. (2020)	9.97	5.25	7.10	47.14	17.37	51.56	34.74	50.38	48.27	46.24

Table 2: Comparative analysis of mortality and readmission prediction on MIMIC-III and MIMIC-IV. Avg denotes the arithmetic mean of F1, Jaccard, AUPRC, and AUROC for each model on the corresponding dataset. Boldface indicates the highest average score under each LLM.

	Model	F1	Jaccard	AUPRC	AUROC
Mortality-III	K2K	18.55	10.22	15.22	61.05
	K2K w/o graph	20.48	11.40	13.18	60.54
	K2K w/o document	16.66	9.09	10.52	55.72
Mortality-IV	K2K Ours	6.61	3.42	2.93	66.50
	K2K w/o graph	4.50	2.30	2.51	60.86
	K2K w/o document	3.57	1.82	2.71	66.41

Table 3: Results of different knowledge sources in K2K

	Model	F1	Jaccard	AUPRC	AUROC
Mortality-III	K2K	18.55	10.22	15.22	61.05
	LLM	4.49	2.29	8.67	55.62
	LLM+Doc	16.00	8.69	11.61	59.40
	LLM+Graph	4.50	2.29	8.67	55.62
	LLM+Doc+Graph	16.00	8.70	11.61	59.41
Readmission-III	K2K	69.31	53.03	62.49	56.64

Table 4: Comparison of Knowledge-Enhanced Models on Mortality and Readmission Prediction (MIMIC-III/IV). LLM refers to Mixtral-7B. LLM+Doc denotes BioMixtral-7B, which is obtained by further training Mixtral-7B on a medical corpus. LLM+Graph refers to Mixtral-7B adapted to graph-based knowledge using LoRA. LLM+Doc+Graph represents BioMixtral-7B further adapted to graph knowledge via LoRA.

	Model	F1	Jaccard	AUPRC	AUROC
Mortality-III	K2K	18.55	10.22	15.22	61.05
	K2K w Euclidean	16.97	9.27	9.67	57.25
	K2K (Mean Only)	12.06	6.42	8.45	52.51
Readmission-III	K2K	69.31	53.03	62.49	56.64
	K2K w Euclidean	63.27	46.28	58.26	53.25
	K2K (Mean Only)	63.98	47.03	54.67	50.92

Table 5: Comparison of K2K with different query construction methods.

Figure 3: K2K performance with different layer knowledge. We used BioMistral-7B, which consists of 32 transformer layers.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: K2K performance with different chunk sizes.

Figure 5: K2K Performance Across Different Top-k Retrieved Knowledge Values on MIMIC-III.

Figure 6: K2K Performance Across Different Top-k Retrieved Knowledge Values on MIMIC-IV.

	Avg	Retrieval Time
KARE	21.11	00:33:52
Prompt-based	22.52	3:26:00
K2K	22.89	0:0:5

Appendix

Table 6: Performance comparison on the MIMIC-III dataset in terms of averaged metrics (Avg), retrieval time.

Memory-Augmented Language Models

Memory-augmented language models combine persistent or temporary memory with neural computation to retain, retrieve, and use information beyond the active context.

See all 291 papers

Future Querying: Can LLMs Serve as Implicit Medical World Models?
Aug 24, 2026Siri Willems, James Butterworth, Lore Goetschalckx +4Medical LLMs Clinical Prediction
imec, AI-labs⋆⋆, Paris, France
ReMedi: Reasoner for Medical Clinical Prediction
May 2, 2026Yushi Cao, Yiming Chen, Hongchao Jiang +2Fine-Tuning Clinical Outcome Prediction
ASUS Intelligent Cloud Services (AICS), Singapore · NTU Artificial Intelligence Center of Research Excellence (NTU AI-CoRE), Taiwan · National University of Singapore, Singapore
ChatHealthAI: Aligning Electronic Health Record Representations with Large Language Models for Grounded Clinical Reasoning
Jun 1, 2026Bo-Hong Wang, Baicheng Peng, Ruilin Wang +3LLM Grounding LLM Alignment
School of Computer Science, McGill University, Montreal, QC, Canada · Mila - Quebec AI Institute, Montreal, QC, Canada

Aug 24, 2026J/K move · Enter open · S save

Future Querying: Can LLMs Serve as Implicit Medical World Models?

Siri Willems, James Butterworth, Lore Goetschalckx, Peter Vrancx+3

imec, AI-labs⋆⋆, Paris, France

Medical LLMs Clinical Prediction Longitudinal EHR Modeling Clinical Outcome Prediction

Open paper AskPDF arXiv

Traditional clinical prediction models rely on task-specific pipelines and curated, structured data, which scale poorly and underutilize unstructured text. To address this, we introduce future querying, a paradigm that probes whether large language models (LLMs) can function as implicit medical world models by evaluating their ability to answer time-indexed clinical queries about a patient's future. Our framework operates on unstructured clinical documentation using endpoint-agnostic training, enabling a single model to answer diverse clinical queries over patient trajectories without manual feature engineering or task-specific retraining. We show that small, locally fine-tuned open-weight models can match or approach larger proprietary systems, making the framework suitable for privacy-preserving, on-premise deployment. Evaluated on a new synthetic medical reports dataset and real ICU notes from the MIMIC-IV dataset, our results provide encouraging evidence that LLMs can capture aspects of clinical dynamics.

Abstract

Memory-Augmented Language Models

Explore similar work

Future Querying: Can LLMs Serve as Implicit Medical World Models?

ReMedi: Reasoner for Medical Clinical Prediction

ChatHealthAI: Aligning Electronic Health Record Representations with Large Language Models for Grounded Clinical Reasoning