Attention Is All You Need: From the Attention Formula to Transformer Architecture and Inference Optimization

The 2017 paper Attention Is All You Need proposed the Transformer. The paper used machine-translation sequence transduction as its main task, replacing the recurrent and convolutional structures common at the time with an attention-based encoder-decoder architecture. Its central claim is that dependencies between sequence positions can be modeled mainly through attention, without relying on recurrent state updates along time. This post explains how the Attention operator inside the Transformer works and why it became the central component of later large-model architectures.

The previous post, Attention Prerequisites: From Sequence Bottlenecks to Dot Products, Softmax, and the GPU Memory Wall, introduced the necessary background: a sequence should not be compressed into one fixed state, tokens need continuous vectors, dot products can turn vector compatibility into scores, Softmax converts a row of scores into non-negative weights that sum to 1, and GPU execution is constrained by HBM traffic. This post starts from those prerequisites and explains why the formula is written as $softmax(QK^T / \sqrt{d_k})V$. The goal is not to memorize the Attention formula, but to understand what each symbol in the formula means.

This post follows two tracks. The first half follows the paper’s core structure: Scaled Dot-Product Attention, masks, Multi-Head Attention, positional encoding, and the Transformer block. The second half discusses later decoder-only LLM practice around the same operator: KV Cache, MQA/GQA, FlashAttention, and PagedAttention. The paper gives the foundational Transformer and Attention form; modern inference optimizations address the execution cost of the same operator under long-context and online-serving workloads.

1. From a Task to One Information Read

The experiments in the Attention Is All You Need paper center on machine translation, but the attention function itself is not a translation-specific module. The same mechanism takes different read patterns in different tasks: in machine translation, the current decoder position reads relevant tokens from the source sentence; in an autoregressive language model, the current position reads relevant tokens from the existing context before predicting the next token. Whether the task is translation or continuation, Attention answers the same question: which positions should the current position read from, and by how much.

Consider this sentence:

Xiao Ming put the apple into the backpack because it was heavy.

When the model processes “it”, that token alone is not enough to determine the referent. The current position needs to look back at “Xiao Ming”, “apple”, “backpack”, and the later description “heavy” to decide which information is most relevant. This process is a soft retrieval operation: the current position asks a question, all candidate positions participate in matching, and each position contributes with a different weight.

The query, key, and value here are not manually annotated labels. They are intermediate vectors produced by projecting the current layer’s input representation through three trainable matrices. During training, backpropagation updates the embedding matrix, $W_Q$, $W_K$, $W_V$, and other parameters together, so these vectors gradually become useful for the final prediction objective. The semantics of QKV come from training, not from rules hard-coded into the architecture.

For this example, the query from “it” can be read as: “I need an earlier object that can be referred to by a pronoun and is compatible with a weight description.” The key exposed by each previous token is like a matching label: Xiao Ming looks like a person and subject, apple looks like an object and food, and backpack looks like an object and container. Query and Key compute compatibility, which decides which positions to read and how much to read from each one.

The matched positions do not contribute their keys; they contribute their values. The value for apple carries semantic content about “apple”, the value for backpack carries semantic content about “backpack”, and the value for Xiao Ming carries semantic content about the person. Key is used for addressing, while Value is used for content transfer; this is the core intuition behind QKV separation.

The model might assign weights like this:

Xiao Ming: 0.05
apple:     0.65
backpack:  0.25
others:    0.05

Then the new representation at “it” is not a hard choice of one word; it is a weighted sum over all values:

$$\begin{aligned} o_{\text{it}} &= 0.05 \cdot v_{\text{Xiao Ming}} + 0.65 \cdot v_{\text{apple}} + 0.25 \cdot v_{\text{backpack}} + 0.05 \cdot v_{\text{others}} \end{aligned}$$

The output is still a vector, but it has absorbed information from the most relevant context positions. In the single-query view, an attention function scores the query against every key, uses Softmax to turn scores into weights, and then takes a weighted sum of the values. Scaled Dot-Product Attention is the matrix form of this information read: $QK^T$ scores all position pairs, $\sqrt{d_k}$ controls scale, Softmax produces weights, and $V$ aggregates content.

