The original Transformer learned to translate one sequence into another. GPT used causal self-attention to learn the probability of text itself. This change made ordinary text a source of training targets: each token supplies the answer for the preceding context.

The central transition is between two probability distributions:

$$p_\theta(y_{1:m}\mid x_{1:n}) =\prod_{t=1}^{m}p_\theta(y_t\mid y_{<t},x_{1:n})$$

and

$$p_\theta(x_{1:T}) =\prod_{t=1}^{T}p_\theta(x_t\mid x_{<t}).$$

The first describes translation conditioned on a source sequence. The second describes an autoregressive language model. Both can use the same basic attention equation.

This article follows that transition through explicit computations, then examines Adam, AdamW, attention optimizations, and KV caching. Original papers and technical reports support the historical claims and named methods. Numerical examples and algebraic derivations are worked examples for this article.

1. Notation: what each matrix contains

A token is an integer index into a vocabulary. It can represent a word, part of a word, punctuation, or another unit chosen by the tokenizer.

Symbol	Meaning
$T$	Number of tokens in a sequence
$d$	Width of a token representation
$d_f$	Width of the intermediate nonlinear representation
$L$	Number of repeated transformations
$h$	Number of query heads
$g$	Number of key/value heads
$d_k,d_v$	Widths of one key and one value
$N_v$	Vocabulary size
$B$	Batch size
$\theta$	All trainable parameters

Use row vectors throughout. Stack one row per token:

$$X=\begin{bmatrix}x_1\\\vdots\\x_T\end{bmatrix} \in\mathbb R^{T\times d}.$$

Here $x_i$ is a representation, while $x_{1:T}$ in a probability expression denotes the token sequence. Context distinguishes these two uses.

Matrix multiplication uses juxtaposition. The symbol $\odot$ means elementwise multiplication. Softmax acts across the keys in each query row. Normalization acts across the $d$ coordinates of each token independently.

Batch dimensions are omitted until the memory calculations. Each depth and each attention operation has its own parameters, even when subscripts are omitted. Unless stated otherwise, equations omit dropout and attention-projection biases to expose the main computation.

2. The original Transformer, expressed as equations

The 2017 Transformer uses an encoder and an autoregressive decoder. Its base configuration has $L=6$ transformations on each side, $d=512$, $h=8$, $d_k=d_v=64$, and $d_f=2048$. It uses sinusoidal positions, ReLU, and normalization after residual addition. These choices belong to the original architecture; later sections introduce different choices. See Attention Is All You Need.

2.1 Token indices become vectors with positions

Let $E\in\mathbb R^{N_v\times d}$ be the embedding matrix. For token index $a_i$, the initial representation is

$$X_i^{(0)}=\sqrt d\,E[a_i]+P_i.$$

For coordinate pairs indexed by $r=0,\ldots,d/2-1$, the original position encoding is

$$P_{i,2r}=\sin\left(\frac{i}{10000^{2r/d}}\right), \qquad P_{i,2r+1}=\cos\left(\frac{i}{10000^{2r/d}}\right).$$

The embedding supplies learned token coordinates. The position vector supplies a fixed function of the token’s location. The sum has the same width $d$.

2.2 Attention is a normalized, input-dependent weighted sum

For one head, compute three projections:

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

where

$$W_Q,W_K\in\mathbb R^{d\times d_k}, \qquad W_V\in\mathbb R^{d\times d_v}.$$

Thus $Q,K\in\mathbb R^{T\times d_k}$ and $V\in\mathbb R^{T\times d_v}$. Each query row $q_i$ is compared with each key row $k_j$:

$$s_{ij}=\frac{q_i k_j^\top}{\sqrt{d_k}}+M_{ij}.$$

The mask $M$ determines which comparisons are allowed. Convert scores into weights, then combine values:

$$a_{ij}=\frac{\exp(s_{ij})}{\sum_{u=1}^{T}\exp(s_{iu})}, \qquad o_i=\sum_{j=1}^{T}a_{ij}v_j.$$

In matrix form:

$$\boxed{ \operatorname{Attn}(Q,K,V;M) =\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}+M\right)V }$$

The intermediate shapes explain the computation:

$$\underbrace{QK^\top}_{T\times T} \quad\Longrightarrow\quad \underbrace{A}_{T\times T} \underbrace{V}_{T\times d_v} =\underbrace{O}_{T\times d_v}.$$

Before dropout, each valid row satisfies $a_{ij}\geq0$ and $\sum_j a_{ij}=1$. A single head therefore returns a convex combination of its projected values. The weights depend on the input through $Q$ and $K$.

For example, take one query whose three scores are $(0,\log 2,0)$. Then

$$a=\left(\frac14,\frac12,\frac14\right).$$

With $v_1=(1,0)$, $v_2=(0,2)$, and $v_3=(3,0)$, the output is

$$o=\frac14(1,0)+\frac12(0,2)+\frac14(3,0)=(1,1).$$

If the third position is masked, its weight becomes zero and the other weights must be renormalized. They become $(1/3,2/3,0)$.

