Quantifying context mixing in transformers

From @mohebbiQuantifyingContextMixing2023.

The Transformer Model

Transformer models in each node, at each layer, mix context information into the representation of each token.
Value Zeroing is a novel approach to quantify the contribution of each context token in determining the final representation of a target token, at each layer of a Transformer.
Value Zeroing is based on the Explaining-by-Removing intuition (Covert et al., 2021) shared by many posthoc interpretability methods, but it takes advantage of a specific feature of Transformers: it zeroes only the value vector of a token $t$ when computing its importance, but leaves the key and query vectors (and thus the pattern of information flow) intact.
This enables the ability to zero-out a token, without removing it from the sentence and thus changing its semantics.

So, inside the Self-Attention Layer of a Transformer, for each attention head $h$, the input vector $x_i$ for the $i^{th}$ token in the sequence is transformed in three distinct vector through the use of three different sets of weights:

• The Query vector ${\mathbf{q}}_i^h$;
• The Key vector ${\mathbf{k}}_i^h$;
• The Value vector ${\mathbf{v}}_i^h$;

Then, the context vector ${\mathbf{z}}_i^h$ for the same $i^{th}$ token in the sequence is generated as a weighted sum over the Value vector:

$${\mathbf{z}}_i^h=\sum _{j=1}^n{\mathbf{a}}_{ij}^h{\mathbf{v}}_j^h$$

where ${\mathbf{a}}_{ij}$ is the raw Attention weight assigned to the $j^{th}$ token in the sequence and computed as a Softmax-normalized dot product between its Query and Key vectors.

Then, the context vector for the $i^{th}$ vector is computed by concatenating all the heads’ outputs, followed by a projection through an output weight matrix ${\mathbf{W}}_O$ and layer normalization:

$${\mathbf{z}}_i=CONCAT(z_i^1,...,z_i^H){\mathbf{W}}_O$$ $${\mathbf{z}}_i=L{N}_{MHA}(z_i+x_i)$$

Finally, the context vector goes through a MLP that first project the input to a space that is traditionally four times larger, goes through a ReLU, and its projected back to the model’s hidden size. This is applied to every ${\mathbf{z}}_i$, producing the representation ${\tilde{\mathbf{x}}}_i$. This is again Layer Normalized:

$${\tilde{\mathbf{x}}}_i=L{N}_{MLP}({\tilde{\mathbf{x}}}_i,{\mathbf{z}}_i)$$

Value Zeroing

In Value Zeroing the weighted sum over the Value vector (The first Equation of the page) is changed by replacing the Value vector associated to $j$ with a zero vector ${\mathbf{v}}_j^h\leftarrow 0,\forall h\in H$ when computing the context vector for the $i^{th}$ token.

This provides a new representation ${\tilde{\mathbf{x}}}_i^{\neg j}$ that has exluded the $j^{th}$ token. Comparing it with the original representation ${\tilde{\mathbf{x}}}_i$ using a pairwise distance metrics (such as the coside distance) we obtain a measure of how much the representation changed by the exlusion of the $j^{th}$ token.

$${\mathbf{C}}_{ij}=cosine\_distance({\tilde{\mathbf{x}}}_i^{\neg j},{\tilde{\mathbf{x}}}_i)$$

Computing this equation for each $i,j$ token in the sequence we obtain a Value Zeroing Matric $\mathbf{C}$ where the value of each cell indicates the degree of saliency to which the $i^{th}$ token is dependent on the $j^{th}$ token to form its contextualized vectorial representation.

In particular, we obtain a $\mathbf{C}$ matrix for each Self-Attention layer in the network.

To obtain a single matrices that represent the co-dependency of the token representation across the network a technique called Rollout is used.

Rollout

From @abnarQuantifyingAttentionFlow2020

Is an attention matrix aggregation technique where each attention matrix are multiplied in order, one after the other, obtaining a final matrix.