Warmup Guide — Model Architecture Mastery

Zero-to-expert primer for Phase 02. Read this before the labs: it builds every concept from first principles — attention, position encodings, normalization, Mixture of Experts, State Space Models, Vision Transformers, and multimodal bridges — assuming only basic Python and linear algebra.

Table of Contents

Chapter 1: Why Architecture Knowledge Is an Optimization Skill

Zero background: a neural network architecture is the fixed computational skeleton — which matrices multiply which activations in what order. Weights change during training; the architecture doesn't.

Why it matters for this curriculum: every optimization decision downstream — which layers to quantize to INT4, how to tile a matmul for an NPU, why accuracy collapses after a graph transformation — is a function of architectural structure. Three concrete examples:

  • GQA shares KV heads → fewer unique weight rows → per-group quantization has fewer groups to calibrate → quantizes more gracefully than MHA.
  • MoE activates 2 of 8 experts per token → 75% of expert weights are cold per token → weight streaming and per-expert quantization become viable; but the router is tiny and precision-sensitive — quantize it last.
  • Mamba carries a recurrent state → no KV cache, O(1) memory per token — but the state is numerically delicate; INT8-quantizing the recurrence accumulates error step by step in a way attention never does.

Misconception to kill now: "architectures are interchangeable backends for the same function." The function may be similar; the hardware cost surface is wildly different, and that surface is your job.

Chapter 2: Attention from Zero

The problem attention solves: process a sequence so each position can use information from any other position, with content-dependent (not fixed) routing.

Mechanism, from nothing: each token's embedding $x_i \in \mathbb{R}^d$ produces three vectors via learned projections: a query $q_i = W_Q x_i$, a key $k_i = W_K x_i$, a value $v_i = W_V x_i$. Token $i$ scores every token $j$ by dot product $q_i \cdot k_j$ — "how relevant is $j$ to me" — then turns scores into a probability distribution and takes the weighted average of values:

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

Every term, justified:

  • $QK^\top$ is an $n \times n$ matrix of all pairwise scores — this is where the $O(n^2)$ cost in both compute and memory comes from.
  • $\sqrt{d_k}$: dot products of random $d_k$-dim vectors have variance $\propto d_k$; unscaled, softmax inputs grow with dimension, saturating it (gradients vanish). Dividing by $\sqrt{d_k}$ keeps variance ≈ 1 regardless of head size.
  • softmax row-wise: each query gets a distribution over keys, so the output is a convex combination of values — bounded, differentiable routing.
  • Causal masking (decoders): set scores $j > i$ to $-\infty$ before softmax, so position $i$ cannot see the future. After softmax those entries are exactly 0.

Multi-head: instead of one attention in $\mathbb{R}^d$, split into $h$ heads of dimension $d/h$, attend independently, concatenate, project. Heads learn different relation types (syntax, coreference, position). Cost is the same FLOPs as single-head at full width; the win is representational.

Production significance: the $n^2$ attention matrix is the reason long-context is hard, the reason FlashAttention exists (Phase 08), and the reason KV caching dominates inference memory (next chapter).

Chapter 3: The KV Cache — Where Inference Cost Lives

The setup: autoregressive generation produces one token at a time. Naively, generating token $t+1$ recomputes attention over all $t$ previous tokens from scratch — $O(t^2)$ work per token, $O(T^3)$ total.

The fix: $K$ and $V$ for past tokens never change (causality). Cache them. Each new token computes only its own $q, k, v$, appends $k, v$ to the cache, and attends: $O(t)$ per token.

The cost you bought: memory. Per token, per layer: $2 \times h_{kv} \times d_{head}$ values. For LLaMA-2-7B (32 layers, 32 KV heads, head dim 128, FP16):

$$32 \times 2 \times 32 \times 128 \times 2,\text{bytes} = 512,\text{KB per token}$$

A 4K-token context costs 2 GB — for batch size 1. This single number explains: why GQA exists (Chapter 4), why batch sizes are memory-bound, why PagedAttention manages KV like virtual memory (Phase 08), and why edge deployment quantizes the KV cache itself.

Misconception: "inference cost = model weights." For long contexts and large batches, KV cache rivals or exceeds weight memory. Always compute both.

Chapter 4: MHA, MQA, GQA — The KV-Head Spectrum

One design axis: how many KV head pairs do $h$ query heads share?

