MedREK: Retrieval-Based Medical Model Editing

• MedREK is a retrieval-based editing framework for medical LLMs that updates specific clinical facts externally, preserving overall model integrity and strict edit locality.
• It employs a shared query-key MLP and an attention-based prompt encoder to mitigate retrieval ambiguity caused by representation overlap in similar medical statements.
• The framework supports both single and batch editing as validated by the MedVersa benchmark, enabling safe and scalable updates in evolving clinical environments.

Searching arXiv for the MedREK paper and key related model-editing references mentioned in the provided data. MedREK is a retrieval-based editing framework for medical LLMs that is designed to update specific clinical facts without full retraining, while preserving strict locality and supporting both single-editing and batch-editing. It is introduced together with MedVersa, an enhanced benchmark with broader coverage of medical subjects and explicit support for evaluating single and batch edits under strict locality constraints. The framework targets two difficulties identified for medical model editing: inaccurate retrieval caused by representation overlap among similar medical statements, and the lack of validated batch-editing methods for realistic medical updates such as guideline changes, contraindications, and retractions (Xia et al., 15 Oct 2025).

1. Problem setting and conceptual position

MedREK is motivated by the observation that medical knowledge changes rapidly and that training corpora may contain outdated errors or hallucinations. In that setting, retraining large medical LLMs is costly and may induce catastrophic forgetting, so model editing is treated as a post-hoc mechanism for updating specific facts while retaining unrelated knowledge. The paper defines editing as

$$f'_\theta = ME(f_\theta, q_e, o_e), \qquad f'_\theta(q_e)=o_e,$$

with efficacy, generality, and locality used as the core evaluation dimensions (Xia et al., 15 Oct 2025).

The framework is positioned against two classes of editors. Parameter-based approaches, including locate-then-edit methods such as ROME, MEMIT, and MedLaSA, directly modify internal parameters identified by causal tracing. In the medical domain, the stated concern is that unintended side effects are unacceptable, and these approaches frequently compromise locality. Retrieval-based editing instead stores updated knowledge externally and conditions generation through retrieved prompts or memories, thereby preserving the base model parameters. MedREK adopts that retrieval-based paradigm but treats retrieval ambiguity as a central medical-domain failure mode because the knowledge space contains many textually similar yet factually distinct items.

The paper frames two obstacles as especially consequential. First, representation overlap can cause inaccurate retrieval and reduce editing accuracy. Competing retrieval-based editors are reported to show high confusion, and cosine-similarity analyses indicate significant overlap among stored keys that increases with batch size; the reported “High-Sim. Samples (%)” reaches 51.09% at larger batches. Second, existing medical editing work is described as being largely restricted to single-sample edits, despite the practical relevance of batch updates.

2. MedVersa benchmark and evaluation protocol

MedVersa is introduced as a new medical editing benchmark intended to enforce strict locality constraints while expanding subject coverage and enabling batch-edit evaluation. It is derived from MedMCQA, a multi-subject medical multiple-choice QA dataset. The construction procedure transforms phrases into clear interrogatives and rephrases them for generalization using Gemini 2.0 Flash, with prompts provided in the appendix of the paper (Xia et al., 15 Oct 2025).

Each edit instance contains three QA pair types. The Efficacy pair $(q_e, o_e)$ measures edit success on the original question. The Generality pair $(q_g, o_e)$ is a rephrased question used to measure generalization of the edit. The Locality pair $(q_l, o_l)$ is an unrelated question from the same medical subject and is designed to ensure that edits do not affect unrelated in-domain knowledge. For each Efficacy item, one wrong option from MedMCQA is selected as the edit target, while the ground truth remains the original correct answer.

A central design decision is strict locality. In MedVersa, Efficacy or Generality questions and Locality questions do not share prompts. Locality items are sampled from different entries within the same subject, which is intended to ensure unrelatedness while controlling for domain. This addresses limitations attributed to MedCF, including prompt overlap between Efficacy and Locality and a narrow pharmacology focus. An intermediate benchmark, MedCF++, is also used; it removes problematic overlaps in MedCF, with 181, 47, and 47 records removed from the train, validation, and test splits respectively, but it remains substantially narrower in subject coverage, with Pharmacology reported at 71.78%.

MedVersa covers 20 medical subjects with a balanced distribution. The subject shares reported in the paper are summarized below.

