Balancing Knowledge Updates: Toward Unified Modular Editing in LLMs

TL;DR

This work tackles the problem of knowledge editing in LLMs by showing that attention (Attn) modules are not only passive processors but also active repositories of factual knowledge, especially in early layers. It introduces IntAttn-Edit, a unified editing framework that extends the linear associative memory paradigm to both MLP and Attn modules, with a causal-contribution–guided balance factor $\alpha$ to allocate updates according to each module's measured knowledge contribution. Through extensive experiments on Mistral-7B and Qwen2.5-7B across ZsRE and WikiData Counterfact, IntAttn-Edit achieves higher edit success, better generalization, and stronger knowledge preservation than prior methods, particularly in batch-editing scenarios. The results demonstrate that jointly editing Attn and MLP, with adaptive balancing, yields robust and scalable knowledge updates for advanced LLMs, reducing residual outdated information and preserving fluency and behavior.

Abstract

Knowledge editing has emerged as an efficient approach for updating factual knowledge in large language models (LLMs). It typically locates knowledge storage modules and then modifies their parameters. However, most existing methods focus on the weights of multilayer perceptron (MLP) modules, which are often identified as the main repositories of factual information. Other components, such as attention (Attn) modules, are often ignored during editing. This imbalance can leave residual outdated knowledge and limit editing effectiveness. We perform comprehensive knowledge localization experiments on advanced LLMs and find that Attn modules play a substantial role in factual knowledge storage and retrieval, especially in earlier layers. Based on these insights, we propose IntAttn-Edit, a method that extends the associative memory paradigm to jointly update both MLP and Attn modules. Our approach uses a knowledge balancing strategy that allocates update magnitudes in proportion to each module's measured contribution to knowledge storage. Experiments on standard benchmarks show that IntAttn-Edit achieves higher edit success, better generalization, and stronger knowledge preservation than prior methods. Further analysis shows that the balancing strategy keeps editing performance within an optimal range across diverse settings.

Figures (8)

• Figure 3: Schematic illustration of the dual-pathway parameter update mechanism in IntAttn-Edit. Both MLP and Attn modules are edited in parallel, with update strengths modulated by the allocation factor $\alpha$ based on the measured causal contributions of each module.
• Figure 4: Effect of the knowledge balancing strategy on evaluation metrics using Qwen2.5-7B and Mistral-7B models on the ZsRE dataset. Optimal values are highlighted for clarity.
• Figure : (a) Probability of MLP over 1000 prompts
• Figure : (a) Logit Difference
• Figure : (a) Probability of MLP over 1000 prompts