The AI Stack: Transformers, Attention & Gradient Flow

A transformer turns token representations into context-dependent representations; attention routes information between positions; and gradient flow determines how the training signal reaches the parameters that produced them. Understanding all three makes the rest of the AI stack easier to reason about: training objectives, model adaptation, inference memory, retrieval, and evaluation.

These are my technical notes as an indie developer, connecting foundational research to practical product decisions. The focus is the decoder-only language-model stack. Encoder-only and encoder–decoder models make different masking and architecture choices. The research milestones here are foundational work through 2023, rather than a survey of the latest models.

1. Map the stack before choosing a tool

There are three useful layers of abstraction. The mathematical layer contains representations, matrix operations, probability distributions, and derivatives. The model layer combines those operations into trainable architectures. The product layer combines a model with data, interfaces, retrieval, tools, and operational constraints.

A library can implement several layers, but its name does not explain which problem it solves. A vector database cannot repair a broken attention mask. A more efficient attention kernel cannot make an unsupported factual claim correct. Diagnose the layer before changing the stack.

2. Tokens, embeddings, and the learning objective

A tokenizer maps text to integer IDs. A token can represent a word, part of a word, punctuation, or bytes; boundaries depend on the tokenizer and language. Looking up an embedding row maps each ID to a vector of width d. For n positions, the representation matrix X has shape n × d. An embedding is a learned coordinate system, not a human-readable definition of a word.

For causal next-token training, the model predicts each target using only the prefix before that target. The loss averages the negative log probability assigned to the correct next token:

pθ(t₁, …, tT) = ∏ₜ pθ(tₜ | t<t)
ℒ(θ) = −(1/T) ∑ₜ log pθ(tₜ | t<t)
θ ← Optimizer(θ, ∇θℒ)

The target is shifted relative to the input. If the model sees the answer token while learning to predict it, a falling loss can conceal data leakage. Padding and any deliberately excluded target positions also need loss masks; a loss mask and an attention mask have different jobs.

During ordinary inference, parameters stay fixed. A prompt changes the computation and the generated distribution, rather than updating the weights. Predicting plausible continuations is also a different objective from verifying facts against a current database.

3. Scaled dot-product attention, with shapes

For a single attention head, learn projections that produce queries, keys, and values. Queries and keys determine routing scores; values contain the vectors that are mixed. Ignoring batch dimensions:

X: n × d
WQ, WK: d × dₖ       WV: d × dᵥ
Q = XWQ: n × dₖ
K = XWK: n × dₖ      V = XWV: n × dᵥ
S = QKᵀ / √dₖ: n × n
A = softmax_rows(S + M): n × n
Y = AV: n × dᵥ

The mask M contains zero for allowed positions and negative infinity for forbidden positions. In causal self-attention, position i can read positions j ≤ i, including itself, and cannot read future positions. Softmax operates across keys separately for each query row. At least one key must remain available in each evaluated row, or the all-masked softmax needs explicit handling.

Why divide by √dₖ? If query and key coordinates are independent, zero mean, and unit variance, their dot product has variance dₖ. Scaling keeps its variance near one under those assumptions. This reduces an initial tendency toward extremely peaked softmax distributions; it is not a guarantee about every trained activation distribution.

Multi-head attention computes several such mixtures, concatenates their outputs, and applies an output projection. Heads can learn different relationships without being assigned human concepts such as grammar or sentiment. The original Transformer paper by Vaswani and colleagues combines these components in an encoder–decoder architecture; decoder-only models retain and adapt its building blocks.

4. A worked attention calculation

Use a query q = [1, 0], three keys, and three values. Treat this as the final position in a three-position causal prefix, so all three keys are visible. These vectors are deliberately invented for arithmetic, rather than extracted from a trained model.

q = [1, 0]
k₁ = [1, 0]    v₁ = [1, 0]
k₂ = [0, 1]    v₂ = [0, 2]
k₃ = [1, 1]    v₃ = [2, 1]

s = [1/√2, 0, 1/√2]
a ≈ [0.401112, 0.197776, 0.401112]
y = a₁v₁ + a₂v₂ + a₃v₃
y ≈ [1.203336, 0.796664]

