Jacobian Scopes: token-level causal attributions in LLMs

TL;DR

Jacobian Scopes introduces gradient-based, token-level causal attribution methods for interpreting large language models, revealing how individual input tokens influence predictions and uncovering biases and mechanisms across various NLP tasks.

Contribution

The paper presents Jacobian Scopes, a novel suite of gradient-based methods for token-level causal attribution in LLMs, grounded in perturbation theory and information geometry.

Findings

Reveals implicit political biases in LLM predictions.

Uncovers word- and phrase-level translation strategies.

Provides insights into mechanisms of in-context learning.

Abstract

Large language models (LLMs) make next-token predictions based on clues present in their context, such as semantic descriptions and in-context examples. Yet, elucidating which prior tokens most strongly influence a given prediction remains challenging due to the proliferation of layers and attention heads in modern architectures. We propose Jacobian Scopes, a suite of gradient-based, token-level causal attribution methods for interpreting LLM predictions. Grounded in perturbation theory and information geometry, Jacobian Scopes quantify how input tokens influence various aspects of a model's prediction, such as specific logits, the full predictive distribution, and model uncertainty (effective temperature). Through case studies spanning instruction understanding, translation, and in-context learning (ICL), we demonstrate how Jacobian Scopes reveal implicit political biases, uncover word- and phrase-level translation strategies, and shed light on recently debated mechanisms underlying in-context time-series forecasting. To facilitate exploration of Jacobian Scopes on custom text, we open-source our implementations and provide a cloud-hosted interactive demo at https://huggingface.co/spaces/Typony/JacobianScopes.