Large Language Models (LLMs) suffer from a critical limitation: they are frozen in time. Once trained, their knowledge is static. If the President of the United States changes, or if a new scientific discovery corrects a previous theory, the model remains ignorant until it undergoes expensive retraining or fine-tuning.

To solve this, researchers developed Model Editing—techniques to surgical update specific facts inside a model without retraining the whole network. One of the most popular methods is ROME (Rank-One Model Editing). It has been hailed as a breakthrough for its ability to locate and edit specific factual associations.

However, there is a catch. While ROME works beautifully for a single edit, it tends to break catastrophe when you try to edit the model repeatedly. This phenomenon is known as Model Collapse.

In this post, we will dive into the paper “Rebuilding ROME” by Gupta et al., which investigates why these collapses happen. It is a fascinating detective story where the culprit turns out to be not the mathematical theory, but a subtle discrepancy in the code implementation. We will explore how the authors diagnosed the problem and proposed r-ROME, a stabilized version that allows for thousands of sequential edits.

The Problem: When “Italy” becomes Everything

Imagine you want to update your model with a sequence of new facts. You edit one fact, then another, then another. Suddenly, after a specific edit, the model loses its mind.

In the original ROME implementation, researchers observed that certain specific edits, dubbed disabling edits, would cause the model to output repetitive, nonsensical text.

Figure 1: A typical generation example after a disabling edit is compared to a normal model edit using ROME. The bold and underlined part in the text is input prompt.

As shown in Figure 1 above, a normal edit results in coherent text (bottom). However, a disabling edit (top) causes the model to enter a degenerate loop, repeating the word “Italy” until it hits the token limit. This isn’t just a minor glitch; it is a total loss of linguistic capability.

This is devastating for Sequential Editing—the realistic scenario where we want to continuously update a model over its lifetime. If just one bad edit can brick the model, the technique is unsafe for production.

Background: How ROME Works

To understand the fix, we first need to understand the mechanism of ROME.

ROME treats the Feed-Forward Networks (FFNs) inside a Transformer as “key-value” memories.

  • The Key (\(k\)): Represents the subject or the query (e.g., “The President of the USA is”).
  • The Value (\(v\)): Represents the target knowledge (e.g., “John Cena”).

When you want to insert a new fact, ROME calculates a specific algebraic update to the weights (\(W\)) of a specific layer. The goal is to maximize the probability of the new target (\(v_e\)) given the prompt (\(k_e\)).

The key vector \(k\) is calculated based on the activations of the neural network at a specific layer. The formal definition for the key at a specific layer \(l^*\) is:

Equation defining the key vector based on layer activations.

Here, the key vector depends on the input \(x\). However, to make sure the edit generalizes (i.e., the model knows the fact regardless of slight phrasing changes), ROME doesn’t just use the raw prompt “The President of the USA is”. Instead, it attaches random prefixes to the prompt (like “I think that…”, “Today is a sunny day…”, etc.) and averages the resulting key vectors.

This averaged key vector, denoted as \(k_e\), is calculated as:

Equation 3: Formula for averaging key vectors over N random prefixes.

By averaging these representations, the edit becomes robust to different contexts.

The Investigation: Why do Models Collapse?

The authors of Rebuilding ROME set out to find why sequential editing led to collapse. They performed thousands of edits on models like GPT-J and GPT2-XL.

They noticed a pattern. The “disabling edits” that broke the model were characterized by a massive spike in the magnitude of the weight update matrix, denoted as \(|\Delta|\). In simple terms, while most edits gently nudged the model’s weights, these specific edits were hitting the weights with a sledgehammer.

This is clearly visible in the analysis below.

Figure 2: Scatter plot comparing the update magnitudes of ROME vs r-ROME.

Look at Plot (a) for the original ROME. You can see two distinct clusters. most edits are normal, but there is a cluster of points on the far left (low entropy, meaning broken generation) and points with oddly high or low \(|\Delta|\) values. These outliers are the model killers.

The “Bug” in the Code

Here is the twist: The authors reviewed the original mathematical derivation of ROME and found it to be sound. Theoretically, these massive updates shouldn’t be happening.

So, they looked at the code.

The update equation for ROME is supposed to look like this:

Equation 6: The theoretical update equation for ROME.

In this equation:

  • \(\Delta\) is the change added to the weights.
  • \(k_e\) is the averaged key vector (from Equation 3).
  • \(v_e\) is the target value vector.
  • \(C_0\) is the covariance matrix of the pre-existing knowledge.

