Attention Mechanism: QKV, Masks, and Multihead Attention

An attention mechanism computes a data-dependent weighted combination of Value vectors. Queries select among permitted Keys, the resulting scores become weights, and those weights determine which Values are mixed into each output. This is best understood as differentiable soft retrieval, not as a literal model of human attention.

The central formula is compact:

Attention(Q, K, V, M) = softmax((QK^T / sqrt(d_k)) + M) V

The difficult engineering work lies around that formula: tensor shapes, mask semantics, numerical stability, positional information, multi-head layout, memory use, inference state, and tests that prove the implementation did not silently attend to the wrong positions.

Key Takeaways

• A Query row produces one distribution over the permitted Key positions; the corresponding output is a weighted sum of Value rows.
• Self-attention, causal self-attention, and cross-attention use the same core operation but differ in where Q/K/V come from and which positions are visible.
• Q, K, and V are learned projections. Labels such as "what I seek" and "what I contain" are useful intuition, not fixed semantic roles.
• Dense attention forms pairwise interactions between target length L and source length S; its score work scales with L * S.
• FlashAttention changes data movement and intermediate storage while preserving exact attention semantics within numeric tolerances. Sparse or linear variants change the operation itself.
• Attention weights are diagnostic evidence about one computation, not a standalone causal explanation of the model's answer.

What Problem Attention Solves

Attention lets a model read from a set of representations without compressing that set into a single fixed vector first. In early neural machine translation, Bahdanau, Cho, and Bengio introduced a differentiable soft alignment so a decoder could use different encoder states for different output words. The 2017 Transformer paper then used scaled dot-product attention as the main mechanism for moving information within and between sequences.

This history matters because "attention" is broader than "Transformer self-attention":

Attention offers a short interaction path between distant positions and supports parallel computation across training tokens. It does not guarantee that every distant dependency is learned, and dense attention pays for all permitted pairs even when most are unhelpful.

QKV and Tensor Shapes

Q, K, and V are interfaces between matching and content transfer. Given target representations X_q and source representations X_kv, learned projections produce:

Q = X_q  W_Q
K = X_kv W_K
V = X_kv W_V

For batched multi-head attention, a useful shape contract is:

B is batch size, H is head count, L is target length, S is source length, and E is the Query/Key head dimension. Ordinary multi-head attention uses the same number of Query and KV heads. Multi-query and grouped-query attention use fewer KV heads and define how they are shared across Query heads.

For one Query row q_i, attention performs three steps:

1. compute compatibility logits q_i dot k_j / sqrt(d_k) for every allowed Key j;
2. normalize those logits across the Key axis;
3. return sum_j(weight_ij * v_j).

The output is a mixture of Values, not a mixture of raw input tokens. Later output projections, residual paths, normalization, feed-forward blocks, and additional layers continue transforming it.

Why divide by sqrt(d_k)

If Query and Key components are roughly zero-centered with unit variance, the variance of their dot product grows with d_k. Large-magnitude logits make softmax distributions extremely sharp and can produce weak gradients. Scaling by 1 / sqrt(d_k) keeps the logit scale more stable under the assumptions used in the original Transformer derivation.

This is an initialization argument, not a universal law about every trained representation. Production kernels may accumulate in a wider dtype, fuse operations, or apply model-specific scaling, so compare against the exact model and runtime.

Masking Is Part of the Contract

A mask controls which Key positions each Query may use. It must be applied to logits before softmax; zeroing weights afterward leaves the remaining row incorrectly normalized.

Three implementation failures are common:

1. Wrong polarity: one API uses true = allowed, another uses true = blocked.
2. Wrong axis: softmax runs over Query positions or heads instead of the Key axis.
3. Fully masked rows: all logits become negative infinity, which can produce NaN instead of a meaningful output.

Current PyTorch scaled_dot_product_attention uses true to mean a position participates for a boolean attn_mask, while MultiheadAttention.key_padding_mask uses true to mean ignored. Its functional SDPA also applies dropout whenever dropout_p > 0, so callers must explicitly pass zero during evaluation. These are API contracts, not properties of the mathematical formula.

A Runnable Go Implementation

The following dependency-free program implements one attention head. It validates shapes, treats true as allowed, rejects fully masked rows, subtracts the row maximum before exponentiation, and exposes weights so tests can verify normalization and masking. It intentionally starts from Q/K/V; a full trainable layer must also learn projections and gradients.

package main

import (
	"errors"
	"fmt"
	"log"
	"math"
)

