LLM Architecture Public Cheat Sheet

A: The Pre-LN structure proceeds as: LayerNorm → Multi-Head Causal Self-Attention → Residual → LayerNorm → FFN → Residual. The causal mask fills the upper triangle of the attention score matrix with $-\infty$, ensuring each token can only attend to itself and preceding tokens. The FFN typically expands the dimension by $4\times$ (or $8/3\times$ for SwiGLU) before projecting back.

Follow-up: What are the pros and cons of Pre-LN vs Post-LN? Why do current LLMs generally use Pre-LN?

Hint: Pre-LN is more stable to train (smoother gradient flow), though some research suggests Post-LN has a slight edge in final performance. Pre-LN became mainstream because Post-LN is unstable in very deep models (>100 layers).

A: Input $X$ is projected by $W_Q, W_K, W_V$ to obtain $Q, K, V$; compute $\text{softmax}(QK^\top / \sqrt{d_k})V$. Time complexity $O(N^2 d)$, space complexity $O(N^2)$ (storing the attention matrix). $N$ is the sequence length, $d$ is the head dimension. The $N^2$ term is the bottleneck for long sequence processing.

Follow-up: What is the role of the $\sqrt{d_k}$ scaling factor? What happens if it is omitted?

Hint: When $d_k$ is large, elements of $QK^\top$ have variance approximately $d_k$, causing softmax saturation and vanishing gradients. Dividing by $\sqrt{d_k}$ normalizes the variance to $\sim 1$.

A: During autoregressive generation, when generating each new token, the K/V vectors of all previous tokens are unchanged. KV Cache stores the already-computed K/V in GPU memory; a new token only needs to compute its own Q and attend to the cache, avoiding redundant computation of historical tokens. This reduces the computation per decode step after prefill from $O(n \cdot d)$ (recomputing all K/V) to $O(d)$ (projecting only the new token) + $O(n \cdot d)$ (attention, but no re-projection needed).

Follow-up: How does KV Cache memory usage scale with sequence length and batch size? What compression methods exist?

Hint: memory = batch_size × seq_len × 2 × L × n_kv_heads × d_k × bytes. Compression methods include MQA/GQA, KV Cache quantization (FP8/INT8), token pruning/eviction, etc.

A:

• Greedy: take the highest-probability token at each step; deterministic, prone to repetition degeneration
• Beam Search: maintain $k$ candidates; suitable for generation with a definite target (e.g. translation); poor diversity in open-ended generation
• Top-$k$: sample from only the top-$k$ tokens by probability; fixed $k$ does not adapt to distribution shape
• Top-$p$: sample from the smallest set with cumulative probability $\geq p$; adaptive candidate set size
• Temperature: $T < 1$ sharpens the distribution (more deterministic), $T > 1$ flattens it (more random); commonly used together with top-$p$

Follow-up: Why does Beam Search perform worse than sampling for open-ended generation?

Hint: Beam Search optimizes sequence probability, tending to select "safe" high-frequency tokens, resulting in outputs that lack diversity and naturalness. Diversity penalties partially mitigate this.

A: In the attention score matrix $S \in \mathbb{R}^{N \times N}$, positions where $j > i$ (the upper triangle) are set to $-\infty$; after softmax these positions have weight 0. This ensures that token $i$, when computing attention, can only see tokens at positions $\leq i$. Decoder-only models train with next-token prediction (computing loss at all positions in parallel); the causal mask ensures that each position's prediction does not leak future information, consistent with autoregressive behavior at inference time.

Follow-up: How does a Prefix LM (e.g. T5 decoder, GLM) mix bidirectional and causal attention within the same sequence?

Hint: Tokens in the prefix portion use bidirectional attention (no mask) among themselves; the generation portion attends bidirectionally to the prefix but uses a causal mask for itself and subsequent tokens.

A: Three core roles:

1. Gradient flow: provides a bypass path around sub-layers, alleviating vanishing gradients in deep networks
2. Feature reuse: $h_l = h_{l-1} + F(h_{l-1})$, the model can selectively preserve or modify information
3. Training stability: combined with Pre-LN, enables stable training of models with hundreds of layers

