Self-Attention

The self-attention mechanism, in its pure form, is permutation invariant. This means that if you shuffle the order of the input tokens, the computed attention weights between any two specific tokens remain identical, as the operation is simply a set computation. This property is why positional encodings must be explicitly added to the input embeddings; without them, the model cannot distinguish the sequence 'dog bites man' from 'man bites dog'. This contrasts sharply with recurrent neural networks, which inherently process data sequentially.

Unlike convolutional neural networks, which build a global view by stacking many local-filter layers, a single self-attention layer provides a global receptive field. Every output token is a direct, weighted combination of every input token. This allows the model to establish long-range dependencies between distant tokens, such as a pronoun at the end of a paragraph and its antecedent at the start, without the information passing through many intermediate steps. This direct path is a key enabler for capturing complex syntactic and semantic structures.

The weights of a self-attention layer are not fixed after training, unlike the static kernels of a convolutional network. Instead, the attention weights are dynamically computed as a function of the input itself via the Query, Key, Value (QKV) projections. This allows the model to flexibly determine which tokens are relevant to each other based on the specific content of the current input sequence. The same model can focus on syntactic relationships in one sentence and semantic relationships in another.

A major practical advantage over recurrent architectures is that self-attention is highly parallelizable. In an RNN, the computation for time step t must wait for the hidden state from time step t-1. In self-attention, the attention scores between all pairs of tokens in a sequence can be computed simultaneously as a series of matrix multiplications. This makes the architecture exceptionally efficient to train on modern GPU and TPU hardware, which are optimized for parallel matrix operations.

The computed attention weight matrix provides a degree of inherent interpretability. By visualizing which tokens a model attended to when generating a specific output, developers can gain insights into the model's reasoning process. For example, in a translation task, attention maps often align with word-to-word correspondences. In a language model, analyzing attention heads can reveal specialized functions, such as heads that focus on syntactic dependencies, coreference links, or semantic similarity.

The fundamental computational limitation of standard self-attention is its quadratic complexity. Both the memory and time required to compute the N x N attention score matrix scale with O(n²), where n is the sequence length. For a sequence of 1,000 tokens, the matrix has 1,000,000 entries. This becomes prohibitive for very long documents, making it the primary bottleneck for processing entire books or high-resolution images. This limitation has driven research into sparse attention and FlashAttention.

The fundamental mathematical operation: Attention(Q,K,V) = softmax(QK^T / √d_k)V. The dot product computes compatibility between Query and Key vectors. Scaling by √d_k prevents the softmax from entering regions of extremely small gradients when dimensionality is large. Without scaling, large dot products push softmax outputs toward one-hot vectors, killing gradient flow.

Instead of computing a single attention function, the model projects Q, K, V into h different subspaces, runs attention in parallel, concatenates outputs, and projects again. This allows the model to jointly attend to information from different representation subspaces at different positions. A single head averages attention; multiple heads can capture syntactic, semantic, and positional relationships simultaneously.

An N × N matrix where entry (i,j) represents the raw compatibility between token i's Query and token j's Key. After softmax normalization along the last dimension, each row becomes a probability distribution over all tokens. This matrix is the core interpretability artifact—visualizing it reveals which tokens the model considers relevant. The matrix is quadratic in sequence length, which is the root of the complexity bottleneck.

A masking mechanism that sets attention scores for future tokens to -∞ before softmax, forcing each position to attend only to itself and preceding tokens. This preserves the autoregressive property essential for text generation. Without causal masking, the model would cheat by peeking at future tokens during training. Implemented via an upper-triangular mask matrix applied to the attention scores.

Standard self-attention requires O(N²) memory and time for sequence length N. For a 32K token context, the attention matrix alone consumes ~4GB at FP16. This is the primary constraint on context window length. Solutions include:

• Sparse attention patterns (local windows, dilated strides)
• FlashAttention (IO-aware tiling)
• Linear attention approximations
• State space models as alternatives

An exact-attention algorithm that avoids materializing the full N×N attention matrix in GPU HBM. It tiles the computation into blocks that fit in SRAM, fusing the softmax reduction across tiles using online statistics. Result: 2-4× speedup and 10-20× memory reduction versus standard implementations. Now the default attention kernel in PyTorch and most LLM training frameworks.