VariantKV headsKV cache vs MHAQuality
MHA$h$ (one per query head)1×best
GQA$g$ groups ($1 < g < h$)$g/h$ (e.g. ⅛ for LLaMA-3)≈MHA with modest training
MQA1$1/h$measurable degradation

Mechanism: in GQA, query heads within a group attend using the same $K, V$ projections. Implementation is just repeat_interleave of KV heads to match query heads (or better: a broadcasted attention kernel that never materializes the repetition).

Why this is the consensus design (LLaMA-3, Mistral, Qwen): KV cache (Chapter 3) is the inference bottleneck; GQA cuts it 4–8× for ~zero quality loss. It also helps quantization: fewer distinct KV projections → fewer calibration distributions.

Lab 01 implements all three behind one config flag — the test checks GQA output matches the naive repeat-expansion within 1e-5, which forces you to get the grouping exactly right.

Chapter 5: Position Encodings — Learned, RoPE, ALiBi

The problem: attention is permutation-invariant — shuffle the tokens and the scores are identical. Order must be injected.

Learned absolute embeddings (GPT-2): a trainable vector per position, added to token embeddings. Simple; cannot extrapolate past training length (position 2049 was never trained); positions don't compose ("distance 5" isn't represented).

RoPE (Rotary Position Embeddings) — used by LLaMA, Qwen, most modern LLMs. Idea: encode position as rotation. Pair up dimensions of $q$ and $k$; rotate each pair $(x_1, x_2)$ at position $m$ by angle $m\theta_i$, with per-pair frequencies $\theta_i = 10000^{-2i/d}$:

$$\begin{pmatrix} \cos m\theta_i & -\sin m\theta_i \ \sin m\theta_i & \cos m\theta_i \end{pmatrix}\begin{pmatrix} x_{2i} \ x_{2i+1} \end{pmatrix}$$

The magic: a rotation by $m\theta$ on $q$ and $n\theta$ on $k$ leaves their dot product depending only on $m - n$ — relative position falls out of absolute rotations. Multi-frequency = a "clock" with many hands: fast hands resolve nearby distances, slow hands resolve far ones. Extrapolation beyond training length works partially and is extended by position interpolation / NTK-aware scaling (scaling $\theta$ at inference).

ALiBi: skip embeddings entirely; subtract a per-head linear penalty $m_h \cdot |i - j|$ from attention logits. Recency bias, baked in. Trains short, generalizes long remarkably well; but no notion of position beyond distance, and weaker at tasks needing precise absolute positions.

Choosing: RoPE is the default; ALiBi when train-short-infer-long matters most. Lab 01's test checks RoPE's relative-position property explicitly.

Chapter 6: Normalization and FFN Variants

LayerNorm from zero: deep nets suffer when activation scales drift across layers. LayerNorm re-standardizes each token's vector: subtract its mean, divide by its std, then apply learned scale/shift. RMSNorm drops the mean subtraction: $\hat{x} = x / \sqrt{\tfrac{1}{d}\sum x_i^2} \cdot \gamma$ — one fewer reduction, ~10-15% faster, and empirically equivalent in transformers. All modern LLMs use RMSNorm.

Pre-norm vs post-norm: normalize before the sublayer (pre-norm, modern) and the residual path stays an identity highway — gradients flow unimpeded, training is stable without warmup tricks. Post-norm (original Transformer) normalizes the sum, which can amplify or attenuate the residual signal — deeper stacks destabilize.

SwiGLU: standard FFN is $W_2,\sigma(W_1 x)$. Gated variants compute $W_3(\sigma(W_1x) \odot W_2x)$ — one path gates the other multiplicatively, letting the FFN model input-dependent feature selection. ~25% more parameters per FFN at the same hidden dim (three matrices, not two); in practice hidden dim is shrunk by ⅔ to compensate. PaLM/LLaMA measurements: better loss per FLOP. SiLU ($x\sigma(x)$) is the usual $\sigma$.