Follow-up: Are there tools for analyzing information flow in the residual stream across Transformer layers?

Hint: The Logit Lens method projects intermediate hidden states through the LM head into the vocabulary space, observing how the predicted distribution changes layer by layer.

A: $d_{\text{ffn}} = 4d_{\text{model}}$ is an empirical setting from the original Transformer paper, without a strict theoretical justification. Subsequent experiments have shown this ratio works well in most settings. The SwiGLU variant uses $\frac{8}{3}d_{\text{model}}$; because it has two gate matrices, the total parameter count is roughly equal to the standard $4\times$ FFN.

Follow-up: What computational role might the FFN play in a Transformer? Is there research that treats the FFN as "knowledge storage"?

Hint: Some research (e.g. the key-value memory perspective) views the first layer of the FFN as a key matrix and the second as a value matrix, analogous to the retrieval mechanism of attention.

A: A single attention head can only learn one type of "attention" pattern. $h$ heads each learn different relational patterns in different subspaces (e.g. syntactic dependencies, coreference, local n-grams, etc.), and are fused by projection through $W^O$ after concatenation. Total parameter count is unchanged (each head has $d_k = d_{\text{model}}/h$), but expressive power is greater.

Follow-up: Do different heads actually learn different patterns? How feasible is head pruning?

Hint: Michel et al. (2019) found that a large number of heads can be removed at inference time with little performance degradation; Voita et al. (2019) analyzed the functional specialization of different heads via head importance analysis.

A:

• MHA: $h$ Q heads each have their own K/V, largest KV Cache
• MQA: all Q heads share 1 set of K/V, KV Cache reduced by $h\times$, but may sacrifice quality
• GQA: $h$ Q heads divided into $g$ groups, each group sharing 1 set of K/V, a tradeoff between MHA and MQA

Converting from MHA to GQA: average the weights of every $h/g$ K/V heads to initialize the new K/V heads, then continue training (uptrain) on a small amount of data (e.g. 5% of the original pretraining volume).

Follow-up: What attention variant does LLaMA-3 use? How is the GQA group count $g$ chosen?

Hint: LLaMA-3 uses GQA. Choosing $g$ is a tradeoff between inference efficiency and model quality, typically determined through small-scale experiments.

A:

• Absolute positional encoding (Sinusoidal/Learned): position information is added directly to the embedding, maximum length is fixed, poor extrapolation
• RoPE: rotations applied to Q/K so that attention scores depend only on relative position $m-n$; supports length extrapolation (PI/NTK/YaRN); used by the LLaMA family
• ALiBi: adds a linear bias $-m_h|i-j|$ to attention scores; no extra parameters, good extrapolation, but prefix caching requires dynamically computing bias offsets, adding implementation complexity (some inference frameworks do not support this)

Follow-up: Does RoPE implement relative encoding by modifying embeddings or by modifying attention scores? What is its mathematical relationship to absolute positional encoding?

Hint: RoPE applies rotations at the embedding level (on Q/K), but the effect is equivalent to adding a relative-position-dependent bias to the attention score. The difference from absolute positional encoding is that it achieves relative positional encoding through rotational invariance.

A: The initial vocabulary consists of all bytes/characters. Iteratively, the most frequent adjacent symbol pair in the corpus is merged into a new symbol and added to the vocabulary, until the target size is reached. Advantages: controllable vocabulary size, can handle OOV. Disadvantages: the same word may have multiple tokenizations (case, spacing variants), lower efficiency for non-English languages. SentencePiece is a language-agnostic BPE implementation operating directly on unicode bytes.

Follow-up: How does tokenizer vocabulary size affect the model? What are the considerations behind GPT-2 (50K) vs LLaMA (32K) vs LLaMA-3 (128K)?

Hint: Effects include embedding layer parameter count ($|V| \times d_{\text{model}}$), information density per token, token efficiency for multilingual text and code, and softmax computation cost. A larger vocabulary is more multilingual-friendly but increases memory.

