Self-attention Mechanism

Sources:

1. Transformer from scratch
2. Attention Is All You Need 2017 paper

Notation

Symbol	Type	Explanation
$x_i$	$\in {\mathbb{R}}^C$	Input vector at position $i$ in the sequence
$y_i$	$\in {\mathbb{R}}^C$	Output vector at position $i$ in the sequence
$T$	$\in \mathbb{N}$	Length of the input and output sequence
$C$	$\in \mathbb{N}$	Dimension of the input and output vectors
$w_{ij}$	$\in \mathbb{R}$	Weight indicating the influence of input vector $x_j$ on $y_i$
$w_{ij}^{\prime }$	$\in \mathbb{R}$	Unnormalized attention score between $x_i$ and $x_j$
$w$	$\in {\mathbb{R}}^{T\times T}$	Weight matrix for the entire sequence
$x$	$\in {\mathbb{R}}^{T\times C}$	Matrix of all input vectors
$y$	$\in {\mathbb{R}}^{T\times C}$	Matrix of all output vectors
$W_q$ (single-head)	$\in {\mathbb{R}}^{C\times C}$	Weight matrix for transforming input vectors into queries in single-head attention
$W_k$ (single-head)	$\in {\mathbb{R}}^{C\times C}$	Weight matrix for transforming input vectors into keys in single-head attention
$W_v$ (single-head)	$\in {\mathbb{R}}^{C\times C}$	Weight matrix for transforming input vectors into values in single-head attention
$q_i$ (single-head)	$\in {\mathbb{R}}^C$	Query vector for the input vector $x_i$ in single-head attention
$k_i$ (single-head)	$\in {\mathbb{R}}^C$	Key vector for the input vector $x_i$ in single-head attention
$v_i$ (single-head)	$\in {\mathbb{R}}^C$	Value vector for the input vector $x_i$ in single-head attention
$W_q$ (multi-head)	$\in {\mathbb{R}}^{(C/h)\times C}$	Weight matrix for transforming input vectors into queries in multi-head attention
$W_k$ (multi-head)	$\in {\mathbb{R}}^{(C/h)\times C}$	Weight matrix for transforming input vectors into keys in multi-head attention
$W_v$ (multi-head)	$\in {\mathbb{R}}^{(C/h)\times C}$	Weight matrix for transforming input vectors into values in multi-head attention
$q_i$ (multi-head)	$\in {\mathbb{R}}^{C/h}$	Query vector for the input vector $x_i$ in multi-head attention
$k_i$ (multi-head)	$\in {\mathbb{R}}^{C/h}$	Key vector for the input vector $x_i$ in multi-head attention
$v_i$ (multi-head)	$\in {\mathbb{R}}^{C/h}$	Value vector for the input vector $x_i$ in multi-head attention
$h$	$\in \mathbb{N}$	Number of heads in multi-head attention
$\mathrm{softmax}$	Function	The softmax function
$\exp$	Function	Exponential function used in softmax calculation
$-\mathrm{\infty }$	Special value	Used to mask future elements in the causality mask

Self-Attention

Self-attention is a sequence-to-sequence operation that transforms an input sequence of vectors $\{x_1,x_2,\dots ,x_T\}$ into an output sequence $\{y_1,y_2,\cdots ,y_T\}$, where each vector has dimension $C$. To compute the output vector $y_i$, self-attention forms a weighted average of the input vectors:

$$y_i=\sum _j{w}_{ij}{x}_j,$$

where $j$ indexes the entire sequence, and the weights $w_{ij}$ sum to one over all $j$. These weights are derived from the input vectors $x_i$ and $x_j$ using a dot product:

$$w_{ij}^{\prime }=x_i^{\mathrm{\top }}{x}_j.$$

The resulting scores are then normalized using the softmax function to form the final attention weights:

$$w_{ij}=\frac{\exp (w_{ij}^{\prime })}{\sum _j\exp (w_{ij}^{\prime })}.$$

