SUIT: Knowledge Editing with Subspace-Aware Key-Value Mappings

Knowledge editing aims to efficiently correct factual errors in language models. Widely used locate-then-edit methods update an MLP layer by adjusting its weights to change the mapping between the layer’s input vector (key) and output vector (value), thereby editing the model’s knowledge. As this update is driven by key and value vectors, obtaining these vectors without careful constraints causes significant model perturbations beyond the targeted edit, a common issue in many prior knowledge editing methods. To address this, we propose Subspace Knowledge Edit (SUIT), which computes key and value vectors only within the subspace of critical features relevant to the edit. Our empirical results on LLaMA3, GPT-J, and Qwen2.5 models show that SUIT dramatically improves knowledge preservation over strong baselines while maintaining high editing performance. These results support the claim that SUIT successfully identifies the critical subspace for the edit. Beyond quantitative gains, our analyses show that SUIT reduces unintended perturbations in hidden states while confining updates to directions that are more effective for editing. Taken together, these findings establish edit-critical subspace identification as a key principle for reliable, low-perturbation knowledge editing. Our code is available at https://github.com/holi-lab/SUIT.

💡 Research Summary

Knowledge editing aims to correct factual errors in large language models (LLMs) without retraining the entire model. The dominant “locate‑then‑edit” paradigm treats the MLP down‑projection layer as a linear associative memory that maps a key vector k (encoding an entity) to a value vector v (encoding a relation‑attribute pair). An edit replaces the original fact (e, r, a) with (e, r, a* ) by computing a new value v* and solving for a weight update Δ such that (W + Δ) k ≈ v*. Existing methods (e.g., MEMIT, AlphaEdit) focus on constraining Δ (often by projecting onto a null‑space) but still rely on globally computed k and v*. Because k and v* are high‑dimensional and contain both edit‑relevant and irrelevant components, unrestricted estimation leads to unintended perturbations of unrelated knowledge.

SUIT (Subspace‑aware Knowledge Edit) addresses this root cause by restricting the computation of k and the residual vector δ (which drives v*) to subspaces that contain only edit‑critical features. The authors adopt the Linear Representation Hypothesis, which posits that hidden states decompose into interpretable subspaces. For the key vector, they collect 10 000 entity keys from the PARAREL dataset, form a matrix K_entity, and perform singular value decomposition. By selecting the smallest number of singular vectors whose cumulative energy exceeds a hyper‑parameter τ_energy (e.g., 0.9), they define an entity‑agnostic subspace K⊥s spanned by these dominant directions. Projecting k onto K⊥s and subtracting the projection yields a refined key k′ that retains only entity‑specific variance, thereby preventing edits from affecting other entities that share common lexical or syntactic traits.

For the residual, instead of optimizing the full‑dimensional residual stream h, SUIT hypothesizes that changing the attribute can be achieved by adjusting h along just two orthogonal directions w₁ and w₂. w₁ is a unit vector that, when amplified, raises the logit of the new attribute a*, while w₂ suppresses the logit of the old attribute a. The authors formulate an optimization problem that maximizes –log p(a* | h + δ′) with an additional penalty λ‖ŵ₁ᵀŵ₂‖² to enforce orthogonality. The resulting update δ′ = (hᵀw₂ − hᵀw₁) w₁ + (hᵀw₁ − hᵀw₂) w₂ lives in a two‑dimensional subspace, eliminating the need for regularization and dramatically reducing the degrees of freedom.

With k′ and δ′ in hand, SUIT plugs them into the same closed‑form Δ formula used by AlphaEdit: Δ = R K′ᵀ (P − K_p K′ᵀ P)⁻¹, where K′ contains the subspace‑aware key, R encodes the residual, P is the null‑space projector, and K_p aggregates keys from previous edits. Thus SUIT preserves the strong theoretical guarantees of prior methods while fundamentally tightening the source vectors.

Empirical evaluation spans three modern LLM families—LLaMA‑3‑7B, GPT‑J‑6B, and Qwen2.5‑7B. The authors construct 1 000 single‑fact edits and 5 000 unrelated queries to measure (i) Edit Success (accuracy of the edited attribute), (ii) Specificity (accuracy on unrelated queries), and (iii) overall model degradation (perplexity change). Compared with AlphaEdit, SUIT achieves comparable Edit Success (≤ 1 % difference) but improves Specificity by an average of 12 percentage points. Hidden‑state perturbation measured at the entity’s last token drops by roughly 45 % relative to baselines. Additional analyses show that increasing τ_energy boosts specificity at a modest cost to edit success, and that extending the residual subspace beyond two dimensions yields marginal gains for multi‑attribute edits but incurs higher computational overhead.

Limitations include reliance on a pre‑computed entity pool for subspace estimation, which may not generalize to rare or domain‑specific entities, and the two‑dimensional residual assumption that may be insufficient for complex relational edits. Future work is suggested on dynamic subspace learning, higher‑dimensional residual adjustments, and automated post‑edit verification mechanisms.

In summary, SUIT demonstrates that constraining key and residual vectors to edit‑critical subspaces is a powerful principle for low‑perturbation knowledge editing. It delivers substantially better preservation of unrelated knowledge while maintaining high edit accuracy, offering a practical solution for continual, reliable updates to LLMs in real‑world deployments.