A: Analogous to OS virtual memory paging: KV Cache is divided into fixed-size blocks, with a block table managing logical-to-physical address mappings, allowing non-contiguous physical memory allocation. This solves the memory fragmentation problem caused by requiring contiguous memory for KV Cache in standard implementations (different request lengths lead to large amounts of internal and external fragmentation), substantially improving GPU memory utilization.

Follow-up: What is the relationship between Continuous Batching and PagedAttention? What layer of problems does each solve?

Hint: Continuous batching solves the scheduling problem (iteration-level scheduling at the request level, so that the batch does not wait for the longest sequence to finish); PagedAttention solves the memory management problem (efficient KV Cache allocation). The two are complementary.

A: Replaces the FFN with $E$ expert FFNs. Each token uses a router (linear layer + softmax + top-$k$ selection) to select $k$ experts (typically $k=2$), and only the selected experts are activated for computation. Total parameters = $E \times$ expert parameters (much larger than dense), but FLOPs per token $\approx (k/E) \times$ dense FLOPs (where "dense" refers to an FFN with the same total parameter count). A load-balancing loss is added to prevent expert collapse (all tokens routing to a few experts).

Follow-up: How do the memory and compute characteristics of MoE models differ between inference and training? What special requirements does MoE impose on tensor parallelism?

Hint: At inference, each token activates only $k$ experts, but all expert parameters must be loaded into memory (high memory requirement, low compute). During training, expert-level parallelism or capacity factor control is needed to manage load.

A: Given a training FLOPs budget $C$, the optimal model parameter count $N$ and training token count $D$ should roughly satisfy $N \propto D$ (both $\propto C^{0.5}$; empirically $D\approx20N$, i.e. ~20 tokens per parameter). That is, many previous large models (e.g. 175B parameters trained on only 300B tokens) were "over-parameterized and under-trained." Training a smaller model on more data yields a better loss per FLOP.

Follow-up: Does the Scaling Law still apply in the post-training (SFT/RLHF) stage?

Hint: Research on scaling laws for post-training is still ongoing. Some work has explored the relationship between SFT data volume and model size, but RLHF scaling behavior is more complex, influenced by reward model quality, KL constraints, and other factors.

A: Formula: $\text{KV Cache} = 2 \times L \times n_{\text{kv}} \times d_k \times \text{seq\_len} \times \text{batch\_size} \times \text{bytes}$

where factor 2 accounts for K and V, $n_{\text{kv}}$ is the number of K/V heads (with GQA, $< h$), $d_k$ is the head dimension. Example with 32 layers, $n_{\text{kv}} = 8$, $d_k = 128$, BF16: approximately 128 KB per token, a sequence of 4096 tokens is approximately 0.5 GB (single sequence), and a batch of 32 is approximately 16 GB.

Follow-up: What is the effect of KV Cache quantization (e.g. FP8/INT8) on generation quality? What implementation approaches exist?

Hint: INT8 quantization generally has very little impact on quality; FP8 reduces it further. Inference frameworks such as vLLM and TensorRT-LLM already include built-in KV Cache quantization support.

A: Each token only attends to tokens within a distance of $\leq W$ (a local window), with complexity $O(NW)$. With multiple stacked layers the receptive field grows linearly with depth (theoretical receptive field $\approx L \times W$ after $L$ layers), analogous to CNN. Mistral-7B uses $W = 4096$.

Advantages: low complexity; at inference, a rolling buffer KV Cache (fixed-size circular buffer) can be used. Disadvantages: a single layer cannot build precise cross-window attention; performs poorly on needle-in-a-haystack retrieval tasks.

Follow-up: How does Mistral's rolling buffer KV cache work?

Hint: The KV Cache uses a circular buffer of fixed size $W$; new tokens overwrite the oldest ones (at position $i \bmod W$), so memory does not grow with sequence length.

A:

$$L(N, D) = \frac{A}{N^\alpha} + \frac{B}{D^\beta} + L_\infty$$