However, the actual implementation in the popular ROME codebase used a slightly different formula:

Equation 5: The implemented update equation showing the asymmetry.

Notice the term in bold: \(\mathbf{k_e^o}\).

The original authors had defined \(k_e^o\) as the key vector for the original prompt only, without any prefixes:

Equation 4: Definition of the original prompt key vector.

The Error: The code was using the averaged key vector (\(k_e\)) in some parts of the fraction, but the original, non-averaged key vector (\(k_e^o\)) in others.

  • Numerator: Uses \(k_e\) (averaged).
  • Denominator: Uses \(k_e^o\) (non-averaged).
  • Residual Calculation: Uses \(k_e^o\) (non-averaged).

This asymmetry created an instability. When the vector space of the averaged key and the original key differed significantly, the denominator in the update equation could become dangerously small or mismatched, resulting in the massive values of \(|\Delta|\) that shattered the model’s weights.

The Solution: r-ROME

The fix proposed by the authors is elegantly simple: Consistency.

They introduced r-ROME (Rebuilt ROME), which adheres strictly to the mathematical derivation. It uses the averaged key vector (\(k_e\)) consistently across the entire update equation.

By replacing the asymmetric usage with a homogeneous one, the “disabling edits” vanish. Referring back to Figure 2(b) (the scatter plot shown earlier), you can see that with r-ROME, the edits form a single, dense cluster. There are no outliers, no massive spikes in update magnitude, and no collapsed low-entropy generations.

An Alternative: p-ROME

The authors also tested a variant called p-ROME, where they used the original prompt key (\(k_e^o\)) consistently everywhere. This effectively removes the prefix-averaging feature entirely. While this also stabilized the model, it slightly reduced the generalization capability (how well the model handles paraphrases of the new fact), confirming that the averaging strategy is useful, provided it is implemented correctly.

Experimental Results: Stability at Scale

To prove that r-ROME solves the problem, the authors simulated sequential editing of up to 5,000 facts on GPT-J (6B parameters).

The Collapse of Original ROME

First, let’s look at what happens with the original implementation.

Figure 3: Sequential editing using original implementation of ROME on GPT-J (6B).

In Figure 3, look at the red line in graph (a). This represents the F1 score on the SST2 task (a standard language benchmark). As the number of edits increases, performance stays high for a while, and then—crash. Around 2,000 edits, the model collapses completely. The graph (b) explains why: the update magnitude \(|\Delta|\) (red line) spikes irregularly.

The Stability of r-ROME

Now, look at the performance of the corrected r-ROME.

Figure 4: Sequential editing with r-ROME on GPT-J.

In Figure 4, the difference is night and day.

  1. Downstream Performance (a): The red line (SST2) and other benchmarks degrade very slowly and gracefully. There is no sudden cliff. The model remains functional even after 5,000 edits.
  2. Update Magnitude (b): The \(|\Delta|\) values are orders of magnitude smaller and grow smoothly. There are no spikes.

Quantitative Metrics

The authors summarized the edit quality in the table below. They measured Efficacy (did the edit work?), Generalization (does it work for paraphrases?), and Locality (did it break other things?).

Table 2: Comparison of Original, r-ROME, and p-ROME performance metrics.

  • Original: High efficacy, but as we saw, it eventually destroys the model.
  • r-ROME: Maintains comparable efficacy (97.92 vs 99.94) but with significantly better stability.
  • Score: The overall score (harmonic mean of metrics) is higher for r-ROME in many contexts, but critically, it actually survives the process.

Conclusion and Implications

The paper “Rebuilding ROME” serves as a crucial reminder in machine learning: implementation details matter. A theoretical derivation can be perfect, but a mismatch in variable usage in the code can lead to catastrophic failure modes like model collapse.

By simply aligning the code with the math—ensuring that the key vectors used in the update equation are consistent—the authors transformed ROME from a fragile method into a robust tool capable of sequential editing at scale.

Key Takeaways:

  1. Disabling Edits occur in ROME due to extremely large weight updates.
  2. These large updates are caused by an asymmetric use of key vectors (mixing averaged keys with original keys) in the original code.
  3. r-ROME fixes this by using averaged keys consistently, eliminating model collapse.
  4. This allows LLMs to be updated sequentially with thousands of new facts without losing their general linguistic abilities.

For students and practitioners working on model editing, this is a green light to revisit sequential editing tasks, armed with a much more stable algorithm.

The r-ROME implementation is available for researchers to use, ensuring that future work on knowledge editing is built on a solid foundation.