Transformers and Self-Supervised Language Models

Definition A tokenizer splits text into tokens from a fixed vocabulary $\mathcal{V}$, and each token is identified by an integer index. A token embedding table is a learned matrix $W_E$ of shape $\lvert \mathcal{V} \rvert \times d$, where $d$ is the model width, and the embedding of a token is the row of $W_E$ at its index. Looking up a token is the same as multiplying its one-hot row vector by $W_E$, and the lookups for an input of $n$ tokens form an $n \times d$ matrix $E$. The index is only a label, and its numerical size means nothing. A positional encoding supplies the position of each token. It is needed because self-attention without a mask is permutation equivariant. Write $f(X) = \operatorname{Attention}(XW_Q, XW_K, XW_V)$ for self-attention without a mask. For every $n \times n$ permutation matrix $\Pi$,

and layers that act on each row separately behave the same way, so without position information a reordering of the input rows only reorders the output rows. An additive encoding sets $H = E + P$, where row $i$ of the $n \times d$ matrix $P$ depends only on position $i$. The rows of $P$ can be learned, one per position up to a maximum length, or fixed, as in the sinusoidal encoding. For even $d$, with positions $i$ and coordinates counted from $0$ and $m = 0, \dots, d/2 - 1$, it is

Other methods put position information into attention itself, for example through relative positions or position-dependent rotations of queries and keys.

Definition Layer normalization acts on each position’s feature vector $\mathbf{h}$ in $\mathbb{R}^d$ separately. It computes the mean and variance over the $d$ features of that one vector,

and returns

where $\varepsilon > 0$ is a small constant, $\boldsymbol{\gamma}$ and $\boldsymbol{\beta}$ in $\mathbb{R}^d$ are a learned scale and shift, and $\odot$ multiplies entry by entry. For an array of shape batch by positions by features it normalizes along the feature axis only. Unlike batch normalization it uses no statistics from other examples or other positions, so it computes the same thing in training and in use. The constant $\varepsilon$ keeps the denominator positive. When all features of a vector are equal, $\sigma^2 = 0$, the normalized part is $\mathbf{0}$, and the result is $\boldsymbol{\beta}$.

Definition A residual connection adds a sublayer’s output to that sublayer’s input, $\mathbf{y} = \mathbf{h} + f(\mathbf{h})$, so $f$ must return vectors of the same width as $\mathbf{h}$. The running sum of the embeddings and all sublayer outputs is called the residual stream. A pre-normalization transformer block maps an $n \times d$ input $H$ to an $n \times d$ output $Z$ in two residual steps,

where $\operatorname{MHA}$ is multi-head self-attention with the block’s mask and its output projection, and $\operatorname{LN}$ acts on each row. The feed-forward network applies the same weights to every row $\mathbf{u}$, written as a row vector,

with $W_1$ of shape $d \times d_{\text{ff}}$, $W_2$ of shape $d_{\text{ff}} \times d$ and an activation $\phi$ such as ReLU. Often $d_{\text{ff}}$ is several times $d$. Attention is the only step that mixes positions. A transformer stacks several such blocks with separate parameters. Other designs normalize after each residual addition instead, add further normalization or use dropout, so the order of operations is part of the architecture and a diagram must match the implementation it describes.

Definition Next-token prediction is the training objective that asks, at every position of a sequence, for the probability of the following token given only the tokens before it. A sequence $x_1, \dots, x_T$ is preceded by a beginning symbol $x_0$, so the inputs are $(x_0, \dots, x_{T-1})$ and the targets are $(x_1, \dots, x_T)$, the inputs shifted by one place. With the prefix $x_{<t} = (x_0, \dots, x_{t-1})$, the token losses and the sequence loss are

measured in nats per token. Over a batch the mean runs over all valid target positions, and padding is left out of both the sum and the count. By the chain rule of probability the product of the conditionals is the model’s probability of $x_1, \dots, x_T$ after $x_0$, so $T \cdot L(\boldsymbol{\theta})$ is the negative log of that probability, and minimizing the summed loss over training sequences is maximum likelihood estimation. The targets come from the text itself, so the objective is self-supervised. A low loss on held-out text means the model gives high probability to text like that text. It does not show that generated statements are true, that its reasoning is reliable or that the training text was permitted for this use.

Definition A decoder-only model is a stack of transformer blocks in which every self-attention layer uses a causal mask, followed by a linear head that maps each final row to $\lvert \mathcal{V} \rvert$ logits and a softmax across the vocabulary. Many designs apply one more layer normalization before the head. Because every other operation acts on each position separately, the output at position $t$ depends only on the tokens at positions up to $t$, so one forward pass computes the next-token distributions for all positions of a training sequence at once. To generate, the model reads a prefix, chooses the next token from the distribution at the last position, appends it and repeats until it produces an end symbol or reaches a length limit. Greedy decoding takes the most probable token, and sampling draws from the distribution, possibly after a temperature or a filter changes it. Other architectures differ in what each position can access. An encoder uses self-attention without a causal mask, so each position can use valid positions on both sides. An encoder-decoder model encodes a source sequence that way, and its decoder uses causal self-attention over the output prefix together with cross-attention whose queries come from the decoder and whose keys and values come from the encoder output.

Definition Masked language modeling hides a random subset of the input tokens, usually by replacing them with a special mask token, and trains a model to predict the original tokens at the hidden positions from all visible tokens on both sides. The loss is the mean of $-\ln P_{\boldsymbol{\theta}}(x_i \mid \text{visible tokens})$ over the hidden positions only. The model is usually an encoder, without a causal mask. Its targets also come from the text itself, so the objective is self-supervised. Its conditional probabilities are not the factors of a left-to-right chain rule, so the sum of its losses is not the negative log likelihood of the sequence, and a masked model does not generate text by appending one token at a time.