2.3 Why divide by the square root of the key width?

Assume, for a scale estimate, that the coordinates of $q$ and $k$ are mutually independent, have zero mean, and have variance one. Then

$$qk^\top=\sum_{r=1}^{d_k}q_rk_r, \qquad \operatorname{Var}(qk^\top)=d_k.$$

Consequently,

$$\operatorname{Var}\left(\frac{qk^\top}{\sqrt{d_k}}\right)=1.$$

This controls the initial score scale as $d_k$ changes. It does not require trained coordinates to remain independent.

The softmax derivative shows why score scale matters:

$$\frac{\partial a_j}{\partial s_u} =a_j(\delta_{ju}-a_u).$$

When one probability is almost one, many of these derivatives approach zero. Scaling reduces the tendency of large dot products to create this saturation at initialization.

2.4 Multiple heads use different learned projections

For head $r$, let

$$O_r=\operatorname{Attn}(XW_{Q,r},XW_{K,r},XW_{V,r};M).$$

Concatenate the results across feature coordinates and project back to width $d$:

$$\operatorname{MHA}(X;M) =[O_1\;O_2\;\cdots\;O_h]W_O, \qquad W_O\in\mathbb R^{hd_v\times d}.$$

The heads can learn different comparisons. No equation assigns a fixed role such as syntax or factual knowledge to a particular head.

For cross-attention, use two sequences with potentially different lengths:

$$Q_r=YW_{Q,r},\qquad K_r=ZW_{K,r},\qquad V_r=ZW_{V,r},$$

where $Y\in\mathbb R^{m\times d}$ and $Z\in\mathbb R^{n\times d}$. Its score matrix has shape $m\times n$. Denote the concatenated, output-projected result by $\operatorname{MHA}_{\mathrm{cross}}(Y,Z)$.

Self-attention and cross-attention differ in the source of $Q,K,V$. Multiple heads can be used in either computation.

2.5 What the two “Linear” operations actually compute

For each token row $x$, the original nonlinear transformation is

$$F(x)=\operatorname{ReLU}(xW_1+b_1)W_2+b_2,$$

with

$$W_1\in\mathbb R^{d\times d_f},\qquad W_2\in\mathbb R^{d_f\times d}.$$

For the original base configuration, this gives

$$\mathbb R^{512} \xrightarrow{xW_1+b_1}\mathbb R^{2048} \xrightarrow{\max(0,\cdot)}\mathbb R^{2048} \xrightarrow{(\cdot)W_2+b_2}\mathbb R^{512}.$$

An operation called “Linear” in a library is usually affine when it includes a bias. The intermediate nonlinearity is essential. Without it,

$$(xW_1+b_1)W_2+b_2 =x(W_1W_2)+(b_1W_2+b_2),$$

which is just one affine transformation.

The same $W_1,W_2$ apply to every token position. This calculation mixes features within a row. Attention supplies direct mixing between rows.

2.6 Residual addition and LayerNorm

For a row $x\in\mathbb R^d$, define

$$\mu(x)=\frac1d\sum_{r=1}^{d}x_r, \qquad \sigma^2(x)=\frac1d\sum_{r=1}^{d}(x_r-\mu(x))^2,$$ $$\operatorname{LN}(x) =\gamma\odot\frac{x-\mu(x)}{\sqrt{\sigma^2(x)+\epsilon}}+\beta.$$

Both $\gamma,\beta\in\mathbb R^d$ are trainable. The scalar $\epsilon>0$ prevents division by zero.

Write one original encoder transformation as

$$U^{(\ell)} =\operatorname{LN}_{\ell,1} \left(X^{(\ell-1)}+\operatorname{MHA}_{\ell}(X^{(\ell-1)};0)\right),$$ $$X^{(\ell)} =\operatorname{LN}_{\ell,2} \left(U^{(\ell)}+F_\ell(U^{(\ell)})\right).$$

Here $M=0$ assumes no padding. Padding keys must be masked in a padded batch. The original training computation also applies dropout to each residual branch before addition.

The order is $\operatorname{LN}(X+F(X))$. This is called Post-LN. Normalization before the residual branch, called Pre-LN, is a later choice discussed in Section 5.1.

2.7 The original decoder has two distinct attention calculations

Let $Z=X^{(L)}$ be the final encoder representations. Initialize the decoder from the shifted target sequence

$$(\mathrm{BOS},y_1,\ldots,y_{m-1}).$$

The representation at position $t$ predicts $y_t$. Its causal mask is

$$M_{ij}=\begin{cases} 0,&j\leq i,\\ -\infty,&j>i. \end{cases}$$

One decoder transformation is

$$U^{(\ell)} =\operatorname{LN}_{\ell,1} \left(Y^{(\ell-1)}+\operatorname{MHA}_{\ell}(Y^{(\ell-1)};M)\right),$$ $$R^{(\ell)} =\operatorname{LN}_{\ell,2} \left(U^{(\ell)}+\operatorname{MHA}_{\mathrm{cross},\ell}(U^{(\ell)},Z)\right),$$ $$Y^{(\ell)} =\operatorname{LN}_{\ell,3} \left(R^{(\ell)}+F_\ell(R^{(\ell)})\right).$$