Quantization note: the elementwise multiply in SwiGLU creates wider activation dynamic range than ReLU FFNs — one reason LLM activations are harder to quantize than CNN activations (foreshadowing Phase 03's SmoothQuant).

Chapter 7: Mixture of Experts

Zero background: a dense FFN applies the same weights to every token. MoE replaces one FFN with $E$ parallel FFNs ("experts") plus a tiny router that picks $K$ experts per token (K=1 Switch, K=2 Mixtral). Total parameters scale with $E$; compute per token scales with $K$ — parameters and FLOPs are decoupled.

Router mechanics: logits $= W_r x$ (an $E$-way linear classifier per token); take top-K; softmax over the selected logits gives combination weights; output = weighted sum of the chosen experts' outputs.

The failure mode — expert collapse: routing is self-reinforcing. An expert that gets more tokens trains more, improves, attracts more tokens; others atrophy. End state: one expert does everything and you paid for $E$.

The fix — auxiliary load-balancing loss (Switch Transformer): with $f_i$ = fraction of tokens routed to expert $i$ and $p_i$ = mean router probability for expert $i$,

$$\mathcal{L}{aux} = E \cdot \sum{i=1}^{E} f_i , p_i$$

This is minimized when both distributions are uniform ($f_i = p_i = 1/E$ gives $\mathcal{L}_{aux}=1$). $f_i$ is non-differentiable (a count); $p_i$ carries the gradient — the product trick lets gradient flow push probability away from overloaded experts.

Production significance: capacity factors (token budgets per expert with overflow dropping), all-to-all communication in expert parallelism, and per-expert quantization. On edge: MoE means most weights are cold per token — pair with layer/expert streaming (Phase 10).

Chapter 8: State Space Models and Mamba

The motivation: attention pays $O(n^2)$ compute and $O(n)$ KV memory. An RNN pays $O(n)$ compute and $O(1)$ memory but can't be parallelized over time during training and forgets. SSMs are the modern reconciliation.

Continuous origin: a linear time-invariant system $h'(t) = Ah(t) + Bx(t),; y(t) = Ch(t) + Dx(t)$. To run on discrete tokens, discretize with step $\Delta$ using Zero-Order Hold (input held constant within a step):

$$\bar{A} = e^{\Delta A}, \qquad \bar{B} = (\bar{A} - I)A^{-1}B$$

giving the recurrence $h_t = \bar{A}h_{t-1} + \bar{B}x_t$, $y_t = Ch_t + Dx_t$. Intuition: $\bar{A}$ entries (< 1) are memory decay coefficients; $\bar{B}$ is input gain. Large $\Delta$ = step fast, forget fast; small $\Delta$ = retain.

S4's contribution: with fixed $A$ (HiPPO-initialized, diagonal), the whole sequence output is a convolution with a precomputable kernel — train in parallel like a CNN, run as an RNN at inference.

Mamba's contribution — selectivity: make $B, C, \Delta$ functions of the input $x_t$. Now the state update is content-dependent: the model can choose, per token, to store or ignore. This breaks the convolution trick (kernel is input-dependent), so Mamba ships a hardware-aware parallel associative scan computed in SRAM with recomputation in backward — an architecture and a kernel co-designed.

Why NPU/edge people care: O(1) inference state (a $d_{state} \times d$ matrix, not a growing KV cache) is the difference between bounded and unbounded memory for streaming workloads. The risk: recurrent state quantization error compounds over time — evaluate long-sequence accuracy, not just per-op SQNR.

Chapter 9: Vision Transformers and DINO

ViT from zero: chop the image into $p \times p$ patches (16×16 typical), flatten each to a vector, project linearly — now you have a "sentence" of patch tokens; run a standard transformer. The projection is implemented as Conv2d with kernel = stride = $p$ (identical math, one fused op — Lab 04 tests that non-overlap property). A learnable CLS token is prepended as the aggregation point; learned position embeddings restore spatial order.

Inductive bias tradeoff: CNNs hard-code locality and translation equivariance; ViTs learn relations from data — worse on small datasets, better at scale, and their global receptive field from layer 1 is why their attention maps are interesting.

DINO — labels-free training: two networks, student and teacher, same architecture. Teacher weights are an EMA of the student (never gradient-trained). Both see different augmented crops of the same image; the student is trained to predict the teacher's output distribution (cross-entropy). Collapse is prevented by two opposing mechanisms: centering (subtract a running mean from teacher logits — kills the one-hot collapse) and sharpening (teacher temperature τ_t ≈ 0.04 < student τ_s ≈ 0.1 — kills the uniform collapse). Multi-crop makes the task "predict global view from local crop", which forces object-level features — and produces the famous segmentation-quality attention maps with zero labels. Lab 04 implements every piece and tests each property in isolation.

Chapter 10: Multimodal Bridges

The design question: you have a pretrained vision encoder (CLIP) and a pretrained LLM. How do images get into the LLM?

LLaVA's answer (Lab 05): project CLIP's patch tokens through a small MLP into the LLM's embedding space and prepend them as if they were text tokens. Train only the projector first (the bridge must learn to "speak embedding-ese" before touching either tower); optionally unfreeze the LLM later for instruction tuning. The loss is computed on answer text positions only — visual positions and the prompt are masked with ignore_index — and there's a one-position causal shift at the modality boundary that is the canonical off-by-one bug (Lab 05 has a test that catches precisely it).

Alternatives: BLIP-2's Q-Former compresses patches into a fixed small set of query tokens (cheaper sequences); Flamingo interleaves gated cross-attention layers (LM untouched, text throughput preserved, but new layers to train). The engineering axis: prepended tokens inflate the KV cache by $N_{patches}$ at every layer — 576 CLIP tokens often dominate a short VQA prompt — which is why token compression matters on edge.

Lab Walkthrough Guidance

Order: Lab 01 → 02 → 03 → 04 → 05 (attention foundations first; MoE/SSM build on it; vision and multimodal reuse your attention implementation mentally).

  • Lab 01 (variants): implement MHA first and get the causal mask + KV cache correct before adding GQA/MQA (the GQA test compares against naive expansion — write that naive expansion yourself to understand it). RoPE before ALiBi.
  • Lab 02 (MoE): router → expert → block → utilization. Verify the aux-loss formula against Chapter 7's math by hand on a 2-expert toy.
  • Lab 03 (SSM): discretize_zoh is pure math — check it against scalar examples ($a=-1, \Delta=0.1 \Rightarrow \bar{a} \approx 0.905$). Then the sequential scan, then selectivity.
  • Lab 04 (DINO): follow the README's order; the no-grad-to-teacher and centering tests are the conceptual heart.
  • Lab 05 (bridge): draw the [v_1..v_N | t_1..t_T] position layout on paper before writing loss().

Success Criteria

You are ready for Phase 03 when you can, from memory:

  1. Compute the KV cache size of any (layers, kv_heads, head_dim, context, dtype) config and state the GQA reduction factor.
  2. Derive why $\sqrt{d_k}$ scaling exists from the variance of random dot products.
  3. Explain RoPE's relative-position property (rotations compose; dot product depends on $m-n$) without notes.
  4. State both MoE collapse modes and how $\sum f_i p_i$ fixes them.
  5. Write the ZOH discretization from the continuous SSM and explain what selectivity changes computationally (convolution trick breaks → scan).
  6. Explain DINO's two anti-collapse mechanisms and what happens if either is removed.
  7. Name the multimodal causal-shift bug and the KV-cache cost of prepended visual tokens.

Interview Q&A

Q: 32K context, LLaMA-3-8B (32 layers, 8 KV heads, head dim 128, FP16). KV cache size? $32 \times 2 \times 8 \times 128 \times 2,\text{B} = 131,\text{KB/token}$; × 32768 tokens ≈ 4.3 GB per sequence. With MHA (32 KV heads) it would be 17 GB — that 4× is GQA's entire reason for existing.

Q: Why does GQA quantize more gracefully than MHA? Fewer distinct KV projection matrices → fewer weight distributions to calibrate, and the shared KV heads see more diverse activation statistics (averaging over query groups), so per-channel ranges are better conditioned. Also the KV cache itself, if quantized, has fewer head-specific distributions.

Q: Your Mamba model passes per-op quantization checks but fails on long sequences. Why? Recurrent state error compounds: each step multiplies state by $\bar{A}$ and adds quantized input contributions; unlike attention (which reads exact cached K/V), small per-step errors accumulate over thousands of steps. Evaluate sequence-level metrics vs length; consider keeping the state and $\Delta$-path in higher precision.

Q: When would you pick ALiBi over RoPE? When the dominant requirement is train-short/infer-long generalization with no inference- time scaling tricks, and the tasks are recency-dominated (streaming ASR post-processing, log modeling). RoPE wins when precise relative structure matters or ecosystem tooling (RoPE-scaling for context extension) is needed.

Q: How would you adapt NPU tiling for an MoE model? Router runs dense and tiny (keep high precision); expert FFNs are the bulk — tile them as independent matmuls with token gathering by expert (sort tokens by expert id to batch); since only K of E experts are hot per token, weights can be streamed/paged per expert, and per-expert quantization scales are natural tile boundaries.

References