The weights sum to one, and the output is a vector. The next stages still apply projections, nonlinear transformations, and more layers. The highest attention weight therefore does not identify a final answer or establish why the complete model produced it. Jain and Wallace provide empirical reasons to be careful when interpreting attention as an explanation.

For numerical stability, compute softmax after subtracting the largest allowed score. Adding a constant to every score in a row leaves the distribution unchanged. The attention experiment above lets you vary a separate temperature: smaller positive temperatures sharpen its synthetic distribution. That is distinct from changing the head dimension or the temperature used to sample final output tokens.

5. Differentiate attention instead of treating it as magic

For one row, write a = softmax(s) and y = Σⱼ aⱼvⱼ. The softmax Jacobian has entries aᵢ(δᵢⱼ − aⱼ). Substituting it into the weighted sum gives a compact derivative. If g is the upstream loss gradient with respect to y:

∂aᵢ/∂sⱼ = aᵢ(δᵢⱼ − aⱼ)
∂y/∂sⱼ = aⱼ(vⱼ − y)
∂ℒ/∂sⱼ = aⱼ gᵀ(vⱼ − y)
∂ℒ/∂vⱼ = aⱼg

sⱼ = qᵀkⱼ / √dₖ
∇qℒ = (1/√dₖ) ∑ⱼ (∂ℒ/∂sⱼ) kⱼ
∇kⱼℒ = (∂ℒ/∂sⱼ) q / √dₖ

This exposes both paths through attention: gradients change what gets routed through queries and keys, and what content gets transmitted through values. Because self-attention derives all three from X, the complete gradient with respect to X includes all three projection paths. The final expression above holds values fixed while differentiating routing scores.

You can check the score derivative with central finite differences. This standalone Python example uses only the standard library, reproduces the worked output, and compares an analytic derivative against a numerical one. The chosen loss is a small linear probe, ℒ = [0.3, −0.2] · y, so the upstream gradient is known exactly.

import math

def softmax(scores):
    peak = max(scores)
    exps = [math.exp(score - peak) for score in scores]
    return [value / sum(exps) for value in exps]

scores = [1 / math.sqrt(2), 0, 1 / math.sqrt(2)]
values = [[1, 0], [0, 2], [2, 1]]
upstream = [0.3, -0.2]

def output(logits):
    weights = softmax(logits)
    return [sum(a * v[d] for a, v in zip(weights, values))
            for d in range(2)]

def loss(logits):
    return sum(g * y for g, y in zip(upstream, output(logits)))

weights, y = softmax(scores), output(scores)
analytic = [a * sum(g * (v[d] - y[d])
                   for d, g in enumerate(upstream))
            for a, v in zip(weights, values)]
epsilon = 1e-5
numeric = []
for j in range(len(scores)):
    plus, minus = scores.copy(), scores.copy()
    plus[j] += epsilon
    minus[j] -= epsilon
    numeric.append((loss(plus) - loss(minus)) / (2 * epsilon))

assert abs(sum(weights) - 1) < 1e-12
assert max(abs(a - b) for a, b in zip(analytic, numeric)) < 1e-7
print("weights:", [round(a, 6) for a in weights])
print("output:", [round(value, 6) for value in y])
print("score gradients:", [round(value, 6) for value in analytic])

A finite-difference check tests a local derivative, not whether an entire trained model is correct. Use double precision, a non-extreme input, and a sensible perturbation: too small a step amplifies rounding error; too large a step measures curvature as well as the derivative.

6. Assemble the transformer block

Attention mixes information across positions. A position-wise MLP transforms each position’s features independently, using shared weights. Residual additions preserve a stream that each sublayer can modify. One pre-normalized block can be written as follows; dropout and bias details are omitted:

U = X + MultiHeadAttention(Norm(X))
Xnext = U + MLP(Norm(U))
MLP(z) = W₂ φ(W₁z + b₁) + b₂

The MLP equation uses a column-vector convention for one token; the earlier X matrix stacks token row vectors. Implementations transpose or arrange their weight tensors accordingly. Keeping the convention explicit prevents dimension errors disguised by familiar notation.