The first attention uses target-prefix representations for $Q,K,V$. The second uses target representations for $Q$ and encoder outputs for $K,V$. All non-padding source positions can be available in this second calculation.

Finally,

$$p_\theta(y_t\mid y_{<t},x_{1:n}) =\operatorname{softmax}(Y_t^{(L)}W_{\mathrm{vocab}}+b_{\mathrm{vocab}}),$$

where $W_{\mathrm{vocab}}\in\mathbb R^{d\times N_v}$. The basic training loss is the sum of negative log probabilities of the correct targets. The original paper additionally uses label smoothing, so its training target distribution is not exactly one-hot.

3. GPT: make the prefix itself the conditioning sequence

GPT removes the separate source encoding and its cross-attention calculation. Causal self-attention remains. Learned absolute positions and GELU are used in the original GPT. Its pretraining objective is the likelihood of the next token. See Improving Language Understanding by Generative Pre-Training.

With token indices $a_{1:T}$, use the input $(\mathrm{BOS},a_1,\ldots,a_{T-1})$ and targets $(a_1,\ldots,a_T)$. A GPT-style prediction has the form

$$H^{(0)}=E[\text{shifted input}]+P_{\mathrm{learned}}, \qquad z_t=H_t^{(L)}W_{\mathrm{vocab}}+b_{\mathrm{vocab}},$$ $$p_{t,v}=\frac{\exp(z_{t,v})}{\sum_{u=1}^{N_v}\exp(z_{t,u})}, \qquad \mathcal L(\theta)=-\frac1T\sum_{t=1}^{T}\log p_{t,a_t}.$$

The recurrence for $H^{(\ell)}$ retains causal self-attention, residual additions, normalization, and a nonlinear feature transformation. There is no cross-attention term dependent on $Z$.

“Decoder-only” therefore describes the causal computation and absence of a separate source encoder. It does not mean that the model stops computing useful representations of its input. It also does not mean that every decoder-only model is a GPT model.

3.1 Parallel training, sequential generation

During training, all target tokens are known. A triangular mask makes it possible to evaluate every next-token loss in one sequence pass. There is no need to run each growing prefix separately.

During generation, the next token is unknown until it is selected:

$$a_{t+1}\sim p_\theta(\cdot\mid a_{1:t}), \qquad p_\theta(\cdot\mid a_{1:t+1}) \text{ requires the selected }a_{t+1}.$$

This dependence across output tokens remains even when all computations within one step use parallel hardware.

3.2 Why this supports general language pretraining

Translation requires paired source and target text. Autoregressive training obtains targets by shifting ordinary text. A task can then be expressed as part of the prefix:

$$p_\theta(\text{answer}\mid\text{instruction},\text{examples},\text{question}).$$

GPT-3 demonstrated task performance from prompts and examples without a separate gradient update for each prompted task. This is an empirical result, not a theorem that scale must produce any particular capability. See Language Models are Few-Shot Learners.

Instruction tuning changes the training distribution to examples of desired responses:

$$\mathcal L_{\mathrm{SFT}} =-\mathbb E_{(q,a)}\sum_t\log p_\theta(a_t\mid q,a_{<t}).$$

Preference-based training can further change which responses receive higher probability. It does not replace next-token probabilities as the generation interface. InstructGPT is an early example of supervised fine-tuning followed by reinforcement learning from human feedback. See Training language models to follow instructions with human feedback.

4. From prediction error to Adam and AdamW

Attention specifies how to calculate predictions. The loss specifies what counts as an error. An optimizer specifies how gradients change the parameters.

For one token with one-hot target vector $y$,

$$\ell=-\sum_{v=1}^{N_v}y_v\log p_v, \qquad p=\operatorname{softmax}(z), \qquad \frac{\partial\ell}{\partial z}=p-y.$$

Backpropagation applies the chain rule from this error to every trainable matrix.

4.1 Attention also receives gradients

For $O=AV$, $A=\operatorname{softmax}(S)$, and $S=QK^\top/\sqrt{d_k}+M$, write $\bar O=\partial\mathcal L/\partial O$. Then

$$\bar V=A^\top\bar O, \qquad \bar A=\bar O V^\top,$$ $$\bar S_{ij} =A_{ij}\left(\bar A_{ij}-\sum_u A_{iu}\bar A_{iu}\right),$$ $$\bar Q=\frac{\bar S K}{\sqrt{d_k}}, \qquad \bar K=\frac{\bar S^\top Q}{\sqrt{d_k}}.$$

For example, $Q=XW_Q$ gives $\bar W_Q=X^\top\bar Q$. Masked entries have zero attention weight and zero score gradient. These equations assume no attention dropout; its derivative must be included when it is active.

4.2 Adam estimates gradient direction and scale

Use $s$ for optimizer steps to distinguish them from token positions. Starting with $m_0=v_0=0$, compute

