Transformer Attention Math: Q/K/V, Softmax Weights, Masks, and KV Cache
AI Learning Project
Transformer Attention Math: Q/K/V, Softmax Weights, Masks, and KV Cache
A three-token hand calculation of scaled dot-product attention, covering Q/K/V, softmax weights, masks, multi-head attention, and KV cache.
AI Learning Project

Transformer Attention Math: Q/K/V, Softmax Weights, Masks, and KV Cache

Reading info

Level: Intermediate Reading time: 14 min
  • Transformer
  • Attention
  • QKV
  • KV Cache
Open knowledge map

The core concept of a Transformer’s Self-Attention mechanism can be intuitively understood as follows: every token in a sequence uses its own Query vector to evaluate the Key vectors of all other tokens. This process determines how “attention” should be distributed across the sentence, and finally, this weight distribution is used to compute a weighted sum of the information vectors (Values). Its mathematical expression is remarkably concise—just one line of code—but it conceals a staggering amount of engineering depth and model training nuances.

In this article, we will start from first principles. We will manually calculate Scaled Dot-Product Attention for a 3-token sequence and dive deep into Q/K/V projections, Softmax saturation, Causal Masking, Multi-head mechanisms, and the notorious inference performance beast known as the KV Cache.

1. The Core Mathematical Formula

This is the foundational equation of the Large Language Model era:

Attention(Q, K, V) = softmax((Q @ K^T) / sqrt(d_k)) @ V

Breaking it down:

  • Q @ K^T produces an attention score matrix of size `[seq_len, seq_len]`. Because it’s a dot product, it measures the “similarity” or “affinity” between pairs of tokens in a high-dimensional space.
  • Why must we divide by sqrt(d_k)? Suppose Q and K have a dimension of `d_k = 4096`, and their elements follow an independent distribution with a mean of 0 and a variance of 1. The variance of their dot product will scale up to `4096`. A massive variance creates extreme score values (e.g., 100 vs -100). When these are passed through a Softmax function, it forces the gradients to near-zero (vanishing gradients), a problem known as “Softmax saturation.”

2. Architectural Diagram: Data Flow and Dimensions


graph TD
    Input[Input Sequence: B, L, d_model] --> WQ(W_q Linear)
    Input --> WK(W_k Linear)
    Input --> WV(W_v Linear)
    
    WQ --> Q[Q: B, h, L, d_k]
    WK --> K[K: B, h, L, d_k]
    WV --> V[V: B, h, L, d_v]
    
    Q --> Dot[Dot Product: Q @ K^T]
    K --> Dot
    Dot --> Scale[Scale by 1/sqrt(d_k)]
    Scale --> Mask[Apply Causal Mask]
    Mask --> Softmax[Softmax along dim L]
    Softmax --> AttentionWeights[Attention Weights: B, h, L, L]
    
    AttentionWeights --> MatMulV[MatMul with V]
    V --> MatMulV
    
    MatMulV --> Context[Context Output: B, h, L, d_v]
    Context --> Concat[Concat Heads: B, L, d_model]
    Concat --> Out[W_o Linear]

3. Practical Demonstration: Self-Attention in NumPy

Formulas can be abstract. Let’s run a highly educational, purely Pythonic NumPy implementation. Imagine our input sequence consists of just 3 tokens (e.g., “AI”, “needs”, “math”) with an embedding dimension of 4:

import numpy as np

# 1. Simulate Q, K, V Matrices (Seq_len=3, d_k=4)
# Representing tokens: "AI", "needs", "math"
Q = np.array([
    [ 1.0,  0.5, -0.2,  0.1],  # AI
    [-0.5,  1.2,  0.8, -0.4],  # needs
    [ 0.2, -0.1,  1.5,  0.9]   # math
])
K = np.array([
    [ 0.8,  0.4, -0.3,  0.0],
    [-0.2,  1.0,  0.5, -0.1],
    [ 0.1, -0.2,  1.1,  0.7]
])
V = np.array([
    [ 1.0,  0.0],
    [ 0.0,  1.0],
    [-1.0, -1.0]
])

d_k = Q.shape[1]

# 2. Calculate Dot-Product Scores and Scale
scores = (Q @ K.T) / np.sqrt(d_k)
print("Scaled Scores:\n", scores)

# 3. Causal Mask
# Mask out future positions to prevent information leakage (cheating)
mask = np.triu(np.ones((3, 3)), k=1)
scores[mask == 1] = -np.inf

# 4. Softmax Normalization
def softmax(x):
    e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
    return e_x / e_x.sum(axis=-1, keepdims=True)