self-attention

A visual illustration of basic self-attention. Note that the softmax operation over the weights is not illustrated.

In matrix form, the self-attention operation is expressed as: $$y=wx,$$

where

$$x=\left[\begin{array}{c}{x}_1^{\mathrm{\top }}\\ ⋮\\ x_T^{\mathrm{\top }}\end{array}\right]\in {\mathbb{R}}^{T\times C},y=\left[\begin{array}{c}{y}_1^{\mathrm{\top }}\\ ⋮\\ y_T^{\mathrm{\top }}\end{array}\right]\in {\mathbb{R}}^{T\times C},w=\left[\begin{array}{cccc}{w}_{11}& w_{12}& \cdots & w_{1T}\\ w_{21}& w_{22}& \cdots & w_{2T}\\ ⋮& ⋮& \ddots & ⋮\\ w_{T1}& w_{T2}& \cdots & w_{TT}\end{array}\right]\in {\mathbb{R}}^{T\times T}.$$

Additional tricks

The actual self-attention used in modern transformers relies on several additional tricks.

Queries, keys, and values

Each input vector $x_i$ in self-attention is used in three ways:

1. to determine weights for its own output (query),
2. to contribute to weights for others' outputs (key), and
3. to participate in the weighted sum (value).

To facilitate these roles, each $x_i$ is linearly transformed into distinct vectors using learnable matrices $W_q$, $W_k$, and $W_v$: $$q_i=W_q{x}_i,k_i=W_k{x}_i,v_i=W_v{x}_i.$$

The attention weights are then computed as:

$$w_{ij}^{\prime }=q_i^{\mathrm{\top }}{k}_j,$$

2. apply a normalization then

$$w_{ij}=\mathrm{softmax}(w_{ij}^{\prime }),$$

and the output vector $y_i$ is: $$y_i=\sum _j{w}_{ij}{v}_j.$$

This mechanism introduces learnable parameters into the self-attention layer, enabling it to adaptively modify the input vectors for their roles.

Another notation is $$\begin{gathered} X=[x_1,x_2,\cdots ,x_T]\in {\mathcal{R}}^{C\times T} \\ Q=W_{q}X\in {\mathcal{R}}^{C_q\times T} \\ K=W_{k}X\in {\mathcal{R}}^{C_q\times T^{\prime }} \\ V=W_{v}X\in {\mathcal{R}}^{C_v\times T^{\prime }} \end{gathered}$$ Thus $$Y=V\text{col-wise-softmax}(K^{T}Q)\in {\mathcal{R}}^{C_v\times T}$$ where $K^{T}Q\in {\mathcal{R}}^{T^{\prime }\times T}$.

$t$-th col in K^T is the attention to the T' searched elements.

Multi-head attention

Rather than applying a single attention function with $C$-dimensional keys, values, and queries, the model applies this function $h$ times independently on different subspaces of dimension $\frac{C}{h}$. The outputs are then concatenated and linearly transformed to form the final output, allowing the model to capture diverse aspects of the input.

Causality mask for autoregressive models

In autoregressive models, where predictions depend solely on preceding elements, a causality mask is applied to prevent attention to future elements.

masked-attention

Masking the self attention, to ensure that elements can only attend to input elements that precede them in the sequence. Note that the multiplication symbol is slightly misleading: we actually set the masked out elements (the white squares) to $-\mathrm{\infty }$.

The mask modifies the dot product scores:

