Mass-Editing Memory in Transformers

• The paper introduces a closed-form weight update that enables simultaneous mass edits of factual associations in transformer MLP submodules while preserving unaltered knowledge.
• The methodology utilizes causal tracing and key–value pair extraction to precisely locate and update memory, achieving high edit success even with large batch sizes.
• Experimental results demonstrate MEMIT’s scalability and precision, outperforming single-edit methods and maintaining locality and robustness under varied conditions.

Mass-editing memory in a transformer (MEMIT) refers to a precise, scalable method for directly modifying the factual knowledge encoded within the parameters of LLMs—particularly within their multi-layer perceptron (MLP) submodules—without retraining or broadly finetuning the entire network. MEMIT stands out for its ability to process thousands of factual edits simultaneously in a single batch, efficiently updating the model’s internal associative memory structures and achieving performance unmatched by prior single-edit approaches.

1. Conceptual Framework and Objective

MEMIT formalizes the model editing task as a constrained weight update problem over critical MLP modules in transformers. The framework views certain MLP feed-forward layers as associative key-value memories, where each association corresponds to a specific factual statement (e.g., a <subject, relation, object> triple). The model editing objective, often referred to as the “preservation-memorization objective,” is to insert a new set of key–value associations corresponding to edited facts, while strictly preserving outputs for a set of keys representing stored but unedited knowledge.

This objective is typically formulated as:

$$\min_{\hat{W}} \lambda \| \hat{W} K_0 - W_0 K_0 \|^2_F + \| \hat{W} K_E - V_E \|^2_F$$

where:

• $W_0$ and $\hat{W}$ are the original and updated layer weight matrices,
• $K_0$ represents keys for knowledge to be preserved,
• $K_E$, $V_E$ are the edited keys–values for insertion,
• $\lambda$ trades off preservation and memorization.

MEMIT solves for a closed-form weight update across (potentially) thousands of factual associations using efficient linear algebra, enabling “mass-editing” in a single batch (Meng et al., 2022, Gupta et al., 21 Mar 2024).

2. Methodology: Locate-then-Edit and Optimization

Layer and Token Localization

MEMIT begins by identifying the critical MLP layers responsible for factual recall. This localization relies on causal tracing techniques, which measure the indirect effect of hidden states on the target output, and typically focuses on hidden activations at the final subject token of each fact prompt (Meng et al., 2022, Gupta et al., 2023).

Key–Value Pair Extraction

For each fact $(s_i, r_i, o_i)$ to be edited, MEMIT computes:

• A key $k_i$—the hidden state at the critical layer and token position, which acts as the address for that fact in the associative memory.
• A target value $v_i$—a hidden representation optimized so that, if substituted at $k_i$’s position, causes the model to output $o_i$.

Formally, for each $i$, the target value is obtained via:

$v_i = h^{l}_{i} + \delta_i$

subject to minimizing

$$\mathbb{E}_{x_j \sim \text{prefixes}}[ -\log P(o_i \mid x_j \oplus p(s_i, r_i)) ]$$

where $\oplus$ denotes concatenation, and $h^l_i$ is the original hidden state at position $k_i$.

Batch Weight Update

The central step spreads the necessary changes over all targeted layers by solving for the update $\Delta$ to the weights $W_0$:

$$\Delta = (V_E - W_0 K_E) K_E^\top ( \lambda C_0 + K_E K_E^\top )^{-1}$$

where $C_0$ is the empirical covariance (outer product sum) over the preserved keys $K_0$, typically approximated with a randomly sampled subset of hidden activations (Ojito et al., 6 Jun 2024).

This closed-form, Batched Least-Squares update distributes modifications efficiently, ensuring the new facts are “memorized” while the behavior for all preserved keys is retained to the extent possible.

3. Experimental Performance and Scaling Properties

MEMIT consistently outperforms single-edit and finetuning-based baselines in both scale and specificity. On benchmarks like zsRE and CounterFact, it maintains high “Edit Success” (whether the target fact is output), “Paraphrase Success” (generalization across paraphrased prompts), and “Neighborhood Success” (locality: unedited nearby facts remain unchanged) even as the number of edits scales into the thousands (Meng et al., 2022).

A comparison table for core editing metrics reported for GPT-J (6B parameters) is as follows:

Method	# Edits	ES (Efficacy)	PS (Paraphrase)	NS (Neighborhood)	Composite Score S
ROME	1	~100%	High	High	High
MEMIT	1–10k	>85% (10k)	High	High	85.8
MEND	1–1k	Drops sharply	Moderate	Low	Lower