2. QKV Role Separation

The prerequisite post used $XX^T$ to show that if a sequence is written as $X \in \mathbb{R}^{N \times D}$, dot products can compute pairwise token relations at once. But using $XX^T$ directly for attention is still too restrictive: the same token representation would have to decide whether two positions are related and also serve as the content sent back after a related position is selected. These two jobs do not have to depend on the same information. A pronoun may use syntactic or semantic cues to match a previous noun, while the matched noun can pass richer entity semantics and contextual information to the next layer. QKV separation lets the model learn “how to decide what to read” separately from “what to send back after it is read”.

This separation is done with three trainable linear projections:

$$Q = XW_Q,\qquad K = XW_K,\qquad V = XW_V$$

Here $W_Q$, $W_K$, and $W_V$ are model parameters. A useful reading is: $Q$ (Query) represents what the current token is looking for, $K$ (Key) represents what kind of query each token can match, and $V$ (Value) represents what content that token contributes when read. Query and Key decide which positions to read and how much to read from them; Value defines what content is read back.

The diagram below shows that $Q$, $K$, and $V$ all come from the same input $X$, but through three different linear projections:

For one sequence, the common shape convention is:

$$X \in \mathbb{R}^{N \times D},\quad Q,K \in \mathbb{R}^{N \times d_k},\quad V \in \mathbb{R}^{N \times d_v}$$

Here $N$ is the number of tokens, $D$ is the input hidden dimension, $d_k$ is the scoring-space dimension, and $d_v$ is the value-space dimension. Linear projection does not change the number of tokens; it changes each token’s representation dimension and role. Q, K, and V are not three new token sequences, but three representations of the same tokens for querying, matching, and content transfer.

3. Deriving Scaled Dot-Product Attention

The standard Scaled Dot-Product Attention formula is:

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

This formula can be decomposed into four steps: compute all token-pair scores with $QK^T$, divide by $\sqrt{d_k}$ to control score scale, apply row-wise Softmax to obtain reading weights, and multiply those weights by $V$ for weighted aggregation. Scaled Dot-Product Attention is matrix computation in four steps: score, scale, normalize, and aggregate. The data flow is shown below.

The first step constructs the relation matrix:

$$S = QK^T,\qquad S_{ij}=q_i \cdot k_j$$

Row $i$ contains the scores produced by token $i$ as a query against all key tokens; column $j$ corresponds to token $j$ as a candidate information source. A larger $S_{ij}$ means position $i$ tends to read more from position $j$. $QK^T$ writes all pairwise matching relations inside the sequence as an $N \times N$ matrix.

The second step is scaling:

$$\tilde{S} = \frac{S}{\sqrt{d_k}}$$

The prerequisite post showed that if the dimensions of $q$ and $k$ are independent, zero-mean, and unit-variance under a simplified assumption, the variance of $q \cdot k$ grows with $d_k$. Without scaling, larger dimensions make score gaps larger, and Softmax more easily approaches one-hot saturation with weaker gradients. Dividing by $\sqrt{d_k}$ controls score scale; it does not change the matching semantics of Attention.

The third step applies Softmax row by row:

$$A = softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)$$

The matrix $A$ is still $N \times N$. Each row is a weight distribution, and $A_{ij}$ means the share of information position $i$ reads from position $j$. These weights are aggregation shares for the current layer and head; they are not a guarantee of human-level causal interpretability. Softmax converts arbitrary real-valued scores into non-negative weights that sum to 1, so those weights can be used in the later weighted sum.

The fourth step multiplies by $V$:

$$O = AV$$

For token $i$, this expands to:

$$o_i = \sum_j A_{ij}v_j$$

The output vector at position $i$ is therefore not only derived from itself; it is a weighted sum over all value vectors. Which positions contribute more is determined by $A_{ij}$, and what content they contribute is determined by $v_j$. The output here is not a new token id; it is the updated hidden state for the same input position $i$. The Attention output is a contextual representation for each position: the position stays the same, but its vector has absorbed sequence information through weights.

Putting the four steps together makes each term’s job explicit: $Q$ asks, $K$ exposes matchable indexes, $QK^T$ builds relation scores, $\sqrt{d_k}$ controls scale, Softmax produces reading weights, and $V$ supplies the content being aggregated. Together, these roles define one differentiable information read.

