PaperlayerView exact revision on arXiv

Source-linked AI summary

$\mu$KE: Matryoshka Unstructured Knowledge Editing of Large Language Models

Zian Su, Ziyang Huang, Kaiyuan Zhang, Xiangyu Zhang

arXiv:2504.01196v2cs.CLcs.AI

TL;DR

AnyEdit significantly improves on existing locate-and-edit methods such as MEMIT and AlphaEdit. µKE uses a Matryoshka-style objective that trains each memory update against subsequent target ranges, with adaptive coefficients informed by target-figure gradient affinity. Up to 12.33% improvement in edit efficacy over AnyEdit; evaluations on two models and five benchmarks also report robustness to diverse edit formats.

  • Problem

    AnyEdit significantly improves on existing locate-and-edit methods such as MEMIT and AlphaEdit.

    Source: §1 · p. 1

  • Method

    µKE uses a Matryoshka-style objective that trains each memory update against subsequent target ranges, with adaptive coefficients informed by target-figure gradient affinity.

    Source: §3.2 · p. 5, §3.3 · p. 6

  • Results

    Up to 12.33% improvement in edit efficacy over AnyEdit; evaluations on two models and five benchmarks also report robustness to diverse edit formats.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The evaluations support µKE as a robust approach to updating large language models through unstructured knowledge editing.

    Source: §8 · p. 10

    Existing unstructured-editing benchmarks lack fine-grained annotations of knowledge units in longer targets, limiting insight into edit efficacy and nuanced evaluation dimensions.

    Source: §8 · p. 10

Abstract

from arXiv · show

Large language models (LLMs) have emerged as powerful knowledge bases yet are limited by static training data, leading to issues such as hallucinations and safety risks. Editing a model's internal knowledge through the locate-and-edit paradigm has proven a cost-effective alternative to retraining, though current unstructured approaches, especially window-based autoregressive methods, often disrupt the causal dependency between early memory updates and later output tokens. In this work, we first theoretically analyze these limitations and then introduce Matryoshka Unstructured Knowledge Editing ($\mu$KE), a novel memory update mechanism that preserves such dependencies via a Matryoshka-style objective and adaptive loss coefficients. Empirical evaluations on two models across four benchmarks demonstrate that $\mu$KE improves edit efficacy by up to 12.33% over state-of-the-art methods, and remains robust when applied to diverse formatted edits, underscoring its potential for effective unstructured knowledge editing in LLMs.

1 Introduction

Locate-and-edit offers a lightweight alternative to retraining, but early methods focused on structured facts and window-based unstructured editing leaves a memory-dependency question. µKE analyzes that issue and introduces a Matryoshka-style update mechanism with adaptive loss coefficients.

1 Introduction

  • Locate-and-edit updates internal model components to incorporate new facts without extensive retraining, while aiming to limit collateral changes to unrelated knowledge.Earlier methods primarily targeted structured (subject, relation, object) triplets.Source: §1 · p. 1
  • AnyEdit chunks long targets into windows and edits them autoregressively, extending locate-and-edit methods to unstructured tasks.The passage contrasts this extension with earlier methods focused on structured facts.Source: §1 · p. 1
  • Window-based autoregressive editing updates successive target windows but overlooks dependencies between earlier memories and later output tokens.The paper frames encoding those dependencies across sequential memory updates as the central unstructured-editing question.Source: §1 · p. 2, Figure 1 · p. 2
  • µKE analyzes missing working-memory dependencies and introduces a Matryoshka-style update mechanism with adaptive loss coefficients for unstructured editing.The method treats early memory updates as condensed shifts that partially cover later edit targets and uses adaptive coefficients to optimize them.Source: §1 · p. 2

2 Preliminaries

The preliminaries describe locate-and-edit as locating memory-bearing layers, optimizing shifts to hidden states, and mapping those updates to model weights. They distinguish a single shift for an entire target from sequential shifts for autoregressive target windows.

