Google

In a multi-head attention mechanism, the queries, keys, and values for the $$j$$-th attention head are obtained by projecting the input representation $$\mathbf{H}$$ into different subspaces via linear transformations. These transformations utilize unique learnable parameter matrices for each head. The projections are defined as follows:

$$\mathbf{Q}^{[j]} = \mathbf{H} \mathbf{W}_j^{q}$$
$$\mathbf{K}^{[j]} = \mathbf{H} \mathbf{W}_j^{k}$$
$$\mathbf{V}^{[j]} = \mathbf{H} \mathbf{W}_j^{v}$$

Here, $$\mathbf{W}_j^{q}$$, $$\mathbf{W}_j^{k}$$, and $$\mathbf{W}_j^{v} \in \mathbb{R}^{d \times \frac{d}{\tau}}$$ denote the parameter matrices of the transformations for the $$j$$-th head.

Query, Key, and Value Projections in Multi-Head Attention

In a multi-head attention mechanism, the key weight matrix for an individual attention head, which can be denoted as $k_h^k$, has a specific shape defined as $d \times \frac{d_h}{M}$. In this formula, $d$ is the dimension of the input representation, $d_h$ is the total dimension of the key projection across all heads, and $M$ is the number of attention heads.

Shape of Key Weight Matrix per Head

In a multi-head attention mechanism, the key weight sub-matrix for an individual attention head, denoted as $W_h^k$, has a shape of $d \times \frac{d}{M}$. This formula applies specifically when the total dimension of the key projection across all heads is equal to the input representation dimension, $d$. In this context, $M$ represents the number of attention heads.

Shape of Key Weight Sub-Matrix per Head

In the context of a multi-head attention mechanism, explain the primary reason for using distinct, learnable weight matrices to project the input representation into separate Query, Key, and Value sets for each individual attention head.

Rationale for Unique Projections in Multi-Head Attention

Based on the mechanism that generates query, key, and value matrices, explain how it is possible for these different heads to learn and attend to such different types of information, even though they all begin with the exact same input representation.

Attention Head Specialization

In the general vector-level formulation of multi-head attention (Eq. 11.5.1), the $$i$$-th attention head output $$\mathbf{h}_i$$ (for $$i = 1, \ldots, h$$) is computed by first projecting a query $$\mathbf{q} \in \mathbb{R}^{d_q}$$, a key $$\mathbf{k} \in \mathbb{R}^{d_k}$$, and a value $$\mathbf{v} \in \mathbb{R}^{d_v}$$ through head-specific learnable weight matrices, and then applying an attention pooling function $$f$$:

$$\mathbf{h}_i = f(\mathbf{W}_i^{(q)} \mathbf{q},\; \mathbf{W}_i^{(k)} \mathbf{k},\; \mathbf{W}_i^{(v)} \mathbf{v}) \in \mathbb{R}^{p_v}$$

Here, $$\mathbf{W}_i^{(q)} \in \mathbb{R}^{p_q 	imes d_q}$$, $$\mathbf{W}_i^{(k)} \in \mathbb{R}^{p_k 	imes d_k}$$, and $$\mathbf{W}_i^{(v)} \in \mathbb{R}^{p_v 	imes d_v}$$ are learnable parameter matrices that project the original representations into subspaces of dimensions $$p_q$$, $$p_k$$, and $$p_v$$ respectively. The function $$f$$ denotes the attention pooling operation, such as additive attention or scaled dot-product attention.

Individual Attention Head Computation (General Vector Form)

In a multi-head attention mechanism, utilizing $$h$$ parallel attention heads could potentially lead to a significant increase in both computational and parametrization costs. To avoid this growth, the dimensionalities of the query, key, and value projections for each individual head (denoted as $$p_q$$, $$p_k$$, and $$p_v$$) are typically constrained to $$p_o / h$$, where $$p_o$$ is the total desired output dimensionality. By ensuring that $$p_q h = p_k h = p_v h = p_o$$, the computations for all $$h$$ heads can be performed in parallel while maintaining overall resource requirements that are comparable to a single-head attention mechanism with dimensionality $$p_o$$.

Parametrization Cost Control in Multi-Head Attention

In a multi-head attention mechanism with 'M' heads, an engineer makes an implementation error. Instead of creating a unique set of learnable weight matrices for the query, key, and value projections for each of the 'M' heads, the same single set of query, key, and value weight matrices is shared across all heads. What is the primary consequence of this error on the model's functionality?

Learn Before

Related