$N$ = parameter count, $D$ = number of training tokens, $L_\infty$ = irreducible loss (the entropy of the data itself). $\alpha$ and $\beta$ measure the rate of diminishing returns from parameters and data, respectively. Over a very wide range, loss maintains a power-law relationship with $N$ and $D$.

Follow-up: Why can emergent abilities (e.g. chain-of-thought reasoning) not be directly predicted from scaling laws?

Hint: Scaling laws describe the smooth decrease in loss (perplexity); emergent abilities are specific task metrics (e.g. accuracy) that suddenly jump from near-random to well-above-random at a certain scale, which may be an artifact of metric choice (Wei et al., 2022; Schaeffer et al., 2023).

A: Embedding layer parameter count = $|V| \times d_{\text{model}}$ (typically weight-tied with the LM head). Large vocabulary: larger share of embedding parameters, higher softmax computation, but higher information density per token and more friendly for multilingual text and code. Small vocabulary: fewer embedding parameters, but Chinese and similar languages require more tokens to represent the same text (longer sequences, slower inference).

Follow-up: Is there research on the scaling relationship between vocabulary size and model capability? How should the optimal vocabulary size be chosen?

Hint: Vocabulary size selection requires balancing multilingual coverage, inference efficiency, and embedding parameter overhead. The main motivation for LLaMA-3's expansion to 128K is to improve token efficiency for multilingual text and code.

A: FlashAttention solves the memory-bound problem of standard attention — needing to write the $N \times N$ attention matrix back to HBM. The core is tiling + online softmax: Q/K/V are loaded in tiles into SRAM, tiled attention computation is performed within SRAM, and via the online softmax algorithm a result mathematically equivalent to standard attention is obtained without storing the full $N \times N$ matrix. FlashAttention is an exact algorithm, not an approximation. IO complexity is reduced from $O(N^2)$ to $O(N^2 d^2 / M)$ ($M$ is SRAM size), and memory from $O(N^2)$ to $O(N)$.

Follow-up: How does online softmax guarantee mathematical equivalence without storing the full score matrix?

Hint: Online softmax maintains running-max and running-sum statistics; when processing each new block, it corrects previous blocks' contributions via rescaling, ultimately producing a result identical to global softmax (requiring an extra rescaling pass).

A:

1. Position Interpolation (PI): linearly scale position indices $m \leftarrow m \cdot L_{\text{train}} / L_{\text{test}}$; simple but loses high-frequency information
2. NTK-aware Scaling: modify the base ($10000 \to \alpha \cdot 10000$), high-frequency dimensions scale less, low-frequency dimensions scale more
3. YaRN: NTK + non-uniform interpolation + attention temperature adjustment; better performance
4. Dynamic NTK: dynamically adjust the base at inference based on actual sequence length; no retraining required

Core challenge: unseen high-frequency rotation angles during training cause attention score distribution shift; extrapolation methods fundamentally balance preserving learned patterns and adapting to new lengths.

Follow-up: What specific improvements does YaRN make over naive PI? Why is attention temperature adjustment needed?

Hint: YaRN uses different interpolation strategies for high-frequency and low-frequency dimensions, and adds an attention temperature factor $1/\sqrt{t}$ to compensate for changes in attention score magnitude caused by interpolation.

A: A small draft model generates $k$ candidate tokens in parallel; a large verifier validates all of them in a single forward pass. Validation uses rejection sampling: acceptance probability $\min(1, p_{\text{verifier}}(x) / p_{\text{draft}}(x))$; on rejection, resample from a corrected distribution $\max(0, p_v - p_d)$ normalized. Mathematically it can be shown that the final token distribution at each position is exactly $p_{\text{verifier}}$.

Throughput gain comes from: the draft model is extremely fast (small model), and verifying $k$ tokens requires only one forward pass (parallel), whereas normal generation requires $k$ passes.

Follow-up: Do the draft model and verifier need to share the same architecture? How does self-speculative decoding work?