4. Mask Constraints

If we only look at $softmax(QK^T / \sqrt{d_k})V$, every token can read every position. Real tasks have boundaries: padding is not real text, and the original paper’s decoder-side masked self-attention cannot read future target tokens. Masks constrain which tokens each position can read, so Attention follows the task rules.

The first kind is the padding mask. During training or batched inference, sequences in one batch may have different lengths, so shorter sequences are padded to a shared length $N$. These padding positions exist only to make tensor shapes rectangular; real tokens should not read from them. A common implementation adds $-\infty$ or a very negative value to scores at padding positions before Softmax. Padding masks solve the problem of “positions that exist in the tensor shape but not in the sequence semantics”.

The second kind is the causal mask, or the look-ahead mask used in the paper’s decoder self-attention. When the machine-translation decoder generates the target sequence, position $i$ can only see itself and previous target tokens, not future target tokens; decoder-only language models use the same visibility constraint for next-token prediction. Therefore row $i$ only allows columns $j \le i$:

allowed positions:

1 0 0 0
1 1 0 0
1 1 1 0
1 1 1 1

This is the lower-triangular mask. Causal masks write the next-token-prediction constraint into the visible range of Attention.

Masks are applied before Softmax, not simply zeroed after Softmax. If illegal positions are blocked before Softmax, the remaining legal positions are automatically renormalized; if weights are zeroed after Softmax, their row sum is no longer 1 unless another normalization step is added. Masks should take effect before Softmax, so weights are redistributed only among legal positions.

Training can also use loss masks, but they are separate from Attention masks. An Attention mask controls which tokens the current position can read; a loss mask controls which positions participate in the loss. In SFT, for example, the user prompt is usually context, while the assistant response is the training target: the response can read the prompt, but prompt positions do not have to contribute to the loss. Attention masks control information flow; loss masks control the training objective.

5. Multi-Head Attention

Single-head Attention has one scoring space and one value space. Language contains many relations at once: local collocations, syntactic dependencies, coreference, long-range semantic links, and formatting boundaries can all matter. If all of them are squeezed into one Attention space, representation capacity is limited. Multi-Head Attention lets the model read different relations in multiple learned subspaces in parallel.

Multi-Head Attention first projects the inputs into each head’s $Q,K,V$ subspaces, then concatenates all head outputs:

$$\begin{aligned} MultiHead(Q,K,V) &= Concat(head_1,\ldots,head_h)W^O \\ head_i &= Attention(QW_i^Q,KW_i^K,VW_i^V) \end{aligned}$$

Each head has its own $W_i^Q$, $W_i^K$, and $W_i^V$, so each head can learn a different matching and content-extraction pattern.

The base Transformer in the paper uses $d_{\text{model}}=512$ and $h=8$, so each head has $d_k=d_v=64$. This does not make the total representation eight times larger; it splits the 512-dimensional representation into 8 parallel subspaces and then mixes the result back to $d_{\text{model}}$ through the output projection. Multi-Head Attention in the paper both increases the number of representation subspaces and keeps each head’s computation small.

Shape-wise, a common setup is $D = h \cdot d_{\text{head}}$. Implementations typically reshape $B \times N \times D$ into a tensor with an explicit head dimension, run Attention in each $d_{\text{head}}$ subspace, and concatenate results back to $D$. The Linear projections, parallel heads, Concat, and output projection in the diagram correspond to this tensor organization. Multi-Head Attention is first an expressiveness design, and it also maps naturally to batched GPU matrix computation.

What individual heads learn should not be over-interpreted as stable human rules. Some heads may attend more locally, some may focus on separators, and some may track longer dependencies, but that is not an explicit semantic assignment in the architecture; it is shaped by data and the training objective. Attention heads can form useful reading patterns, but a single head’s weights are not a reliable natural-language explanation.

6. Positional Encoding and RoPE

