Attention Flows for General Transformers

TL;DR

This paper introduces a method to quantify input token influence in Transformer models using flow networks and maxflow algorithms, providing insights into model decision processes.

Contribution

It formalizes a flow network approach to compute token influence in general Transformer architectures, including auto-regressive decoders, and offers a visualization library.

Findings

Flow network construction yields Shapley values for token influence

Method applies to various Transformer architectures and tasks

Provides a visualization tool for attention flow analysis

Abstract

In this paper, we study the computation of how much an input token in a Transformer model influences its prediction. We formalize a method to construct a flow network out of the attention values of encoder-only Transformer models and extend it to general Transformer architectures including an auto-regressive decoder. We show that running a maxflow algorithm on the flow network construction yields Shapley values, which determine the impact of a player in cooperative game theory. By interpreting the input tokens in the flow network as players, we can compute their influence on the total attention flow leading to the decoder's decision. Additionally, we provide a library that computes and visualizes the attention flow of arbitrary Transformer models. We show the usefulness of our implementation on various models trained on natural language processing and reasoning tasks.