$$g_s=\nabla_\theta\mathcal L_s(\theta_{s-1}),$$ $$m_s=\beta_1m_{s-1}+(1-\beta_1)g_s, \qquad v_s=\beta_2v_{s-1}+(1-\beta_2)g_s\odot g_s.$$

Correct the zero-initialization bias:

$$\hat m_s=\frac{m_s}{1-\beta_1^s}, \qquad \hat v_s=\frac{v_s}{1-\beta_2^s}.$$

The update is

$$\boxed{ \theta_s=\theta_{s-1} -\eta_s\frac{\hat m_s}{\sqrt{\hat v_s}+\epsilon} }$$

All vector operations in this update are elementwise. The numerator estimates recent gradient direction. The denominator adjusts the step for each coordinate’s recent squared-gradient scale. The quantity $v_s$ estimates a second raw moment, not a centered variance or a Hessian. See Adam: A Method for Stochastic Optimization.

For a constant gradient $g$, the recurrence gives $m_s=(1-\beta_1^s)g$. Bias correction recovers $\hat m_s=g$. The same argument gives $\hat v_s=g\odot g$.

4.3 AdamW separates weight decay from gradient adaptation

An L2 penalty changes the objective to

$$\mathcal J(\theta)=\mathcal L(\theta)+\frac\lambda2\|\theta\|_2^2, \qquad \nabla\mathcal J=\nabla\mathcal L+\lambda\theta.$$

If this combined gradient enters Adam, the penalty also enters both moment estimates. Its effect therefore depends on gradient history.

AdamW instead forms moments from the data-loss gradient and applies decay separately:

$$\boxed{ \theta_s=(1-\eta_s\lambda)\theta_{s-1} -\eta_s\frac{\hat m_s}{\sqrt{\hat v_s}+\epsilon} }$$

Here $\lambda$ uses the convention in which decay is scaled by the learning rate. A schedule can equivalently specify the decay coefficient directly. For plain SGD, an L2 gradient gives this multiplicative decay. For Adam, the gradient transformation breaks that equivalence. See Decoupled Weight Decay Regularization.

Parameter groups can use different $\lambda$, including zero. Thus decay need not apply to biases or normalization parameters. Adam and AdamW are training algorithms; ordinary generation does not run these updates.

4.4 Learning-rate schedules are a separate choice

The original Transformer uses an Adam learning-rate schedule of the form

$$\eta_s=d^{-1/2}\min\left(s^{-1/2},\;s\,w^{-3/2}\right),$$

where $w$ is the warmup duration in optimizer steps. The rate first rises linearly, then decreases as $s^{-1/2}$. This schedule is a training choice from the original Transformer paper, not part of Adam’s definition.

5. Change the equations to train deeper models

The next changes concern representations and gradient flow. Their usefulness is empirical; none is a universal requirement for an LLM.

5.1 Pre-LN moves normalization before each residual branch

Compare

$$y=\operatorname{LN}(x+F(x)) \qquad\text{with}\qquad y=x+F(\operatorname{LN}(x)).$$

Using column-style Jacobian notation for this derivative only, the second form gives

$$J_y=I+J_FJ_{\operatorname{LN}}.$$

The identity term is a direct gradient path through the residual addition. In the first form, the derivative is $J_{\operatorname{LN}}(I+J_F)$. This difference helps explain the training behavior studied in On Layer Normalization in the Transformer Architecture. It does not guarantee that gradients remain bounded at arbitrary depth.

GPT-2 moved normalization to the inputs of its residual branches and added final normalization. See Language Models are Unsupervised Multitask Learners. A Pre-LN recurrence should therefore not be labeled as the original 2017 ordering.

5.2 RMSNorm removes centering

RMSNorm uses

$$\operatorname{RMSNorm}(x) =\gamma\odot\frac{x}{\sqrt{\frac1d\sum_{r=1}^{d}x_r^2+\epsilon}}.$$

It scales the row without subtracting its mean. The standard form above has a learned scale and no additive bias. See Root Mean Square Layer Normalization.

For $x=(2,4,6,8)$, with $\gamma=1$, $\beta=0$, and $\epsilon$ ignored,

$$\operatorname{LN}(x) =\frac{(-3,-1,1,3)}{\sqrt5}, \qquad \operatorname{RMSNorm}(x) =\frac{(2,4,6,8)}{\sqrt{30}}.$$

RMSNorm does not force a zero mean. It also does not directly bound the magnitude of the unnormalized residual stream.

5.3 SwiGLU introduces multiplication between learned features

GELU can be written as $\operatorname{GELU}(u)=u\Phi(u)$, where $\Phi$ is the standard normal cumulative distribution function. A GELU transformation is

$$F_{\mathrm{GELU}}(x)=\operatorname{GELU}(xW_1+b_1)W_2+b_2.$$

A common bias-free SwiGLU form is

$$\operatorname{SiLU}(u)=u\sigma(u), \qquad \sigma(u)=\frac1{1+e^{-u}},$$ $$\boxed{ F_{\mathrm{SwiGLU}}(x) =\left[\operatorname{SiLU}(xW_g)\odot(xW_u)\right]W_d }$$

