Self-attention and multi-head in transformers

Attention Mechanisms in Transformers: Self-Attention, Cross-Attention, and Multi-Head

Attention mechanisms enable AI models to dynamically weigh the importance of elements in an input sequence. For each token, weights are computed using Query, Key, and Value matrices, creating contextualized representations. Self-attention examines connections within a single sequence, cross-attention links the encoder and decoder, and multi-head attention parallelizes the process to capture different data aspects.

Self-Attention Theory

The input sequence is represented by matrix $X \in \mathbb{R}^{n \times d}$, where $n$ is the number of tokens and $d$ is the embedding dimension. From $X$, we project matrices $Q = XW^Q$, $K = XW^K$, $V = XW^V$ with weights $W^{Q,K,V} \in \mathbb{R}^{d \times d_k}$.

The attention scores matrix $A_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}}$ measures the relevance of the $j$-th token to the $i$-th. Scaling by $\sqrt{d_k}$ stabilizes gradients, preventing vanishing gradients in softmax.

The normalized matrix $N = \mathrm{softmax}(A)$ across rows provides a probabilistic distribution. Output: $\mathrm{Attention}(Q, K, V) = N V$.

Self-Attention Example

Consider the sentence "Karina is going to the store," tokenized as ["karina", "is", "to", "the", "store"] with $n=5$. Embeddings (simplified):

| Token | Embedding |

|--------|-----------|

| karina | [1, 0] |

| is | [0, 1] |

| to | [1, 1] |

| the | [1, 1] |

| store | [0.5, 1] |

Matrix $X_{5\times2} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \\ 1 & 1 \\ 0.5 & 1 \end{bmatrix}$. For simplicity, $W^Q = W^K = W^V = I$, so $Q=K=V=X$.

$QK^T = \begin{bmatrix} 1 & 0 & 1 & 1 & 0.5 \\ 0 & 1 & 1 & 1 & 1 \\ 1 & 1 & 2 & 2 & 1.5 \\ 1 & 1 & 2 & 2 & 1.5 \\ 0.5 & 1 & 1.5 & 1.5 & 1.25 \end{bmatrix}$, divided by $\sqrt{2} \approx 1.41$:

$A \approx \begin{bmatrix} 0.71 & 0 & 0.71 & 0.71 & 0.35 \\ 0 & 0.71 & 0.71 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.71 & 0.71 & 1.42 & 1.42 & 1.06 \\ 0.35 & 0.71 & 1.06 & 1.06 & 0.89 \end{bmatrix}$.

After softmax per row: $N \approx \begin{bmatrix} 0.25 & 0.15 & 0.25 & 0.25 & 0.10 \\ 0.17 & 0.23 & 0.23 & 0.23 & 0.23 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.15 & 0.15 & 0.30 & 0.30 & 0.10 \\ 0.14 & 0.23 & 0.28 & 0.28 & 0.07 \end{bmatrix}$.

Output $N \times V = \begin{bmatrix} 0.73 & 0.68 \\ 0.57 & 0.87 \\ 0.68 & 0.81 \\ 0.68 & 0.81 \\ 0.63 & 0.84 \end{bmatrix}$. The new embeddings now incorporate context: the vector for "is" reflects human movement rather than just an abstract tense.

import torch  
import torch.nn as nn  

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x):
        Q = self.q_linear(x)
        K = self.k_linear(x)
        V = self.v_linear(x)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

Cross-Attention: Linking Sequences

In cross-attention, Query comes from the decoder ($Q = X_{dec}W^Q$), while Key and Value come from the encoder ($K = X_{enc}W^K$, $V = X_{enc}W^V$). Computation is identical: $\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left( \frac{QK^T}{\sqrt{d_k}} \right) V$.

Key differences:

| Component | Self-Attention | Cross-Attention |

|-----------|----------------|-----------------|

| Q | X | X_dec |

| K | X | X_enc |

| V | X | X_enc |

This lets the decoder focus on relevant parts of the encoder's context.

import torch  
import torch.nn as nn  

class CrossAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.embed_dim = embed_dim
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x_dec, x_enc):
        Q = self.q_linear(x_dec)
        K = self.k_linear(x_enc)
        V = self.v_linear(x_enc)
        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.bmm(attn, V)
        return out

Multi-Head Attention: Parallel Perspectives

Multi-head attention runs $h$ independent attention computations with projections $Q^{(i)} = X W_Q^{(i)}$, $K^{(i)} = X W_K^{(i)}$, $V^{(i)} = X W_V^{(i)}$. Each head: $\mathrm{head}_i = \mathrm{Attention}(Q^{(i)}, K^{(i)}, V^{(i)})$.

Results are concatenated and projected: $\mathrm{MultiHead}(X) = \mathrm{Concat}(head_1, \dots, head_h) W^O$. Head size $d_k = d / h$, where $d$ is embed_dim.

import torch
import torch.nn as nn

class MultiHeadSelfAttention(nn.Module):  
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.q = nn.Linear(embed_dim, embed_dim)
        self.k = nn.Linear(embed_dim, embed_dim)
        self.v = nn.Linear(embed_dim, embed_dim)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        B, T, D = x.shape
        Q = self.q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.k(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.v(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, V).transpose(1, 2).contiguous().view(B, T, D)
        return self.out(out)

Key Takeaways

• Self-attention contextualizes tokens: Each output vector depends on the entire sequence.
• Scaling by $\sqrt{d_k}$ is essential: It stabilizes softmax and gradients.
• Cross-attention for seq2seq: Links encoder and decoder in transformers.
• Multi-head captures diverse aspects: Grammar, semantics, long-range dependencies.
• PyTorch implementations use Linear and bmm: For batch efficiency.

— Editorial Team