Hint: They do not need the same architecture, but the vocabulary must be the same. Self-speculative decoding uses a portion of the same model's layers (early exit) or skips some layers (layer skipping) as the draft, avoiding the need for a separate model.

A:

1. Positional encoding extrapolation: PI/NTK/YaRN modify RoPE, requiring a small amount of continued training
2. Sliding Window Attention: local window $O(NW)$; stacking layers expands the receptive field
3. Sparse Attention: Longformer/BigBird mixed global + local + random attention
4. Ring Attention: splits long sequences across multiple devices and passes KV through all-to-all communication; theoretically supports unlimited length
5. Attention sink: retains KV of the first few tokens (StreamingLLM), addressing the "attention sink" phenomenon in sliding window attention

Follow-up: What bandwidth requirements does Ring Attention have? How does it differ from Sequence Parallelism?

Hint: Ring Attention requires high-bandwidth interconnects between devices (e.g. NVLink); communication volume is proportional to window size. The difference from traditional Sequence Parallelism is that Ring Attention's partitioning is a pipeline-style transfer along the sequence dimension, rather than simple data parallelism.

A: Each expert has a capacity limit within a batch (capacity factor × tokens/expert). Tokens that exceed the limit are dropped — that expert is skipped and the residual or zero output is used directly. During training the capacity factor is typically set to 1.0–1.25; setting it too large wastes computation, setting it too small causes token loss.

Token dropping introduces training noise and degrades gradient quality. Mitigation approaches include: increasing the capacity factor (at the cost of wasted computation), improving router design (e.g. load-balancing loss), and using expert choice routing (experts select tokens rather than tokens selecting experts).

Follow-up: What is the difference between expert choice routing and token choice routing? What are the pros and cons of each?

Hint: Token choice (standard top-k): each token selects $k$ experts, which may cause load imbalance. Expert choice: each expert selects the top-$k$ tokens, naturally load-balanced, but some tokens may not be selected by any expert.

A: During training: the full sequence is processed in parallel; an $N \times N$ causal mask matrix is constructed and Q/K/V at all positions are computed in one pass. During inference (with KV Cache): the new token's Q is a $(1, d)$ vector, K/V Cache is an $(n, d)$ matrix ($n$ = current sequence length), attention computation is $(1, d) \times (d, n) = (1, n)$, and no explicit causal mask is needed (since K/V does not contain future tokens). The prefill stage (processing the prompt) is similar to training, using a causal mask and parallel computation.

Follow-up: During prefill, if the prompt is very long (e.g. 100K tokens), where is the compute bottleneck? What optimizations exist?

Hint: The prefill bottleneck is $O(N^2)$ attention computation and $O(N)$ KV Cache writes. Optimizations include chunked prefill (processing in chunks interleaved with decode requests), FlashAttention to reduce memory and compute, and prefix caching (reusing KV Cache for requests with the same prefix).

A: Dense models: increasing parameter count = increasing compute (FLOPs/token $\propto N$). MoE models: total parameters $N_{\text{total}} \gg N_{\text{active}}$ (active parameters), FLOPs/token $\propto N_{\text{active}}$. This means MoE can have a larger "knowledge capacity" under the same compute budget.

From a scaling law perspective: MoE loss is mainly determined by active parameters $N_{\text{active}}$ and data volume $D$ (similar to dense), but increasing the number of experts still yields gains when data is abundant (different experts can specialize in different domains of knowledge). However, MoE scaling efficiency is constrained by router quality, load balance, and expert utilization.

Follow-up: What innovations does DeepSeek's MoE design (e.g. DeepSeek-V2/V3) introduce? What are the benefits of separating shared experts from routed experts?

Hint: DeepSeek-V2/V3 introduces a design with shared experts (all tokens pass through) + routed experts (selected by the router); shared experts handle general capabilities, routed experts handle domain knowledge. There are also fine-grained expert splitting (more but smaller experts) and other design choices.

