Transformers and Pretraining

Left-to-right language modeling (also called causal, autoregressive or decoder): given input tokens, predict output tokens one by one, conditioning on the prior context.

It is embedded (row of the embedding matrix \(E\)) plus a positional encoding, passed up through \(L\) transformer blocks (each adding attention and feedforward outputs to its residual stream), then the language modeling head (unembedding + softmax) outputs a distribution over next tokens.

The view of a token’s processing as one stream of \(d\)-dimensional vectors: the input embedding is passed straight up, and each component (attention, feedforward) adds its output to it. At any layer the stream is the sum of the original embedding and all previous outputs.

Multi-head (self-)attention. It builds contextual representations of a token by integrating information from surrounding tokens, and it moves information from other tokens’ residual streams into the current one.

The final unembedding matrix \(U\) plus softmax that map the last block’s output embedding to a probability distribution over next tokens.

From about 12 to 96 or more blocks. Over a context window it maps input vectors \((x_1, \ldots, x_n)\) to output vectors \((h_1, \ldots, h_n)\) of the same length.

A static vector is the same in every context, but it refers to the chicken in “…because it was too tired” and to the road in “…because it was too wide”. Its meaning depends on context.

Pronouns (it → chicken or road), agreement (“The keys to the cabinet are on the table”: the verb agrees with a distant subject), and word senses (bank as riverside because of pond).

Representations of a token’s meaning in its context. Layer by layer, a transformer combines a token’s representation from the previous layer with information from neighboring tokens, giving richer contextual representations at each layer.

The mechanism that weighs and combines the representations of other tokens in the context at layer \(k\) to build the representation of a token at layer \(k+1\).

Mostly to chicken and road: at that point it could corefer with either, so its representation should draw on both.

The chicken didn’t cross the road because it

darker = more attention from “it” (illustrative)

Its own input \(x_i\) and the representations of all prior tokens \(x_1 \ldots x_{i-1}\) in the context window, but no tokens after \(i\). It maps an input sequence to an output sequence of the same length, computing every position in parallel.

The output is a weighted sum of the prior inputs, \(a_i = \sum_{j \le i}\alpha_{ij}x_j\), with weights from a softmax over dot-product similarities: \(\alpha_{ij} = \text{softmax}(x_i \cdot x_j)\) for \(j \le i\).

\(x_i\) itself, since a vector is very similar to itself and gets a high dot product; other similar context tokens also get some weight.

Query: the current element being compared to preceding inputs. Key: a preceding input being compared to the current element to get a similarity weight. Value: a preceding element’s content that gets weighted and summed into the output.

By projecting each input with learned matrices: \(q_i = x_iW^Q\), \(k_i = x_iW^K\), \(v_i = x_iW^V\).

Dot products can be arbitrarily large, and exponentiating large values causes numerical issues and loss of gradients in training. Scaling by the square root of the query/key dimensionality keeps scores in a reasonable range.

\(\text{score}(x_i, x_j) = \frac{q_i \cdot k_j}{\sqrt{d_k}}\); \(\alpha_{ij} = \text{softmax}(\text{score}(x_i, x_j))\) for \(j \le i\); \(\text{head}_i = \sum_{j \le i}\alpha_{ij}v_j\); \(a_i = \text{head}_i W^O\).

It reshapes the head output back to the model dimension. \(W^Q, W^K\): \([d \times d_k]\); \(W^V\): \([d \times d_v]\); the head output is \([1 \times d_v]\); \(W^O\): \([d_v \times d]\), so input and output are both \([1 \times d]\).

It makes the transformer very modular: blocks and components can be stacked and combined freely.

Each head can attend to the context for a different purpose, e.g. specializing in different linguistic relationships or patterns. A heads run in parallel at the same depth, each with its own \(W^{Q_c}, W^{K_c}, W^{V_c}\).

The \(A\) head outputs (each \([1 \times d_v]\)) are concatenated into \([1 \times Ad_v]\), then projected by \(W^O\) of shape \([Ad_v \times d]\): \(a_i = (\text{head}^1 \oplus \cdots \oplus \text{head}^A)W^O\).

The head dimensionality is set to \(d_v = d/A\), so the heads together have the size of one full-width head.

A multi-head self-attention layer, a feedforward layer, residual connections, and layer normalization (“layer norm”).

Input → layer norm → attention, whose output is added to the input; the sum → layer norm → feedforward, whose output is added again, giving \(h_i\).

A fully connected 2-layer network (one hidden layer) applied to each token position independently. The same weights are shared across positions within a layer, but differ between layers.

Usually \(d_{ff} \approx 4d\) (e.g. 2048 for \(d = 512\)). The feedforward networks thus hold most of a transformer’s parameters and seem to store most of its factual knowledge.

