Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Massive Editing for Large Language Models via Meta Learning

Chenmien Tan, Ge Zhang, Jie Fu

TL;DR

This work tackles the challenge of updating knowledge in large language models without catastrophic forgetting or excessive retraining. It introduces MALMEN, a hyper-network-based editor that aggregates parameter shifts through a least-squares formulation solved via the normal equation, enabling mass edits with constrained memory. The method decouples hyper-network and LM computations to support arbitrary batch sizes and demonstrates superior scalability across BERT-base, GPT-2, T5-XL, and GPT-J on FEVER and zsRE, outperforming strong baselines like MEND and MEMIT. Ablation studies validate design choices, showing robust editing performance and improved memory efficiency as the number of edits grows. The approach offers a practical path toward maintaining up-to-date, factually accurate LMs in real-world deployments, with explicit limitations and avenues for future improvements.

Abstract

While large language models (LLMs) have enabled learning knowledge from the pre-training corpora, the acquired knowledge may be fundamentally incorrect or outdated over time, which necessitates rectifying the knowledge of the language model (LM) after the training. A promising approach involves employing a hyper-network to generate parameter shift, whereas existing hyper-networks suffer from inferior scalability in synchronous editing operation amount. To mitigate the problem, we propose the MAssive Language Model Editing Network (MALMEN), which formulates the parameter shift aggregation as the least square problem, subsequently updating the LM parameters using the normal equation. To accommodate editing multiple facts simultaneously with limited memory budgets, we separate the computation on the hyper-network and LM, enabling arbitrary batch size on both neural networks. Our method is evaluated by editing up to thousands of facts on LMs with different architectures, i.e., BERT-base, GPT-2, T5-XL (2.8B), and GPT-J (6B), across various knowledge-intensive NLP tasks, i.e., closed book fact-checking and question answering. Remarkably, MALMEN is capable of editing hundreds of times more facts than strong baselines with the identical hyper-network architecture and outperforms editor specifically designed for GPT. Our code is available at https://github.com/ChenmienTan/malmen.

Massive Editing for Large Language Models via Meta Learning

TL;DR

This work tackles the challenge of updating knowledge in large language models without catastrophic forgetting or excessive retraining. It introduces MALMEN, a hyper-network-based editor that aggregates parameter shifts through a least-squares formulation solved via the normal equation, enabling mass edits with constrained memory. The method decouples hyper-network and LM computations to support arbitrary batch sizes and demonstrates superior scalability across BERT-base, GPT-2, T5-XL, and GPT-J on FEVER and zsRE, outperforming strong baselines like MEND and MEMIT. Ablation studies validate design choices, showing robust editing performance and improved memory efficiency as the number of edits grows. The approach offers a practical path toward maintaining up-to-date, factually accurate LMs in real-world deployments, with explicit limitations and avenues for future improvements.

Abstract

While large language models (LLMs) have enabled learning knowledge from the pre-training corpora, the acquired knowledge may be fundamentally incorrect or outdated over time, which necessitates rectifying the knowledge of the language model (LM) after the training. A promising approach involves employing a hyper-network to generate parameter shift, whereas existing hyper-networks suffer from inferior scalability in synchronous editing operation amount. To mitigate the problem, we propose the MAssive Language Model Editing Network (MALMEN), which formulates the parameter shift aggregation as the least square problem, subsequently updating the LM parameters using the normal equation. To accommodate editing multiple facts simultaneously with limited memory budgets, we separate the computation on the hyper-network and LM, enabling arbitrary batch size on both neural networks. Our method is evaluated by editing up to thousands of facts on LMs with different architectures, i.e., BERT-base, GPT-2, T5-XL (2.8B), and GPT-J (6B), across various knowledge-intensive NLP tasks, i.e., closed book fact-checking and question answering. Remarkably, MALMEN is capable of editing hundreds of times more facts than strong baselines with the identical hyper-network architecture and outperforms editor specifically designed for GPT. Our code is available at https://github.com/ChenmienTan/malmen.
Paper Structure (24 sections, 2 theorems, 3 equations, 11 figures, 5 tables, 2 algorithms)