As batch size and edit count increase, MEMIT’s locality degrades more gracefully than earlier methods, making it well-suited for real-world batch updates.

4. Limitations, Batch Structure, and Key Collisions

While MEMIT’s batch update mechanism excels at large edit sets, two important limitations have been identified:

1. Edit Batch Size Degradation: As batch sizes increase (especially above ~1024), performance on paraphrase and neighborhood metrics sharply decreases. Very large “one-shot” edits may introduce more interference and unanticipated generalization (Yoon et al., 1 May 2024). Sequentially applying smaller batches or “sequential-batch” editing achieves higher stability than single massive updates. The relevant update matrix

$$\Delta = (V_E - W_0 K_E) K_E^\top (\lambda C_0 + K_E K_E^\top)^{-1}$$

becomes less well-conditioned with very large $|K_E|$.

1. Key Collisions in Same-Subject Editing: If multiple facts in a batch share the same subject, the corresponding keys $k_i$ are nearly identical. Standard MEMIT’s mechanism cannot correctly map a single key to several different values, resulting in “key collisions” and drastically reduced edit success rates (often falling below 50% at large batch sizes for repeated subjects). MEMIT-Merge mitigates this by merging the value computation for each subject group, achieving >90% edit success in same-subject batch editing (Dong et al., 11 Feb 2025).

5. Robustness, Long-Form Editing, and Failure Modes

Robustness to Context and Prefix: MEMIT’s efficacy can degrade in context-rich or noisy settings due to “embedding collisions”—where hidden representations of different facts become indistinguishable. NAMET, a noise-aware extension, introduces left-padding with [unk] tokens to decorrelate hidden keys and values, improving efficacy (up to 15% higher than MEMIT in scale tests), generalization, and resistance against long prefixes (Dai et al., 17 May 2025).

Long-Form Generation and Locality: MEMIT, like several “rank-one update” approaches, can cause “factual drift” in long-form generations. While the edited fact is inserted successfully for the subject passage, ground-truth properties not meant for editing may be altered unpredictably, resulting in reduced locality and consistency (Rosati et al., 14 Feb 2024). This over-editing is rooted in the global nature of the weight update.

6. Precomputation Efficiency and Practical Considerations

Original MEMIT procedures required extensive layerwise precomputation—estimating covariance matrices with up to 44 million hidden vectors, demanding >36 GPU-hours on a 6B-parameter model. Theoretical analysis shows that the number of precomputed vectors can be reduced to less than 0.3% of the original requirement (e.g., from 44M to ~32k for GPT-J) if a sufficient dynamic multiplier is chosen, drastically reducing time and compute overhead without sacrificing efficacy (Gupta et al., 4 Jun 2025).

In practice, this reduction is achieved by sampling only $d_m \cdot d_k$ hidden states, where $d_k$ is the key vector dimensionality and $d_m$ a small multiplier (e.g., 2–3), ensuring the key covariance matrix is full-rank and well-conditioned.

7. Applications, Extensions, and Future Directions

MEMIT’s ability to perform mass, precise, and efficient memory edits unlocks dynamic knowledge management in LLMs:

• Updating Factual Knowledge: Incorporating breaking news or correcting errors without re-training.
• Customizing Domain Knowledge: Industry-specific or user-specific updates in deployed LLMs.
• Privacy and Unlearning: As shown in (Li et al., 26 May 2025), applying MEMIT with a refusal or “empty-set” (∅) response enables targeted unlearning (removal) as a constrained case of editing.
• Cross-Lingual and Attention-Integrated Extensions: MEMAT, as an extension of MEMIT, demonstrates that supplementing MLP edits with selective attention head modifications not only enhances confidence but improves multilingual propagation of edits and magnitude metrics by ~10% (Tamayo et al., 4 Feb 2025).

Prospective research is focused on mitigating catastrophic forgetting after many sequential edits (Gupta et al., 15 Jan 2024), refining key–value modeling to avoid collisions (Duan et al., 8 Feb 2025), introducing adaptive batch strategies (Yoon et al., 1 May 2024), and reducing long-form factual drift.

---

In conclusion, MEMIT provides a theoretically principled, empirically validated approach for scalable transformer memory editing. Its combination of closed-form batch updates, efficient precomputation, and extensibility to address key collision and robust generalization has established it as a foundation for contemporary research in transformer knowledge control, update, and unlearning.