Subject	Share
Anatomy	11.22%
Microbiology	10.18%
Physiology	10.08%
Surgery	10.02%
Social & Preventive Medicine	9.97%
Gynaecology & Obstetrics	8.14%
Ophthalmology	7.94%
Forensic Medicine	6.78%
Pediatrics	6.61%
ENT	6.02%
Medicine	2.71%
Pathology	2.62%
Pharmacology	2.05%
Biochemistry	1.98%
Orthopaedics	0.94%
Radiology	0.82%
Psychiatry	0.64%
Anaesthesia	0.47%
Dental	0.44%
Skin	0.37%

The benchmark supports single-editing and batch-editing with configurations of 10, 50, and 100 edits per batch. The evaluation metrics are formally defined as

$$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$

$$\text{Generality} = \mathbb{E}_{(q_g,o_e)} \mathbb{I}\{f'_\theta(q_g)=o_e\},$$

$$\text{Locality} = \mathbb{E}_{(q_l,o_l)} \mathbb{I}\{f'_\theta(q_l)=f_\theta(q_l)=o_l\},$$

$\text{Fluency} = \sum_k f_n(k)\log_2 f_n(k),$

and

$$\text{Average} = \frac{\frac{\text{Eff.}+\text{Gen.}}{2}+\frac{\sum_{m\in \text{Loc.}} m}{|\text{Loc.}|}}{2}.$$

On MedCF++, the locality sub-metrics are TD, EM, SS, and TS, defined in the appendix.

3. Architecture: shared query-key retrieval and key-aware prompting

MedREK consists of a knowledge base of key-value pairs, a shared query-key representation module, an attention-based prompt encoder, and a prototype-gated retrieval rule. The end-to-end workflow is described as follows. First, a key is formed from the knowledge key, specifically the subject-relation pair $(s_i, r_i)$, with the object excluded to reduce noise. The value is a key-aware continuous prompt token sequence computed from the full triplet $(q_e, o_e)$0. At inference time, the user query is encoded into the same shared query-key space, the most similar key is retrieved if it exceeds a learned prototype threshold, and the associated prompt is appended to the LLM input embedding. If retrieval fails the gate, no prompt is used and the frozen model is run unchanged (Xia et al., 15 Oct 2025).

The input encoder tokenizes text with the RoBERTa tokenizer, extracts last-layer hidden states $(q_e, o_e)$1 and the CLS embedding $(q_e, o_e)$2, and pools them into

$(q_e, o_e)$3

Queries and keys are mapped with a shared residual MLP:

$(q_e, o_e)$4

The same encoder is applied to knowledge keys to obtain $(q_e, o_e)$5, thereby enforcing a unified query-key embedding space. The paper attributes retrieval precision largely to this shared-space alignment and notes that no orthogonality loss is introduced; separation is instead pursued via contrastive objectives and shared-space matching.

Triplet values are encoded by a separate MLP,

$(q_e, o_e)$6

and $(q_e, o_e)$7 is used by the prompt encoder. The prompt encoder generates a continuous prompt of length $(q_e, o_e)$8 by forming

$(q_e, o_e)$9

followed by multi-head attention:

$(q_g, o_e)$0

The attention uses standard scaled dot-product attention,

$(q_g, o_e)$1

Stored prompt tensors are denoted $(q_g, o_e)$2 and are later concatenated to the LLM input embedding.

Retrieval is governed by a learnable prototype vector $(q_g, o_e)$3 that serves as a dynamic threshold. For a query embedding $(q_g, o_e)$4 and the top-1 key $(q_g, o_e)$5, the retrieval rule is

$(q_g, o_e)$6

Scoring uses dot-product similarity. Editing is therefore realized by prefixing the retrieved continuous prompt to the query embeddings while keeping the base model frozen.

4. Optimization objectives and batch-editing mechanism

The total objective is

$(q_g, o_e)$7

The editing loss is averaged over a batch of size $(q_g, o_e)$8:

$(q_g, o_e)$9

Its three components are

$(q_l, o_l)$0

$(q_l, o_l)$1

and

$(q_l, o_l)$2

Here $(q_l, o_l)$3 is the base LLM, $(q_l, o_l)$4 is $(q_l, o_l)$5 with the embedding layer removed, $(q_l, o_l)$6 denotes the embedding layer, and $(q_l, o_l)$7 denotes prompt concatenation (Xia et al., 15 Oct 2025).

The contrastive objective is defined to support prompt learning and selective retrieval. Positive alignment for editing is

$(q_l, o_l)$8

while the selective overflow control term is

$(q_l, o_l)$9

The batch-averaged contrastive loss is

$$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$0

with InfoNCE

$$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$1

where $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$2 and $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$3 defaults to 1.

Single-editing adds or updates one $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$4 pair, whereas batch-editing constructs a knowledge base containing many pairs simultaneously. The shared modules $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$5, $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$6, the prompt encoder, and the prototype $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$7 are trained jointly. At inference, each query independently retrieves its top-1 key and either receives the associated prompt or falls back to the unmodified model. This design is presented as the mechanism by which MedREK avoids parameter conflicts under batch edits. The paper further states that the external knowledge base grows with edits and that, in naive implementations, retrieval scales linearly with knowledge-base size.

5. Empirical results, ablations, and qualitative behavior

The empirical evaluation uses two backbones: Meditron-7B, described as an LLaMA2-based medical model, and HuatuoGPT-o1-8B, described as an LLaMA3-based medical complex reasoning model. Baselines include MEND, MEMIT, MedLaSA, and the retrieval-based editor RECIPE. Benchmarks are MedCF++ and MedVersa, with metrics reported for Efficacy, Generality, Locality, Fluency, and the weighted Average (Xia et al., 15 Oct 2025).

On Meditron-7B, MedREK is reported to outperform RECIPE on both single-editing and batch-editing. For single-editing on MedCF++, MedREK obtains Eff 78.50, Gen 80.61, Loc TD/EM/SS/TS 99.42/98.96/99.34/98.67, Flu 587.00, and Avg 89.33, whereas RECIPE obtains Eff 72.66, Gen 77.12, Loc 92.64/99.80/90.29/90.59, Flu 586.68, and Avg 84.11. On MedVersa single-editing, MedREK reports Eff 74.49, Gen 70.46, Loc 100.00, Flu 579.63, and Avg 86.24; RECIPE reports Eff 57.89, Gen 57.44, Loc 99.03, Flu 599.66, and Avg 78.35.

The batch-editing results are central to the paper’s claims. At 100 edits on MedCF++, MedREK reports Eff 77.96, Gen 79.88, Loc TD/EM/SS/TS 97.76/98.20/97.57/97.19, Flu 588.01, and Avg 88.30. At the same scale, RECIPE reports Eff 68.30, Gen 70.74, Loc 72.74/79.73/72.48/71.58, Flu 588.21, and Avg 71.83. On MedVersa with 100 edits, MedREK reports Eff 74.49, Gen 70.46, Loc 99.45–100.00 across settings, Flu approximately 598–600, and Avg approximately 85.96–86.24, whereas RECIPE reports Eff 57.89, Gen approximately 57.3–57.4, Loc 79.55–96.72, Flu approximately 600–663, and Avg approximately 65.49–77.18.

On HuatuoGPT-o1-8B, the same pattern is reported. For single-editing on MedCF++, MedREK reports Eff 77.05, Gen 78.66, Loc TD/EM/SS/TS 99.60/97.90/98.95/98.24, Flu 653.11, and Avg 88.26, compared with RECIPE Avg 84.58. For single-editing on MedVersa, MedREK reports Eff 69.47, Gen 63.43, Loc 100.00, Flu 634.05, and Avg 83.22, compared with RECIPE Avg 75.26. For 100-edit batches on MedCF++, MedREK reports Eff 76.42, Gen 78.15, Loc approximately 96.6 on the four locality sub-metrics, and Avg 86.83, compared with RECIPE Avg 73.38. On MedVersa at 100 edits, MedREK reports Eff 69.47, Gen 63.28, Loc 99.06, Flu 663.02, and Avg 82.72, compared with RECIPE Avg 65.49.

The ablations attribute distinct roles to the two architectural components. Without the shared MLP, MedREK reports Eff 56.74, Gen 59.28, locality in the approximate range 81–84, and Avg 69.91 on MedCF++ with 100 edits. Without the attention-based prompt encoder, it reports Eff 73.17, Gen 75.63, locality approximately 94–97, and Avg 85.24. The full system reports Eff 77.96, Gen 79.88, locality approximately 97–98, and Avg 88.30. The paper’s interpretation is that the shared MLP is critical for retrieval precision, whereas the attention encoder improves edit quality.

Retrieval analyses are used to explain these numbers. MedREK’s query-key distributions are described as tightly aligned for Efficacy and Generality, whereas RECIPE shows scatter and distance. Prototype-gating ratios are reported in two forms: for Efficacy and Generality, $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$8, where $$\text{Efficacy} = \mathbb{E}_{(q_e,o_e)} \mathbb{I}\{f'_\theta(q_e)=o_e\},$$9 indicates correct retrieval; for Locality, $$\text{Generality} = \mathbb{E}_{(q_g,o_e)} \mathbb{I}\{f'_\theta(q_g)=o_e\},$$0, where $$\text{Generality} = \mathbb{E}_{(q_g,o_e)} \mathbb{I}\{f'_\theta(q_g)=o_e\},$$1 indicates no retrieval, which is the desired behavior. The paper states that MedREK yields more correct retrievals for edit queries and fewer undesired retrievals for locality queries.

A qualitative example from MedVersa uses the question “What is the treatment for multiple carboxylase deficiency?” with ground truth Biotin and counterfactual edit target Thiamine. The rephrased Generality question is “What is the therapeutic management for multiple carboxylase deficiency?” The Locality question, from the same subject area Pediatrics but a different fact, is “At what age does purposeful movement typically start in infants?” with ground truth 6 months. The intended behavior is that the key formed from a subject-relation representation such as $$\text{Generality} = \mathbb{E}_{(q_g,o_e)} \mathbb{I}\{f'_\theta(q_g)=o_e\},$$2 retrieves the appropriate prompt for the target fact, while the prototype gate blocks retrieval for the unrelated locality question.

6. Limitations, implementation details, and broader significance

The paper states several limitations. Because MedVersa uses Gemini 2.0 Flash to rewrite and rephrase source items, subtle changes in phrasing may bias difficulty. Reliance on a RoBERTa tokenizer and encoder plus a shared MLP does not fully eliminate all representation overlaps, and no formal guarantee of zero-conflict retrieval is claimed. Prototype gating is learned but heuristic; locality may be challenged by adversarially phrased queries that closely mimic edited keys. Batch-editing is evaluated at 10, 50, and 100 edits, but extreme scales and lifelong sequential edits are not exhaustively studied. MedVersa covers 20 subjects, yet some remain small-share categories, so rare or long-tail biomedical facts may need further coverage (Xia et al., 15 Oct 2025).

The implementation details reported in the paper are specific. Code and data are released at the stated repository. Training and inference use NVIDIA RTX 5090 GPUs. The learning rate is $$\text{Generality} = \mathbb{E}_{(q_g,o_e)} \mathbb{I}\{f'_\theta(q_g)=o_e\},$$3, batch size is 8, and both representation dimension and hidden size are 4096. The number of prompt tokens is 3 for MedCF++ and 8 for MedVersa. The number of knowledge prototype tokens is 10. Training proceeds to 200 epochs, with early signs of loss increase observed after approximately 150 epochs for both MedREK and RECIPE; final checkpoints are selected by the smallest loss. RECIPE is tuned to use the same number of prompt tokens.

Efficiency is treated as a practical advantage of the frozen-parameter design. For a single edit on Meditron-7B and MedVersa, the reported edit time is 0.012 s for MedREK, compared with 0.006 s for RECIPE, while parameter-modifying methods are stated to be orders slower, with MEMIT at 18.239 s and MedLaSA at 16.787 s due to per-edit finetuning. The memory footprint is described as modest because only prompts and keys are stored while the base LLM remains frozen.

In broader methodological terms, MedREK is positioned as a retrieval-based alternative to parameter modification for safety-critical medical editing. Its distinctive contributions are the shared query-key MLP, the attention-based prompt encoder conditioned on the full knowledge triplet, and the dynamic prototype gate that reduces over-retrieval. The paper characterizes this combination as the first validated solution for batch-editing in medical LLMs under strict locality constraints. A plausible implication is that MedREK is most relevant in settings where high locality is a hard requirement and updates arrive as batches of related but non-identical facts, rather than as isolated corrections.