$$w_{ij}^{\prime }=\{\begin{array}{ll}{q}_i^{\mathrm{\top }}{k}_j,& \text{if }j\le i,\\ -\mathrm{\infty },& \text{if }j>i.\end{array}$$

This ensures that during the softmax calculation:

$$w_{ij}=\frac{\exp (w_{ij}^{\prime })}{\sum _{k=1}^i\exp (w_{ik}^{\prime })}$$

Future elements ($j>i$) are effectively ignored, as $\exp (-\mathrm{\infty })=0$, maintaining the autoregressive property and ensuring the model does not "look ahead" in the sequence.

Code

Single self-attention head

class AttentionHead(nn.Module):
    """
    A single attention head for self-attention in transformer models.

    Attributes:
        head_size (int): Dimension of the out put of the attention head.
        embd_dim (int): Dimension of the input embeddings.
        context_len (int): Maximum sequence length for attention.
        to_query, to_key, to_value (nn.Linear): Layers to compute query, key, and value vectors.
        tril (torch.Tensor): Lower triangular matrix for causal masking.
    """

    def __init__(self, head_dim, embd_dim, context_len):
        """
        Initializes the AttentionHead.

        Args:
            head_dim (int): Dimension of the out put of the attention head.
            embd_dim (int): Dimension of the input embeddings.
            context_len (int): Maximum sequence length.
        """
        super().__init__()
        
        self.head_size, self.embd_dim, self.context_len = head_dim, embd_dim, context_len
        self.to_query = nn.Linear(embd_dim, head_dim)
        self.to_key = nn.Linear(embd_dim, head_dim)
        self.to_value = nn.Linear(embd_dim, head_dim)
        
        # Causal mask for future tokens
        self.register_buffer('tril', torch.tril(torch.ones(self.context_len, self.context_len)))
        
    def forward(self, x, causality_mask=False):
        """
        Computes attention scores.

        Args:
            x (torch.Tensor): Input of shape [B, T, C].
            causality_mask (bool): If True, applies a mask to prevent attending to future tokens.

        Returns:
            torch.Tensor: Output of shape [B, T, head_dim].
        """
        B, T, C = x.shape
        k = self.to_key(x) # [B, T, head_dim]
        q = self.to_query(x) # [B, T, head_dim]
        v = self.to_value(x) # [B, T, head_dim]
        
        weight = q @ k.transpose(-2,-1) # [B, T, T]
        
        if causality_mask:
            # Masking future tokens
            weight = weight.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # [B, T, T]
            
        weight = F.softmax(weight, dim=-1)
        y = weight @ v # [B, T, head_dim]
        return y

Muti-head attention block

class MultiHeadAttention(nn.Module):
    """
    Multi-head self-attention mechanism.

    This class implements multiple heads of self-attention in parallel, 
    allowing the model to attend to information from different representation 
    subspaces at different positions.

    Attributes:
        head_num (int): Number of attention heads.
        embd_dim (int): Dimension of the input embeddings.
        context_len (int): Maximum sequence length for attention.
        head_size (int): Dimension of each attention head.
        heads (nn.ModuleList): List of AttentionHead modules.
        proj (nn.Linear): Linear layer to combine the outputs from all heads.
    """

    def __init__(self, head_num, embd_dim, context_len):
        """
        Initializes the MultiHeadAttention.

        Args:
            head_num (int): Number of attention heads.
            embd_dim (int): Dimension of the input embeddings.
            context_len (int): Maximum sequence length.
        """
        super().__init__()
        self.head_num, self.embd_dim, self.context_len = head_num, embd_dim, context_len
        self.head_size = embd_dim // head_num
        self.heads = nn.ModuleList([AttentionHead(self.head_size, self.embd_dim, self.context_len) for _ in range(head_num)])
        self.proj = nn.Linear(self.embd_dim, self.embd_dim)  # Mix the data output by the heads back together.
        

    def forward(self, x):
        """
        Applies multi-head attention to the input.

        Args:
            x (torch.Tensor): Input tensor of shape [B, T, C].

        Returns:
            torch.Tensor: Output tensor of shape [B, T, embd_dim].
        """
        outs = torch.cat([head(x) for head in self.heads], dim=-1)  # Concatenate outputs from all heads
        x = self.proj(outs)  # Project back to the original embedding dimension
        return x