Transformer Self Attention and Key Value Cache Linear Algebra

The Core Problem of Sequence Modeling

Before the introduction of the Transformer architecture by Vaswani and colleagues in 2017, sequence processing relied on recurrent neural networks (RNNs) and long short term memory networks (LSTMs). These architectures processed text step by step, updating a hidden state vector $h_t$ at each time step.

While intuitive, sequential recurrence created two major bottlenecks:

1. Information loss over long distances: In long sequences, early tokens become diluted through repeated matrix multiplications. Even with gating mechanisms, context beyond a few hundred tokens degraded quickly.
2. Sequential execution prevents hardware scaling: Computing step $t$ strictly required the completion of step $t - 1$. GPUs, which achieve their immense speed by running thousands of arithmetic cores concurrently, sat idle waiting for sequential feedback loops.

The Transformer eliminated recurrence entirely, replacing it with an operation that computes pairwise relationships between all tokens across the sequence at once: Multi Head Self Attention.

Scaled Dot Product Attention Mathematical Formulation

At the core of the Transformer lies the scaled dot product attention equation:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V$$

Let us examine each component of this equation through linear algebra.

Step 1: Linear Projections for Queries, Keys, and Values

Suppose we have an input sequence of $T$ tokens, where each token is represented by a hidden embedding vector of dimension $d_{\text{model}}$. We can stack these vectors into an input matrix $X$ of shape $(T, d_{\text{model}})$.

To allow the model to extract different contextual features, we multiply $X$ by three learned parameter projection matrices:

• Query Projection Matrix $W^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$: Produces queries $Q = X W^Q$ with shape $(T, d_k)$. Queries represent what each token is looking for in other tokens.
• Key Projection Matrix $W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$: Produces keys $K = X W^K$ with shape $(T, d_k)$. Keys represent what information each token offers to others.
• Value Projection Matrix $W^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$: Produces values $V = X W^V$ with shape $(T, d_v)$. Values represent the actual content transmitted when a query and key match.

In standard architectures, $d_k = d_v = d_{\text{model}} / h$, where $h$ is the number of parallel attention heads.

Step 2: The Attention Score Matrix

Next, we compute the dot product between every query and every key:

$$S = Q K^T$$

Because $Q$ has dimensions $(T, d_k)$ and $K^T$ has dimensions $(d_k, T)$, the resulting matrix $S$ has dimensions $(T, T)$. The scalar at coordinate $(i, j)$ represents the raw compatibility score between the $i$-th token and the $j$-th token.

Step 3: Variance Scaling by $\sqrt{d_k}$

Why do we divide the dot product matrix by $\sqrt{d_k}$?

Assume that the components of $q_i$ and $k_j$ are independent random variables with zero mean and unit variance ($\sigma^2 = 1$). Their dot product is a sum of $d_k$ independent products:

$$q_i \cdot k_j = \sum_{m=1}^{d_k} q_{i,m} k_{j,m}$$

The mean of this sum remains zero, but its variance grows directly with dimension:

$$\text{Var}(q_i \cdot k_j) = d_k \times \text{Var}(q_{i,m}) \times \text{Var}(k_{j,m}) = d_k$$

In modern models where $d_k$ is 64 or 128, the magnitude of these dot products can grow very large (values of 20 to 50). When large values enter the softmax function, the outputs push into regions with near zero gradients, causing gradient vanishing during training.

Dividing by $\sqrt{d_k}$ normalizes the variance back to 1.0, preserving healthy gradient flow throughout backpropagation.

Step 4: Softmax and Causal Masking

In decoder models like GPT and Llama, future tokens must be concealed so that token $i$ cannot look ahead at token $i+1$. We achieve this by adding a causal mask matrix $M$ where upper triangular positions are set to negative infinity ($-\infty$):

$$A = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} + M\right)$$