This paper contains 24 sections, 2 theorems, 3 equations, 11 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Editing with External Memory
  4. Locate-and-edit
  5. Meta-learning
  6. Problem Formulation
  7. Method
  8. Aggregating Parameter Shifts using the Normal Equation
  9. Memory Economic Training
  10. Empirical Evaluation
  11. Scaling Curve
  12. Ablation Study
  13. Discussion
  14. Limitations
  15. Extended Discussion
  16. ...and 9 more sections

Key Result

Theorem 1

Suppose that $U\in\mathbb R^{d\times n}$, $f:\mathbb R^{d'\times n}\to\mathbb R^{d'\times d},D\mapsto DU^T(UU^T+\lambda I)^{-1}$, and $g:\mathbb R^{d'\times d}\to\mathbb R$ is a differentiable function. Then, $\nabla (g\circ f)=\nabla g\cdot(UU^T+\lambda I)^{-1}U$

Figures (11)

  • Figure 1: The procedure to compute $\nabla_\theta L_{\mathrm{meta}}$ in mitchell2022mend. The cached tuples are fed into the hyper-network to generate the updated LM parameter, which is differentiable w.r.t. the hyper-network parameter. After feeding the equivalent and unrelated prompt-answer tuples into the LM, the meta loss is back propagated along the red arrow.
  • Figure 2: The overall procedure to compute the meta gradient $\nabla_\theta L_{\mathrm{meta}}$: (i) The edit prompt-answer tuples $(x_i,y_i)_{i=1}^m$ are fed into the LM $p_\mathcal{W}$, where keys $(u_{\ell,j})_{\ell\in\mathcal{L},j\in[n]}$ are cached. (ii) Back-propagate the standard fine-tuning loss $L$ and cache the gradient with respect to values, i.e., $(\nabla_{v_{\ell,j}} L)_{\ell\in\mathcal{L},j\in[n]}$. (iii) Feed cached tuples $(u_{\ell,j},\nabla_{v_{\ell,j}} L)_{\ell\in\mathcal{L},j\in[n]}$ into the hyper-network $g_\theta$ to infer the updated LM parameter $\tilde{\mathcal{W}}$ without caching any intermediate variable. (iv) Feed the equivalent and unrelated tuples $(x_i^e,y_i^e)_{i=1}^m,(x_i^u,y_i^u)_{i=1}^m$ into the post-edit LM $p_{\tilde{\mathcal{W}}}$. (v) Back-propagate the meta loss $L_{\mathrm{meta}}$ on linear layers $\ell\in\mathcal{L}$ and compute $(\nabla_{D_\ell}L_{\mathrm{meta}})_{\ell\in\mathcal{L}}$ and $(dL_{\mathrm{meta}}/d\lambda_\ell)_{\ell\in\mathcal{L}}$. (vi) Again, feed the cache tuples $(u_{\ell,j},\nabla_{v_{\ell,j}}L)_{\ell\in\mathcal{L},j\in[n]}$ into the hyper-network $g_\theta$ to generate the value difference matrices $(D_\ell)_{\ell\in\mathcal{L}}$ but with the training mode. (vii) Back propagate the proxy loss $\tilde{L}=\sum_{\ell\in\mathcal{L},j\in[n]}\nabla_{d_{\ell,j}}L_{\mathrm{meta}}^T d_{\ell,j}$. All inputs to neural networks, including $(x_i,y_i,x_i^e,y_i^e,x_i^u,y_i^u)_{i=1}^m$ and $(u_{\ell,j},\nabla_{v_{\ell,j}}L)_{\ell\in\mathcal{L},j\in[n]}$, can be split into batches, where the gradients of the meta loss $L_{\mathrm{meta}}$ and proxy loss $\tilde{L}$ are accumulated.
  • Figure 3: Hooked module in GPT-J. The module of interest is highlighted with a yellow pentagram.
  • Figure 4: Scaling curve of BERT-base (110M)
  • Figure 5: Scaling curve of GPT-2 (124M)
  • ...and 6 more figures

Theorems & Definitions (4)

  • Theorem 1
  • proof
  • Theorem 2
  • proof