Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Empirical Study on Updating Key-Value Memories in Transformer Feed-forward Layers

Zihan Qiu, Zeyu Huang, Youcheng Huang, Jie Fu

TL;DR

This work treats transformer FFNs as key-value memories, formalized as $FFN(h) = f(h K^T) V$, and empirically compares updating keys versus updating values across knowledge editing and tuning tasks. Through back-propagation updates (and LoRA variants) in GPT-J (6B), GPT2-xl, and Llama2-7B under 4-bit quantization, the study shows updating keys generally yields better generalization, locality, and efficiency than updating values. The results suggest that altering the mechanism controlling knowledge (keys) is often more effective than directly editing stored content (values), with practical implications for efficient, robust model editing.

Abstract

The feed-forward networks (FFNs) in transformers are recognized as a group of key-value neural memories to restore abstract high-level knowledge. In this work, we conduct an empirical ablation study on updating keys (the 1st layer in the FFNs layer) or values (the 2nd layer in the FFNs layer). We compare those two methods in various knowledge editing and fine-tuning tasks of large language models to draw insights to understand FFNs further. Code is available at $\href{https://github.com/qiuzh20/Tuning-keys-v.s.-values}{this\,repo}$.

Empirical Study on Updating Key-Value Memories in Transformer Feed-forward Layers

TL;DR

This work treats transformer FFNs as key-value memories, formalized as , and empirically compares updating keys versus updating values across knowledge editing and tuning tasks. Through back-propagation updates (and LoRA variants) in GPT-J (6B), GPT2-xl, and Llama2-7B under 4-bit quantization, the study shows updating keys generally yields better generalization, locality, and efficiency than updating values. The results suggest that altering the mechanism controlling knowledge (keys) is often more effective than directly editing stored content (values), with practical implications for efficient, robust model editing.

Abstract

The feed-forward networks (FFNs) in transformers are recognized as a group of key-value neural memories to restore abstract high-level knowledge. In this work, we conduct an empirical ablation study on updating keys (the 1st layer in the FFNs layer) or values (the 2nd layer in the FFNs layer). We compare those two methods in various knowledge editing and fine-tuning tasks of large language models to draw insights to understand FFNs further. Code is available at .
Paper Structure (8 sections, 1 figure, 4 tables)

This paper contains 8 sections, 1 figure, 4 tables.

Table of Contents

  1. Introduction
  2. Experiments
  3. Conclusion
  4. URM statement
  5. Appendix
  6. More Details for Knowledge Editing and LoRA Tuning
  7. More experimental results about the Knowledge Editing
  8. More experimental results about the LoRA Tuning

Figures (1)

  • Figure 1: Left: FFNs operation is conceptualized as a key-value memory. The input hidden states interact with keys (rows of $K$) through an inner product to obtain activation values. These activations then serve as weights for summing values (rows of $V$). Right: To update these key-value memories, one can either directly modify the values that store relevant information or alter the keys to adjust the weights of existing values.