All you need for attention

This blog provides a detailed explanation of the attention mechanism, which is a core component of many modern neural network architectures, especially in natural language processing tasks.

Input and Output

The basic inputs and usually of the attention mechanism is token, which is usually a 1D vector representing a word or subword in a sequence (in vision, it could be a patch of an image). Given a sequence of tokens, the attention aims to learn the relationships between them and produce a new sequence of tokens that captures these relationships. Thus, the input and output shape are usulally $[B,N,C]$ where $B$ is the batch size, $N$ is the number of tokens, and $C$ is the number of channels (or features).

Equations of attention

The basic attention mechanism can be described below (also known as scaled dot-product attention):

\begin{equation} \operatorname{Attention}(Q,K,V) = \operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V \end{equation}

Here is what each value means:

• $Q$ — queries: what each token is looking for.
• $K$ — keys: what each token offers or represents for matching.
• $V$ — values: the information retrieved after matching.
• $QK^\top$ — similarity scores between every query and key.
• $d_k$ — dimension of each key/query vector. $d_k = c$ if there is no input projection.
• $\operatorname{softmax}$ — turns each row of scores into attention weights that sum to 1.
• Multiplication by $V$ — produces a weighted mixture of the value vectors.

Rescale factor

In the above equation, the similarity score is divided by $\sqrt{d_k}$ to prevent scores from becoming too large since the dot product grows with the dimension of the vectors. The reason of choosing $\sqrt{d_k}$ is that the standard deviation of the dot product of two random vectors of dimension $d_k$ increases by $\sqrt{d_k}$ times, and thus dividing by $\sqrt{d_k}$ helps normalize the scores and keep them in a reasonable range for the softmax function.

import numpy as np
n=10
d=100
a=np.random.randn(n,d)
b=np.random.randn(n,d)
c=a@b.T
np.std(c/np.sqrt(d))

1.0391848482914605

If the normalization is missing, the softmax function may heavily favor one token over others, leading to a loss of information and poor performance in the model. For example, softmax values of vector $[1,2]$ and $[10,20]$ are totally different, with the latter being almost a one-hot vector:

from scipy.special import softmax
print(softmax([1,2]), softmax([10,20]))

[0.26894142 0.73105858] [4.53978687e-05 9.99954602e-01]

Below is a simple implementation of the attention mechanism in PyTorch:

from torch import nn
import torch,math

class ScaleDotProductAttention(nn.Module):
    def __init__(self, d_k:int):
        super().__init__()
        self.d_k = d_k

    def forward(self, q, k, v):
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attention_weights = nn.functional.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, v)
        return output

Learnable parameters, self- and cross-attention

In the above implementation, there is no learnable parameters and the attention is a simple function of the input tokens representing an interaction between them. In practice, the attention mechanism is usually implemented with learnable parameters, which are used to transform the input tokens into queries, keys, and values. The learnable parameters are usually implemented as linear layers for the input and output tokens. Below is a simple implementation:

class Attention(nn.Module ):
    def __init__(self,
                dim_in:int,
                dim_hidden:int,
                dim_out:int,
                 ):
        super().__init__()
        self.w_q = nn.Linear(dim_in, dim_hidden)
        self.w_k = nn.Linear(dim_in, dim_hidden)
        self.w_v = nn.Linear(dim_in, dim_hidden)
        self.attention = ScaleDotProductAttention(dim_hidden)
        self.output_map = nn.Linear(dim_hidden, dim_out)

    def forward(self, q, k, v):
        q=self.w_q(q)
        k=self.w_k(k)
        v=self.w_v(v)
        return self.output_map(self.attention(q, k, v))

From practical perspective, there are two types of attention mechanisms: self-attention and cross-attention. In self-attention, the queries, keys, and values are all derived from the same input sequence, allowing the model to capture relationships within the same sequence. In cross-attention, the queries come from one sequence while the keys and values come from another sequence (thus queries must have same dimension as keys/values but can have different number of tokens), enabling the model to learn relationships between different sequences. Although the above implementation can be used for both self- and cross-attention by changing the input tensors, it is more common and efficient to only keep a single input learnable linear layer for self-attention:

class SelfAttention(nn.Module):
    def __init__(self,
                dim_in:int,
                dim_hidden:int,
                dim_out:int,
                 ):
        super().__init__()
        self.w_qkv = nn.Linear(dim_in, 3*dim_hidden)
        self.attention = ScaleDotProductAttention(dim_hidden)
        self.output_map = nn.Linear(dim_hidden, dim_out)

    def forward(self, x):
        qkv = self.w_qkv(x)
        q, k, v = torch.chunk(qkv, 3, dim=-1)
        return self.output_map(self.attention(q, k, v))

Multi-head attention

Multi-head attention is a commonly used extension of the attention mechanism that allows the model to jointly attend to information from different representation subspaces at different positions. It first splits the queries, keys, and values into multiple heads, applies the attention mechanism to each head in parallel, and then concatenates the results. The key idea behind multi-head attention is that each head can learn to focus on different parts of the input sequence, capturing various aspects of the relationships between tokens. This allows the model to capture richer and more diverse information compared to a single attention head. Below is a simple implementation of multi-head attention in PyTorch:

from einops import rearrange

class MultiheadCrossAttention(nn.Module):

    def __init__(self,
                dim_in:int,
                dim_hidden:int,
                dim_out:int,
                num_heads:int):
        super().__init__()
        assert dim_hidden % num_heads == 0, "dim_hidden must be divisible by num_heads"
        self.num_heads = num_heads
        self.dim_per_head = dim_hidden // num_heads
        self.w_q = nn.Linear(dim_in, dim_hidden)
        self.w_k = nn.Linear(dim_in, dim_hidden)
        self.w_v = nn.Linear(dim_in, dim_hidden)
        self.attention = ScaleDotProductAttention(self.dim_per_head)
        self.output_map = nn.Linear(dim_hidden, dim_out)

    def get_heads(self, tokens, linear_layer):
        b, n, _ = tokens.size()
        tokens = linear_layer(tokens)
        tokens = tokens.view(b, n, self.num_heads, self.dim_per_head)
        # move the head dimension to the second position so that we can perform attention for each head in parallel
        tokens = tokens.transpose(1, 2)
        # or use rearrange:
        # tokens = rearrange(tokens, 'b n (h d) -> b h n d', h=self.num_heads)
        return tokens

    def forward(self, q, k, v):
        q = self.get_heads(q, self.w_q)
        k = self.get_heads(k, self.w_k)
        v = self.get_heads(v, self.w_v)

        out = self.attention(q, k, v)
        # move the head dimension back
        # since the `.view()` operation requires the tensor to be contiguous in memory and transpose operations can make it non-contiguous, we need to call `.contiguous()` after the transpose operation
        out= out.transpose(1,2).contiguous()
        # merge the head dimension back into the feature dimension
        out = out.view(q.size(0), -1, self.num_heads * self.dim_per_head)
        # or use rearrange:
        # out = rearrange(out, 'b h n d -> b n (h d)')
        return self.output_map(out)        

Attention weight, attention mask and causal attention

In attention mechanism, the output of the softmax function has a very clear interpretation: it is the attention wight, which indicates how much each token should attend to every other token. The attention wight is a matrix of shape $[N_q,N_k]$ where $N_q$ is the number of queries and $N_k$ is the number of keys. Value at entry $(i,j)$ means: “How much does query token $i$ attend to key token $j$?” Thus, it is also a common practice to visualize the attention wight matrix to understand how the model is attending to different tokens in the input sequence.

Besides, the clear meaning of the attention wight also allows us to apply an attention mask to control which tokens can attend to which other tokens. For example, in causal attention, we want to prevent a token from attending to future tokens in the sequence (i.e., we should only keep the lower triangular part of the attention wight matrix). This is achieved by applying a mask that sets the attention wights for future tokens to negative infinity before applying the softmax function. This way, the softmax will assign zero attention weight to future tokens, effectively preventing them from being attended to.

Below is a simple implementation of multi-head self-attention with masking:

from typing import Optional
from einops import rearrange
class MultiHeadSelfAttention(nn.Module):

    def __init__(self,
                 in_dim:int,
                 hidden_dim:int,
                 out_dim:int,
                 num_heads:int,
                 causal:bool=False
                 ):
        super().__init__()
        assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads"
        self.num_heads = num_heads
        self.dim_per_head = hidden_dim // num_heads
        self.w_qkv = nn.Linear(in_dim, 3*hidden_dim)
        self.output_map = nn.Linear(hidden_dim, out_dim)
        self.causal = causal

    def forward(self, x, 
                mask:Optional[torch.Tensor]=None):
        # mask: a tensor with shape of [N, N] or can be broadcasted to shape [B, H, N, N]. True means the position is allowed.
        n = x.shape[1]
        qkv = self.w_qkv(x)
        qkv = rearrange(qkv, "b n (3 h d) -> b h n d 3", h=self.num_heads, d=self.dim_per_head)
        q,k,v = qkv[...,0], qkv[...,1], qkv[...,2]
        score = q @ k.transpose(-2, -1) / math.sqrt(self.dim_per_head)
        if self.causal:
            causal_mask = torch.ones(
                n, n, device=x.device, dtype=torch.bool
            ).triu(1)
            score = score.masked_fill(causal_mask, float("-inf"))

        if mask is not None:
            score = score.masked_fill(~mask.bool(), float("-inf"))

        weights = nn.functional.softmax(score, dim=-1)
        out = rearrange(weights @ v, "b h n d -> b n (h d)")
        return self.output_map(out)

Computational and memory complexity

In this section, we will analyze the computational complexity of the attention mechanism. The computational complexity of the attention mechanism is easy to analyze, as most of the operations are matrix multiplications. For matrix multiplication, the computational complexity is $O(n m p)$ for multiplying a matrix of shape $[n,m]$ with a matrix of shape $[m,p]$. The memory complexity is also easy to analyze, as most of the operations are linear transformations and the attention wight matrix. For linear transformation, the memory complexity is $O(n d)$ for storing the output of shape $[n,d]$. For attention wight matrix, the memory complexity is $O(n^2)$ for storing the attention wight of shape $[n,n]$.

Given a multi-head attention with $h$ heads, $d_h$ dimension per head, and an input of $n$ tokens with dimension $d$, the computational complexity can be broken down as follows:

Here’s the clean breakdown.

A compact way to say the computational complexity is:

and under the standard assumption that $d$ is fixed and $h d_h = d$, this is usually simplified to:

For memory complexity, a similar analysis yields:

In llms, since the number of tokens is usually much larger than the number of channels, the $O(n^2)$ term dominates both the computational and memory complexity. This is why in most blog or research papers, the attention mechanism is often described as having $O(n^2)$ complexity. However, it is important to note that the actual complexity can vary depending on the specific implementation and application. For example, in some cases, the dimension of the tokens may be larger than the number of tokens, in which case the $O(n d^2)$ term may dominate. Additionally, some implementations may use optimizations such as sparse attention or low-rank approximations to reduce the complexity of the attention mechanism. Thus, it is important to consider the specific context and implementation when analyzing the complexity of the attention mechanism.

Transformer block and position embeddings

In pratice, the attention mechanism does not work alone, but is usually combined with other components to form a transformer block. A transformer block typically consists of a multi-head attention layer followed by a feed-forward neural network, with residual connections and layer normalization applied to both components. The transformer block can be stacked multiple times to form a deep transformer model. Below is a simple implementation of a transformer block with self-attention:

class FeedForward(nn.Module):
    def __init__(self, in_dim:int, hidden_dim:int, out_dim:int):
        super().__init__()
        self.linear1 = nn.Linear(in_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, out_dim)
        self.activation = nn.ReLU()

    def forward(self, x):
        return self.linear2(self.activation(self.linear1(x)))


class TransformerBlock(nn.Module):
    def __init__(self,
                 dim:int,
                 hidden_dim:int,
                 num_heads:int,
                 causal:bool=False):
        super().__init__()
        # the dimension usually remains the same due to the requirement of residual connection, so we can use the same dimension for input and output
        self.attention = MultiHeadSelfAttention(dim, dim, dim, num_heads, causal)
        self.feed_forward = FeedForward(dim, hidden_dim, dim)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x):
        attn_out = self.attention(x)
        x = self.norm1(x + attn_out)
        ff_out = self.feed_forward(x)
        x = self.norm2(x + ff_out)
        return x