weights = softmax(scores)
print("Attention Weights:\n", np.round(weights, 3))

# 5. Value Weighting (Context Output)
context = weights @ V
print("Context Output:\n", context)

If you run this code, you’ll observe that the first row (the word “AI”) assigns attention weights exclusively to itself. The third row (“math”) distributes its attention across the preceding two tokens. This perfectly demonstrates the essence of autoregressive models: they must synthesize historical context without peaking into the future.

4. What Does the Causal Mask Actually Change?

As demonstrated in the code above, in autoregressive generation tasks, if the model is currently predicting the 3rd token, it absolutely cannot “see” tokens 4 or 5. Right before the Softmax operation, we forcefully overwrite the upper triangular matrix of the attention scores to negative infinity (-inf). After passing through the Softmax, these specific weights are mathematically crushed to exactly 0. Therefore, the Mask does not delete tokens; rather, it performs a probabilistic cutoff, ensuring that illegal attention allocation is impossible.

5. An Engineer’s Perspective: The VRAM Killer and KV Cache

Real-World Insight: In textbooks, you see elegant matrix multiplication. But in industrial LLM deployment, what you see is a relentless stream of OOM (Out of Memory) exceptions.

During inference, large language models operate token-by-token. When generating token $t+1$, the K and V matrices for the previous $t$ tokens remain completely identical! If we were to naively multiply the full L x d_model matrices over and over, the computational waste would be catastrophic.

The KV Cache is the ultimate space-for-time tradeoff.

  • We allocate a massive contiguous block in the GPU VRAM to cache historically generated K and V tensors.
  • For every new token generated, we only compute $Q_{new}, K_{new}, V_{new}$ for that single token, and strictly append $K_{new}$ into the VRAM cache block.
  • **The Cost is Staggering:** For slightly longer contexts (even just 10K tokens), a single batch’s KV Cache footprint can easily exceed the memory required to load the model weights themselves! This is exactly why the industry invented PagedAttention (the core of vLLM), MQA (Multi-Query Attention), and GQA (Grouped-Query Attention)—they are all desperate engineering hacks designed to shrink the KV Cache memory footprint.

6. Shape and Mask Checks

The easiest attention bugs are the ones that still produce a tensor. First, verify that batched attention uses batch x heads x tokens x dim and transposes only the final two dimensions for Q @ K^T. Then verify that the causal mask describes query-token to key-token visibility and broadcasts over batch and head. Finally, confirm that softmax is applied over the key dimension, not over the query dimension.

For the three-token toy example, every attention row should sum to approximately 1, future positions should be zero after masking, and the context output should have the same final dimension as V. A heatmap is useful for debugging, but it is not causal proof of why a model predicted a token; it only shows how the current weighted read used Value vectors.

7. Attention Verification Matrix

Self-attention often fails with code that runs but has the wrong semantics. Use the matrix below to audit the NumPy toy example here and to transfer the checks to batched, multi-head, or cached inference implementations.

Check Correct evidence Common mistake
Score shape Q @ K.T produces a query-token by key-token matrix. Transposing the wrong dimensions and mixing batch or head into attention.
Scaling and softmax Scores are divided by sqrt(d_k); each row sums to about 1 over keys. Normalizing over queries, or skipping scaling and saturating attention early.
Causal mask Future positions are near 0 after softmax while historical positions remain visible. Reversing the mask so the current token sees the future but not the past.
KV cache Each new token appends only K_new and V_new; history is not recomputed. Recomputing all K/V every step, or letting cache length drift from position encoding.

8. Visualizations and Data Flow Summary

Scaled dot-product attention weight heatmap for three tokens
Each row represents the attention distribution of a Query token over all historical Key tokens. This is the classic Attention Heatmap. The reason the model “understands” human language is encoded entirely within these shifting gradient weights.

This mechanism may seem like simple linear algebra, but it currently supports the absolute frontier of global AI research. The next time your Transformer script crashes, your first instinct should always be: Print the shape of every single tensor, and trace the matrix multiplication on a piece of scrap paper.

Internal links

Related tutorials

Review foundation Convolution and Receptive Field Math

Compute convolution output size, receptive fields, channel mixing, and im2col layout.

 Continue next LLM, RAG, and Agent Security

Separate instructions from data and enforce tool permissions against indirect prompt injection.

 Learning path AI Learning Path

Read AI foundations, machine learning, CNN work, traffic defense, and AI security in order.

 Resources Resource Library

Download companion code, datasets, diagrams, model files, and notes.

 Topic hub Student AI Projects

Connect handwritten digits, browser demos, and CNN projects into a portfolio path.

