Understanding the Collapse of LLMs in Model Editing

TL;DR

This paper investigates why large language models often collapse after editing, identifies key causes related to key handling and token distribution, and proposes a simple method to prevent collapse while preserving edit effectiveness.

Contribution

It uncovers root causes of model collapse in LLM editing and introduces a straightforward approach to mitigate this issue effectively.

Findings

The identified causes include key handling inconsistencies and token distribution differences.

The proposed method prevents model collapse during editing.

The approach maintains the effectiveness of model edits.

Abstract

Despite significant progress in model editing methods, their application in real-world scenarios remains challenging as they often cause large language models (LLMs) to collapse. Among them, ROME is particularly concerning, as it could disrupt LLMs with only a single edit. In this paper, we study the root causes of such collapse. Through extensive analysis, we identify two primary factors that contribute to the collapse: i) inconsistent handling of prefixed and unprefixed keys in the parameter update equation may result in very small denominators, causing excessively large parameter updates; ii) the subject of collapse cases is usually the first token, whose unprefixed key distribution significantly differs from the prefixed key distribution in autoregressive transformers, causing the aforementioned issue to materialize. To validate our findings, we propose a simple yet effective approach: uniformly using prefixed keys during editing phase and adding prefixes during testing phase to ensure the consistency between training and testing. The experimental results show that the proposed solution can prevent model collapse while maintaining the effectiveness of the edits.