Representation Shattering in Transformers: A Synthetic Study with Knowledge Editing

We thank the reviewer for appreciating our empirical analysis and mechanism of representation shattering. Please see our detailed responses to specific comments below.

---

#1: Multi-Step Editing

Thank you for this suggestion! We note we already provide an analysis of multi-step and layerwise edits with real LLMs (Llama 3 and Mamba in Sec. 4.5 and App. G.1) via MEMIT:

For each counterfactual edit distance $d$, we have $(\text{twelve months}) \times (\text{three offsets}) = 36$ total edits—we apply these sequentially in batches of size $4$ using MEMIT, using consistent sample orders and batch orders for fair comparisons.

This was done to magnify performance degradations and representation shattering, as singular edits yielded minute effects initially. This qualitatively suggests shattering amplifies in multi-step editing, but a rigorous comparison of Batch vs. Sequential vs. Sequential-Batched editing (as in [1]) is warranted. We will incorporate such comparisons into future work.

[1] Yoon et al., 2024

---

#2: Layerwise Analysis

Current KE methods (ROME/MEMIT) specify early layers (typically early-layer MLPs) as edit sites based on causal tracing. Altering this would require formulating and validating a new KE approach, beyond our scope. We do note that we include analyses in App. F.3 and G.2 visualizing representations from various model sites, hence showing the effects of KE across layers.

---

#3: Connections to Catastrophic Forgetting Mitigation

Thank you for highlighting connections to forgetting mitigation strategies! Indeed, our framework can naturally evaluate such strategies: effective forgetting mitigation methods should preserve global representation geometry when instilling new knowledge (minimizing $R(D_*)$). Recent work in fact examines regularization-based KE methods (e.g., [1]), suggesting compatibility with classical forgetting mitigation literature [2]. A thorough analysis in the context of our work is however best deferred to a future paper.

[1] Gu et al., 2024

[2] https://arxiv.org/abs/1612.00796

---

#4: More Fine-Grained Tasks Beyond Month Ordering

We appreciate your acknowledgment that the months-of-the-year task was intentionally chosen due to its clear interpretability and structured cyclic geometry. This choice enabled controlled experiments directly linking representation geometry and performance.