2 Preliminaries

  • Locate-and-edit locates likely knowledge-bearing layers, optimizes a hidden-state shift against the edit target, then maps memory updates to constrained model-weight shifts.The hidden-state update is introduced by adding a bias term to a selected state; weight updates may include locality constraints.Source: §2 · p. 3
  • One-for-All places a single memory shift before the first target token and optimizes it for the full target.This assumes one working-memory update can edit all later generations.Source: §2 · p. 3
  • Window-by-Window splits a long target into windows and updates a separate working memory for each window autoregressively, conditioned on preceding contexts.Each shift is optimized to minimize the target window's negative log likelihood given previous contexts.Source: §2 · pp. 3–4

3 Matryoshka Unstructured Knowledge Editing

Section 3 diagnoses missing cross-window memory dependency in window-by-window updates and introduces µKE’s Matryoshka objective with affinity-informed adaptive coefficients.

3.1 Window-by-Window Update Overlooks Memory Dependency

  • Full-target optimization adds gradients from later windows to earlier memory updates, exposing the dependency that the window-by-window objective ignores.Parallel optimization performs less well than sequential updating, potentially because of a larger parameter space and design details such as normalization and batch updates.Source: §3.1 · p. 4

3.2 A Matryoshka-Style Working Memory Update Objective

  • µKE averages weighted NLL losses over nested target spans beginning at each window, allowing later-token gradients to update earlier working memories.Early windows receive larger coefficients because their memories get no second chance for editing; later windows can be addressed in subsequent updates.Source: §3.2 · p. 5, Figure 2 · p. 5

3.3 Adaptive Coefficients Informed by Affinity Between Edit Target Figures

  • With a static objective, edit efficacy varies across target lengths; µKE uses adaptive coefficients negatively correlated with each target figure’s affinity to the first figure.This compensates for less-affine figures that may otherwise be neglected; the coefficients are adjusted per sample.Source: §3.2 · p. 5, §3.3 · p. 6
  • Affinity computation is efficient because locate-and-edit updates use a much smaller gradient space than LLM pretraining or fine-tuning.Affinity is cosine similarity between target-figure gradient directions.Source: §3.3 · p. 6

4 Experiment Setup

Experiments evaluate µKE on two instruction-tuned LLMs against editing baselines, using edit-quality, hallucination, and locality benchmarks.

Models and Baselines

  • Experiments use Qwen2.5-7B-Instruct and Llama3-8B-Instruct, comparing against original models, MEMIT, AlphaEdit, UnKE, AnyEdit, and AnyEdit∗.Source: §4 · p. 6

Benchmarks

  • Evaluation covers UnKEBench, AKEW’s CounterFact and MQuAKE subsets, EditEverything, SelfCheckGPT, MMLU, and IFEval.SelfCheckGPT assesses hallucination reduction, while MMLU and IFEval assess preservation of general capabilities.Source: §4 · p. 6

Metrics

  • Editing efficacy is measured with BLEU, ROUGE-L, and BERTScore; each experiment runs three trials and reports mean values.Standard deviations are reported in Appendix D.4.Source: §4 · p. 6

5 Results

On long-form edits, µKE outperforms AnyEdit on BLEU and is competitive on ROUGE-L; it also outperforms AnyEdit across diverse domains. The section evaluates whether editing preserves general capabilities using MMLU and IFEval.

5.1 Long-Form Knowledge Editing Evaluation

  • µKE improves over AnyEdit in BLEU by up to 12.33% on original and 10.5% on paraphrased questions, with higher or competitive ROUGE-L in most cases.µKE* reaches scores as high as 99.996% across BLEU, ROUGE-L, and BERTScore.Source: Table 1 · p. 7, §5.1 · p. 7
  • On UnKEBench with Qwen2.5-7B-Instruct, µKE exceeds UnKE on paraphrased questions by +18.65% BLEU, +8.50% BERTScore, and +23.37% Rouge-L.Across most evaluation settings, µKE or µKE* achieves the strongest generalization performance.Source: §5.1 · p. 7

5.2 Diverse-Formatted Knowledge Editing Evaluation

  • µKE outperforms AnyEdit across EditEverything subdomains and shows robust, consistent performance across domains, while AnyEdit drops on Poetry for both evaluated LLMs.The domains include mathematics, poetry, news, programming, and chemistry.Source: §5.2 · pp. 7–8, Figure 4 · p. 8