A gated feedforward activation used by most LLMs: a component-wise product of two linear transforms of the input, one acting as the value and one, passed through the Swish function \(\text{Swish}_\beta(x) = x\,\sigma(\beta x)\), as a gate: \(\text{FFN} = (\text{Swish}(xW_1 + b_1) \otimes (xW_2 + b_2))W_3 + b_3\).

A normalization that keeps hidden values in a range good for gradient-based training. Despite its name it is applied to a single token’s embedding vector, not a whole layer.

Compute the mean \(\mu = \frac{1}{d}\sum_j x[j]\) and standard deviation \(\sigma = \sqrt{\frac{1}{d}\sum_j (x[j] - \mu)^2}\), normalize, and apply a learnable gain and offset: \(\text{LayerNorm}(x) = \gamma\frac{x - \mu}{\sigma} + \beta\).

A simpler norm used by many modern networks: it rescales the vector by its root-mean-square, skipping the mean-centering step.

\(t^1_i = \text{LayerNorm}(x_i)\); \(t^2_i = \text{MultiHeadAttention}(t^1_i, [t^1_1, \ldots, t^1_{i-1}])\); \(t^3_i = t^2_i + x_i\); \(t^4_i = \text{LayerNorm}(t^3_i)\); \(t^5_i = \text{FFN}(t^4_i)\); \(h_i = t^5_i + t^3_i\).

Multi-head attention: it is the only component that takes input from other tokens’ residual streams, moving information from neighboring streams into the current one (in different subspaces).

At early blocks it mainly represents the current token; at the highest blocks it usually represents the following token, since the model is trained to predict the next token.

Prenorm (most common today) applies layer norm before attention and the FFN, and needs one extra layer norm at the very top. The original postnorm transformer applied it after them; prenorm works better.

The block’s computation for each position is independent of the others given the inputs, so all \(N\) input embeddings can be packed into a matrix \(X\) of shape \([N \times d]\) and processed with matrix multiplications.

\(Q = XW^Q\), \(K = XW^K\), \(V = XW^V\); then \(\text{head} = \text{softmax}\big(\text{mask}\big(\frac{QK^\top}{\sqrt{d_k}}\big)\big)V\) and \(A = \text{head}\,W^O\). \(QK^\top\) is an \([N \times N]\) matrix of all query–key comparisons.

\(QK^\top\) compares each query with every key, including future ones, which would let the model see the word it must predict. The upper triangle (\(j > i\)) is set to \(-\infty\), which the softmax turns into 0.

The softmax exponentiates scores: \(e^{-\infty} = 0\) removes a position entirely, while a score of 0 would still give it weight \(e^0 = 1\) before normalization.

Quadratic: each layer computes dot products between every pair of tokens (the \(N \times N\) matrix), which makes very long documents expensive.

Each head \(c\) computes \(Q^c, K^c, V^c\) from \(X\) and its own masked softmax attention, giving \([N \times d_v]\); the \(A\) outputs are concatenated into \([N \times Ad_v]\) and projected by \(W^O\) to \([N \times d]\).

\(O = X + \text{MultiHeadAttention}(\text{LayerNorm}(X))\) and \(H = O + \text{FFN}(\text{LayerNorm}(O))\), where the FFN and layer norm are applied to each row in parallel. Input and output are both \([N \times d]\), so layers stack.

For each of the \(N\) tokens, look up its token embedding (row of \(E\), shape \([|V| \times d]\), selected by token id or by a one-hot vector) and add a positional embedding for its position.

Token embeddings don’t depend on position, and attention by itself is order-insensitive, so word order must be represented explicitly.

Randomly initialized embeddings for each position up to a maximum length, learned along with the other parameters (a matrix \(E_{pos}\)), added to the token embedding.

Later positions have fewer training examples, so their embeddings may be poorly trained and generalize badly. Alternatives: sinusoidal embeddings (static sine/cosine functions of position, as in the original transformer), or relative position methods applied inside attention at each layer, most popularly RoPE (Rotary Position Embedding).

They may capture relationships among positions, e.g. that position 4 is closer to position 5 than to position 17.

Additional neural circuitry added on top of the basic transformer to apply it to a task; the LM head is the circuitry needed for language modeling.

An n-gram model conditions on \(n-1\) words; a transformer on its whole context window, which can be hundreds of thousands of tokens or more.

From the last layer’s output for token \(N\), \(h^L_N\) \([1 \times d]\): logits \(u = h^L_N E^\top\) of shape \([1 \times |V|]\), then probabilities \(y = \text{softmax}(u)\). In training it predicts at every position, in inference just at the last.

Using the same weights for two matrices: the embedding matrix \(E\) maps one-hot tokens to embeddings at the input, and its transpose \(E^\top\) maps the final embedding back to scores over the vocabulary. Training makes \(E\) good at both mappings.

They correspond roughly to the decoder half of the original encoder-decoder transformer (built for machine translation); the causal LM paradigm later used only that part.