Search questions

FAQ

Who is this article for?

This article is for readers who want an intermediate-level guide to Transformer Attention Math. It takes about 14 min and focuses on Transformer, Attention, QKV, KV Cache.

What should I read next?

The recommended next step is LLM, RAG, and Agent Security, so the article connects into a longer learning route instead of ending as an isolated note.

Does this article include runnable code or companion resources?

Yes. Use the run notes, resource cards, and download links on the page to reproduce the example or inspect the companion files.

How does this article fit into the larger site?

It is connected to the article context block, learning routes, resources, and project timeline so readers can move from concept to implementation.

Article context

AI Learning Project

A practical route from AI concepts to machine learning workflow, evaluation, neural networks, Python practice, handwritten digits, a CIFAR-10 CNN, adversarial traffic-defense notes, and AI security.

Level: Intermediate Reading time: 14 min
  • Transformer
  • Attention
  • QKV
  • KV Cache

Your next step

Continue: LLM Visualizer
Review the foundation Open resource
Other language version Transformer Attention 数学：Q/K/V、Softmax 权重、Mask 与 KV Cache
Share summary Transformer Attention Math

Hand-calculate Q/K/V scores, softmax weights, masks, multi-head structure, and KV cache.

Download share card Open share center

Companion resources

Leave a Reply

Project timeline

Published posts

  1. AI Basics Learning Roadmap Separate AI, machine learning, and deep learning before going into implementation details.
  2. Machine Learning Workflow Follow the practical path from data and features to training, prediction, and evaluation.
  3. Model Training and Evaluation Understand loss, overfitting, train/test splits, accuracy, recall, and F1.
  4. Neural Network Basics Move from perceptrons to activation, forward propagation, backpropagation, and training loops.
  5. Matrix Calculus for Neural Networks Derive dL/dW for y = Wx + b and verify it with finite differences.
  6. Backpropagation as a Computation Graph Trace local gradients through ReLU and softmax cross-entropy in a two-layer MLP.
  7. Gradient Descent and Optimizer Geometry Compare gradient descent, momentum, and Adam on a visible quadratic loss surface.
  8. Convolution and Receptive Field Math Compute convolution output size, receptive fields, channel mixing, and im2col layout.
  9. Transformer Attention Math Hand-calculate Q/K/V scores, softmax weights, masks, multi-head structure, and KV cache.
  10. Python AI Mini Practice Run a small scikit-learn classification task and read the experiment output.
  11. Handwritten Digit Dataset Basics Read train.csv, test.csv, labels, and the flattened 28 by 28 pixel layout before training the classifier.
  12. Handwritten Digit Softmax in C Follow the C implementation from logits and softmax probabilities to confusion matrices and submission export.
  13. Handwritten Digit Playground Notes See how the offline classifier was adapted into a browser demo with drawing input and probability output.
  14. CIFAR-10 Tiny CNN Tutorial in C Build and train a small convolutional neural network for CIFAR-10 image classification, then read its loss and accuracy output.
  15. High-Entropy Traffic Defense Notes Study encrypted metadata leaks, entropy, traffic classifiers, and a defensive Python chaffing prototype.
  16. AI Security Threat Modeling Build a defense map with NIST adversarial ML, MITRE ATLAS, and OWASP LLM risks.
  17. Adversarial Examples and Robust Evaluation Evaluate clean and perturbed accuracy with an FGSM-style digits experiment.
  18. Data Poisoning and Backdoor Defense Study poison rate, trigger behavior, attack success rate, and training pipeline controls.
  19. Model Privacy and Extraction Defense Measure membership inference signal and surrogate fidelity against a local toy model.
  20. LLM, RAG, and Agent Security Separate instructions from data and enforce tool permissions against indirect prompt injection.