5.3 General Capability Preservation Evaluation

  • The general-capability evaluation assesses whether edited LLMs preserve natural language understanding and instruction following using MMLU and IFEval.The supplied passage describes the evaluation focus but does not report its results.Source: §5.3 · p. 8

6 Analysis

The analysis finds adaptive Lµ more robust to target-length variation and more stable during optimization than static objectives, while reported performance changes little across tested window sizes.

General capability evaluation

  • The analysis includes pre- and post-edit evaluations on MMLU and IFEval, but the supplied passage does not report the outcomes.Source: Table 2 · p. 8

Length and objective analysis

  • Adaptive Lµ shows the best robustness to target length; BLEU varies greatly for static objectives such as AnyEdit’s window objective and Matryoshka-Vanilla.The comparison used 100 samples with a window size of 20, learning rate 0.5, and 25 optimization steps.Source: §6 · p. 9

Optimization stability

  • Parallel memory-update optimization is harder than sequential; both Matryoshka methods outperform sequential-NLL, and dynamic-coefficient Lµ is more stable than static coefficients.Dynamic Lµ increases smoothly, especially for paraphrased questions.Source: §6 · p. 9, Figure 6 · p. 9

Window-size analysis

  • Across window sizes from 10 to 100 tokens, µKE’s metric variations are typically within 3%, and reported results do not depend on tuning an optimal size.The authors identify input-dependent window selection as a future direction.Source: §6 · pp. 9–10, Figure 7 · p. 9

7 Related Work

Related work spans structured factual editing, newer unstructured methods for free-form knowledge, and several alternative or complementary model-editing strategies.

Unstructured Editing of LLMs

Model Editing Methods

  • Beyond locate-then-edit, related approaches include retrieval, meta-learning, adapters, minimal weight updates, multi-hop editing, and task arithmetics.The passage also notes refinements to locate-then-edit, including RECT (Gu et al., 2024) and AlphaEdit (Fang et al., 2024).Source: §7 · p. 10

8 Conclusion

The conclusion presents μKE as a response to disrupted causal dependencies in window-based editing, reporting stronger editing efficacy and identifying optimization, benchmark, and knowledge-unit challenges for future work.

8 Conclusion

  • Evaluations on two models across five benchmarks show μKE improves editing efficacy over state-of-the-art techniques and provides a robust approach to updating LLMs.The paper's conclusion describes the evaluations as comprehensive.Source: §8 · p. 10
  • The locate-and-edit paradigm remains tightly coupled to its optimization workflow, although μKE's sample-wise adaptive loss alleviates this issue.The authors identify more effective and generalizable optimization strategies as future work.Source: §8 · p. 10
  • Existing unstructured-editing benchmarks lack fine-grained annotations of knowledge units in long targets, limiting insight into edit efficacy and nuanced behaviors.The authors call for broader benchmarks and metrics, including evaluation of interdependent edits that preserve causal memory dependencies.Source: §8 · p. 10
  • Knowledge units in long targets need not match fixed-window lengths, motivating future work on dynamic localization and finer-grained models of memory dependency and causality.The authors say this work would likely require a more intrinsic and fine-grained model of memory dependency and causality in autoregressive language models.Source: §8 · p. 10

A.1 Evaluation Details

The evaluation details specify general-ability tests, their prompting and scoring settings, and the experimental resources and baseline-reproduction procedure.

A.1 Evaluation Details

  • MMLU tests general understanding across subjects, while IFEval measures instruction following; both are used to assess preservation of general ability after editing.MMLU uses 5-shot multi-turn prompts with batch size 8; IFEval uses 0-shot prompts with batch size 16 and strict or loose prompt-level accuracy.Source: §A.1 · p. 14
  • Experiments used four NVIDIA A100 SXM4 GPUs, bfloat16 models, chat templates, and ten repetitions, with means and standard errors reweighted by inverse-variance weighting.The compute environment was an AWS EC2 p4d.24xlarge instance; the two model families used different system prompts.Source: §A.1 · p. 14
  • Baseline replication followed AnyEdit settings, including method-specific edited layers, chunk sizes, optimization steps, and three-run mean reporting.The MEMIT-based baseline used up to 25 steps per chunk; the UnKE-based method used 50 optimization steps, and key-value updates were limited to 25 steps.Source: §A.1 · p. 14

