MedREK: Retrieval-Based Editing for Medical LLMs with Key-Aware Prompts

TL;DR

This work tackles the rapid evolution and potential inaccuracies in medical knowledge encoded in LLMs by introducing MedVersa, a batch-editing medical knowledge benchmark, and MedREK, a retrieval-based editing framework. MedREK combines a shared query-key MLP for precise knowledge retrieval with an attention-based prompt encoder that generates knowledge-specific prompts to guide edits, thereby preserving locality while updating facts. Empirical results show state-of-the-art performance across Efficacy, Generality, and Locality metrics in both single- and batch-edit settings, and demonstrate robust retrieval effectiveness and practicality for real-world medical knowledge maintenance. The authors also provide a dataset (MedVersa) and code to facilitate reproducibility and further research in medical LLM knowledge editing.

Abstract

LLMs hold great promise for healthcare applications, but the rapid evolution of medical knowledge and errors in training data often cause them to generate outdated or inaccurate information, limiting their applicability in high-stakes clinical practice. Model editing has emerged as a potential remedy without full retraining. While parameter-based editing often compromises locality and is thus ill-suited for the medical domain, retrieval-based editing offers a more viable alternative. However, it still faces two critical challenges: (1) representation overlap within the medical knowledge space often causes inaccurate retrieval and reduces editing accuracy; (2) existing methods are restricted to single-sample edits, while batch-editing remains largely unexplored despite its importance for real-world medical applications. To address these challenges, we first construct MedVersa, an enhanced benchmark with broader coverage of medical subjects, designed to evaluate both single and batch edits under strict locality constraints. We then propose MedREK, a retrieval-based editing framework that integrates a shared query-key module for precise matching with an attention-based prompt encoder for informative guidance. Experimental results on various medical benchmarks demonstrate that our MedREK achieves superior performance across different core metrics and provides the first validated solution for batch-editing in medical LLMs. Our code and dataset are available at https://github.com/mylittleriver/MedREK.

Figures (5)

• Figure 1: Pipeline of our MedREK. Left: Construction of the knowledge base by encoding medical knowledge into key-value pairs. Right: Inference process where different types of queries are encoded to retrieve relevant knowledge and generate attention-based prompts for precise model editing.
• Figure 2: (a) Visualization of key representations in retrieval-based editing methods. Key overlap raises retrieval errors, decreasing editing accuracy. (b) Example of incorrect retrieval caused by key overlap. (c) Cosine similarity measurements quantifying degree of overlap among key representations.
• Figure 3: Comparison of MedVersa and MedCF++ in terms of medical subjects and dataset statistics. We provide the full distribution of medical subjects in the two datasets in Appendix.
• Figure 4: Distribution of query and corresponding key representations (i.e., the keys of the $k$–$v$ pairs in the knowledge base) under a batch of 100 edits using Meditron-7B on the MedCF++ dataset.
• Figure 5: (a-b) Distribution of relative similarity between the query, the ground-truth knowledge, the most similar knowledge and the prototype for test samples in 100-edit batch-editing using Meditron-7B on MedCF++. For Eff. and Gen., x values < 1 indicate incorrect or no retrieval; values > 1 indicate correct retrieval. For Loc., x values = 1 indicate no retrieval (locality preserved); values > 1 imply unintended retrieval. (c) Retrieval statistics across Eff., Gen., and Loc.