Disentangling Knowledge Representations for Large Language Model Editing

1. Problem Statement

Large Language Models (LLMs) contain vast amounts of knowledge that often become outdated or inaccurate. Knowledge Editing aims to update specific facts within an LLM without retraining the entire model. While existing methods (e.g., ROME, MEMIT, AlphaEdit) are effective at injecting new knowledge and preserving coarse-grained irrelevant knowledge (facts with completely different subjects), they struggle with fine-grained irrelevant knowledge.

• The Core Challenge: Fine-grained irrelevant knowledge consists of facts that share the same subject as the edited fact but differ in relation and object (e.g., editing "The President of USA is Biden" to "Trump" should not alter "The Capital of USA is Washington").
• Root Cause: In current LLM architectures, the representation of a subject (e.g., "USA") inherently encodes multiple attributes (relations and objects). When editing methods modify the subject representation to inject new knowledge, the target knowledge and the fine-grained irrelevant knowledge become entangled in the parameter space. Consequently, updating the target often inadvertently corrupts the unrelated facts sharing the same subject.

2. Methodology: DiKE

The authors propose DiKE (Disentangling Knowledge representations for LLM Editing), a novel locate-then-edit framework. The approach operates in two distinct phases:

A. Knowledge Representation Disentanglement (KRD)

Before editing, a lightweight module is trained to decompose the subject representation into two distinct components:

1. Target-Knowledge-Related ($z^r_e$): Encodes information specific to the fact being edited (subject + relation + object).
2. Target-Knowledge-Unrelated ($z^u_e$): Encodes all other attributes of the subject (fine-grained irrelevant knowledge).

Mechanism:

• Architecture: Consists of a Disentangler (two projection networks) and a Recomposer (linear combination).
• Training Objectives:
  • Contrastive Loss: Maximizes mutual information between the disentangled components and the original subject representation while pushing the related and unrelated components apart in the vector space.
  • Constraint Loss: Ensures the unrelated component ($z^u_e$) can still correctly predict fine-grained irrelevant facts when fed back into the model.
  • Reconstruction Loss: Ensures the recomposed representation matches the original subject representation.

B. Disentanglement-based Knowledge Edit (DKE)

Once the representations are disentangled, the editing process is performed with strict constraints:

1. Target Injection: The new knowledge is injected only into the target-related component ($z^r_e$).
2. Preservation Constraint: The target-unrelated component ($z^u_e$) is explicitly constrained to remain invariant (unchanged) during the update.
3. Closed-Form Solution: Leveraging matrix theory, the authors derive a rank-one parameter update formula. This update satisfies three simultaneous constraints:
   • Inject the new key-value pair (Target Editing).
   • Preserve existing coarse-grained knowledge (Standard MEMIT constraint).
   • Preserve the disentangled unrelated representation (Fine-grained constraint).

The final weight update $\hat{W}$ is derived as:
$$\hat{W} = W + (W_3^\top W_3 + E)^{-1} \Delta_{MEMIT}$$
Where $W_3$ is the projection matrix from the Disentangler, and $\Delta_{MEMIT}$ is the standard update derived from MEMIT. This ensures the edit is minimally invasive to the unrelated knowledge.

3. Key Contributions

1. Problem Identification: The paper formally categorizes irrelevant knowledge into coarse-grained and fine-grained types, identifying the latter as a critical failure point in current editing methods due to representation entanglement.
2. Novel Framework (DiKE): Introduces a two-stage approach (KRD + DKE) that explicitly disentangles subject representations before editing, allowing for precise updates without collateral damage to related facts.
3. Theoretical Derivation: Derives a closed-form, rank-one parameter update that mathematically guarantees the preservation of the disentangled unrelated components while injecting new knowledge.
4. New Benchmark (FINE-KED): Constructs a new dataset specifically designed to evaluate fine-grained irrelevant knowledge preservation. It categorizes test instances into Easy, Middle, and Hard levels based on the semantic similarity between the edited relation and the preserved relation.

4. Experimental Results

The authors evaluated DiKE on GPT2-XL (1.5B), GPT-J (6B), and LLaMA-3 (8B) against baselines like FT-L, MEND, ROME, MEMIT, and AlphaEdit.

• FINE-KED Performance:
  • DiKE significantly outperforms all baselines in Relational Locality (preserving fine-grained knowledge), achieving up to 8.3% improvement over the best baseline on average.
  • It maintains high Efficacy (success rate of the edit), demonstrating that preservation does not come at the cost of editing capability.
  • In "Hard" scenarios (high semantic similarity between edited and preserved relations), DiKE shows the most significant advantage, whereas methods like ROME and MEMIT suffer severe degradation.
• General Benchmarks (COUNTERFACT & MQUAKE):
  • DiKE achieves competitive performance on standard editing metrics (Efficacy, Paraphrase, Neighborhood scores) and multi-hop reasoning tasks (MQUAKE), proving it does not sacrifice general editing ability.
• Batch Editing:
  • In subject-consistent batch editing (editing multiple facts about the same subject simultaneously), DiKE maintains high stability, while other methods (especially ROME) degrade sharply due to conflicting updates on entangled representations.
• Ablation Studies:
  • Removing the disentanglement module or the preservation constraints leads to significant drops in Relational Locality, confirming the necessity of both components.

5. Significance

• Precision in Editing: DiKE addresses a fundamental limitation in current LLM editing: the inability to update specific attributes of an entity without corrupting other attributes of the same entity. This is crucial for maintaining the integrity of complex knowledge graphs within models.
• Efficiency: By deriving a closed-form solution, DiKE avoids the need for iterative optimization or expensive retraining, making it scalable for large-scale knowledge updates.
• Trustworthiness: The method enhances the reliability of LLMs in dynamic environments (e.g., news updates, factual corrections) by ensuring that updates do not introduce subtle hallucinations or factual errors in related domains.
• Future Direction: The work establishes a new paradigm for knowledge editing based on representation disentanglement, suggesting that separating semantic attributes within the latent space is a more robust strategy than simply constraining parameter updates.

In summary, DiKE represents a significant advancement in knowledge editing by mathematically isolating target knowledge from fine-grained irrelevant knowledge, thereby solving the "ripple effect" problem that has plagued previous editing techniques.