Self-Attention in LLMs Cheatsheet

What is Self-Attention?

Self-attention is the core mechanism that lets a transformer model decide, for every token in a sequence, how much to “attend to” (weigh the information from) every other token in that same sequence — including itself — when building its representation. It’s what replaced the recurrence of RNNs and the fixed local windows of convolutions: instead of processing tokens strictly in order or only looking at nearby neighbors, self-attention lets any token directly look at any other token, regardless of distance, and learn how relevant each one is.

The word “self” distinguishes it from the earlier “cross-attention” used in encoder-decoder models, where one sequence (the decoder) attends to a different sequence (the encoder’s output). In self-attention, queries, keys, and values all come from the same sequence — the sequence is attending to itself.

This mechanism is the central idea introduced in the 2017 paper “Attention Is All You Need,” and it’s the computational backbone of essentially every modern LLM (GPT, Llama, Claude, Gemini, and the rest).

The Core Intuition

For each token, self-attention asks three questions, encoded as three learned vectors:

• Query (Q): “What am I looking for?” — represents the current token’s request for relevant information.
• Key (K): “What do I contain?” — represents what each token (including itself) has to offer, as a label that can be matched against queries.
• Value (V): “What information do I actually pass along?” — the actual content that gets mixed into the output once a token is deemed relevant.

A token’s query is compared against every token’s key to produce a relevance score (via a dot product). Those scores are turned into weights (via softmax), and the output for that token is a weighted sum of all the value vectors, weighted by how relevant each other token’s key was to this token’s query. Every token in the sequence does this simultaneously, so the whole sequence transforms into a new sequence of representations, each one now informed by (a weighted mixture of) every other token.

A commonly used analogy: think of it like a library search. The query is your search terms, the keys are the titles/tags on every book, the values are the actual book contents. You compare your search terms against every book’s tags to decide how much attention each book deserves, then you read a blend of the books weighted by that relevance.

The Math: Scaled Dot-Product Attention

Given an input sequence represented as a matrix of token embeddings, three learned weight matrices (W_Q, W_K, W_V) project each token’s embedding into its Query, Key, and Value vectors:

Q = X · W_Q
K = X · W_K
V = X · W_V

Attention is then computed as:

Attention(Q, K, V) = softmax( (Q · Kᵀ) / √d_k ) · V

Breaking this down step by step:

1. Q · Kᵀ — dot product every query against every key, producing a matrix of raw relevance scores (one row per token, one column per token it’s comparing against).
2. / √d_k — scale down by the square root of the key dimension. Without this, dot products of high-dimensional vectors can grow very large, which pushes softmax into regions with extremely small gradients. This single division is what makes it “scaled” dot-product attention rather than plain dot-product attention.
3. softmax(...) — normalize each row of scores into a probability distribution that sums to 1, so the result is interpretable as “how much weight does each other token get.”
4. ... · V — use those weights to compute a weighted sum of the value vectors, producing the final output representation for each token.

Multi-Head Attention

Rather than computing a single attention pass, transformers split Q, K, and V into several smaller “heads” that each run scaled dot-product attention independently, with their own learned projection matrices, and then concatenate the results back together and pass them through one more learned projection.

MultiHead(Q, K, V) = Concat(head_1, ..., head_h) · W_O
where head_i = Attention(Q·W_Q_i, K·W_K_i, V·W_V_i)

The point of doing this in parallel, smaller subspaces rather than one large one is that different heads can specialize — empirically, different heads in a trained transformer tend to learn different kinds of relationships (e.g. one head might track syntactic dependencies like subject-verb agreement, another might track coreference like which noun a pronoun refers to, another might attend mostly to the immediately preceding token). A single large attention computation can’t easily represent several distinct relational patterns at once in this way.

Causal (Masked) Self-Attention

Autoregressive, decoder-only LLMs (the GPT/Llama/Claude style architecture that generates text one token at a time) use a specific variant called causal self-attention, or masked self-attention. Because the model must only ever use information available “so far” when predicting the next token, each token is prevented from attending to any token that comes after it in the sequence — only to itself and tokens before it.

This is implemented by adding a mask to the raw attention scores before the softmax step: positions that should be invisible get set to negative infinity, so after softmax their weight becomes exactly zero. This produces the characteristic lower-triangular attention pattern visualized as a staircase: token 1 can only see token 1; token 2 can see tokens 1-2; token 5 can see tokens 1-5; and so on.

