--- date: '2024-02-07' description: Attention, next-token training, decoding, and particle-based steering in transformer language models. id: Transformers layout: technical modified: 2026-10-23 16:58:14 GMT-04:00 seealso: - '[[thoughts/LLMs|LLMs]]' - '[[thoughts/Embedding|embedding]]' - '[[thoughts/PD disaggregated serving|disaggregated serving]]' - '[[thoughts/KV compression|KV compression]]' - '[[thoughts/KV offloading|KV offloading]]' - '[[thoughts/mechanistic interpretability|mechanistic interpretability]]' - '[[thoughts/Speculative decoding|Speculative decoding]]' - '[[/thoughts/tsfm|toronto school of foundational modelling]]' tags: - ml - technical title: Transformers created: '2024-02-07' published: '2024-02-07' pageLayout: technical slug: thoughts/Transformers permalink: https://aarnphm.xyz/thoughts/Transformers.md generator: quartz: v4.6.0 hostedProvider: Cloudflare baseUrl: aarnphm.xyz full: https://aarnphm.xyz/llms-full.txt --- A transformer combines [[thoughts/Attention#Multi-head Attention|attention]] with position-wise feed-forward layers, residual connections, and normalization. Attention lets each position read weighted combinations of other positions’ value vectors. The weights come from query–key scores computed from the current representations using learned projections. A high weight says where information enters that computation; it does not identify a high-entropy token or guarantee semantic importance. \[@vaswani2023attentionneed\] The architecture and the training objective are separate choices. The original transformer had an encoder and a decoder for translation. Here I focus on a decoder-only [[thoughts/Autoregressive models|autoregressive]] language model: read a prefix, produce a distribution for the next token, append a token, repeat. [1706.03762](https://arxiv.org/html/1706.03762v7#S3) \[@vaswani2023attentionneed\] . There is also a more recent [blog post](https://magazine.sebastianraschka.com/p/beyond-standard-llms) that goes beyond just a simple GPT-like architecture (probably a bit more updated with Kimi, Gated Delta Rule, et al.) ## internals See also: [transformers from scratch](https://e2eml.school/transformers.html), @geva2021transformerfeedforwardlayerskeyvalue Token IDs select rows from an embedding table. Position information enters through positional embeddings or an operation such as rotary position encoding. Each block updates those representations through attention and a feed-forward network. A final projection produces one logit per vocabulary item. [GPT-2 report, §2.3](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) ```mermaid flowchart LR A["token IDs"] --> B["token embeddings + position information"] B --> C["causal transformer blocks"] C --> D["hidden states"] D --> E["vocabulary logits"] E --> F["next-token distribution"] ``` For one attention head, with queries $Q$, keys $K$, values $V$, and key width $d_k$: $$ \operatorname{Attention}(Q,K,V) =\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}+C\right)V, \qquad C_{t,j}=\begin{cases}0&j\leq t,\\-\infty&j>t.\end{cases} $$ The softmax runs across key positions. The causal mask $C$ prevents position $t$ from reading future tokens. Multiple heads use separate projections, then combine their outputs. [1706.03762](https://arxiv.org/html/1706.03762v7#S3.SS2) \[@vaswani2023attentionneed\] For {{sidenotes[^a batch]}} of $B$ sequences with $T$ input positions and vocabulary $\mathcal V$, logits have shape $B\times T\times |\mathcal V|$. Let $y_{b,t}=x_{b,t+1}$ be the target after input position $t$, and let $m_{b,t}$ equal one exactly where that target should contribute to training. With $N_{\mathrm{sup}}=\sum_{b,t}m_{b,t}>0$: {{sidenotes[a batch]}}: A batch can contain padding or targets excluded from supervision, so the denominator counts supervised positions. For example, sequences with three and one supervised targets contribute four terms in total. Using the padded batch size would change the loss scale when padding changes. Training supplies the actual prefix at every position. The causal mask permits all these next-token predictions in one forward pass. During generation, each sampled token becomes part of the next prefix. $$ \mathcal L(\theta) =-\frac{1}{N_{\mathrm{sup}}} \sum_{b=1}^{B}\sum_{t=1}^{T} m_{b,t}\log p_\theta(y_{b,t}\mid x_{b,\leq t}). $$ ## memory limitations. See also: [[thoughts/state-space models|state-space models]], [Karpathy’s linked thread](https://x.com/karpathy/status/1691571869051445433). A transformer uses memory for weights, temporary activations, and, during cached decoding, previous keys and values. Dense attention over a sequence of length $T$ computes $T^2$ query–key scores per head. Implementations can avoid storing the complete score matrix; that changes memory traffic and storage without removing all those interactions. [2009.06732](https://arxiv.org/abs/2009.06732) \[@tay2022efficienttransformerssurvey\] For an ordinary uncompressed KV cache with $L$ layers, $B$ sequences of length $T$, $h_{\mathrm{KV}}$ key/value heads, equal key and value width $d_h$, and $s$ bytes per element: $$ \operatorname{KVBytes}=2BLTh_{\mathrm{KV}}d_hs. $$ The factor of two counts keys and values. Under these assumptions, doubling the context doubles cache storage. Grouped-query attention reduces the number of KV heads; compression and offloading change what is stored or where it lives. The serving bottleneck also depends on how quickly hardware can move those bytes. [2403.14123](https://arxiv.org/abs/2403.14123) \[@gholami2024aimemorywall\] ## inference. Prefill processes a prompt and builds its KV cache. Decode extends the sequence, usually one token per iteration. [[thoughts/PD disaggregated serving#prefill/decode|Disaggregated serving]] runs those stages on separate workers, which requires transferring or sharing the cache between them. ### sampling Given finite logits $z_i$ over vocabulary $\mathcal V$, the base distribution is $$ p_i=\frac{e^{z_i}}{\sum_{j\in\mathcal V}e^{z_j}}. $$ A sampler changes this distribution before drawing a token. The following definitions describe individual operations. Their composition needs an explicit order. #### temperature For $\tau>0$: $$ p_i^{(\tau)}=\frac{e^{z_i/\tau}}{\sum_j e^{z_j/\tau}}. $$ Let $A=\{i:z_i=\max_j z_j\}$ be the set of tied maxima. Then $$ \lim_{\tau\to0^+}p_i^{(\tau)}=\frac{\mathbf1[i\in A]}{|A|}, \qquad \lim_{\tau\to\infty}p_i^{(\tau)}=\frac{1}{|\mathcal V|}. $$ The low-temperature limit is uniform over exact ties. An implementation’s greedy mode can choose one tied token through a separate tie-breaking rule. For fixed finite logits, entropy is nondecreasing with temperature: $$ H(p^{(\tau)})=-\sum_i p_i^{(\tau)}\log p_i^{(\tau)}, \qquad \frac{dH}{d\tau}=\frac{\operatorname{Var}_{p^{(\tau)}}(z)}{\tau^3}\geq0. $$ This follows by differentiating the softmax partition function. It concerns the next-token distribution before any truncation. It does not establish the quality or diversity of complete answers. #### top-k Choose a deterministic ordering for probability ties, and let $S_k$ contain the first $k$ tokens in descending probability order. Keep those tokens and renormalize: $$ p_i^{(k)}=\frac{p_i\mathbf1[i\in S_k]}{\sum_{j\in S_k}p_j}. $$ Exactly $k$ candidates survive under this definition. A threshold implementation that keeps every token tied with the $k$th token can retain more. The parameter fixes a count even when probability is distributed unevenly across those candidates. #### top-p (nucleus) Write the descending probabilities as $p_{(1)},\ldots,p_{(|\mathcal V|)}$. For a mass threshold $\rho\in(0,1]$, choose $$ k_\rho=\min\left\{k:\sum_{j=1}^{k}p_{(j)}\geq\rho\right\}. $$ Retain the first $k_\rho$ tokens, including the token that crosses the threshold, and renormalize. The number retained changes with the distribution. A direct implementation sorts the vocabulary; this is an implementation choice, not a requirement that every sampler use a full sort. \[@holtzman2020curiouscaseneuraltext\] [1904.09751](https://arxiv.org/abs/1904.09751) \[@holtzman2020curiouscaseneuraltext\] #### min-p For a relative threshold $m\in[0,1]$, retain $$ S_m=\{i:p_i\geq m\max_j p_j\}. $$ Then renormalize on $S_m$. At $m=0$, every token survives; at $m=1$, only the maxima survive. For positive $m$ after temperature scaling, the same condition can be written without computing probabilities: $$ z_i-z_{\max}\geq\tau\log m. $$ This makes the interaction with temperature explicit. Increasing $\tau$ widens the allowed logit gap for fixed $m$. Nguyen et al. report gains from min-p, including at high temperatures. \[@nguyen2025turningheatminpsampling\] A subsequent [2506.13681](https://arxiv.org/html/2506.13681v1) \[@schaeffer2025minpmaxexaggerationcritical\] disputes the claimed advantage after checking the evaluations and matching hyperparameter searches. The relative probability floor defines the method; superiority over top-p remains an empirical claim to evaluate for the model, task, and tuning budget. #### pipeline order Filters can change the probabilities that later filters receive. Take $$ p=(0.60,0.25,0.15),\qquad k=2,\qquad\rho=0.65. $$ Top-p alone retains the first two tokens. If top-k runs first and its output is renormalized, the first token has probability $0.60/0.85\approx0.706$. Top-p then retains only that token. Reversing these operations retains two tokens. A serving library’s sampler determines its order, renormalization, and tie rules. Check the deployed version before assuming a pipeline. The [vLLM sampler source](https://github.com/vllm-project/vllm/blob/main/vllm/v1/sample/sampler.py) is an implementation reference whose contents can change. [[thoughts/Speculative decoding#von Neumann acceptance rejection|See also: speculative decoding]]. ### KV A KV cache stores the key and value representations from earlier positions at each cached layer. In a causal model, appending a token leaves those earlier representations unchanged, so the next query can reuse them. Caching avoids recomputing the prefix; the query still has to attend over the available keys and values. Mooncake organizes serving around this cache and [[thoughts/PD disaggregated serving|prefill and decode disaggregation]]. Its cache transfers and placement policies are system-level choices. \[@qin2024mooncakekvcachecentricdisaggregatedarchitecture\] ```jsx imports={Zoomable,MooncakeFlow} ``` Prefill streams KV one layer at a time. Before layer $\ell$ runs attention, the node waits for that layer’s cached prefix to arrive and starts loading layer $\ell+1$; after attention, the messenger ships layer $\ell$ to the decode node. Overlap hides much of the transfer time while each layer has enough computation to cover it. The first load and final transfer still add latency. A large cached prefix leaves less compute to hide communication; congestion or a slower link can make transfers dominate. [2407.00079](https://arxiv.org/html/2407.00079v5#S5.SS2) \[@qin2024mooncakekvcachecentricdisaggregatedarchitecture\] . ```jsx imports={Zoomable,MooncakeLayerStream} ``` Rate source: [Mooncake Transfer Engine](https://github.com/kvcache-ai/Mooncake#mooncake-transfer-engine). The benchmark reports its payload as 40 GB. ### next-token prediction. Ordinary autoregressive sampling draws one token from the distribution conditioned on the existing prefix: $$ x_{t+1}\sim p_\theta(\cdot\mid x_{\leq t}), \qquad p_\theta(x_{1:T})=\prod_{t=1}^{T}p_\theta(x_t\mid x_{