However, we fully recognize your interest in finer-grained setups. We anticipate more realistic KE tasks—typically involving densely interconnected entities and facts—would exhibit even more severe representation shattering. Since greater interconnectedness implies more overlapping representations and deeper manifold entanglement (elaborated in our response to reviewer (see rVM5, #4), we expect edits in realistic scenarios to cause geometric distortions at least as severe as our cyclic-month edits. Our current findings represent a conservative estimate of the impact representation shattering can have in practical, large-scale knowledge editing scenarios; we plan to explicitly explore these richer domains as an immediate follow-up.

---

#5: Generalization Beyond Cyclic Structures

We acknowledge your concern regarding our emphasis on cyclic structures, and fully agree that exploring other structures (e.g., trees or hierarchies) is crucial. Our current focus on cycles is deliberate: cycles are the simplest nontrivial graphs exhibiting global geometric constraints, ideal for precisely measuring and visualizing representation shattering. Their symmetric nature mitigates extraneous considerations like hierarchical relations, equivalence classes, and edge cases for root/leaf entities.

As you note, our preliminary experiments on tree-structured KGs (App. G.3) align closely with our cyclic results: edits causing larger geometric displacement yield greater representation shattering (e.g., moving "Paris" to UK vs. Spain). Thus, we strongly believe the core phenomenon generalizes beyond cycles. Still, more systematic evaluations with alternative relational structures is needed—this is a key next step we intend to pursue in future work.

---

#6: Code Release

We are committed to releasing our code for the camera-ready version for reproducibility (see rVM5, #6).

---

---

Summary: We again thank the reviewer for their insightful feedback that has helped us emphasize the existing multi-step editing and layerwise analyses in our work, discuss catastrophic forgetting literature, and affirmed our intent to explore more complex setups. We hope our rebuttals address the reviewer's questions, and that they will continue to champion our paper's acceptance!

We thank the reviewer for appreciating our experiments and analysis of representation shattering. Please see our responses below.

---

#1: Generality and Tree-Structured KGs

We acknowledge that our tree-based experiment in App. G.3 is preliminary, but believe further exploration of tree-shaped KGs is best reserved for follow-up work. Specifically, our paper aims to establish representation shattering, define evaluations to detect/visualize it, and demonstrate correspondence between synthetic and large-scale models. Cyclic graphs are uniquely suited for this: they are the simplest non-trivial graphs with global geometric constraints where post-edit effects can be exhaustively elicited. By contrast, trees lack this symmetry and introduce complexities around hierarchical relations, equivalence classes, and edge cases for root/leaf entities.

We agree that tree-shaped KGs are an important next step, and we intend to pursue this direction. For this work, however, we prioritize simplicity—our cyclic results isolate the essence of representation shattering for the community to build on.

---

#2: Mechanistically Causal vs. Correlational Evidence

While several of our experiments elicit correlational evidence, we emphasize our counterfactual editing experiments suggest a causal role for representation geometry in model capabilities (see Sec. 4.4/4.5, esp. Tab. 2, Fig. 6, Fig. 7). We summarize these results below.

Summary of counterfactual experiments. We systematically vary counterfactual edit distance (CE distance) while holding all other variables fixed. For a given subject entity, edits with larger CE distances imply greater displacement in the representation manifold and higher $R(D_*)$. In reverse, higher $R(D_*)$ implies a larger CE distance. This manipulation approximates an intervention on representation geometry itself.

Across all cases, higher CE distance leads to greater representation shattering and larger degradation in direct recall, logical inference, and compositional inference. While we cannot rule out models with scattered but functional representations, it is unclear how to obtain such a model to test this. We argue that the manipulation of CE distance—which directly impacts geometry—offers strong evidence for a causal link between shattering and performance loss.

We will clarify that our experiments test for behaviors consistent with a mechanistic explanation, though we do not claim formal proof.

---

#3: Realism of the "Months" Task

We agree the months-of-the-year task is simpler than real-world KE scenarios involving entity attributes. We selected it to isolate and quantify representation shattering in pretrained models.

Though somewhat artificial, the fact that shattering arises in this clean cyclic setting suggests that richer domains—where knowledge is denser and more interdependent—would experience more severe effects. That is, if shattering already causes degradation in this minimal case, it likely plays a role in realistic tasks too. While we don't claim it explains all KE failures, our results show shattering is a major contributor. Exploring broader domains remains an important future direction (Sec. 5).

---

#4: Layer Choice & Key-Value Assumption (ROME/MEMIT)

In all synthetic experiments, we use a 2-layer nanoGPT Transformer (explicitly stated in Sec. 3.3) Given this architecture, layer 1 is the only editing site compatible with the early-site assumptions of ROME. Further, causal tracing and MEMIT use a sliding window of layers: in a 2-layer model, any non-1 window size would encompass the whole model. We thus used layer 1 to avoid confusion, but early sanity checks with a singleton window re-implementation of causal tracing also confirmed layer 1 as the primary factual recall site. If the reviewer believes these results should be included, we can prepare them for the final version.

We also qualitatively verified the key-value assumptions in exploratory stages. Prompts about the same entity yield clusters in the input space of the layer 1 MLP; prompts with different subjects resolving to the same object yield clusters in its output space. These patterns support the key-value framework underlying ROME/MEMIT as compatible with our synthetic setup. If it will help address reviewer's concerns, we are happy to add visualizations confirming these assumptions.

---

#5: Code Release

We will release our code for the camera-ready version for reproducibility (see rVM5, #6).

---

---

Summary: We again thank the reviewer for their thoughtful feedback that helped us clarify the rationale behind our synthetic design, show how our experiments suggest a causal link between representation geometry and model performance, acknowledge the simplicity of our Llama task as deliberate, and justify our editing layer choices and key-value assumptions. We hope our rebuttals address the reviewer's questions and that they will champion our paper's acceptance!

We thank the reviewer for appreciating our experiments and analysis of representation shattering. Please see our responses below.

---

#1: $R(D_*)$ and Permutation Invariance

Great question! We intentionally made our distortion metric $R(D_*)$ to not be permutation-invariant, i.e., $R(D_*) > 0$ even if $D_*$ is a permutation of $D_{\varnothing}$. This is because $R(D_*)$ tracks how much each entity's position in the representation space, as identified by its token, has changed—not whether the manifold is isomorphic under another labeling.

For example, in your suggested swap of entity 2 with entity 3, the manifold may remain consistent, but the label-respecting configuration changes, and so do the model's semantics. A zero value would suggest no parameter shift even if outputs differ.

A permutation-invariant version (e.g., via minimum matching) would conflate isomorphic configurations and assign zero cost to degenerate cases—like total entity permutation—despite incoherence. For our purposes, permutation invariance is a feature. We will clarify this in the final version of the paper.

---

#2: Shape of $R(D_*)$

$D$ is an $n \times n$ matrix over $n$ entities, where $D_{ij}$ is the Euclidean distance between representations of $x_i$ and $x_j$. We exclude relation tokens. We'll state this in the final version.

---

#3: Isomap vs. PCA

Isomap preserves geodesic distances along the manifold, which is essential for visualizing non-linear geometries (e.g., rings, torii) [1,2]. Meanwhile, PCA, being linear, can collapse or stretch such structures, obscuring effects like representation shattering. This makes us prefer Isomap for our analysis, since faithful projections of the original topology turn out to be critical for interpreting how KE affects global geometry. We use PCA as well when linearity suffices (see App. F.2).

[1] https://www.nature.com/articles/s41593-019-0460-x

[2] https://www.nature.com/articles/s41583-022-00642-0

---

#4: Held-Out Tokens and Logical Inference

We note that "held-out" refers to specific facts, not entire relation tokens. Specifically, when generating the training dataset, we drop sequences that state one direction of a pair of conjugate facts with fixed probability $p$. That is, for any entity $x_i$ and relations $r, r'$, suppose the fact $(x_i, r, x_j)$ always implies $(x_j, r', x_i)$ (i.e. $r=$I_C1 and $r'=$I_A1). Our DGP may drop one of these facts from the training data (with probability $p$). Even if $(x_j, r', x_i)$ is absent, one can still learn the conjugacy of $r$ and $r'$ and the existence of $x_i$ and $x_j$ via other examples; however, failure to infer this relation indicates the model has rote memorized relations rather than understanding the global structure.

This discussion is detailed in App. B, and we'll clarify this further in the final version.

---

#5: Bulk Edits and Avoiding Shattering

In theory, we agree that a fully self-consistent batch of edits could reduce shattering. In practice, we observe:

• KE methods like ROME and MEMIT maximize the likelihood of edited facts without explicit constraints on representation geometry. Like task-specific fine-tuning, this can induce interference and catastrophic forgetting.

• Batch sizes are small because of compute limits. Edits are not independent or additive, so parallel updates can introduce inconsistencies, with degradation compounding over edits [1].

• Enumerating all semantically connected facts is intractable. Editing "Eiffel Tower is in Paris Rome" would require accounting for every related event or entity; a self-consistent closure is undefined.

Thus, while comprehensive edits may reduce shattering, they do not eliminate it.

[1] Yoon et al., 2024

---

#6: Prompt Format for Inference

• Compositional: The model is given a chain of relations (e.g., $x_i\, r_1\, r_2$) and must produce $x_k$, with $(x_i, r_1, x_j)$ and $(x_j, r_2, x_k)$ seen separately in training but not composed. Input: $\text{ctx}\, x_i\, \vec{r}$ where $\vec{r} = r_1r_2$.

• Logical: The model is asked about a fact $(x_j, r', x_i)$ where $(x_i, r, x_j)$ was seen but $r'$ was withheld. If the model understands relation symmetry, it can infer the inverse.

We'll further clarify these formats in the final version.

---

#7: Tree-Structured KG (App. G.3)

Thank you for the suggestion to add a forward reference—we will do so in Sec. 4.5.

---

#8: Code Release

We will release our code for the camera-ready version (see response to rVM5, #6).

---

#9: Typos

Thank you! We'll fix these typos in the final version.

---

---

Summary: We again thank the reviewer for their detailed review, which helped emphasize the design of our distortion metric $R(D_*)$, justify our visualization strategy, address inference task design, and explain the limits of bulk editing. We hope our rebuttals address the reviewer's questions and that they will champion our paper's acceptance!

We thank the reviewer for appreciating our experiments and analysis of representation shattering. Please see our responses below.

---

#1: Scope of Models and Methods Tested

We emphasize our evaluation already spans multiple architectures and scales: small Transformer models trained from scratch on synthetic knowledge graphs, pretrained decoder-only models without fine-tuning (GPT-2), pretrained decoder-only models with instruction fine-tuning (Llama 3 8B Instruct), and pretrained structured state-space models without fine-tuning (Mamba 2.8B). All consistently exhibit representation shattering that scales with counterfactual edit distance (see Sec. 4.5 and App. G.1).

Regarding the reviewer's suggested settings, we note the following.

• BERT: Modern KE methods exclusively focus on causal decoder-only models (ROME, MEMIT, PMET, AlphaEdit), due to which bidirectional encoder models like BERT are not amenable to analysis under the purview of our work.

• GPT-4: We note GPT4 is a closed-source model. This precludes parameter-space editing and replicating our methods.

• Mixtral: We note our Llama 3 8B experiments already feature similar scale, decoder-only construction, and instruction tuning as a Mixtral model. While MoE architecture effects are interesting, there is no precedent for such an experiment, so this lies outside our scope for now.

• RAG and WISE: RAG, RETRO, and KE protocols like WISE operate under different assumptions (e.g., inference-time components). Thus, they extend beyond our current scope of analyzing direct parameter-editing methods.

Nevertheless, we agree the directions above are promising. We'll include these in an expanded "Future Directions" section, citing WISE.

---

#2: Layerwise Analysis of Representation Shattering

Current KE methods use causal tracing to identify the layer to edit, often focusing on early MLP layers. Modifying this step would require a new KE protocol, which is beyond our scope. Nevertheless, our manuscript includes visualizations from different sub-layers and depths (see App. F.3 and G.2) that show the spatial extent of shattering. We are open to extending these analyses if helpful (e.g., as outlined in response to reviewer kjUD, #2).

---

#3: Visualization Methods

We refer the reviewer to our comment justifying the use of Isomap in response to reviewer gDBJ, #3. In brief, we use Isomap due to its ability to preserve geodesic distances, hence observing how cyclical manifolds distort under KE. UMAP/t-SNE show clustering, but Isomap's focus on global geometry better captures topological collapses from KE.

---

#4: Theoretical Explanation for Transformer Fragility under KE

Though mainly empirical, our study aligns with a broader theoretical framework.

• Transformers store factual associations in key-value pairs within MLP layers [1, 2] corroborated by our synthetic model analysis (see also response to rub2, #4).
• Parameter sharing and superposition cluster unrelated facts in overlapping subspaces [1], making them vulnerable to unintentional interference.
• Entities and relations often form structured manifolds (e.g., cycles, hierarchies), which aid compositional inference [4].
• KE methods (ROME, MEMIT, etc.) enact local weight updates that deform these manifolds, causing representation shattering for unedited facts residing in shared sub-regions.

Hence, we believe fragility arises from the entangled, compressed nature of factual storage, rather than from simply retaining knowledge. We will include this discussion in the final version of the paper.

[1] Geva et al., 2020

[2] Meng et al., 2022a

[3] https://transformer-circuits.pub/2023/toy-double-descent/index.html

[4] Engels et al., 2024

---

#5: Evaluation on Tree-Structured Knowledge

App. G.3 explores KE on a tree-structured city–country graph using GPT-2 XL. While trees are more challenging to visualize than cyclical manifolds, smaller-distance edits (e.g., relocating "Paris: to "Spain:) exhibit less shattering than those with larger distance (e.g., "UK:). This supports our hypothesis that shattering severity scales with the manifold distance of the edit.

---

#6: Supplemental Materials and Code Release

See Apps. A/B/C for documentation on pseudocode, hyperparameters, data, and architectures; we are happy to provide more!

We also highlight App. A.1, where we state our commitment to publicly releasing the source code for our experiments—covering both synthetic and naturalistic settings—in the camera-ready version of this paper. This ensures transparency and reproducibility.

---

---

Summary: We again thank the reviewer for their insightful comments that have helped us clarify our design choices, the constraints of current KE layer selection methods, our use of Isomap, and the theoretical underpinnings of Transformer fragility. We hope our rebuttals address the reviewer's questions and that they will champion our paper's acceptance!