💬 NLP

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

The paper introduces Keys to Knowledge (K2K), a novel framework that enhances healthcare prediction by replacing computationally intensive external retrieval with efficient, activation-guided internal key-value memory access, thereby achieving state-of-the-art performance with minimal inference latency.

Original authors: Mingchen Li, Jiatan Huang, Zonghai Yao, Hong yu

Published 2026-04-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mingchen Li, Jiatan Huang, Zonghai Yao, Hong yu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a brilliant doctor who has read every medical textbook ever written. You are incredibly smart, but when a patient walks in with a complex, rare set of symptoms, you sometimes get stuck. You might "hallucinate" (make up a fact) because you can't quite remember the specific details, or you might take too long searching through your massive mental library to find the right answer. In a hospital emergency room, you don't have time to flip through thousands of books; you need the answer now.

This paper introduces a new system called K2K (Keys-to-Knowledge) that solves this problem by turning the doctor's brain itself into a super-fast, searchable library, without needing to look outside.

Here is how it works, broken down into simple concepts:

1. The Problem: The "External Library" is Too Slow

Usually, when AI doctors need extra info, they use a method called RAG (Retrieval-Augmented Generation). Think of this like a doctor running out of the room to a giant, dusty library to find a specific book, read it, and then run back to the patient.

The Issue: This takes too long (high latency). In an emergency, seconds count. Also, carrying a huge book into the room makes the doctor's head feel heavy and confused (computational cost).

2. The Solution: The "Internal Filing Cabinet"

The researchers realized that the AI model (the "doctor") already stores a massive amount of medical knowledge inside its own brain (its mathematical parameters). It's like the doctor has a filing cabinet built right into their skull, but they've never learned how to open the right drawer quickly.

K2K teaches the AI how to open those drawers instantly. Instead of running to an external library, it pulls the exact fact it needs from its own internal memory.

3. How K2K Works (The Three Magic Steps)

The paper describes three clever tricks to make this internal search work perfectly:

A. Building the Memory (The "Internal Filing Cabinet")

First, they take the AI and teach it specific medical facts (like disease relationships) by "injecting" this knowledge directly into its brain using a technique called LoRA.

Analogy: Imagine taking a general encyclopedia and secretly writing specific notes about heart attacks and diabetes directly into the margins of the pages the AI reads most often. Now, the knowledge is part of the AI's DNA, not just a book on a shelf.

B. The "Smart Flashlight" (Activation-Guided Probe)

When the AI looks at a patient's history, it needs to know what to search for. Standard AI just looks at the whole sentence and gets a blurry idea. K2K uses a "Smart Flashlight."

Analogy: If a patient says, "I have chest pain and shortness of breath," a normal search might look at the whole sentence equally. K2K's flashlight zooms in on the most important words (the "outliers" or rare symptoms) and ignores the boring filler words. It uses a special math trick (Mahalanobis distance) to figure out which words are the "golden keys" that will unlock the right memory.

C. The "Triage Nurse" (Cross-Attentive Reranking)

Once the AI finds a few relevant facts from its internal memory, it might find too many or slightly wrong ones. It needs a way to pick the best one.

Analogy: Imagine the AI finds three different medical notes about "chest pain." A human triage nurse would look at the patient's specific situation and say, "Note A is about a heart attack, Note B is about indigestion, and Note C is about anxiety. For this patient, Note A is the most critical." K2K acts as this triage nurse, instantly weighing the options and combining the best facts to make the final prediction.

4. The Results: Faster and Smarter

The researchers tested this on real hospital data (predicting if a patient will die or be readmitted).

The Outcome: K2K was faster and more accurate than the old methods. It didn't need to run to the external library. It just looked inside its own brain, used its "Smart Flashlight" to find the key, and let the "Triage Nurse" pick the best answer.
Why it matters: In healthcare, speed saves lives. By removing the delay of searching external databases, K2K makes AI doctors ready for real-time, high-stakes decisions.

Summary

Think of K2K as upgrading a doctor from someone who has to run to the library for every answer, to a doctor who has memorized the library but also has a magical, instant-search system built into their mind. It's efficient, it's fast, and it stops the AI from making things up because it can actually "remember" the facts it needs.

1. Problem Statement

Large Language Models (LLMs) show promise in healthcare but face critical barriers in high-stakes clinical settings:

Hallucinations & Context Gaps: LLMs often lack granular, patient-specific medical context and may generate factually incorrect information.
Latency in Standard RAG: Traditional Retrieval-Augmented Generation (RAG) relies on searching massive external knowledge bases (structured graphs or unstructured documents). This introduces significant computational overhead and high latency, making it impractical for time-sensitive clinical decision-making.
Prompt Limitations: Injecting external knowledge via long input prompts expands context length, escalating inference costs and limiting scalability.
Retrieval Quality: Existing internal retrieval methods often use raw queries that lack discriminative power, leading to the retrieval of irrelevant keys, and fail to account for per-dimension variance in embeddings.

