Source-linked AI summary
On Identifiability in Transformers
Gino Brunner, Yang Liu, Damián Pascual, Oliver Richter, Massimiliano Ciaramita, Roger Wattenhofer
TL;DR
The paper asks whether attention weights and token identities in Transformer representations are identifiable, a question tied to interpreting model predictions. It analyzes attention and contextual embeddings theoretically and experimentally, finding non-identifiable attention weights, persistent but depth-decreasing token identity, and strong contextual mixing. The results motivate effective attention and gradient-based attribution as tools for understanding Transformer representations.
Problem
The paper investigates whether attention weights and token embeddings provide stable, interpretable mappings in self-attention models.
Method
The paper combines structural identifiability analysis with experiments on BERT token embeddings and Hidden Token Attribution based on gradient attribution.
Results
Attention weights are not identifiable when sequence length exceeds the relevant rank condition; token identity remains largely recoverable, while contextual embeddings strongly mix token and context information.
Takeaways & Limitations
Attention distributions are not directly interpretable on their own, so effective attention and token-attribution analyses provide complementary tools for examining Transformer representations.
Takeaways & Limitations
Strong information mixing leaves it unclear whether token identifiability alone is sufficient to equate hidden embeddings with their input words.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper studies identifiability in self-attention models, covering attention weights, token identity across layers, and context-token mixing in contextual embeddings.
- Identifiability concerns whether neural models learn stable representations, affecting the replicability and interpretability of their predictions.
- Attention weights are not identifiable under mild sequence-length and head-dimension conditions, so infinitely many weights can yield the same output.The authors propose effective attention to remove weight components that do not influence predictions.
- Experiments find that contextual word embeddings generally maintain token identity across transformer layers, although this pattern does not always hold in later layers.Identity information is largely encoded in embedding angles and can be recovered through a learned linear mapping followed by nearest-neighbour lookup.
- Hidden Token Attribution measures how input tokens contribute to contextual embeddings and finds strong mixing of token and context information.Token contribution decreases monotonically with depth, while the corresponding token typically remains the largest individual contributor; aggregated context is mostly local.
2 BACKGROUND ON TRANSFORMERS
Transformers use multi-head self-attention to let each input token incorporate information from every other token, producing contextual word embeddings across layers.
- Each self-attention layer projects token embeddings into query, key, and value vectors before computing an attention-head output.The paper denotes the query, key, and value matrices as Q, K, and V, with sequence length d_s and query/value dimensions d_q and d_v.
- The attention matrix assigns influence weights from every sequence position to the hidden embedding computed at each position.Its output is represented as Attention(Q, K, V) = A · V, where A is produced by softmax.
- Self-attention is non-local because each token can attend to all other tokens regardless of their distance in the input.
- Successive self-attention layers aggregate contextual information into embeddings associated with the input words.
- The experiments use BERT, with input tokens represented by sums of token, segment, and position embeddings and contextual embeddings indexed by layer.All experiments use the pre-trained uncased BERT-Base model.
3 ATTENTION IDENTIFIABILITY
The paper shows that self-attention weights generally cannot be uniquely recovered from a head’s output, because value-vector combinations have a non-trivial null space. It proposes effective attention to isolate the attention component that affects outputs and reports empirical divergence from raw attention.
- 3.2 THE NULL SPACE OF T: The identifiability problem arises because attention weights combine T vectors, whose rank is bounded by the head dimension.When the number of weighting coefficients exceeds rank(T), free variables remain.
- 3.2 THE NULL SPACE OF T: Attention weights are not identifiable when sequence length exceeds the attention head dimension, allowing infinitely many weights to produce the same output.The null-space dimension increases with sequence length.
- 3.3 THE NULL SPACE WITH PROBABILITY CONSTRAINTS: Even under probability-simplex constraints, alternative attention distributions can exist when the augmented null space is non-trivial.The construction requires preserving the output, row sums, and non-negativity.
- 3.4 EFFECTIVE ATTENTION: Effective attention is the component of attention weights that affects the model output after removing the null-space component.It is computed by projecting attention weights onto the orthogonal complement of the left null space.
- 3.4 EFFECTIVE ATTENTION: The effective-attention formulation is not guaranteed to be a probability distribution because it uses the unconstrained null space of T.The authors leave constrained and alternative formulations for future research.
- 3.4 EFFECTIVE ATTENTION: Effective and raw attention diverge increasingly with sequence length, as their Pearson correlation decreases in the reported experiments.The authors interpret this pattern as evidence that more raw attention lies in the null space for longer sequences.
4 TOKEN IDENTIFIABILITY
This section tests whether contextual embeddings preserve the identity of their input tokens across Transformer layers. It finds that identity remains largely recoverable, although it changes with depth and does not by itself reveal how much contextual mixing occurred.
- 4.1 TOKEN IDENTIFIABILITY: Token identifiability is defined by whether a mapping can recover each contextual embedding’s corresponding input token.Evaluation uses nearest-neighbour recovery within the same input sentence.
- 4.2 EXPERIMENTAL RESULTS AND DISCUSSION: Contextual embeddings remain close to their input embeddings through layer 4, after which token identifiability decreases linearly.This pattern holds for both cosine and L2 similarity measures.
- 4.2 EXPERIMENTAL RESULTS AND DISCUSSION: 93% of last-layer tokens are recovered using a learned linear projection and cosine-distance nearest-neighbour lookup.The result indicates that substantial identity information remains in the final contextual embeddings.
- 4.2 EXPERIMENTAL RESULTS AND DISCUSSION: Cosine distance recovers identity more effectively than L2 distance, suggesting that identity information is largely encoded in embedding angles.Tokens remain mostly identifiable despite decreasing identifiability with depth.
- 4.2 EXPERIMENTAL RESULTS AND DISCUSSION: A probe trained on one layer can generalize across layers, but layer-specific transfer patterns reveal that token identity changes through the model.Training at layer 1 falls below 70% when tested at layers 11 and 12, whereas training at layer 12 reaches 94% on layer 12 and transfers increasingly well to earlier layers.
- 4.2 EXPERIMENTAL RESULTS AND DISCUSSION: Token identifiability alone cannot establish that hidden embeddings represent only their input words because self-attention may strongly mix information across layers.The section therefore motivates measuring how much of each input token remains in contextual embeddings.
5 ATTRIBUTION ANALYSIS TO IDENTIFY CONTEXT CONTRIBUTION
This section introduces Hidden Token Attribution to quantify how input tokens contribute to intermediate contextual embeddings. The analysis finds strong, depth-dependent mixing, with local context generally dominating despite increasing long-range contributions.
- 5.1 TOKEN MIXING: CONTRIBUTION OF INPUT TOKENS: Hidden Token Attribution uses gradient attribution to measure each input token’s contribution to an intermediate hidden embedding.Unlike standard attribution, it tracks contributions to hidden representations rather than explaining the model’s final decision.
- 5.1 TOKEN MIXING: CONTRIBUTION OF INPUT TOKENS: The original token’s median contribution falls from 30.6% after the first layer to 14.4% at layer 6 and 10.7% after the last layer.Contribution decreases monotonically with depth.
- 5.1 TOKEN MIXING: CONTRIBUTION OF INPUT TOKENS: The original input token is always the largest contributor during the first three layers, while the percentage for which it is not largest reaches 18% at layer 6 and 30% in the last two layers.The corresponding token generally remains the largest individual contributor even as mixing increases.
- 5.1 TOKEN MIXING: CONTRIBUTION OF INPUT TOKENS: From layer three onward, self-attention strongly mixes input information, yet the resulting mixtures often still permit recovery of token identity.This connects the attribution results with the decreasing identifiability rate observed after layer three.
- 5.2 CONTEXT AGGREGATION: Nearest neighbours contribute more in early layers, whereas tokens 11 positions or more away contribute most in deeper layers.Despite increasing long-range dependencies, contextual contributions remain mostly local on average.
- 5.2 CONTEXT AGGREGATION: The results suggest that BERT learns local operators from data, although exceptions such as the [CLS] token occur.This is notable because self-attention itself is a non-local operator.
6 RELATED WORK
Related work has questioned whether attention distributions uniquely determine model behaviour. This paper extends that concern to multi-head, multi-layer self-attention through theoretical and empirical analyses of attention and token mixing.
- ATTENTION INTERPRETABILITY: Earlier studies constructed alternative attention distributions that leave LSTM encoder-decoder decisions unchanged and found limited correlation with importance measures.These findings motivated interpreting attention distributions with care.
- ATTENTION INTERPRETABILITY: This paper provides a theoretical proof that attention weights are non-identifiable and studies multi-head, multi-layer Transformer self-attention rather than single-layer RNN attention.Its token classification and token mixing experiments show that non-identifiable tokens increase with depth.
- TRANSFORMER REPRESENTATIONS: Findings that only a small number of BERT heads materially affect outputs are consistent with the paper’s conclusion that many attention weights do not affect downstream components.The comparison links head importance results with attention non-identifiability.
7 CONCLUSION
The paper uses identifiability to analyze attention weights, token representations, and context mixing in Transformers. It finds non-unique attention distributions, largely retained token identity, and strong information mixing, while identifying open questions about their relationship.
- Attention identifiability: Attention weights can be non-identifiable, so infinitely many attention distributions may produce the same representation and output.The paper proposes effective attention by projecting out the null space to improve interpretability.
- Token identity and mixing: Tokens remain largely identifiable through learned linear transformations and cosine-similarity nearest-neighbor lookup, although identifiability decreases in later layers.The paper also reports that tokens mix heavily with context inside Transformers.
- Open questions: The paper identifies the relationship between token mixing and retained token identity as an open question for future research.It notes that a token can gather contextual information while retaining the essence of the original input word.
A.2 ADDITIONAL RESULTS OF THE EFFECTIVE ATTENTION VS. RAW ATTENTION RESULTS
Additional results compare raw, effective, and null attention, showing that raw attention can contain substantial null-space mass and that effective attention can expose different interactions. A single-head example illustrates this divergence but does not support broad generalization.
- Aggregate comparison: Most raw attention toward [SEP] lies effectively in the null space, so raw and effective attention differ substantially.The comparison recreates attention patterns using average raw and effective attention.
- Single-head example: Raw attention concentrates on [CLS], [SEP], and near-diagonal alignments, whereas effective attention is more uniform and removes [CLS] and [SEP] attention.The effective pattern exposes a non-diagonal relation between “make” and co-referential “them” 48 positions away.
- Null component: Null attention is more uniform than raw attention and appears to carry the mass associated with [CLS] and [SEP].The null and effective components show both redundancy and complementary elements.
- Interpretive scope: The example is a proof of concept that raw and effective attention can diverge qualitatively, but conclusions should not be extrapolated from one observation.The authors suggest that null-space attention may obscure interactions recoverable through effective attention.
B.2 GENERALIZATION ERROR
Additional experiments examine generalization, layerwise token identity, and recovery of neighboring tokens. They broadly support persistent but depth-dependent token information, with weaker recovery for neighboring positions and asymmetric left-right behavior.
- Generalization error: 0.74, 0.8, 0.81, 0.82 are the MLP test token identifiability rates using 50%, 80%, 90%, and 100% of training data.The result indicates little benefit from adding more data in this experiment.
- Generalization error: The linear model and MLP have similar generalization error, while cosine-trained models have nearly equal linear and nonlinear performance.The cosine-distance generalization error is not reported because the two perceptrons perform almost equally.
- Generalization across layers: Linear and MLP transformations trained on selected layers generalize to later layers with the same qualitative trends reported in the main text.The additional figures cover linear L2 and cosine losses, plus MLP L2 and cosine losses.
- Single-layer identifiability: Single-layer results show that BERT applies small changes to contextual embeddings, with angular information affected less than information measured by L2 distance.Naive baselines perform well across individual layers.
- Single-layer identifiability: Token identifiability generally decreases in later layers, and fine-tuning comparisons on MRPC and CoLA test whether task-dependent adaptation contributes to this pattern.The figures compare non-fine-tuned and fine-tuned BERT across layers.
- Context identifiability: Neighboring-token recovery declines with distance and depth, and left and right neighbors show different behavior despite bidirectional pretraining.Neighboring tokens are recovered to a much lower degree than same-position tokens.
- Context identifiability: Relative context contributions are computed by averaging left and right neighbors within distance groups and normalizing each group’s contributions.The groups include positions 1–5, 6–10, and 11 onward.
C.3 CONTEXT CONTRIBUTION TO CLS TOKEN
Hidden Token Attribution shows that [CLS] aggregates context from across the sequence, especially near [SEP], while token contributions evolve across layers and differ by part of speech.
- Context contribution to [CLS]: Hidden Token Attribution analyzes how context contributes to the [CLS] token used for BERT classification.[CLS] is especially informative because BERT bases classification decisions on its output.
- Context contribution to [CLS]: [CLS] aggregates context from all tokens but most strongly from tokens around the first [SEP].The two tokens with the highest contribution are directly before and after [SEP], and the second sentence contributes more than the first.
- Context aggregation examples: Context aggregation is illustrated as mostly local for individual words, with later layers incorporating syntactic, semantic, and cross-sentence context.Examples include “fit” drawing from “strategy,” and “long” drawing from “term” and corresponding tokens in the second sentence.
- Token contribution across layers: Most token types are already equally mixed by layer 1, while [CLS] retains over 40% median original contribution.The original contribution to [CLS] and [SEP] falls significantly by layer 4.
- Token contribution across layers: Across layers, nouns, verbs, adjectives, and adverbs retain more contribution from their corresponding input embeddings than function words and symbols.This pattern becomes clear by layer 4 and remains evident through layer 12.
- Additional datasets: The appendix reports qualitatively similar conclusions across CoLA and MNLI, supporting generalization across data domains.The token-identifiability findings are described as qualitatively in line with MRPC, with no meaningful discrepancy for MNLI.
D.2.1 COLA EXPERIMENTS
On CoLA, token mixing and context aggregation follow the main experiments: input-token contribution is slightly larger, while context remains mostly local before longer-range dependencies grow later.
- Token mixing: CoLA shows the same general token-mixing behavior as MRPC, with slightly larger original-token contributions across layers.The percentage of tokens that are not maximum contributors is also slightly larger, explained by CoLA’s shorter sequences.
- Context aggregation: CoLA context is aggregated mostly locally, with long-range dependencies increasing in later layers.The average sequence length of 11 tokens explains the smaller relative contribution from tokens beyond the 10th neighbour.
D.2.2 MNLI EXPERIMENTS
MNLI-matched reproduces the main token-identifiability and token-mixing patterns, with no meaningful discrepancy from the main-text results.
- Token mixing: MNLI-matched reproduces the main token-mixing analysis without a meaningful discrepancy.The cited results compare the matched MNLI dataset with the main-text experiments.