Here $W_g,W_u\in\mathbb R^{d\times d_f}$ and $W_d\in\mathbb R^{d_f\times d}$. This form adds a second input projection and an elementwise product. See GLU Variants Improve Transformer.

The SiLU output can be negative or greater than one. The word “gate” therefore does not imply a probability restricted to $[0,1]$.

Ignoring biases, a two-projection transformation has $2dd_f$ parameters. SwiGLU has $3dd_f$. To match the parameter count of a two-projection transformation with width $4d$, choose approximately

$$3dd_f=2d(4d) \quad\Longrightarrow\quad d_f=\frac83d,$$

then adjust the width for hardware requirements. Comparing both forms at the same $d_f$ changes their parameter budgets.

6. Optimize attention from projections to weighted sums

The attention calculation has several costs: projections, pairwise scores, normalization, weighted sums, and memory traffic. Changes at these points solve different problems.

6.1 Fuse projections without changing their result

The three projection matrices can be concatenated:

$$[Q\;K\;V]=X[W_Q\;W_K\;W_V].$$

This identity allows one larger matrix multiplication in place of three separate calls. It preserves the real-arithmetic result. Speed depends on implementation and hardware. Floating-point operation ordering can produce small numerical differences.

For $d_k=d_v=d_h$, $h$ query heads, and $g$ key/value heads, the combined projection width is $(h+2g)d_h$.

6.2 RoPE puts position differences into query-key products

For a two-coordinate row vector, define

$$R(\phi)=\begin{bmatrix} \cos\phi&-\sin\phi\\ \sin\phi&\cos\phi \end{bmatrix}.$$

Use the transpose when applying this conventional rotation matrix to row vectors:

$$\widetilde q_i=q_iR(i\omega)^\top, \qquad \widetilde k_j=k_jR(j\omega)^\top.$$

Their dot product is

$$\widetilde q_i\widetilde k_j^\top =q_iR(i\omega)^\top R(j\omega)k_j^\top =q_iR((j-i)\omega)k_j^\top.$$

The explicit rotation depends on the relative offset $j-i$. In higher dimensions, apply separate rotations to coordinate pairs, for example with frequencies

$$\omega_r=10000^{-2r/d_k},\qquad r=0,\ldots,d_k/2-1.$$

This is the central construction in RoFormer. The remaining content dependence still comes from $q_i,k_j$. Standard RoPE rotates queries and keys; values remain unrotated. The relative-position identity does not guarantee useful behavior beyond the model’s trained context length.

6.3 MHA, GQA, and MQA differ in which heads share keys and values

Assume $g$ divides $h$. Map query head $r\in\{1,\ldots,h\}$ to key/value head

$$\pi(r)=1+\left\lfloor\frac{r-1}{h/g}\right\rfloor.$$

Then

$$O_r=\operatorname{softmax}\left( \frac{Q_rK_{\pi(r)}^\top}{\sqrt{d_h}}+M \right)V_{\pi(r)}.$$

MHA has $g=h$. Multi-query attention, or MQA, has $g=1$. Grouped-query attention, or GQA, uses an intermediate number of key/value heads. See the MQA paper and GQA paper.

For $h=32$, these choices might use $g=32$, $g=8$, and $g=1$. All retain 32 query heads. Each query head still produces its own attention weights. Sharing $K,V$ does not make the score matrices identical.

Key/value projection width and cache size scale with $g$. Pairwise score and value-aggregation arithmetic still scales with $h$ for a direct implementation. GQA therefore does not reduce every attention cost by $h/g$.

These choices change parameter sharing. They usually require appropriate training or adaptation; they are not exact replacements for an arbitrary trained MHA checkpoint.

6.4 Stable softmax subtracts a row maximum

For finite scores $s_1,\ldots,s_T$, choose $m=\max_j s_j$. Then

$$\frac{e^{s_j}}{\sum_u e^{s_u}} =\frac{e^{s_j-m}}{\sum_u e^{s_u-m}}.$$

Subtracting $m$ preserves the result and prevents overflow in the exponentials. An implementation must handle a fully masked row separately: subtracting $-\infty$ from $-\infty$ is undefined.

6.5 FlashAttention avoids storing the full score and probability matrices

For one query, attention can be expressed as a ratio:

$$o=\frac{\sum_j e^{s_j}v_j}{\sum_j e^{s_j}}.$$

For a set of valid keys $\mathcal A$, define a maximum, a scalar sum, and a vector sum:

$$m_{\mathcal A}=\max_{j\in\mathcal A}s_j, \quad \ell_{\mathcal A}=\sum_{j\in\mathcal A}e^{s_j-m_{\mathcal A}}, \quad u_{\mathcal A}=\sum_{j\in\mathcal A}e^{s_j-m_{\mathcal A}}v_j.$$

For a disjoint set $\mathcal B$, merge the statistics using

