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

arXiv: https://arxiv.org/abs/2604.07659 · HTML: https://mlsanity.com/paper/2604.07659/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 33 references that are in the map.
- First posted on arXiv 8 Apr 2026.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling](https://mlsanity.com/paper/1803.01271/) (2018)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Unsupervised Dense Information Retrieval with Contrastive Learning](https://mlsanity.com/paper/2112.09118/) (2021)
- [RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism](https://mlsanity.com/paper/1608.05745/) (2016)
- [Deepr: A Convolutional Net for Medical Records](https://mlsanity.com/paper/1607.07519/) (2016)
- [GraphCare: Enhancing Healthcare Predictions with Personalized Knowledge Graphs](https://mlsanity.com/paper/2305.12788/) (2023)
- [AdaCare: Explainable Clinical Health Status Representation Learning via Scale-Adaptive Feature Extraction and Recalibration](https://mlsanity.com/paper/1911.12205/) (2019)
- [StageNet: Stage-Aware Neural Networks for Health Risk Prediction](https://mlsanity.com/paper/2001.10054/) (2020)
- [Reasoning-Enhanced Healthcare Predictions with Knowledge Graph Community Retrieval](https://mlsanity.com/paper/2410.04585/) (2024)
- [BiomedRAG: A Retrieval Augmented Large Language Model for Biomedicine](https://mlsanity.com/paper/2405.00465/) (2024)
- [To Generate or to Retrieve? On the Effectiveness of Artificial Contexts for Medical Open-Domain Question Answering](https://mlsanity.com/paper/2403.01924/) (2024)
- [CancerLLM: A Large Language Model in Cancer Domain](https://mlsanity.com/paper/2406.10459/) (2024)
- [Semantic Structure based Query Graph Prediction for Question Answering over Knowledge Graph](https://mlsanity.com/paper/2204.10194/) (2022)
- [Understand the Dynamic World: An End-to-End Knowledge Informed Framework for Open Domain Entity State Tracking](https://mlsanity.com/paper/2304.13854/) (2023)
- [Activation-aware Probe-Query: Effective Key-Value Retrieval for Long-Context LLMs Inference](https://mlsanity.com/paper/2502.13542/) (2025)

## Built on by (most cited first, 1 of 1)

- [RIVO: Lightweight Adaptive Rerouting for Deep Research Agents](https://mlsanity.com/paper/2610.15910/) (2026)

## Abstract (arXiv)

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.