func matrixWidth(name string, matrix [][]float64) (int, error) {
	if len(matrix) == 0 || len(matrix[0]) == 0 {
		return 0, fmt.Errorf("%s must not be empty", name)
	}
	width := len(matrix[0])
	for i, row := range matrix {
		if len(row) != width {
			return 0, fmt.Errorf("%s row %d has width %d, want %d", name, i, len(row), width)
		}
	}
	return width, nil
}

func scaledDotProductAttention(
	q, k, v [][]float64,
	allowed [][]bool,
) ([][]float64, [][]float64, error) {
	qWidth, err := matrixWidth("q", q)
	if err != nil {
		return nil, nil, err
	}
	kWidth, err := matrixWidth("k", k)
	if err != nil {
		return nil, nil, err
	}
	vWidth, err := matrixWidth("v", v)
	if err != nil {
		return nil, nil, err
	}
	if qWidth != kWidth {
		return nil, nil, errors.New("q and k head dimensions must match")
	}
	if len(k) != len(v) {
		return nil, nil, errors.New("k and v must have the same source length")
	}
	if len(allowed) != len(q) {
		return nil, nil, errors.New("mask must have one row per query")
	}

	scale := math.Sqrt(float64(qWidth))
	output := make([][]float64, len(q))
	weights := make([][]float64, len(q))

	for i := range q {
		if len(allowed[i]) != len(k) {
			return nil, nil, fmt.Errorf("mask row %d has the wrong source length", i)
		}

		logits := make([]float64, len(k))
		maxLogit := math.Inf(-1)
		for j := range k {
			if !allowed[i][j] {
				logits[j] = math.Inf(-1)
				continue
			}
			var dot float64
			for d := range q[i] {
				dot += q[i][d] * k[j][d]
			}
			logits[j] = dot / scale
			if logits[j] > maxLogit {
				maxLogit = logits[j]
			}
		}
		if math.IsInf(maxLogit, -1) {
			return nil, nil, fmt.Errorf("query row %d is fully masked", i)
		}

		weights[i] = make([]float64, len(k))
		var denominator float64
		for j := range logits {
			if allowed[i][j] {
				weights[i][j] = math.Exp(logits[j] - maxLogit)
				denominator += weights[i][j]
			}
		}
		output[i] = make([]float64, vWidth)
		for j := range weights[i] {
			weights[i][j] /= denominator
			for d := range v[j] {
				output[i][d] += weights[i][j] * v[j][d]
			}
		}
	}
	return output, weights, nil
}

func rowSum(row []float64) float64 {
	var total float64
	for _, value := range row {
		total += value
	}
	return total
}

func main() {
	q := [][]float64{{1, 0}, {0, 1}}
	k := [][]float64{{1, 0}, {0, 1}}
	v := [][]float64{{2, 1}, {0, 3}}
	causal := [][]bool{{true, false}, {true, true}}

	output, weights, err := scaledDotProductAttention(q, k, v, causal)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("row sums: %.1f %.1f\n", rowSum(weights[0]), rowSum(weights[1]))
	fmt.Printf("masked future weight: %.1f\n", weights[0][1])
	fmt.Printf("output shape: %d x %d\n", len(output), len(output[0]))
}

// row sums: 1.0 1.0
// masked future weight: 0.0
// output shape: 2 x 2

This code is a teaching reference, not a replacement for a fused framework kernel. It demonstrates the invariants that production tests should keep:

• every unmasked row sums to approximately one;
• blocked positions receive zero weight;
• output shape is [query length, value dimension];
• mismatched dimensions and fully masked rows fail explicitly;
• large logits do not overflow the exponential.

Self-Attention, Cross-Attention, and Causal Attention

These names describe data flow and visibility, not different scoring formulas.

Self-attention uses one representation stream for Q, K, and V. Encoder-style self-attention often allows every real token to read every other real token. The operation is permutation-equivariant unless the model introduces position through positional embeddings, rotary transforms, relative bias, or another positional mechanism.

Causal self-attention is self-attention with a rule that position t cannot read positions after t. During training, a triangular mask allows all target positions to be processed in parallel without leaking future targets. During autoregressive inference, the model processes newly generated positions while a KV cache retains prior Keys and Values.

Cross-attention takes Queries from one stream and Keys/Values from another. In encoder-decoder translation, decoder positions query encoder outputs. In multimodal models, text tokens may query image or audio representations. Cross-attention therefore has target length L and source length S, which need not match.

The broader Transformer architecture guide explains how attention interacts with residual connections, normalization, feed-forward blocks, and positional information.

Multi-Head, Multi-Query, and Grouped-Query Attention

Multi-head attention applies separate learned projections per head, computes attention in parallel, concatenates head outputs, and mixes them through an output projection. Heads provide multiple representation subspaces, but no architecture rule assigns one head to syntax and another to coreference. Specialization, redundancy, and interactions must be measured after training.