$$m=\max(m_{\mathcal A},m_{\mathcal B}),$$ $$\ell=e^{m_{\mathcal A}-m}\ell_{\mathcal A} +e^{m_{\mathcal B}-m}\ell_{\mathcal B}, \qquad u=e^{m_{\mathcal A}-m}u_{\mathcal A} +e^{m_{\mathcal B}-m}u_{\mathcal B}.$$

After all valid keys have been processed, $o=u/\ell$. Empty or fully masked sets are skipped or handled explicitly.

This derivation permits an exact attention result in real arithmetic while processing scores in tiles. FlashAttention uses this principle with fast on-chip memory and reduced transfers to GPU high-bandwidth memory. It also uses recomputation during the backward pass. See FlashAttention.

A conventional materialized attention matrix requires $O(BhT^2)$ elements. FlashAttention avoids keeping that full matrix in high-bandwidth memory. Dense pairwise arithmetic remains $O(BhT^2d_h)$. It improves memory use and data movement without turning dense attention into a linear-time calculation.

6.6 A local mask changes the attention function

If attention is restricted to a causal window of $w$ tokens, define

$$M_{ij}=\begin{cases} 0,&\max(1,i-w+1)\leq j\leq i,\\ -\infty,&\text{otherwise}. \end{cases}$$

The allowed pair count becomes $O(Tw)$ instead of $O(T^2)$. To realize the saving, the implementation must skip disallowed pairs. This restriction removes direct access to older positions within that attention calculation. It changes the function, whereas FlashAttention preserves the chosen dense or masked function.

7. KV cache: derive the reusable state

Consider generation in a causal decoder with fixed parameters and dropout disabled. A representation at position $i$ depends only on tokens at positions $1,\ldots,i$. Appending future tokens does not change that prefix.

This statement holds by induction over depth. It is true for token and fixed-position embeddings. Causal attention only reads earlier or current positions. Per-token normalization and feature transformations do not introduce a dependence on future positions.

Thus, for an unchanged prefix and position convention,

$$k_i^{(\ell)}\text{ and }v_i^{(\ell)} \quad\text{remain unchanged when a token is appended.}$$

7.1 Compute one new row and reuse the previous rows

For one head at one depth, let $u_t$ be the current input to its projections:

$$q_t=u_tW_Q,\qquad k_t=u_tW_K,\qquad v_t=u_tW_V.$$

Append the new key and value:

$$K_{1:t}=\begin{bmatrix}K_{1:t-1}\\k_t\end{bmatrix}, \qquad V_{1:t}=\begin{bmatrix}V_{1:t-1}\\v_t\end{bmatrix}.$$

The new attention result is

$$\boxed{ o_t=\operatorname{softmax}\left( \frac{q_tK_{1:t}^\top}{\sqrt{d_k}} \right)V_{1:t} }$$

The score shape is now $1\times t$. Each depth has its own cache. The new token must pass through every depth before the next-token probabilities are available.

For RoPE, cache the rotated keys using their correct positions. Rotate the new query and key at the new position. A change in position indexing or position-scaling rules can invalidate a cache.

7.2 Why cache K and V, but not old Q?

The expression for $o_t$ needs the current query $q_t$, all keys $k_{1:t}$, and all values $v_{1:t}$. It does not contain $q_1,\ldots,q_{t-1}$.

Old queries were needed to calculate old outputs. They are not needed to calculate the next output. This is why standard autoregressive inference retains K and V. Training has different requirements because backward computation needs additional activations or their recomputation.

Caching is exact for the same causal model under these assumptions, up to numerical effects. Quantized or truncated caches introduce additional approximations.

7.3 Prefill and decode have different cost shapes

Prefill processes a prompt of length $P$ and creates its caches. Dense attention has $O(P^2)$ query-key pairs. The final prompt position supplies probabilities for the first generated token.

Each subsequent decode step processes one newly selected token. With $t$ cached positions, its attention work is $O(thd_h)$ per depth. Its projections and per-token nonlinear transformations also have a cost, independent of $t$ when model widths are fixed.

For generation from a short initial prefix to length $N$, a naive method that repeatedly recomputes the full prefix incurs pairwise attention work proportional to

$$\sum_{t=1}^{N}t^2=\frac{N(N+1)(2N+1)}6=O(N^3).$$

With cached keys and values, the pairwise work is proportional to

$$\sum_{t=1}^{N}t=\frac{N(N+1)}2=O(N^2).$$

These are attention-only totals for that comparison. KV caching does not make a full dense-attention generation run linear in output length. Each new query still reads the retained context.

7.4 Calculate the cache size

With equal key and value widths $d_h$ and $b$ bytes per stored element, an unsharded cache for equal-length sequences uses

$$\boxed{M_{\mathrm{KV}}=2BLTgd_hb\text{ bytes}.}$$

The factor two counts keys and values. More generally, the per-depth factor is $g(d_k+d_v)$. This formula excludes weights, activations, allocator overhead, and quantization metadata.

Take $B=1$, $L=32$, $T=8192$, $h=32$, $d_h=128$, and $b=2$:

Attention choice	KV heads $g$	Cache calculation	Cache size
MHA	32	$2\cdot32\cdot8192\cdot32\cdot128\cdot2$	4 GiB
GQA	8	$2\cdot32\cdot8192\cdot8\cdot128\cdot2$	1 GiB
MQA	1	$2\cdot32\cdot8192\cdot1\cdot128\cdot2$	128 MiB

These are worked examples, not measurements of named models. For a batch of eight equal-length sequences, multiply each value by eight. Lower cache size can reduce memory traffic, but it does not imply the same factor of end-to-end speedup.

8. MLA: cache a latent representation instead of expanded K and V

DeepSeek-V2 introduces multi-head latent attention, or MLA, and DeepSeek-V3 retains it. The following discussion concerns those documented architectures. See the DeepSeek-V2 report and DeepSeek-V3 report.

8.1 A low-rank factorization exposes a smaller state

First omit positional rotations to isolate the algebra. Compress a token representation into a latent row:

$$c_j=x_jW_{DKV}\in\mathbb R^{d_c}.$$

For head $r$, define

$$k_{j,r}^{C}=c_jU_{K,r},\qquad v_{j,r}=c_jU_{V,r},$$

where $U_{K,r}\in\mathbb R^{d_c\times d_k}$ and $U_{V,r}\in\mathbb R^{d_c\times d_v}$. The superscript $C$ denotes the content part of a key.

Associativity gives

$$q_{i,r}^{C}(k_{j,r}^{C})^\top =q_{i,r}^{C}U_{K,r}^\top c_j^\top =\widehat q_{i,r}c_j^\top, \qquad \widehat q_{i,r}=q_{i,r}^{C}U_{K,r}^\top.$$

The value sum also factors:

$$\sum_j a_{ij,r}v_{j,r} =\left(\sum_j a_{ij,r}c_j\right)U_{V,r}.$$

These identities let an implementation attend over cached latent rows without repeatedly expanding historical keys and values. They are exact for the specified factorized weights. They do not prove that an arbitrary MHA model can be compressed to any chosen rank without error.

8.2 Separate the positional key contribution

If a position-dependent rotation acts on the full expanded content key, the score contains

$$q_iR_i^\top R_jU_K^\top c_j^\top.$$

The query-side expression now depends on key position $j$. A single fixed absorbed query projection cannot remove that dependence.

A simplified expression for the decoupled position treatment is

$$s_{ij,r} =\frac{ q_{i,r}^{C}(k_{j,r}^{C})^\top +\widetilde q_{i,r}^{R}(\widetilde k_j^{R})^\top }{\sqrt{d_k+d_R}}+M_{ij},$$

where the $R$ terms carry RoPE and the positional key is shared across heads. The cache stores both $c_j$ and $\widetilde k_j^R$. The V2 equations also include query compression; it is omitted here because it does not add historical KV state.

The resulting idealized cache size is

$$M_{\mathrm{MLA}}=BLT(d_c+d_R)b.$$

This is a different storage strategy from choosing fewer KV heads. Its benefit depends on $d_c+d_R$ compared with the expanded key/value widths. Actual speed also depends on the kernels and memory layout.

9. Store the cache efficiently: paging and quantization

9.1 Paging changes allocation, not attention weights

For page capacity $p$ tokens and sequence length $T$, the required number of pages is

$$n_{\mathrm{pages}}=\left\lceil\frac Tp\right\rceil.$$

Logical position $j\in\{1,\ldots,T\}$ maps to

$$\text{logical page}=\left\lfloor\frac{j-1}{p}\right\rfloor, \qquad \text{offset}=(j-1)\bmod p.$$

A page table maps logical pages to physical storage. This allows cache growth without one large contiguous allocation. The unused capacity in the final page is less than $p$ tokens per sequence. Additional allocator and page-table costs still exist. PagedAttention applies this principle to LLM serving. See Efficient Memory Management for Large Language Model Serving with PagedAttention.

Prefix sharing can reuse pages only when tokens, model parameters, positional treatment, and other computation-affecting settings match. A common textual suffix with different earlier context does not generally produce the same keys and values.

9.2 Quantization changes the numerical representation

A simple symmetric quantizer for a cache group is

$$\widehat x=s\,\operatorname{clip} \left(\operatorname{round}(x/s),-q_{\max},q_{\max}\right),$$

where $s>0$ is a scale and $q_{\max}$ is the largest positive integer magnitude. With $q$ bits per element, the idealized payload size becomes

$$M_{\mathrm{KV,quant}}\approx 2BLTgd_h\frac q8.$$

Scales, packing, and other metadata add storage. Quantization errors in K change the scores; errors in V change the weighted sums. Their effect on output quality must be measured. Paging preserves the stored values; quantization generally changes them.

10. Put a modern causal decoder into one recurrence

A representative dense decoder using Pre-RMSNorm, RoPE, GQA, and SwiGLU can be written as follows. This is a worked composition of the preceding equations, not a specification for every GPT or every LLM. Related choices appear in the LLaMA paper, although individual model versions use different head-sharing choices.

