On Identifiability in Transformers

TL;DR

This work analyzes identifiability in Transformer architectures, focusing on self-attention and token embeddings. It proves that attention weights are non-identifiable when sequence length exceeds the attention head dimension and introduces effective attention as a diagnostic tool. Through token identifiability and gradient-based Hidden Token Attribution, the paper shows that input tokens largely retain identity via angular information and that context is heavily mixed into hidden embeddings, though aggregation remains mostly local. The findings challenge naive interpretations of attention distributions and provide practical techniques to better understand and investigate Transformer models.

Abstract

In this paper we delve deep in the Transformer architecture by investigating two of its core components: self-attention and contextual embeddings. In particular, we study the identifiability of attention weights and token embeddings, and the aggregation of context into hidden tokens. We show that, for sequences longer than the attention head dimension, attention weights are not identifiable. We propose effective attention as a complementary tool for improving explanatory interpretations based on attention. Furthermore, we show that input tokens retain to a large degree their identity across the model. We also find evidence suggesting that identity information is mainly encoded in the angle of the embeddings and gradually decreases with depth. Finally, we demonstrate strong mixing of input information in the generation of contextual embeddings by means of a novel quantification method based on gradient attribution. Overall, we show that self-attention distributions are not directly interpretable and present tools to better understand and further investigate Transformer models.

Figures (45)

• Figure 1: (a) Each point represents the Pearson correlation coefficient of effective attention and raw attention as a function of token length. (b) Raw attention vs. (c) effective attention, where each point represents the average (effective) attention of a given head to a token type.
• Figure 2: (a) Identifiability of contextual word embeddings at different layers. Here, $\hat{g}$ is trained and tested on the same layer. (b) $g_{cos,l}^{lin}$ trained on layer $l$ and tested on all layers.
• Figure 3: (a) Contribution of the input token to the embedding at the same position. The orange line represents the median value and outliers are not shown. (b) Percentage of tokens $\tilde{P}$ that are not the main contributors to their corresponding contextual embedding at each layer.
• Figure 4: (a) Relative contribution per layer of neighbours at different positions. (b) Total contribution per neighbour for the first, middle and last layers.
• Figure 5: Effective attention (a) vs. raw attention (b). (a) Each point represents the average effective attention from a token type to a token type. Solid lines are the average effective attention of corresponding points in each layer. (b) is the corresponding figures using raw attention weights.