Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards Interpreting Language Models: A Case Study in Multi-Hop Reasoning

Mansi Sakarvadia

TL;DR

Attention Lens is an open source tool that translates the outputs of attention heads into vocabulary tokens via learned transformations called lenses and demonstrates the use of lenses to reveal how a model arrives at its answer and uses them to localize sources of model failures such as in the case of biased and malicious language generation.

Abstract

Answering multi-hop reasoning questions requires retrieving and synthesizing information from diverse sources. Language models (LMs) struggle to perform such reasoning consistently. We propose an approach to pinpoint and rectify multi-hop reasoning failures through targeted memory injections on LM attention heads. First, we analyze the per-layer activations of GPT-2 models in response to single- and multi-hop prompts. We then propose a mechanism that allows users to inject relevant prompt-specific information, which we refer to as "memories," at critical LM locations during inference. By thus enabling the LM to incorporate additional relevant information during inference, we enhance the quality of multi-hop prompt completions. We empirically show that a simple, efficient, and targeted memory injection into a key attention layer often increases the probability of the desired next token in multi-hop tasks, by up to 424%. We observe that small subsets of attention heads can significantly impact the model prediction during multi-hop reasoning. To more faithfully interpret these heads, we develop Attention Lens: an open source tool that translates the outputs of attention heads into vocabulary tokens via learned transformations called lenses. We demonstrate the use of lenses to reveal how a model arrives at its answer and use them to localize sources of model failures such as in the case of biased and malicious language generation.

Towards Interpreting Language Models: A Case Study in Multi-Hop Reasoning

TL;DR

Attention Lens is an open source tool that translates the outputs of attention heads into vocabulary tokens via learned transformations called lenses and demonstrates the use of lenses to reveal how a model arrives at its answer and uses them to localize sources of model failures such as in the case of biased and malicious language generation.

Abstract

Answering multi-hop reasoning questions requires retrieving and synthesizing information from diverse sources. Language models (LMs) struggle to perform such reasoning consistently. We propose an approach to pinpoint and rectify multi-hop reasoning failures through targeted memory injections on LM attention heads. First, we analyze the per-layer activations of GPT-2 models in response to single- and multi-hop prompts. We then propose a mechanism that allows users to inject relevant prompt-specific information, which we refer to as "memories," at critical LM locations during inference. By thus enabling the LM to incorporate additional relevant information during inference, we enhance the quality of multi-hop prompt completions. We empirically show that a simple, efficient, and targeted memory injection into a key attention layer often increases the probability of the desired next token in multi-hop tasks, by up to 424%. We observe that small subsets of attention heads can significantly impact the model prediction during multi-hop reasoning. To more faithfully interpret these heads, we develop Attention Lens: an open source tool that translates the outputs of attention heads into vocabulary tokens via learned transformations called lenses. We demonstrate the use of lenses to reveal how a model arrives at its answer and use them to localize sources of model failures such as in the case of biased and malicious language generation.

Paper Structure

This paper contains 29 sections, 8 equations, 2 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Probing Models
  4. Activation Engineering
  5. Model Editing
  6. Circuit Discovery
  7. Knowledge Extraction
  8. Memorization
  9. Language Model Reasoning
  10. Retrieval Augmented Generation
  11. Memory Injections
  12. Introduction
  13. Background & Notation
  14. Multi-hop vs. single-hop prompts
  15. Transformer Architecture
  16. ...and 14 more sections

Figures (2)

  • Figure 1: A multi-hop prompt vs. two analogous single-hop prompts. The outputs are from GPT2-Small.
  • Figure 2: Diagram of language model reasoning. Highest ranked attention outputs of GPT2-Small at layer $\ell=9$, head $h=8$ when projected into vocabulary space (via the GPT2-Small embedding matrix) for a single-hop prompt (green) and its multi-hop counterpart (red).