Contrast this with the bidirectional self-attention used in encoder-only models like BERT, where every token can attend to every other token in both directions — appropriate for a model whose job is to build a representation of a complete, already-known sequence (e.g. for classification or masked-word prediction), rather than to generate text one token at a time.

Why Self-Attention Beats RNNs and CNNs for This

• Long-range dependencies: An RNN has to carry information through every intermediate timestep sequentially, so relevant information from far earlier in a sequence can fade or get overwritten (the classic vanishing-gradient problem). Self-attention connects any two tokens directly in a single step, regardless of distance.
• Parallelization: RNNs must be processed one timestep at a time because each step depends on the previous one’s hidden state, making them inherently sequential and slow to train. Self-attention computes all tokens’ representations simultaneously (it’s just matrix multiplications), which parallelizes extremely well on GPUs/TPUs.
• Flexible receptive field: CNNs need many stacked layers to let information from distant positions reach each other (since each convolution only sees a local window). Self-attention gives every layer a full, direct view of the whole sequence from the start.

The tradeoff is cost: naive self-attention computes a full N×N matrix of pairwise scores for a sequence of length N, so its compute and memory cost grows quadratically with sequence length — this is the central reason long context windows are expensive, and the motivation behind the efficiency techniques below.

Self-Attention in the Transformer Block

Self-attention doesn’t operate alone — each transformer layer wraps it with normalization and a feed-forward network, plus residual (skip) connections that let gradients flow more easily through many stacked layers:

x = x + MultiHeadSelfAttention(LayerNorm(x))
x = x + FeedForward(LayerNorm(x))

(Exact ordering of normalization relative to the sub-layers — “pre-norm” vs. “post-norm” — varies by model family, but the residual-connection-around-attention pattern is close to universal.)

Positional Information

A subtle but important point: the scaled dot-product formula above treats the input as an unordered set of tokens — nothing in the attention computation itself encodes word order. Position information has to be injected separately, commonly through:

• Absolute positional embeddings — a learned or fixed vector added to each token’s embedding based on its position index (original Transformer, early GPT models).
• Rotary Position Embeddings (RoPE) — rotates the Query and Key vectors by an angle that depends on token position, so the dot product between two tokens’ Q and K naturally encodes their relative distance. Used by Llama, Mistral, and most current open-weight LLMs.
• ALiBi (Attention with Linear Biases) — adds a fixed penalty to attention scores that grows linearly with the distance between two tokens, biasing attention toward nearby tokens without any learned position embedding at all.

Efficiency Variants for Long Context

Because standard self-attention is quadratic in sequence length, a number of techniques have emerged to make long-context inference and training more tractable:

The KV Cache (Inference-Time Detail)

During autoregressive generation, recomputing Key and Value vectors for every previous token at every new generation step would be wasteful, since those tokens’ Keys and Values don’t change once computed (causal masking means earlier tokens never see later ones). So inference engines cache the Key and Value vectors for all previously generated tokens (the “KV cache”) and only compute Q, K, V for the newest token at each step, reusing the cached K/V for everything before it. This is a major reason generation speed and memory usage scale with context length, and it’s the direct motivation behind MQA and GQA above, which shrink how much K/V data needs to be cached per token.

Quick Reference: Minimal PyTorch-Style Pseudocode

import torch
import torch.nn.functional as F

def self_attention(x, W_q, W_k, W_v, causal=True):
    Q = x @ W_q          # (seq_len, d_k)
    K = x @ W_k          # (seq_len, d_k)
    V = x @ W_v          # (seq_len, d_v)

    d_k = Q.shape[-1]
    scores = (Q @ K.T) / (d_k ** 0.5)   # (seq_len, seq_len)

    if causal:
        mask = torch.triu(torch.ones_like(scores), diagonal=1).bool()
        scores = scores.masked_fill(mask, float('-inf'))

    weights = F.softmax(scores, dim=-1)
    return weights @ V   # (seq_len, d_v)

Created using Anthropic Claude. This covers the standard self-attention mechanism as broadly implemented across current LLMs; exact architectural choices (positional encoding scheme, attention variant, normalization placement) differ by model family — verify specifics against a given model’s technical report before relying on them.