head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)
MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O

Decoder inference introduced a second head-count question: how many distinct Key/Value heads must be stored and read?

The MQA paper targets the memory-bandwidth cost of incremental decoding. The GQA paper defines the intermediate design and evaluates uptraining under its stated model and compute conditions. Neither is a lossless serving switch for an arbitrary MHA checkpoint; the architecture and weights must support it.

Complexity and Efficient Attention

Dense attention does not have one meaningful complexity number without dimensions. With target length L, source length S, model width d_model, and head width d_k:

• Q/K/V projections include terms proportional to sequence length times d_model squared;
• score and Value aggregation include terms proportional to L * S * d_k per head;
• materializing all score or probability tensors costs O(B * H * L * S) memory;
• causal autoregressive decoding reuses historical K/V projections but still reads an expanding cache.

This separates four optimization families that are often conflated:

FlashAttention is IO-aware exact attention: tiling and recomputation reduce reads and writes between GPU memory levels. FlashAttention-2 improves GPU work partitioning. These papers report results for specific hardware, models, shapes, and baselines; they do not justify a universal speedup claim. A framework may also fall back to another backend when dtype, shape, mask, or hardware requirements are not met.

For long-context architecture choices, distinguish exact kernels from algorithms that alter the visible attention pattern. The long-context decision guide covers when more model context is actually the right system choice.

Attention Weights Are Diagnostics, Not Causal Proof

An attention heatmap shows the normalized routing coefficients for one head, layer, input, and forward pass. It does not include:

• the magnitude and direction of Value vectors;
• the head's output projection;
• residual-stream information that bypasses the head;
• interactions with other heads and later layers;
• whether changing the attended token changes the prediction.

The paper Attention is not Explanation found weak relationships between attention weights and some feature-importance measures, and showed that substantially different attention distributions could produce similar predictions in its evaluated NLP models. The responsible conclusion is narrow: raw weights alone do not establish causal importance.

Use heatmaps to find hypotheses, mask errors, positional patterns, or unusual concentration. Test explanations with interventions such as token or head ablation, activation patching, controlled counterfactuals, and task-level output measurements.

How to Test an Attention Implementation

An implementation should pass semantic tests before performance tuning:

1. Shape contract: verify output and score shapes for L != S, multiple batches, and multiple heads.
2. Mask contract: assert that blocked probabilities are zero and allowed rows sum to one.
3. Causality: change a future token and prove earlier causal outputs do not change.
4. Cross-attention provenance: change the source stream and verify only the intended path changes.
5. Numerical stability: test large logits, low-precision inputs, and fully masked rows.
6. Reference parity: compare a small deterministic case against a trusted framework implementation.
7. Gradient check: for training code, compare analytical and finite-difference gradients on tiny tensors.
8. Backend audit: record which fused kernel actually ran and why fallback occurred.
9. Workload benchmark: report latency and peak memory by batch size, sequence shape, dtype, mask type, and hardware.
10. Task quality: confirm that a faster kernel or changed attention pattern preserves the target metric.

Do not use a random heatmap as proof that the model learned a linguistic relationship. Do not report a kernel benchmark without tensor shapes and hardware. Do not change mask conventions at an API boundary without an explicit conversion test.

Common Questions

Is attention the entire Transformer?

No. A Transformer block also includes learned projections, output mixing, residual connections, normalization, feed-forward layers, positional mechanisms, and often architecture-specific routing or gating. Attention is the token-interaction mechanism inside that larger system.

Why can self-attention train in parallel but generation remains sequential?

During teacher-forced training, the complete target sequence is available and a causal mask prevents future leakage, so all positions can be evaluated together. During autoregressive generation, token t+1 depends on the sampled or selected token t, so output positions still arrive sequentially even when each step uses parallel matrix operations.

How should I choose the number of heads?

Treat head count as a model and kernel design parameter. It must satisfy the implementation's divisibility and layout constraints, and it changes head dimension, memory behavior, parallelism, and learned capacity. Compare candidate architectures under the same training budget and serving workload rather than assuming more heads are better.

Does FlashAttention make dense attention linear-time?

No. It avoids materializing the full score matrix in high-bandwidth memory and reduces IO, but exact dense attention still evaluates pairwise Query-Key interactions. Its memory and wall-clock behavior improve without turning the arithmetic into a universally linear-time algorithm.

Where should I continue learning?

Read Attention Mechanism for a compact definition, Transformer Architecture for the complete block, and KV Cache for autoregressive inference state. The original Bahdanau attention paper, Transformer paper, and current framework documentation should remain the source of truth for historical and API claims.