The core idea of Online Softmax is that softmax can be computed incrementally by maintaining a running maximum and a correction factor. After dividing Q and K into blocks (tiles), local $QK^\top$ scores are computed block by block; each block yields a local max $m_{\text{local}}$. Comparing with the current global max $m_{\text{global}}$, the previously accumulated $\sum \exp$ is multiplied by a correction factor $e^{m_{\text{global}} - m_{\text{new}}}$, and $m_{\text{global}}$ is then updated. In this way, the partial sum of each block can be correctly rescaled, and the final result is element-wise identical to computing the full softmax at once. The output $O$ accumulates in a similarly online weighted manner. Throughout the entire process, the $N \times N$ matrix never fully exists in HBM; only local tiles are processed in SRAM.

Follow-up: If Online Softmax needs to retroactively correct previously accumulated $O$ when processing each tile, is the final output $O$ written back in one shot or corrected multiple times? How does FlashAttention avoid extra memory writes from these corrections?

In Post-LN, LayerNorm comes after the residual: $\text{output} = \text{LN}(x + \text{SubLayer}(x))$. In deep networks, gradients on the residual path must pass through the LN at every layer; the LN Jacobian depends on the activation statistics of that layer, causing inter-layer gradient scale coupling. As depth increases, gradients in early layers tend to grow or decay exponentially (a variant of gradient explosion/vanishing), requiring careful warmup strategies. In Pre-LN, $\text{output} = x + \text{SubLayer}(\text{LN}(x))$, LN is at the sub-layer input side, and the residual path has no LN blockage — gradients can flow directly along the residual (similar to ResNet's identity shortcut), allowing deep models to train stably without warmup. The tradeoff is that Pre-LN may yield slightly lower final performance than carefully tuned Post-LN.

Follow-up: DeepNorm (a variant proposed by Microsoft) stabilizes training on top of Post-LN by adjusting the residual scaling coefficient $\alpha$ and initialization. What is its core mathematical principle?

Router collapse takes several forms: (1) complete collapse — almost all tokens select only one or two experts, leaving other experts with no gradient updates; (2) partial imbalance — most experts are utilized but a few "star experts" are overloaded; (3) oscillatory cycling — expert utilization alternates across training steps without converging. The auxiliary loss ($\mathcal{L}_{\text{aux}} = \alpha \cdot E \sum f_i p_i$) is only the most basic approach. Other strategies include: Expert-Choice Routing — each expert actively selects its top-$k$ tokens (rather than tokens selecting experts), naturally guaranteeing load balance; Random routing with noise — adding tunable noise to router logits to prevent deterministic collapse; dynamic capacity factor adjustment — adaptive capacity based on training progress; and using a smaller learning rate or separate optimizer for router parameters to avoid excessive fluctuation in routing decisions.

Follow-up: In Expert-Choice Routing, some tokens may be selected by zero or multiple experts simultaneously. What are the effects on model training and inference respectively, and how are they handled?

PagedAttention maintains a block table (logical-to-physical page mapping) for each request. When multiple requests need to share the same KV Cache prefix (e.g. a system prompt, or multiple branches of the same parent sequence in beam search), they can share the same physical block, avoiding duplicate storage. In this case the block's reference count > 1. When a request needs to modify (append new tokens to) a shared block's content, Copy-on-Write is triggered: a new physical block is allocated, the original block's content is copied, the original block's reference count is decremented, and the new data is written. This mechanism is most beneficial in: (1) beam search, where multiple beams share a prefix; (2) parallel sampling (generating multiple responses for the same prompt); (3) prefix caching (reusing the system prompt's KV across multi-turn conversations).

Follow-up: How should PagedAttention's block size be chosen? What problems arise from a block size that is too small or too large? What is the typical block size in production systems?

Main approaches include: (1) KV Cache quantization — quantizing K/V from FP16 to INT8/INT4 or lower precision, reducing memory but requiring care about attention precision loss, with V quantization having a larger impact on output quality; (2) Token eviction — discarding KV of less important historical tokens based on a policy, e.g. H2O (Heavy Hitter Oracle) retains tokens with high attention scores, StreamingLLM retains "attention sinks" (first few tokens) + a sliding window; (3) KV merging — merging K/V vectors of adjacent tokens into one (e.g. via weighted average or PCA), trading precision for space; (4) cross-layer KV sharing — adjacent layers share a single set of K/V, reducing total cache volume. Each method is a three-way tradeoff of precision, memory, and compute overhead: quantization has the lowest compute overhead but is precision-bounded; eviction strategies are simple but may discard critical information; merging strategies are flexible but introduce additional computation.

Follow-up: StreamingLLM observes the attention sink phenomenon — the first few tokens of the model receive high attention scores regardless of semantic relevance. What do you think is the underlying reason for this?

In the standard ReLU FFN $\text{FFN}(x) = W_2 \cdot \text{ReLU}(W_1 x)$, ReLU is a hard element-wise gate (output is 0 or linear), with a crude information pathway. SwiGLU introduces multiplicative gating: $\text{SwiGLU}(x) = (\text{Swish}(xW_1) \odot xW_3) W_2$, where $xW_3$ acts as the gate and $\odot$ is element-wise multiplication. This multiplicative interaction gives the FFN richer feature combination capability — the element-wise product of two linear transformations is equivalent to a bilinear operation, able to model more complex feature dependencies. Swish ($x \cdot \sigma(x)$) is itself a smooth, non-monotonic activation that provides smoother gradient flow than the sparsity of ReLU. Empirically, when total parameter count is held constant (reducing $d_{\text{ffn}}$ from $4d$ to approximately $\frac{8}{3}d$ to account for the additional gate matrix $W_3$), SwiGLU consistently outperforms ReLU/GELU FFN on multiple benchmarks.

Follow-up: The GLU family (including ReGLU, GeGLU, SwiGLU) shares the same gating framework. Why was SwiGLU chosen as the mainstream variant (e.g. LLaMA, Mistral family) rather than the others?

Speculative Decoding's speedup depends on the acceptance rate (the proportion of draft model predictions accepted by the large model). Scenarios where speedup degrades include: (1) highly divergent distributions — when the draft model and verifier distributions differ greatly (e.g. different architecture families), many tokens are rejected and speculation is nearly useless; (2) highly random generation — at higher temperature, the probability distribution over correct tokens is flat, making it harder for the draft model to match the verifier's sampling; (3) structured/constrained output — e.g. in JSON format generation, certain positions have only a few legal tokens and the draft model may not cover them; (4) draft model too small — quality is too low, causing $\mathbb{E}[\text{accepted tokens}]$ to approach 1. Self-Speculative Decoding (e.g. Medusa, EAGLE) avoids using an independent draft model and instead makes lightweight predictions on top of the large model itself: for example adding extra prediction heads after the final layers (Medusa), or using a small draft decoder that reuses the large model's hidden states (EAGLE), eliminating the memory and loading overhead of a separate small model.

Follow-up: The Medusa method adds multiple prediction heads in parallel on top of the large model, each responsible for predicting the $k$-th future token. How are these heads trained? Why can a single head not simply be shared to predict multiple positions?

The Chinchilla Scaling Law optimizes training loss given a FLOPs budget. But in actual deployment, training is a one-time cost while inference is an ongoing cost — a model may serve billions of queries. A smaller model has lower FLOPs/decode token at inference, higher throughput, lower latency, and lower deployment cost. Therefore, if a smaller model is trained on far more tokens than the Chinchilla optimum, even though its training loss is worse than a larger model trained with the same FLOPs, the smaller model's cost-efficiency at inference may be significantly better — i.e. it can serve more requests under the same inference budget with limited performance loss. The LLaMA family pioneered this approach: models obtained with fewer parameters and more training tokens outperform larger models trained at the Chinchilla-optimal ratio in inference efficiency. The core logic is to shift the optimization objective from "minimize training loss" to "achieve target performance at minimum inference cost."

Follow-up: If inference cost is also in the optimization objective, how would you define an "inference-optimal" model size choice from the perspective of training token count, model parameter count, and inference hardware constraints?