Reproduce “Backward Lens: Projecting Language Model Gradients into the Vocabulary Space” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Shahar Katz, Yonatan Belinkov, Mor Geva, Lior Wolf

The 'Backward Lens' technique extends vocabulary projection to the backward pass of Transformer language models, revealing how knowledge is encoded during training. This method uncovers an 'imprint and shift' mechanism for MLP layers and allows for the development of an efficient knowledge editing approach called 'Forward Pass Shifting', which requires only a single forward pass.

Key insights

  • Gradient matrices in Transformer sequence models are inherently low-rank, with the final MLP layer exhibiting a consistent rank-1 structure.
  • Knowledge acquisition in MLP layers occurs via an 'imprint and shift' mechanism: FF1 updates 'imprint' input features, while FF2 updates 'shift' outputs toward target embeddings.
  • Gradient contributions are localized and sparse, with major updates concentrated at subject tokens in early layers and the last token in middle layers.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.