A.2 Implementation Details

This material covers implementation choices and locate-and-edit mechanics, followed by supplementary comparisons on hallucination reduction, lifelong baselines, batch editing, and reliability.

A.2 Implementation Details

  • For λ_i,j, the implementation sets T_aff to 3 and the learning rate to 0.5; μKE is otherwise optimized for 25 steps for baseline fairness.The reported window sizes are 30 tokens for Llama3-8B-Instruct and 20 for Qwen2.5-7B-Instruct, without overlap.Source: §A.2 · p. 14

B Broader Definition of “Working Memory”

  • The paper uses “working memory” for a short-lived, capacity-limited store that is actively manipulated to affect later outcomes, linking the concept to human cognition.In μKE, a specific transformer layer's activations serve to keep relevant information online for selective updating and subsequent processing.Source: B Broader Definition of “Working Memory” · p. 14

C.3 Updating Working Memories

  • Locate-and-edit updates a selected hidden state by gradient descent; normalization projects each update into the hooked hidden state's L2 ball to preserve locality.Removing normalization empirically improves edit success but harms locality, and the authors hypothesize that optimization difficulty helps explain the need for sequential updates.Source: §C.3 · p. 16

C.4 Batch Update

  • The batch-update stage maps optimized working-memory shifts into model weight shifts; for W_down methods it uses constrained least squares, while UnKE uses gradient descent.The appendix also describes combining updated memories for unstructured edits as a memory-modifying set.Source: §C.4 · pp. 16–17

D.1 Unstructured Editing for Hallucination Reduction

  • On SelfCheckGPT with Llama3-8B-Instruct, μKE performs better than AnyEdit and reduces hallucinations relative to the pre-edit model.Table 4 reports these results.Source: §D.1 · p. 17, Table 4 · p. 17

D.2 Comparison with Lifelong Editing Baselines

  • Under default settings, WISE shows incomparable and unstable performance relative to AnyEdit and μKE across four benchmarks and two models.Tables 5–8 report WISE results for Llama3-8B-Instruct and Qwen2.5-7B-Instruct.Source: §D.2 · p. 17, Table 5 · p. 18, Table 6 · p. 18

D.3 Batch Editing

  • μKE consistently outperforms AnyEdit across batch sizes in Figure 8, supporting the importance of preserving memory dependency in unstructured batch editing.The figure reports batch-editing performance on three benchmarks.Source: Figure 8 · p. 18, §D.3 · p. 18

D.4 Statistical Reliability

  • Table 9 reports μKE and μKE* results on UnKEBench, CounterFact, and AKEW-MQuAKE, including means and standard deviations from three runs.Paraphrased metrics are not applicable for AKEW-MQuAKE.Source: §D.4 · p. 19, Table 9 · p. 19

D.5 Ablation Study on Adaptive Coefficients

The ablation reports a generality benefit from adaptive coefficients; the following material quantifies computation overhead and presents a poetry-edit output example.

D.5 Ablation Study on Adaptive Coefficients

  • On Qwen2.5-7B-Instruct with 100 samples, both Matryoshka-Vanilla and adaptive coefficients are effective for edit success, while adaptive coefficients improve paraphrased-question generality.The ablation compares all-one coefficients with adaptive coefficients computed using different optimization steps.Source: §D.5 · p. 19

D.6 Computational Overhead

  • μKE increases elapsed time by 4.2% with cache and 33.8% without cache, while peak VRAM usage increases negligibly.The reported elapsed time without cache includes I/O costs and intermediate inference, among other factors.Source: §D.6 · p. 19, Table 11 · p. 19

D.7.1 A failing example of AnyEdit’s performance on EditEverything-Poetry

  • The appendix shows μKE and AnyEdit outputs for a rain-poem edit using Qwen2.5-7B-Instruct, omitting special tokens.The edit target asks for a poem inspired by the sound of rain.Source: §D.7.1 · p. 20

Related papers

Loading 2504.01196v2…