The preceding sections mostly used self-attention as the running case. Self-attention means that $Q,K,V$ are all projected from the same token sequence: $Q=XW^Q$, $K=XW^K$, and $V=XW^V$; encoder self-attention and decoder masked self-attention both fit this pattern, while cross-attention usually takes $Q$ from the decoder and $K,V$ from the encoder. Self-attention is a use case of the Attention operator, not a separate formula.

The Attention operator here means the scoring, scaling, normalization, and weighted-sum computation $softmax(QK^T/\sqrt{d_k})V$. It receives numeric $Q,K,V$ matrices; a row index in the matrix certainly corresponds to a token position in the program, but that row index is not itself a numeric feature inside $QK^T$. The computation preserves row order, but “which position” and “how far apart” do not automatically become information the model can use.

If a permutation matrix $P$ reorders all input rows, pure self-attention reorders its output in the same way: $SelfAttention(PX)=P\,SelfAttention(X)$. This is permutation equivariance: the program knows the row order, while the operator itself computes weights only from matches between row vectors. Positional mechanisms turn external row indices or relative distances into numeric signals and inject them into token representations or Attention scores.

The original Attention Is All You Need uses sinusoidal positional encoding: it generates a positional vector $p_t$ for position $t$ and adds it to the token embedding $e_t$:

$$x_t = e_t + p_t$$

The same token therefore enters later layers with a different representation at different positions. Absolute positional encoding writes “where this token is” directly into the input representation.

$p_t$ is a $d_{\text{model}}$-dimensional vector; $PE(pos,2j)$ and $PE(pos,2j+1)$ below are only two scalar components of that vector, at dimensions $2j$ and $2j+1$:

$$\begin{aligned} PE(pos,2j) &= \sin\left(\frac{pos}{10000^{2j/d_{\text{model}}}}\right) \\ PE(pos,2j+1) &= \cos\left(\frac{pos}{10000^{2j/d_{\text{model}}}}\right) \end{aligned}$$

Stacking the scalar components across all dimensions gives the full positional vector $p_{pos}$. Here $pos$ is the position and $j$ indexes a pair of dimensions; even dimensions use sine, odd dimensions use cosine, and different dimensions correspond to different wavelengths. Because the function can be evaluated at any $pos$, the model does not need to learn a separate parameter for every maximum position; the paper also chose this form in part to support extrapolation beyond training lengths. The key point of sinusoidal positional encoding is that fixed functions turn position into a vector that can be added to token representations.

Modern decoder-only LLMs often use RoPE (Rotary Position Embedding). RoPE is not the positional encoding from the original Transformer paper, but it is useful to explain alongside Attention because it acts on $Q$ and $K$, rather than being simply added to input embeddings. A two-dimensional rotation is:

$$R_\theta = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix}$$

Inside each query vector and each key vector, RoPE groups adjacent dimensions into two-dimensional planes: $(0,1)$, $(2,3)$, and so on. The query at position $m$ rotates these two-dimensional components by $m\theta$, while the key at position $n$ rotates them by $n\theta$. When they are dotted together, the shared absolute-position part cancels out, and the positional relation is controlled by the angle difference $(n-m)\theta$. The key intuition of RoPE is that after position becomes rotation angle, the Attention score naturally contains the relative distance between query and key.

Long context touches this mechanism directly: when sequence length goes beyond the common training range, query-key relative distances grow, and the corresponding rotation angles move into less familiar regimes for the model. Attention scores come from $QK^T$, while RoPE directly changes the geometry of $Q$ and $K$, so discussions of extrapolation, context extension, and KV Cache often return to RoPE.

7. Attention Inside the Transformer Block

The original Transformer is a sequence-to-sequence architecture. The encoder reads source tokens and produces contextual representations; the decoder reads the target sequence shifted right by one position, uses a causal mask so each position can only see itself and earlier tokens, and uses encoder-decoder attention to read encoder outputs. The encoder turns the input sequence into readable representations, while the decoder generates the output sequence conditioned on the generated prefix and those input representations.

Add & Norm in the diagram means residual connection followed by layer normalization. For a sublayer $Sublayer$, the paper uses the post-LN form: Add & Norm first adds the sublayer output back to the original representation, then normalizes the sum.

$$LayerNorm(x + Sublayer(x))$$

