Self-Attention

Self-attention lets every position in a sequence look at every other position and build a context-aware representation by taking a weighted average of value vectors, where the weights come from the compatibility between queries and keys.

Scaled Dot-Product Attention

Each token is projected into three vectors: a query $q$, a key $k$, and a value $v$. Stacking these over a sequence of length $T$ gives matrices $Q,K\in {\mathbb{R}}^{T\times d_k}$ and $V\in {\mathbb{R}}^{T\times d_v}$. Attention is defined as:

$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$

The score matrix $Q{K}^{\top }\in {\mathbb{R}}^{T\times T}$ holds a dot product for every query-key pair. Softmax over the last axis turns each row into a probability distribution, and multiplying by $V$ produces a convex combination of value vectors.

Why divide by $\sqrt{d_k}$?

If the components of $q$ and $k$ are independent with mean 0 and variance 1, then $q\cdot k={\sum }_{i=1}^{d_k}{q}_i{k}_i$ has variance $d_k$. Large-magnitude scores push softmax into saturated regions where gradients vanish. Scaling by $\sqrt{d_k}$ restores unit variance and keeps the softmax well conditioned.

Causal Masking

For autoregressive decoding a token must not attend to future positions. We add a mask $M$ to the scores before softmax, setting disallowed entries to $-\infty$ so their softmax weight becomes 0:

${\text{scores}}_{ij}=\frac{q_i\cdot k_j}{\sqrt{d_k}}+M_{ij},M_{ij}=\{\begin{array}{ll}0& j\le i\\ -\infty & j>i\end{array}$

Complexity

The score matrix and the softmax are both $O(T^2)$ in memory, and forming scores plus the context costs $O(T^{2}d)$ in compute. This quadratic dependence on sequence length is the main scaling bottleneck of attention and the motivation behind the KV Cache at inference time and many efficient-attention variants.

Multi-Head Attention

Rather than one attention with dimension $d_{\text{model}}$, we run $h$ heads in parallel, each of width $d_k=d_{\text{model}}/h$. Heads can specialize (some track syntax, some track long-range coreference), and their outputs are concatenated and projected:

$\text{MHA}(X)=[{\text{head}}_1;\dots ;{\text{head}}_h]W^O,{\text{head}}_i=\text{Attention}(X{W}_i^Q,X{W}_i^K,X{W}_i^V)$

Code: einsum and einops

import torch
import torch.nn.functional as F
from einops import rearrange
 
def scaled_dot_product_attention(q, k, v, causal=False):
    # q, k, v: (batch, seq, d_k)
    d_k = q.shape[-1]
    scores = torch.einsum('bqd,bkd->bqk', q, k) / (d_k ** 0.5)
    if causal:
        T = scores.shape[-1]
        mask = torch.triu(torch.ones(T, T, device=q.device, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(mask, float('-inf'))
    attn = F.softmax(scores, dim=-1)
    context = torch.einsum('bqk,bkd->bqd', attn, v)
    return context, attn

The multi-head version uses einops.rearrange to split the model dimension into heads, run attention per head, then merge back:

import torch.nn as nn
 
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)
 
    def forward(self, x, causal=False):
        # x: (b, t, d_model)
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
        # split heads: (b, t, (h d)) -> (b, h, t, d)
        q, k, v = (rearrange(t, 'b t (h d) -> b h t d', h=self.n_heads) for t in (q, k, v))
        d = q.shape[-1]
        scores = torch.einsum('bhqd,bhkd->bhqk', q, k) / (d ** 0.5)
        if causal:
            T = scores.shape[-1]
            mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
            scores = scores.masked_fill(mask, float('-inf'))
        attn = torch.softmax(scores, dim=-1)
        ctx = torch.einsum('bhqk,bhkd->bhqd', attn, v)
        # merge heads back: (b, h, t, d) -> (b, t, (h d))
        ctx = rearrange(ctx, 'b h t d -> b t (h d)')
        return self.out(ctx)

Shape discipline

Keeping named axes with einops (b h t d) makes the head split/merge unambiguous and removes error-prone transpose/view chains.

Related

Self-attention is the core operation of Transformers. Because it is permutation-equivariant, order information must be injected via Positional Encoding. Masked self-attention drives Decoder-Only Transformers, and caching past keys and values via the KV Cache avoids recomputing attention over the whole prefix at every generation step.