Published resources

  1. Python AI practice code guide The article includes a runnable scikit-learn classification script.
  2. digit_softmax_classifier.c The C source for the handwritten digit softmax classifier.
  3. train.csv.zip Compressed handwritten digit training set with 42000 labeled samples.
  4. test.csv.zip Compressed handwritten digit test set with 28000 unlabeled samples.
  5. sample_submission.csv The official submission format example for checking the final output columns.
  6. submission.csv The prediction file generated by the current C project.
  7. digit-playground-model.json The compact softmax demo model and sample set used by the browser playground.
  8. digit-sample-grid.svg A small handwritten digit preview grid extracted from the training set.
  9. Handwritten digit project bundle Contains the source file, compressed datasets, submission files, browser model, and preview grid.
  10. cifar10_tiny_cnn.c source Single-file C tiny CNN with CIFAR-10 loading, convolution, pooling, softmax, and backpropagation.
  11. model_weights.bin sample weights Model weights generated by one local small-sample run.
  12. test_predictions.csv sample predictions Sample test prediction output from the CIFAR-10 tiny CNN.
  13. CNN project explanation PDF Companion explanation material for the CNN project.
  14. Virtual Mirror redacted code skeleton A redacted mld_chaffing_v2.py control-flow skeleton with secrets, node topology, and target lists removed.
  15. Virtual Mirror stress-test template A redacted CSV template for CPU, memory, peak threads, pulse rate, latency, and error measurements.
  16. Virtual Mirror classifier-evaluation template A CSV template for TP, FN, FP, TN, accuracy, precision, recall, F1, ROC-AUC, entropy, and JS divergence.
  17. Virtual Mirror resource notes Notes explaining why the public resources include only redacted code, test templates, and architecture context.
  18. AI Security Lab README Setup, safety boundaries, and quick-run commands for the AI Security series.
  19. AI Security Lab full bundle Includes safe toy scripts, result CSVs, risk register, attack-defense matrix, and architecture diagram.
  20. AI security risk register CSV risk register template for AI threat modeling and release review.
  21. AI attack-defense matrix Maps attack surface, toy demo, metric, and defensive control into one CSV table.
  22. AI Security Lab architecture diagram Shows threat modeling, robustness, data integrity, model privacy, and RAG guardrails.
  23. FGSM digits robustness script FGSM-style perturbation and accuracy-drop experiment for a local digits classifier.
  24. Data poisoning and backdoor toy script Demonstrates poison rate, trigger behavior, and attack success rate on digits.
  25. Model privacy and extraction toy script Outputs membership AUC, target accuracy, surrogate fidelity, and surrogate accuracy.
  26. RAG prompt injection guard toy script Uses a deterministic toy agent to demonstrate external-data demotion and tool-policy blocking.
  27. Deep Learning Math Lab README Setup commands, script entry points, generated outputs, and figure notes for the math series.
  28. Deep learning math full lab bundle Bundles NumPy scripts, CSV outputs, formula diagrams, loss contours, convolution figures, and attention heatmaps.
  29. Gradient check results CSV Stores MSE analytic gradients, finite-difference gradients, and error norms.
  30. Optimizer path CSV Step-by-step coordinates and loss for gradient descent, momentum, and Adam on a 2D quadratic.
  31. Attention weights CSV Scores, softmax weights, and context vectors for a three-token scaled dot-product attention example.
  32. Deep learning math figure set Includes matrix shapes, computation graphs, loss contours, convolution scans, and attention heatmaps.
  33. Deep learning math interactive visualizer Browser modules for gradient checking, optimizer paths, convolution output size, and attention heatmaps.
  34. Deep Learning topic share card A 1200x630 SVG card for sharing the Deep Learning / CNN topic hub.
  35. Machine Learning From Scratch share card A 1200x630 SVG card for the K-means, Iris, and ML workflow topic hub.
  36. Student AI Projects share card A 1200x630 SVG card for handwritten digits, C classifiers, and browser demos.
  37. CNN convolution scan animation An 8-second Remotion animation showing how a 3x3 convolution kernel scans an input and builds a feature map.

Current route

  1. AI Basics Learning Roadmap Learning path step
  2. Machine Learning Workflow Learning path step
  3. Model Training and Evaluation Learning path step
  4. Neural Network Basics Learning path step
  5. Matrix Calculus for Neural Networks Learning path step
  6. Backpropagation as a Computation Graph Learning path step
  7. Gradient Descent and Optimizer Geometry Learning path step
  8. Convolution and Receptive Field Math Learning path step
  9. Transformer Attention Math Learning path step
  10. LLM Visualizer Learning path step
  11. Python AI Mini Practice Learning path step
  12. Handwritten Digit Dataset Basics Learning path step
  13. Handwritten Digit Softmax in C Learning path step
  14. Handwritten Digit Playground Notes Learning path step
  15. CIFAR-10 Tiny CNN Tutorial in C Learning path step
  16. High-Entropy Traffic Defense Notes Learning path step
  17. AI Security Threat Modeling Learning path step
  18. Adversarial Examples and Robust Evaluation Learning path step
  19. Data Poisoning and Backdoor Defense Learning path step
  20. Model Privacy and Extraction Defense Learning path step
  21. LLM, RAG, and Agent Security Learning path step

Next notes

  1. Add more image-classification and error-analysis cases
  2. Turn common metrics into a quick reference
  3. Add more AI security defense experiment notes
Scroll down