Position must also enter the system. In RoPE, query and key coordinate pairs are rotated as a function of position. Their inner product depends on the relative rotation, since (Rᵢq)ᵀ(Rⱼk) = qᵀRᵢᵀRⱼk. This supplies positional structure without simply adding a position vector to every input. It does not promise correct behavior at arbitrarily long contexts. See Su and colleagues’ RoFormer paper.

LayerNorm centers and rescales features. RMSNorm rescales using their root mean square without subtracting the mean. The RMSNorm paper studies that simplification. The exact normalization, activation, position scheme, and head layout vary across architectures; a diagram labeled transformer does not uniquely specify them.

7. Gradient flow through depth

Backpropagation applies the chain rule in reverse. Let Jₗ be the Jacobian of layer l’s output with respect to its input, and let gₗ be the loss gradient at that input. Repeated matrix multiplication transports the signal back through depth:

hₗ₊₁ = fₗ(hₗ)
Jₗ = ∂hₗ₊₁/∂hₗ
gₗ = Jₗᵀgₗ₊₁
g₀ = J₀ᵀJ₁ᵀ … J_(L−1)ᵀ g_L

For a scalar chain with derivative a at every layer, the gradient gain is aᴸ. With a = 0.9 and L = 100, its magnitude is about 0.0000266. With a = 1.1, it is about 13,781. Deep matrix networks involve directions and singular values, so the scalar example illustrates amplification and attenuation without being a complete model of real training.

A residual layer hₗ₊₁ = hₗ + Fₗ(hₗ) changes its Jacobian to I + JFₗ. The identity term creates a direct derivative path. This is the structural insight behind deep residual learning, later used throughout transformer stacks.

Residual layer: hₗ₊₁ = hₗ + Fₗ(hₗ)
Gradient: gₗ = (I + JFₗ)ᵀgₗ₊₁
Scalar toy: direct gain = |a|ᴸ
Scalar toy: residual gain = |1 + a|ᴸ

An identity term does not guarantee preservation: JFₗ = −I cancels it, and positive amplification can still explode. The interactive gradient plot intentionally allows both outcomes. Real stability depends on the residual branch, normalization, initialization, precision, and optimizer. In a pre-normalized block, JF includes the derivative of the normalization inside that branch.

The original Transformer used post-normalization. Xiong and colleagues analyze how moving normalization inside the residual branches changes gradient behavior at initialization. Their result is a reason to inspect normalization placement, rather than treating pre-normalization as a universal promise that warm-up or careful tuning is unnecessary.

8. The training loop and its failure modes

A training step selects a batch, runs the forward computation, measures a loss, computes derivatives, and applies an optimizer update. Gradient accumulation combines several microbatches before that update; it requires loss scaling that matches the intended effective batch average.

Gradient clipping limits an excessive update signal; it cannot restore a vanished gradient or fix a wrong objective. Lower-precision training changes numerical behavior and memory use; the choice of format and accumulation precision matters. Activation checkpointing trades extra recomputation for reduced saved-activation memory. These choices interact, so a single change should be evaluated with both loss curves and the target task.

9. Breakthroughs that changed what could be built

Modern language models emerged from several complementary advances. This short timeline names the constraint each addressed, rather than implying that one paper invented the whole stack.

The Chinchilla result concerns allocating a training compute budget between parameters and data. Deployment may favor a different allocation when inference volume dominates lifetime cost. There is no universal token-to-parameter ratio that settles every product’s economics.

Instruction tuning and preference optimization change model behavior after pretraining. They do not independently establish factual accuracy. In DPO, a useful quantity is the log probability ratio between the trainable policy and a reference policy. The objective increases the chosen answer’s relative advantage over a rejected answer for the same prompt. The quality and coverage of those preferences remain part of the system’s limitations.

10. Inference: prefill, decode, and KV memory

During prefill, the model processes a prompt and constructs key and value states for its layers. During decode, a generated token adds new states and attends to the available prefix. A KV cache reuses previous key and value projections, avoiding their recomputation at each generated token. It stores sequence-specific activations; it does not update the model’s knowledge or weights.

For a uniform transformer stack with fixed batch and sequence length, an approximate KV payload is:

KV bytes ≈ 2 × B × L × T × Hkv × dh × b
B: batch size       L: layers       T: cached tokens
Hkv: KV heads       dh: head width  b: bytes per element

B=1, L=32, T=8192, Hkv=8, dh=128, b=2
KV payload = 1,073,741,824 bytes = 1 GiB

That example excludes model weights, allocator overhead, temporary buffers, and other runtime state. Changing to 32 KV heads makes this payload 4 GiB under the same assumptions. Grouped-query attention shares keys and values across groups of query heads, reducing this term; it is an architectural choice with quality and training considerations.

FlashAttention changes how exact attention is computed through tiling and online softmax, reducing memory transfers and avoiding a fully materialized n × n attention matrix in high-bandwidth memory. Dense prefill still has quadratic attention arithmetic. A cached decode step attends over a growing prefix, while projections and MLPs add their own costs. See the original FlashAttention paper and the PyTorch SDPA documentation for an optimized interface and its mask and dropout semantics.

Longer prompts can increase both waiting time and memory pressure. Measure time to first token separately from output-token throughput. Quantization can reduce weight or cache storage, but quality and numerical effects need measurement for the specific model, format, and task.

11. Prompting, retrieval, and LoRA solve different problems

The original RAG work combined parametric generation and a retrievable document index. A typical inference-time RAG application searches documents, adds excerpts to context, and generates an answer without backpropagating through those retrieved documents. Retrieval quality, document permissions, and whether the answer actually follows the evidence all need separate evaluation.

For LoRA, freeze a base matrix W₀ and train a low-rank update. With W₀ shaped d_out × d_in, A is r × d_in and B is d_out × r:

W = W₀ + (α/r)BA
Trainable parameters = r(d_in + d_out)
Full matrix parameters = d_in × d_out

This reduces the trainable parameter count when r is small. It does not eliminate the base model, its activations, or every training-memory cost. The LoRA paper studies adaptation with this parameterization. Use retrieval for source-dependent evidence and tuning for learned behavior; neither removes the need to evaluate the complete application.

12. A design exercise: a Shopify support assistant

Suppose a shop wants help answering product and delivery questions. Start with the source of truth: catalog facts, policy documents, and an authenticated order-status API. A language model can interpret a question and compose an answer, while the application controls access and executes permitted tools.

For a policy question, retrieve the relevant policy passages and include document identifiers. For an order question, verify the user’s access and fetch structured order data. Have the model explain the returned state; do not make the model invent it. Keep retrieved text separate from application instructions, because documents can contain misleading or adversarial directions.

Build an evaluation set before adding more infrastructure. Include answered and unanswerable questions, stale or conflicting policy text, missing retrieval results, and attempts to access another customer’s order. Score answer support, correct tool selection, authorization behavior, latency, and cost separately. A good next-token loss cannot stand in for these application checks.

This exercise illustrates the boundary between a capable model and a useful product. The architecture follows the task and the evidence, with an explicit path for missing information and escalation.

13. Checks that connect the math to the implementation

1. Causal isolation: change a future token and confirm earlier outputs remain unchanged in evaluation mode with dropout disabled.
2. Probability invariants: allowed attention weights sum to one, masked weights are zero, and adding one constant to a row’s logits preserves its distribution.
3. Derivative agreement: compare local analytic or automatic derivatives with double-precision finite differences on small, non-extreme examples.
4. Shape checks: distinguish query-head count from KV-head count and keep row-vector and column-vector conventions explicit.
5. Memory accounting: estimate KV payload, then compare measured peak memory with all the additional runtime costs included.
6. Task evaluation: hold out realistic requests, track failure categories, and compare end-to-end quality and latency before changing the stack.

These checks connect theory to observable behavior. They also make a technical discussion more precise: which quantity changed, which assumption failed, and what evidence would distinguish an architecture problem from a data or product problem?

14. Continue from the foundations

The most useful next step is to reproduce the tiny attention calculation, verify its derivative, and then inspect one transformer block with real tensor shapes. Trace the same example forward and backward before moving on to larger models. The original research linked throughout this note is the starting point for each breakthrough; the toy calculations and design exercise connect those ideas to concrete questions.

For the product side of that exploration, browse the independent work or read about the developer behind GradFlowLab.