Meanwhile, it is also very important to note that the attention mechanism is permutation-equivariant, meaning that it does not take into account the order of the tokens in the input sequence. This means that the attention mechanism treats the input tokens as a set, rather than a sequence, and does not capture any information about the relative positions of the tokens. This can be problematic for tasks that require an understanding of the order of the tokens, such as language modeling or machine translation.

To address this, position embeddings are usually added to the input tokens to provide information about their positions in the sequence. Position embeddings can be learned or fixed, and they are typically added to the input tokens before they are fed into the attention mechanism. Two common parameter-free ways to introduce order are sinusoidal positional embeddings and rotary positional embeddings (RoPE).

Sinusoidal positional embeddings

Sinusoidal positional embeddings is the default method used in the original Transformer paper. It uses a combination of sine and cosine functions of different frequencies to encode the position of each token in the sequence. The idea is that each dimension of the positional embedding corresponds to a sinusoid of different wavelength, allowing the model to learn to attend to tokens based on their relative positions.

For $p$th token in the sequence, the positional embedding is a vector of the same dimension as the token embedding, where each dimension is computed using a sine or cosine function of the position $p$ and a frequency that depends on the dimension index. The even-numbered and odd-numbered terms can be written as

where $\omega_i=10000^{-2i/d}$. It is normally added once, before the first attention block. It is important to note that the sinusoidal positional embeddings can also be used to compute relative positions between tokens. This is because the sine and cosine functions are periodic, and thus the difference between two positions can be represented as a linear combination of the sine and cosine functions. Specifically, for two tokens at positions $p$ and $p+\Delta$, we have $PE(p+\Delta)=A_\Delta PE(p)$ means that a relative shift can be represented by a position-independent linear transformation.

Below is a simple implementation of sinusoidal positional embeddings in PyTorch and a visualization of the embeddings:

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 4096, base: float = 10_000.0):
        super().__init__()
        position = torch.arange(max_seq_len, dtype=torch.float32)[:, None]
        inv_freq = base ** (-torch.arange(0, dim, 2, dtype=torch.float32) / dim)
        angles = position * inv_freq[None, :]

        pe = torch.zeros(max_seq_len, dim)
        pe[:, 0::2] = angles.sin()
        pe[:, 1::2] = angles[:, :pe[:, 1::2].shape[1]].cos()
        self.register_buffer("pe", pe, persistent=False)

    def forward(self, x):
        # x: [batch, sequence, dim]
        if x.size(1) > self.pe.size(0):
            raise ValueError("sequence is longer than max_seq_len")
        return x + self.pe[:x.size(1)].to(dtype=x.dtype, device=x.device)

embedding = SinusoidalPositionalEncoding(dim=128, max_seq_len=100)
embedding=embedding.pe.numpy()
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 5))
plt.imshow(embedding, aspect='auto', cmap='coolwarm')
plt.colorbar()
plt.xlabel('Embedding Dimension')
plt.ylabel('Position')
plt.show()

The above plot also indicates an important property of the sinusoidal positional embeddings: small dimension index captures the short-term dependencies while large dimension index captures the long-term dependencies. This is because the sine and cosine functions with small frequencies (large dimension index) vary slowly, while those with large frequencies (small dimension index) vary rapidly. For example, when the dimension index is small, a small change in position will result in a large change in the embedding, allowing the model to capture fine-grained differences between nearby tokens. However, since the sine and cosine functions are periodic, the embeddings will eventually repeat for large enough positions, which may limit the ability of the model to capture long-term dependencies. While for large dimension index, the embeddings only show large changes for large changes in position, allowing the model to capture long-term dependencies.

Rotary positional embeddings (RoPE)

RoPE is a more recent method that encodes the position of each token by rotating the query and key vectors in the attention mechanism. More and more recent transformer models have adopted RoPE as their default positional encoding method.

RoPE does not add anything to the input. After projecting the input to queries and keys, it groups each attention head into two-dimensional coordinate pairs and rotates each pair by an angle determined by its position. Suppose one attention head produces a query vector with head dimension $d_h=8$:

RoPE groups these eight numbers into four 2D vectors:

Each pair can be treated mathematically like an $(x,y)$ coordinate:

RoPE then rotates this point according to token position $p$:

For example, if one pair is

and the rotation angle is $90^\circ$, it becomes

Its length is unchanged:

So RoPE does not add new dimensions or change the vector size. It only mixes the two feature values within every pair.

Different pairs use different rotation speeds:

and so forth. Some pairs rotate quickly and capture short positional differences, and others rotate slowly and retain information over longer distances, similar to sinusoidal embeddings. If we define a rotation matrix $R$ for $p$th query and $s$th key vector, we can write the above as

Because rotation matrices satisfy $R_p^T R_s=R_{s-p}$ (can be derived from trigonometric transformations or definition of rotations), the attention dot product becomes

Thus, although each vector is rotated using an absolute position, the positional part of the attention wight depends naturally on the relative displacement $s-p$. RoPE is applied to $Q$ and $K$ in every attention layer, but normally not to $V$.

Below is a simple implementation of RoPE in PyTorch:

def rope_cos_sin(seq_len: int, head_dim: int, device, base: float = 10_000.0):
    # Compute trigonometric values in float32 for numerical stability.
    assert head_dim % 2 == 0, "RoPE requires an even head dimension"
    inv_freq = base ** (-torch.arange(0, head_dim, 2, device=device).float() / head_dim)
    positions = torch.arange(seq_len, device=device).float()
    angles = positions[:, None] * inv_freq[None, :]
    return angles.cos(), angles.sin()


def apply_rope(x, cos, sin):
    # x: [batch, heads, sequence, head_dim]; rotate adjacent channel pairs.
    original_dtype = x.dtype
    pairs = x.float().reshape(*x.shape[:-1], -1, 2)
    even, odd = pairs[..., 0], pairs[..., 1]
    cos, sin = cos[None, None, :, :], sin[None, None, :, :]
    rotated = torch.stack((even * cos - odd * sin,
                           even * sin + odd * cos), dim=-1)
    return rotated.flatten(-2).to(original_dtype)


class PositionalMultiHeadSelfAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int, causal: bool = False, use_rope: bool = False):
        super().__init__()
        assert dim % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.causal = causal
        self.use_rope = use_rope
        if use_rope:
            assert self.head_dim % 2 == 0
        self.qkv = nn.Linear(dim, 3 * dim)
        self.output = nn.Linear(dim, dim)

    def forward(self, x, mask=None):
        batch, seq_len, dim = x.shape
        qkv = self.qkv(x).reshape(batch, seq_len, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)
        q, k, v = (t.transpose(1, 2) for t in (q, k, v))

        if self.use_rope:
            cos, sin = rope_cos_sin(seq_len, self.head_dim, x.device)
            q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        if self.causal:
            causal_mask = torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool).triu(1)
            scores = scores.masked_fill(causal_mask, float("-inf"))
        if mask is not None:
            scores = scores.masked_fill(~mask.bool(), float("-inf"))

        attended = scores.softmax(dim=-1) @ v
        attended = attended.transpose(1, 2).contiguous().reshape(batch, seq_len, dim)
        return self.output(attended)


class PositionalTransformerBlock(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, num_heads: int,
                 position_type: str = "rope", max_seq_len: int = 4096,
                 causal: bool = False):
        super().__init__()
        if position_type not in {"sinusoidal", "rope"}:
            raise ValueError("position_type must be 'sinusoidal' or 'rope'")

        self.position = (SinusoidalPositionalEncoding(dim, max_seq_len)
                         if position_type == "sinusoidal" else nn.Identity())
        self.attention = PositionalMultiHeadSelfAttention(
            dim, num_heads, causal=causal, use_rope=(position_type == "rope")
        )
        self.feed_forward = FeedForward(dim, hidden_dim, dim)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x, mask=None):
        x = self.position(x)
        x = self.norm1(x + self.attention(x, mask))
        x = self.norm2(x + self.feed_forward(x))
        return x

Modern Transformers often prefer RoPE because it makes relative position directly affect attention through the rotation of query and key vectors, while sinusoidal embeddings only add absolute-position information to token representations and the model must learn how to extract useful relative information. But RoPE may not be suitable for all tasks, and sinusoidal embeddings are still widely used in many models.