Here Add means the sublayer output is added back to the original representation as a correction, and Norm normalizes the resulting hidden vector. Add & Norm lets each sublayer update the representation without discarding it, while keeping numerical scale more stable across stacked layers.

The position-wise feed-forward network applies the same two-layer MLP independently to every position: It is a per-token nonlinear transform, not the part that exchanges information across positions.

$$FFN(x)=\max(0,xW_1+b_1)W_2+b_2$$

In the base Transformer, $d_{\text{model}}=512$ and the inner dimension is $d_{ff}=2048$. This sublayer does not directly mix different tokens; it adds nonlinear transformation capacity to each token’s own hidden vector. Attention handles cross-position communication, while the position-wise FFN processes the representation already gathered at each position.

Linear + Softmax at the top of the diagram is the decoder output layer. The final decoder hidden vector is projected to one logit per vocabulary token, then Softmax turns those logits into a next-token probability distribution. Embeddings plus positional encodings turn discrete tokens into input representations, while Linear + Softmax turns output representations back into vocabulary distributions.

One core reason the paper put self-attention inside the encoder and decoder is shorter dependency paths and higher parallelism. With sequence length $n$ and representation dimension $d$, self-attention has per-layer complexity $O(n^2d)$, but the number of sequential operations is $O(1)$ and the maximum path length between any two positions is also $O(1)$; RNNs have $O(n)$ sequential operations and $O(n)$ maximum path length. The advantage of Attention is not only global visibility; it turns sequence dependencies from time recursion into parallel matrix computation.

Layer type	Complexity per layer	Sequential operations	Maximum path length
Self-attention	$O(n^2d)$	$O(1)$	$O(1)$
Recurrent	$O(nd^2)$	$O(n)$	$O(n)$
Convolutional	$O(knd^2)$	$O(1)$	$O(\log_k n)$

8. Autoregressive Inference and KV Cache

KV Cache is not a model structure introduced by the original paper; it is an implementation technique for the same Attention operator during modern autoregressive inference. Training can process an entire sequence in parallel, but inference generates autoregressively: predict the next token from the existing context, append it, then predict the following token. Inference is therefore usually split into prefill and decode. The value of KV Cache only becomes clear after separating prefill from decode.

Prefill processes the user’s prompt. The model sends the whole prompt through every Attention layer at once; under the causal mask, each position can only read itself and earlier positions, and each layer computes and stores $K,V$ for every prompt token. When predicting the first new token, the output at the last prompt position produces the logits; that output has already used the last prompt position’s $Q$ to attend over the visible $K,V$ inside the prompt. Prefill first processes the existing context and builds the historical $K,V$ cache that later decode steps will read repeatedly.

Decode starts after the first new token has been generated. Each step takes the token generated by the previous step as the current input and computes its $Q,K,V$ in each layer; the new $K,V$ are appended to this request’s KV Cache, while the new $Q$ matches against all $K$ from the prompt through the current token and then aggregates the corresponding $V$ by those weights. Decode does not recompute historical tokens; it lets the current token’s Query read the context already cached for the same request.

The diagram below shows only one decode step. The current token computes its own $Q,K,V$; $K,V$ are written into cache, and $Q$ attends over the cached $K,V$.

Without KV Cache, step $t$ would recompute keys and values for the previous $t-1$ history tokens. But with fixed model parameters and the already processed context, those historical $K$ and $V$ tensors have already been computed and do not need to be regenerated from scratch at every step. KV Cache avoids repeated computation of historical Key/Value tensors during autoregressive decode.

For each decode step, KV Cache works in three steps:

1. Compute $Q,K,V$ for the current input token.
2. Append the current token’s $K,V$ to this request’s own KV Cache.
3. Use the current token’s $Q$ to compute Attention weights over all cached $K$, then use those weights to aggregate the corresponding $V$.

This is a space-time trade-off: store historical key/value tensors in HBM to avoid repeated projection and historical computation later. KV Cache makes decode faster, but its cost is GPU memory that grows with batch, layer count, and context length.

A rough size estimate for KV Cache is:

$$KV\ Cache \approx 2 \times B \times L \times N \times h_{kv} \times d_{\text{head}} \times bytes$$