Start from token embeddings $H^{(0)}=E[a_{1:T}]$. At depth $\ell$, let

$$X=\operatorname{RMSNorm}_{\ell,1}(H^{(\ell-1)}).$$

For each query head $r$ and KV head $u$, compute

$$Q_r=XW_{Q,r},\quad K_u=XW_{K,u},\quad V_u=XW_{V,u},$$ $$\widetilde Q_{i,r}=Q_{i,r}R_i^\top, \qquad \widetilde K_{j,u}=K_{j,u}R_j^\top,$$ $$O_r=\operatorname{softmax}\left( \frac{\widetilde Q_r\widetilde K_{\pi(r)}^\top}{\sqrt{d_h}}+M \right)V_{\pi(r)}.$$

Then apply the two residual additions:

$$U=H^{(\ell-1)}+[O_1\;\cdots\;O_h]W_O, \qquad Z=\operatorname{RMSNorm}_{\ell,2}(U),$$ $$\boxed{ H^{(\ell)} =U+\left[\operatorname{SiLU}(ZW_g)\odot(ZW_u)\right]W_d }$$

The final row predicts the next token:

$$p_\theta(a_{T+1}\mid a_{1:T}) =\operatorname{softmax}\left( \operatorname{RMSNorm}_{\mathrm{final}}(H_T^{(L)})W_{\mathrm{vocab}} \right).$$

During training, shifted targets supply cross-entropy gradients and AdamW can update the parameters. During generation, fixed parameters and causal dependence permit KV reuse. FlashAttention determines how the attention sums are evaluated. Paging determines where the cache lives. Each choice acts on a specific equation or its execution.

Change	Mathematical or execution effect
Translation objective to causal language modeling	Learn $\prod_t p_\theta(a_t\mid a_{<t})$ from text
Adam to AdamW	Separate multiplicative parameter decay from gradient moments
Post-LN to Pre-LN	Move normalization inside the residual branch
LayerNorm to RMSNorm	Remove mean subtraction
ReLU/GELU transformation to SwiGLU	Add a learned projection and featurewise multiplication
Absolute additive positions to RoPE	Rotate Q and K so their product contains position differences
MHA to GQA/MQA	Share K and V across query heads
Materialized attention to FlashAttention	Evaluate the same attention through tiled sums and less memory traffic
Full-prefix recomputation to KV caching	Reuse unchanged historical K and V
Expanded KV storage to MLA latent storage	Exploit a trained factorization and retain a separate positional key
Contiguous allocation to paging	Change the address mapping of retained cache values
Full-precision cache to quantized cache	Reduce bytes per element while introducing numerical error

The connection from Transformer to LLM is visible in these equations: causal probabilities define the task, learned matrix products and nonlinear functions define the model, gradient updates define learning, and reusable prefix state makes sequential generation practical.

References

The following original papers and technical reports are cited in the text. Years identify their first public release; conference or journal publication can be later.

1. Vaswani, A., et al. (2017). Attention Is All You Need. arXiv:1706.03762. Original encoder–decoder Transformer, scaled dot-product attention, sinusoidal positions, and training schedule.
2. Radford, A., et al. (2018). Improving Language Understanding by Generative Pre-Training. GPT and causal language pretraining.
3. Radford, A., et al. (2019). Language Models are Unsupervised Multitask Learners. GPT-2 and its normalization placement.
4. Brown, T. B., et al. (2020). Language Models are Few-Shot Learners. arXiv:2005.14165. GPT-3 and learning tasks from examples in the prompt.
5. Ouyang, L., et al. (2022). Training language models to follow instructions with human feedback. arXiv:2203.02155. InstructGPT, supervised fine-tuning, and reinforcement learning from human feedback.
6. Kingma, D. P., and Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv:1412.6980. Gradient moment estimates, bias correction, and adaptive parameter updates.
7. Loshchilov, I., and Hutter, F. (2017). Decoupled Weight Decay Regularization. arXiv:1711.05101. AdamW and the distinction between L2 regularization and weight decay.
8. Xiong, R., et al. (2020). On Layer Normalization in the Transformer Architecture. arXiv:2002.04745. Post-LN, Pre-LN, and gradient behavior at initialization.
9. Zhang, B., and Sennrich, R. (2019). Root Mean Square Layer Normalization. arXiv:1910.07467. RMSNorm and normalization without mean subtraction.
10. Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv:2002.05202. Gated nonlinear transformations, including SwiGLU.
11. Su, J., et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864. RoPE and relative position in query-key products.
12. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150. Multi-query attention and reduced key/value memory traffic during decoding.
13. Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245. Grouped-query attention and adaptation from MHA checkpoints.
14. Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135. Tiled exact attention, online softmax, and reduced memory traffic.
15. DeepSeek-AI, et al. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434. MLA, latent KV compression, and decoupled RoPE.
16. DeepSeek-AI, et al. (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437. Use of MLA in DeepSeek-V3.
17. Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180. Paged KV allocation and cache sharing in vLLM.
18. Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971. A causal language model combining pre-normalization, RMSNorm, SwiGLU, and RoPE.