Google

Within a multi-head attention mechanism, the output of each distinct attention head, denoted as $$\mathrm{head}_j$$, is computed by applying the Query-Key-Value (QKV) attention function to a specific sub-space of the model's representation. The operation utilizes the corresponding query, key, and value matrices for that particular head—$$\mathbf{Q}^{[j]}$$, $$\mathbf{K}^{[j]}$$, and $$\mathbf{V}^{[j]}$$—resulting in the following equation: $$\mathrm{head}_j = \mathrm{Att}_{\mathrm{qkv}}(\mathbf{Q}^{[j]}, \mathbf{K}^{[j]}, \mathbf{V}^{[j]})$$

Individual Attention Head Formula

Given a representation matrix $$\mathbf{H} \in \mathbb{R}^{m \times d}$$, the multi-head self-attention function computes its output by concatenating the results from multiple individual attention heads. This relationship is formalized as:

$$F(\mathbf{H}) = \mathrm{Merge}(\mathrm{head}_1, \dots, \mathrm{head}_\tau) \mathbf{W}^{\mathrm{head}}$$

In this equation, $$\mathrm{Merge}(\cdot)$$ signifies the concatenation of its inputs. Each element $$\mathrm{head}_j$$ represents the output derived from applying Query-Key-Value (QKV) attention to a specific sub-space of the initial representation. Finally, the concatenated results are projected via multiplication with a parameter matrix $$\mathbf{W}^{\mathrm{head}} \in \mathbb{R}^{d \times d}$$ to yield the final sequence representation.

Multi-Head Attention Output Calculation

In a causal multi-head attention mechanism, the output for a single head `j` at a specific token position `i` is computed using the standard Query-Key-Value (QKV) attention function. This calculation is restricted to the current and preceding tokens to maintain the autoregressive property. The formula is: $$\text{head}_j = \text{Att}_{\text{qkv}}(\mathbf{q}_i^{[j]}, \mathbf{K}_{\leq i}^{[j]}, \mathbf{V}_{\leq i}^{[j]})$$ Here, $\mathbf{q}_i^{[j]}$ is the query vector for the i-th token projected for head `j`, while $\mathbf{K}_{\leq i}^{[j]}$ and $\mathbf{V}_{\leq i}^{[j]}$ are the key and value matrices for head `j`, containing information from tokens 0 up to `i`.

Causal Attention Output for a Single Head and Token

In a multi-head attention mechanism, each individual attention head computes its output using its own unique Query, Key, and Value matrices, which are distinct linear projections of the same input. What is the primary functional consequence of this design choice?

Based on the formula for an individual attention head, `head_j = Att_qkv(Q^[j], K^[j], V^[j])`, what is the most probable cause of the issue described in the case study?

Debugging an Attention Head

In a multi-head attention mechanism, the output of each individual attention head, denoted as $\text{head}_j$, is a vector. This vector belongs to a real-valued vector space of dimension $d_h$, which is represented by the notation: $$\text{head}_j \in \mathbb{R}^{d_h}$$

Dimensionality of an Attention Head Output

You are examining the computation for a single attention head within a multi-head attention layer. Arrange the following steps in the correct chronological order to produce the output for this individual head.

During text generation, the output of the $$j$$-th attention head at step $$i$$ is computed by applying the Query-Key-Value (QKV) attention function to its specific feature sub-space. This operation utilizes the current token's query vector, $$\mathbf{q}_{i}^{[j]}$$, along with the cached key and value matrices for all tokens up to step $$i$$, denoted as $$\mathbf{K}_{\le i}^{[j]}$$ and $$\mathbf{V}_{\le i}^{[j]}$$. By projecting these representations onto the $$j$$-th sub-space, the model can be interpreted as performing attention on a group of independent feature sub-spaces in parallel. The calculation is formalized as: $$\mathrm{head}_j = \mathrm{Att}_{\mathrm{qkv}}(\mathbf{q}_{i}^{[j]}, \mathbf{K}_{\le i}^{[j]}, \mathbf{V}_{\le i}^{[j]})$$.

Learn Before

Related