Applying softmax across each row normalizes all scores into positive probabilities that sum to 1.0. For any position with $-\infty$, $e^{-\infty} = 0$, completely blocking information leakage from subsequent tokens.

Step 5: Value Aggregation

Finally, we multiply the normalized attention matrix $A$ of shape $(T, T)$ by the value matrix $V$ of shape $(T, d_v)$:

$$\text{Output} = A V$$

The result has shape $(T, d_v)$. Each row in this output is a weighted convex combination of all value vectors in the sequence, weighted by how relevant each past token was to the current query.

Multi Head Attention and Tensor Concatenation

Single head attention allows tokens to focus on only one type of relationship at a time. Multi Head Attention (MHA) splits the query, key, and value vectors into $h$ independent subspaces.

For each head $i \in {1, \dots, h}$:

$$\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$$

The outputs of all $h$ heads are concatenated horizontally:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O$$

A final linear projection matrix $W^O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$ blends the diverse perspectives (such as grammatical agreement, coreference, factual retrieval, and syntax) back into the shared model representation.

Autoregressive Generation and Key Value Cache Mechanics

During training, large language models process full prompt sequences in parallel using causal masking. But during generation (inference), the model produces text one token at a time:

1. Token $T$ generates token $T+1$.
2. In the next step, tokens $1$ through $T+1$ must be evaluated to predict token $T+2$.

Notice that for all past tokens $1$ to $T$, their keys and values never change. Recomputing $Q, K, V$ for every preceding token on every single generated word would result in $O(N^2)$ computational waste for an $N$ token response.

To eliminate this redundant work, inference engines store past key and value vectors in memory. This data structure is known as the Key Value Cache (abbreviated as KV Cache).

KV Cache Memory Footprint Formula

The memory required to store the KV cache for a running model can be calculated using this formula:

$$\text{Memory}{\text{KV}} = 2 \times n{\text{layers}} \times n_{\text{heads_kv}} \times d_{\text{head}} \times \text{seq_len} \times \text{batch_size} \times \text{bytes_per_param}$$

Let us calculate the memory requirement for a standard 70B parameter model with FP16 precision (2 bytes per value):

• Layers ($n_{\text{layers}}$): 80
• Key Value Heads ($n_{\text{heads_kv}}$): 8 (using Grouped Query Attention)
• Dimension per head ($d_{\text{head}}$): 128
• Context length ($\text{seq_len}$): 8,192 tokens
• Batch size: 1

$$\text{Memory} = 2 \times 80 \times 8 \times 128 \times 8,192 \times 1 \times 2 \text{ bytes} \approx 2.68 \text{ Gigabytes}$$

For a batch of 16 concurrent users with 32,000 tokens of context, the KV cache alone requires over 170 GB of video RAM, exceeding the memory needed to store the model weights themselves.

Modern Optimizations: Multi Query and Grouped Query Attention

To reduce this immense memory footprint, modern model architectures avoid storing independent key and value heads for every query head:

• Multi Head Attention (MHA): 32 query heads, 32 key heads, 32 value heads. Maximum capacity, highest memory cost.
• Multi Query Attention (MQA): 32 query heads, but only 1 shared key head and 1 shared value head. Reduces KV cache memory by 96%, but slightly diminishes model nuance.
• Grouped Query Attention (GQA): Divides query heads into groups (for example, 8 groups of 4 query heads) where each group shares one key and value head. Used in Llama 3, DeepSeek, and Mistral, GQA delivers 99% of MHA accuracy with an 87% reduction in KV cache memory.

Summary and Practical Takeaway

Self attention provides the expressive mathematical engine that allows modern transformers to capture contextual relationships across thousands of tokens. By decomposing inputs into queries, keys, and values, normalizing with $\sqrt{d_k}$, and caching past states during generation, language models deliver fast, coherent text synthesis.

You can observe these 3D tensor flows, embedding spaces, and layer normalization transformations in real time inside our WebGL 3D Transformer Visualizer.