Choosing which token to generate based on the model’s probabilities. Repeatedly choosing the next token conditioned on previous choices, left to right, is causal or autoregressive generation.

Always generating the most probable token: \(\hat{w}_t = \arg\max_{w \in V} P(w \mid w_{<t})\). A greedy algorithm makes locally optimal choices.

Its choices are extremely predictable, so the text is generic and repetitive; it is even deterministic (same context → same output). An extension, beam search, works for constrained tasks like MT.

Repeatedly choosing the next token at random according to its probability under the model, conditioned on previous choices, until the end-of-sequence token: \(w_i \sim p(w_i \mid w_{<i})\).

The long tail of odd, low-probability tokens adds up to a non-trivial share of the probability mass, so weird tokens get chosen often enough to produce strange sentences.

Quality vs diversity: emphasizing the most probable words gives more accurate, coherent, factual but boring and repetitive text; giving more weight to middle-probability words gives more creative, diverse but less factual, less coherent text.

Divide the logits by a temperature \(\tau\) before the softmax: \(y = \text{softmax}(u/\tau)\). With \(\tau < 1\), high-probability tokens become more likely and rare ones less likely; as \(\tau \to 0\) it approaches greedy decoding; \(\tau > 1\) flattens the distribution toward uniform.

It makes the scores larger in magnitude, and the softmax pushes larger differences further toward 1 for the top value and 0 for the rest.

Thermodynamics: a system at high temperature is flexible and explores many states; at low temperature it stays in a subset of low-energy (better) states.

Keep only the \(k\) most probable tokens, renormalize their probabilities, and sample from them. \(k = 1\) is greedy decoding; larger \(k\) gives more diverse but still good text.

A fixed \(k\) ignores the shape of the distribution: sometimes the top 10 words hold almost all the mass, sometimes only a small part. Top-p keeps the smallest set of top words whose probabilities sum to at least \(p\), so the pool grows and shrinks with the model’s confidence.

Top-k first as a hard filter to eliminate garbage, then top-p to adapt to the model’s confidence.

No gold labels need to be added: the natural sequence of words supervises itself, as the model is trained to predict the true next token at every position.

Cross-entropy between the predicted distribution and the one-hot true next token, \(-\sum_{w}y_t[w]\log\hat{y}_t[w]\), which reduces to \(-\log\hat{y}_t[w_{t+1}]\): the negative log probability of the correct next token.

Always feeding the model the correct history \(w_{1:t}\) to predict the next token, rather than its own previous predictions.

The average over positions, \(\frac{1}{T}\sum_t -\log\hat{y}_t[w_{t+1}]\), minimized by gradient descent with backpropagation. All weights are updated, including the embedding matrix \(E\).

The desired outputs are known in advance and the causal mask prevents each position from attending to its own target, so all \(N\) positions are scored against their true next tokens in one pass: \(N\) training examples at once.

Completely, with text; shorter documents are packed together, separated by a special end-of-text token. Batch sizes are very large (millions of tokens).

\(\text{Perplexity}(w_{1:T}) = P(w_{1:T})^{-1/T}\), the inverse probability normalized by length. It is the exponent of the mean cross-entropy loss (in nats, since log is natural log).

It depends on the number of tokens in a text, so it is very sensitive to the tokenization algorithm. It is best used to compare models with the same tokenizer.

Rotary position embeddings (RoPE), RMSNorm, and gated feedforward layers with SwiGLU.

An inference-time optimization: since tokens are generated one at a time, the key and value vectors of all previous tokens are stored in memory and reused, so only the new token’s query, key and value must be computed.

At inference, tokens are generated iteratively one by one; the future tokens don’t exist yet, so each new token’s attention must be computed when it is produced.

Mixture-of-Experts models, Grouped-Query Attention (GQA), and quantization.

Number of layers \(L\), model dimension \(d\), number of attention heads \(A\), feedforward dimension \(d_{ff}\), vocabulary size \(|V|\), context length, and total parameter count.

Jelinek’s team used it for the whole set of linguistic influences on word-sequence probabilities (grammar, semantics, discourse, speaker characteristics), contrasted with the acoustic model.

Discriminatively trained self-supervised word predictors (Rosenfeld’s maxent LMs at IBM) and pretrained embeddings (LSA). Bengio’s network learned embeddings as part of predicting the next word.

Probably the RNN language model of Mikolov et al. (2010, 2011).

It pretrained a biLSTM with a language-model objective on a large corpus, then fine-tuned and added task heads, making it clear that language models could be a general solution for NLP problems.

That a transformer-based autoregressive LM could perform many NLP tasks, like summarization and question answering, zero-shot.

A general term for using large language model technology across domains (vision, speech, genetics…), where the elements need not be words.

Attention (encoder-decoder attention from Graves 2013 and Bahdanau et al. 2015, extended to self-attention) and memory networks, which contributed the key–query–value terminology for matching a query against keys in an associative memory.