2. Methodology: Keys-to-Knowledge (K2K)

K2K is a novel framework that bypasses external retrieval by facilitating internal, key-based knowledge access directly within the LLM's parameter space. It treats the Feed-Forward Network (FFN) layers of the transformer as an associative key-value memory system.

The framework operates in three core stages:

A. Internal Memory Construction

Instead of relying on external databases, K2K transforms the pre-trained LLM into a retrieval memory:

Mechanism: It leverages the insight that FFN keys ( $W_1$ ) implicitly store factual knowledge.
Knowledge Infusion: For domain-specific medical knowledge absent in pre-training, the authors use LoRA (Low-Rank Adaptation).
- Document-Level Memory ( $K_{doc}$ ): A domain-adapted model is used where FFN keys represent unstructured clinical text.
- Graph-Level Memory ( $K_{graph}$ ): Structured medical knowledge graphs are linearized into text triples and infused into the model via LoRA adapters. The resulting adapter matrices serve as structured knowledge sources.

B. Activation-Guided Probe Construction

To ensure accurate retrieval, the framework constructs a high-quality "probe query" from the input context, addressing the issue that raw queries often yield similar, non-discriminative keys.

Contextual Activation Weighting: Instead of simple mean pooling, the method identifies critical tokens and outlier features.
Mahalanobis Distance Approximation: It computes a diagonal approximation of the Mahalanobis distance to weight tokens. This accounts for per-dimension variance in the embedding space, emphasizing directions with low variance (high sensitivity) rather than treating all dimensions equally (as Euclidean distance does).
Result: This produces an enhanced probe vector ( $Q^+$ ) that better captures semantic "anchors" for precise key retrieval.

C. Cross-Attentive Reranking

To integrate diverse retrieved knowledge dynamically:

Retrieval: The enhanced probe queries retrieve top- $k$ keys from both document and graph internal memories.
Cross-Attention (CA): A cross-attention mechanism reranks and refines these retrieved vectors based on the specific clinical query.
Fusion: The retrieved document and graph knowledge are pooled, concatenated with the original input representation, and fed into an MLP for the final binary prediction (e.g., mortality or readmission).

3. Key Contributions

Internal Retrieval Framework: Proposes K2K, which eliminates the latency of external search by retrieving knowledge directly from the LLM's internal FFN key space.
LoRA-Based Knowledge Injection: Demonstrates an efficient method to infuse structured (graph) and unstructured (document) medical knowledge into the model parameters without full fine-tuning.
Advanced Query Construction: Introduces an Activation-Guided Probe using Mahalanobis distance approximation to improve retrieval precision by accounting for embedding variance.
Dynamic Integration: Develops a Cross-Attentive Reranking module that adaptively reweights multi-source internal knowledge based on task requirements.

4. Experimental Results

The framework was evaluated on MIMIC-III and MIMIC-IV datasets for two tasks: Mortality Prediction and 30-Day Readmission Prediction.

Performance: K2K achieved State-of-the-Art (SOTA) performance across four benchmark datasets, outperforming:
- Traditional sequential models (GRU, RETAIN, TCN).
- Existing retrieval-based models (KARE, Standard RAG).
- Prompt-based generative retrieval methods.
- Example: On MIMIC-III Mortality, K2K achieved an average score of 26.26, surpassing the next best method (KARE at 23.57) and significantly outperforming the base LLM without retrieval (23.92).
Efficiency: K2K demonstrated substantially lower retrieval latency compared to KARE and standard RAG, as it avoids external database queries and long context processing.
Ablation Studies:
- Knowledge Sources: Removing either document or graph knowledge degraded performance, confirming the value of multi-source internal memory.
- Query Strategy: The Mahalanobis-guided query construction significantly outperformed Euclidean distance and mean pooling strategies.
- Layer Analysis: Valuable knowledge was found not only in deep layers but also in shallow layers (e.g., layers 5, 8, 10), suggesting a non-linear distribution of knowledge.

5. Significance

Clinical Viability: By removing external search latency, K2K makes LLM-based decision support feasible for time-critical clinical environments.
Theoretical Insight: The work validates that LLM parameters are not just static weights but can be explicitly accessed as a retrievable knowledge base, challenging the notion that external RAG is the only path to grounded generation.
Scalability: The approach offers a scalable solution for knowledge-intensive tasks by leveraging the model's own capacity, reducing the dependency on massive, heterogeneous external knowledge bases.

In conclusion, K2K represents a paradigm shift from external retrieval to internal memory retrieval, offering a faster, more accurate, and computationally efficient framework for healthcare prediction.