The factor 2 stands for K and V; $B$ is batch size, $L$ is number of layers, $N$ is cached token count, $h_{kv}$ is the number of KV heads, $d_{\text{head}}$ is the dimension per head, and $bytes$ is the number of bytes per value. KV Cache cost is not a constant; it scales with concurrency, context length, and model architecture.

This explains why inference services are often constrained by memory capacity and bandwidth. Long context, high concurrency, many layers, and many heads all enlarge KV Cache; even though each decode step adds only one token, it repeatedly reads historical K/V. The bottleneck in autoregressive inference is often not “can we compute the current token,” but “can we store and read the history fast enough”.

9. KV Cache Compression and Memory Optimization

MQA, GQA, FlashAttention, and PagedAttention in this section are not part of the 2017 paper; they are later engineering developments that address memory bottlenecks when Attention is used for long-context and online inference workloads. In standard Multi-Head Attention, every query head has its own K/V head; with many heads, KV Cache becomes large. MQA (Multi-Query Attention) lets multiple query heads share one K/V set, while GQA (Grouped-Query Attention) groups query heads and shares one K/V set within each group. The core goal of MQA/GQA is to trade off model quality, KV Cache size, and inference throughput.

Let $h$ be the number of query heads and $h_{kv}$ be the number of KV heads. Standard MHA usually has $h_{kv}=h$; MQA is close to $h_{kv}=1$; GQA sits between the two. Smaller $h_{kv}$ reduces the corresponding factor in the KV Cache formula and lowers memory and bandwidth pressure, but sharing K/V can reduce expressiveness. GQA is common because it keeps more representational capacity than MQA while using less KV Cache than standard MHA.

FlashAttention addresses a different problem. A naive implementation of standard Attention creates an $N \times N$ score matrix and an $N \times N$ probability matrix, and long sequences make those intermediate matrices expensive to read and write through HBM. FlashAttention does not change the mathematical result; it tiles $Q$, $K$, and $V$ closer to fast on-chip storage and uses online softmax to avoid writing the full attention matrix back to HBM. FlashAttention optimizes the data movement path; it is not an approximation of Attention.

The intuition behind online softmax is that Softmax needs the maximum and normalization denominator for a row, but tiled computation cannot see the whole row at once. FlashAttention maintains the current maximum and normalization statistics for each row; when a new block arrives, it updates those statistics and rescales previous partial results. This lets tiled computation produce the same result as full Softmax while reducing HBM round trips.

PagedAttention focuses on KV Cache management in inference serving. In online serving, requests have different lengths, arrive at different times, and finish at different times; reserving one contiguous memory block per request can cause fragmentation and waste. PagedAttention borrows the idea of virtual memory paging: split KV Cache into fixed-size blocks, and use a block table to map logical sequence positions to physical GPU memory blocks. FlashAttention optimizes I/O inside one Attention computation, while PagedAttention optimizes KV Cache memory management under high-concurrency inference.

These techniques look separate, but they all center on the same problem: make the limited GPU memory hierarchy serve more useful tokens. MQA/GQA changes the number of KV heads, FlashAttention changes how intermediate matrices are read and written, and PagedAttention changes how KV Cache is allocated and reused. The main line of modern Attention optimization is reducing repeated computation, memory footprint, and HBM traffic without changing the core semantics.

Wrapping Up

Attention can now be understood in three layers. At the mathematical layer, $QK^T$ scores matches, $\sqrt{d_k}$ controls scale, Softmax normalizes into reading weights, and $V$ supplies the content being aggregated. Scaled Dot-Product Attention is differentiable information retrieval, not a disconnected string of matrix symbols.

At the architectural layer, masks constrain visibility, Multi-Head Attention provides multiple matching subspaces, positional mechanisms inject order, and the Transformer block uses Residual Connection, LayerNorm, and MLPs to stack Attention into a deep network. Attention is the core operator, but Transformer capability comes from how Attention combines with the other sublayers.

At the systems layer, KV Cache supports autoregressive inference, MQA/GQA reduce KV cost, FlashAttention reduces HBM round trips, and PagedAttention improves memory management under high-concurrency serving. Attention in large models must both express dependencies mathematically and fit the GPU memory hierarchy and online inference workload.