Self-Attention

The self-attention operation itself is permutation-invariant; it treats an input sequence as an unordered set. To recover sequential order, positional encodings (either fixed sinusoidal or learned embeddings) are added to the input tokens. This allows the model to be aware of the absolute and relative positions of tokens, which is crucial for understanding language structure and image patch layout.

• Absolute Positional Encoding: Adds a unique vector for each position index.
• Relative Positional Encoding: Encodes the distance between token pairs, often more effective for longer sequences.

The standard self-attention mechanism has O(n²) time and memory complexity, where n is the sequence length. This is because it computes a compatibility score for every pair of tokens in the sequence to create an n x n attention matrix. This becomes a significant bottleneck for long sequences (e.g., long documents, high-resolution images).

This limitation has driven research into efficient attention variants like:

• Sparse Attention: Only computes scores for a subset of token pairs.
• Linearized Attention: Reformulates the operation to approximate softmax attention with linear complexity.
• FlashAttention: An IO-aware algorithm that optimizes memory access for dramatic speedups on GPU hardware.

Unlike recurrent neural networks (RNNs) which process data sequentially, self-attention provides direct, single-step connections between any two tokens in a sequence, regardless of distance. This allows it to capture long-range dependencies without the vanishing/exploding gradient problems common in RNNs.

For example, in the sentence "The award, which was given for lifetime achievement, sat proudly on the shelf," self-attention can directly link the verb "sat" to the distant subject "award," bypassing the intervening relative clause. In images, a patch in the top-left corner can directly attend to a patch in the bottom-right.

The attention weights are not fixed parameters but are dynamically computed based on the content of the input sequence for each forward pass. This is known as content-based addressing. The model learns to assign high attention weights to tokens that are semantically relevant to the current token being processed.

• Query, Key, Value Vectors: Each token is projected into three vectors. The compatibility between a token's Query and another token's Key determines the attention weight applied to that token's Value.
• This allows the model to focus on different parts of the input for different contexts, enabling nuanced, context-aware representations.

All attention scores for a sequence can be computed simultaneously in a single matrix multiplication step. This makes self-attention highly parallelizable and well-suited for modern hardware accelerators like GPUs and TPUs, which excel at batched matrix operations.

This contrasts sharply with the inherent sequentiality of RNNs, where the computation for timestep t depends on the hidden state from timestep t-1, limiting parallel processing. The parallel nature of self-attention is a primary reason for the transformer architecture's superior training efficiency on large datasets.

Instead of performing a single attention function, the transformer uses Multi-Head Attention. This mechanism runs multiple, independent self-attention operations ("heads") in parallel, each with its own learned projection matrices.

• Each head can learn to focus on different types of relationships (e.g., syntactic, semantic, long-distance, local).
• One head might attend to a verb's direct object, while another attends to coreferential pronouns.
• The outputs of all heads are concatenated and linearly projected to form the final output, allowing the model to jointly attend to information from different representation subspaces.

In sequence-to-sequence models like the original Transformer, self-attention allows the encoder to create a context-rich representation of the source sentence. Each word's representation becomes a weighted sum of all other words, directly capturing syntactic and semantic relationships (e.g., subject-verb agreement across long distances). The decoder uses masked self-attention on its own output to ensure predictions are based only on previously generated words, enabling autoregressive generation.

Decoder-only models like GPT use causal self-attention (or masked self-attention). This mechanism allows each token in a sequence to attend only to previous tokens, preventing information flow from the future. This is critical for autoregressive text generation. The model builds a contextual representation for the next word by analyzing all preceding words, capturing narrative flow, topic coherence, and long-range dependencies within paragraphs or documents.

When applied to images, self-attention operates on a sequence of image patches. A Vision Transformer (ViT) treats each 16x16 pixel patch as a token. The self-attention mechanism computes relationships between all patches, allowing the model to understand global image structure. For example, it can learn that a patch containing a 'wheel' is highly relevant to a patch containing a 'car body', regardless of their spatial separation in the image, enabling holistic scene understanding.

While self-attention operates within a single modality, its extension, cross-attention, is key for multimodal models. In architectures like Flamingo or BLIP, a language model's decoder uses cross-attention layers where the query vectors come from the text tokens, and the key/value vectors come from encoded visual features. This allows the language model to 'attend to' and ground its reasoning in specific regions of an image, enabling tasks like Visual Question Answering (VQA) and detailed image captioning.

AlphaFold2's Evoformer module heavily utilizes self-attention and cross-attention to reason about protein sequences and multiple sequence alignments (MSAs). Self-attention operates on the residue representations, allowing the model to infer long-range interactions between amino acids that may be far apart in the linear sequence but close in the folded 3D structure. This direct modeling of all-pair interactions was a breakthrough over previous iterative methods.

Models like Codex apply self-attention to source code, treating it as a sequence of tokens. The mechanism learns complex, long-range dependencies inherent in code, such as:

• The connection between a function definition and its call site many lines later.
• The matching of opening and closing brackets.
• The scope of variables. By attending to all relevant parts of the code context, the model can generate syntactically correct and semantically plausible code completions and translations.

Cross-attention (or encoder-decoder attention) is a mechanism where the queries come from one sequence, and the keys and values come from another. This is fundamental to sequence-to-sequence tasks. In a transformer decoder, the self-attention layer is masked to prevent attending to future tokens, while a subsequent cross-attention layer allows the decoder to attend to the entire encoded input sequence from the encoder.

• Critical for machine translation, summarization, and multimodal tasks like image captioning.
• In vision-language models, it allows language tokens to attend to visual features.

Scaled dot-product attention is the specific mathematical operation that computes self-attention in the original transformer. It is defined as: Attention(Q, K, V) = softmax((QK^T) / sqrt(d_k)) V. The scaling factor (1 / sqrt(d_k), where d_k is the key dimension) is crucial. It prevents the dot products from growing large in magnitude (which pushes the softmax into regions with extremely small gradients), ensuring stable training.

• The core computational unit for attention.
• The scaling is a key detail for effective optimization.

Causal attention (or masked self-attention) is a variant where each position in a sequence can only attend to previous positions and itself. This is enforced by applying a mask (typically an upper-triangular matrix of -inf) to the attention logits before the softmax. It ensures the autoregressive property necessary for generative models, preventing the model from 'cheating' by looking at future tokens during training or generation.

• Foundational for decoder-only models like GPT.
• Enables sequential, left-to-right generation of text, code, or other sequences.

Positional encoding is the method of injecting information about the order of tokens into a transformer model. Since self-attention is inherently permutation invariant—it treats a sequence as a set—positional encodings are added to the token embeddings to give the model a sense of sequence order. The original transformer uses fixed, sinusoidal functions. Alternatives include learned positional embeddings or relative position biases.

• Without it, the model cannot distinguish between 'dog bites man' and 'man bites dog'.
• A critical component for processing sequential data with self-attention.

A long-range dependency is a relationship between elements in a sequence that are far apart. Capturing these is a key challenge for sequence models. Recurrent Neural Networks (RNNs) struggle due to vanishing gradients. Self-attention fundamentally solves this by computing relationships between all pairs of positions in a single layer with a constant number of operations, regardless of distance. This allows a token at the start of a paragraph to directly influence a token at the end.

• Enables modeling of complex grammatical structures and coreference resolution (e.g., linking a pronoun to a noun many sentences earlier).
• A primary motivation for the transformer architecture.