Transformer
A transformer combines attention layers with feedforward networks, residual connections, and normalization. For text generation, we train it to predict the next token from the tokens that come before it. We will assemble these pieces and follow the computation through training and generation.
Transformer Blocks
A single attention head gives each query one distribution over keys. Yet a token may need several relationships at once: a subject, a nearby modifier, and a matching delimiter. Multi-head attention creates several routing patterns in parallel.
Let \(\mathbf X\in\mathbb R^{n\times d}\) contain \(n\) token representations of width \(d\). With \(h\) heads, head \(r\) uses its own projections into a narrower space, typically of width \(d_k=d_v=d/h\): \[ \operatorname{head}^{(r)} =\operatorname{softmax}\!\left( \frac{\mathbf Q^{(r)}(\mathbf K^{(r)})^\top}{\sqrt{d_k}} \right)\mathbf V^{(r)}. \] Concatenate the \(h\) outputs and mix them with an output projection: \[ \operatorname{MultiHead}(\mathbf X) =\left[\operatorname{head}^{(1)}\mid\cdots\mid\operatorname{head}^{(h)}\right] \mathbf W_O. \] Each head has shape \(n\times d_v\), their concatenation has shape \(n\times hd_v\), and \(\mathbf W_O\in\mathbb R^{hd_v\times d}\) returns it to the model width. Different heads can learn different notions of relevance because they have different query and key projections. The output projection lets information from their separate subspaces interact.
Attention routes values across positions, but a model also needs to transform the information arriving at each position. A transformer therefore follows attention with the same MLP at every token: \[ \operatorname{MLP}(\mathbf u) =\mathbf W^{(2)}a(\mathbf W^{(1)}\mathbf u+\mathbf b^{(1)})+\mathbf b^{(2)}. \] Here \(a\) is an elementwise activation, the hidden width is \(m\), \(\mathbf W^{(1)}\in\mathbb R^{m\times d}\), and \(\mathbf W^{(2)}\in\mathbb R^{d\times m}\), with matching bias vectors. This is the column-vector convention from neural networks. Applied to \(\mathbf X\), the MLP acts on each row’s feature vector separately. Attention mixes across positions. The MLP mixes and nonlinearly transforms features within one position.
Deep stacks need stable paths for signals and gradients. A common pre-normalized transformer block is \[ \mathbf X\leftarrow \mathbf X+\operatorname{MultiHead}(\operatorname{LN}(\mathbf X)), \] followed by \[ \mathbf X\leftarrow \mathbf X+\operatorname{MLP}(\operatorname{LN}(\mathbf X)). \] Layer normalization keeps each token’s features at a controlled scale. Residual addition means neither sublayer replaces the running representation. Each contributes a correction to a shared residual stream. An embedding layer, a stack of these blocks, and a final map to vocabulary logits form the core model.
Heads decide where to read, the multilayer perceptron transforms what was read, and the residual stream preserves a route through the stack. Without a visibility rule, every head can use later tokens that will not exist during generation.
Next-Token Training
Language modeling asks the representation at position \(i\) to predict the next token using only tokens available through position \(i\). Ordinary self-attention would allow query \(i\) to inspect later keys and copy the answer during training. That shortcut produces a low training loss but cannot work during generation, when future tokens do not exist.
A causal mask sets every future score to \(-\infty\) before softmax, with query and key indices \(i,j\in\{1,\ldots,n\}\): \[ s_{ij}\leftarrow-\infty\qquad\text{when }j>i. \] Since \(e^{-\infty}=0\), future positions receive exactly zero weight, while the remaining entries are renormalized: \[ a_{ij}= \begin{cases} \dfrac{e^{s_{ij}}}{\sum_{\ell=1}^{i}e^{s_{i\ell}}},&j\leq i,\\[1em] 0,&j>i. \end{cases} \]
A causal mask permits a query to read only its own position and earlier positions. Its target is the next observed token.
The mask must be applied before softmax so the visible weights are renormalized to sum to one.
With a triangular mask, one forward pass over a training sequence creates a valid next-token prediction at every position simultaneously. Each position has a different visible prefix, but matrix operations evaluate all of them in parallel. Cross-entropy compares each output distribution with the actual next token, and backpropagation trains the whole stack.
Let the vocabulary contain \(K\) items, indexed by \(1,\ldots,K\). For a training sequence \(c_1,\ldots,c_T\) of integer length \(T\geq2\), each \(c_i\) is a vocabulary index. With model parameters \(\boldsymbol\theta\), let \(p_{ij}=\Pr_{\boldsymbol\theta}(Y=j\mid c_1,\ldots,c_i)\) be the predicted probability of vocabulary item \(j\in\{1,\ldots,K\}\) after prefix \(i\in\{1,\ldots,T-1\}\). Cross-entropy from logistic regression gives the objective \[ \mathcal L(\boldsymbol\theta) =\frac1{T-1}\sum_{i=1}^{T-1}\ell_i, \qquad \ell_i=-\log p_{i,c_{i+1}}. \] Every position predicts the next observed token, so one training pass supplies several labeled predictions. The target is shifted by one position. It is never part of its own visible prefix.
The same restriction must survive the entire stack rather than one attention row.
Claim: Suppose the input representation at position \(i\) depends only on the token at \(i\) and fixed information for that position. In a stack whose every self-attention layer uses the causal mask and whose normalization and MLP act separately at each position, the representation at position \(i\) depends only on tokens at positions \(1,\ldots,i\).
Proof
A computation path is a chain of dependencies by which an input token can affect an output. At query position \(i\), masked attention uses only the query at \(i\) and keys and values at positions \(j\leq i\). Its softmax denominator also contains only these allowed positions, so even the attention weights have no dependence on a future representation. Layer normalization, the MLP, and residual addition act separately at each position.
Every connection that can transmit information from position \(j\) to position \(i\) therefore satisfies \(j\leq i\). Along a path through any number of blocks, position indices can only stay the same or increase. Because the input at position \(j\) depends only on token \(j\) and fixed position information, no token after \(i\) has a path to the final representation or logits at \(i\). This proves causality for the entire stack.
The mask therefore defines a directed information graph for the entire model. Training can evaluate all next-token losses in parallel because each prediction has no computational path to its future target.
All-pairs routing requires score and weight matrices with one entry per token pair, so attention work grows quadratically with sequence length. This quadratic cost matters when comparing attention with alternative routing mechanisms.
Generation
After training, the final token’s representation is mapped to logits \(\mathbf z\in\mathbb R^K\), one for each vocabulary item, and softmax gives a probability distribution over the next token. Generation repeats three steps:
- Run the transformer on the available prefix.
- Choose a next token from the final distribution.
- Append that token and run again.
The causal mask guarantees consistency: the representation used for a prediction depends on exactly the kind of prefix that will be available at generation time.
The generation view below reveals one additional token at a time. A fixed illustrative sequence shows the information available at each step.
Generation computes a distribution from the available prefix, chooses a token, appends it, and repeats.
Each chosen token becomes part of the next prefix, so these generation steps must occur in order.
How we choose from the distribution changes generation without changing the trained model. Greedy decoding takes the largest logit, while sampling draws from the distribution. A temperature \(\tau>0\) rescales logit gaps: \[ p_j(\tau)=\frac{e^{z_j/\tau}}{\sum_{\ell=1}^K e^{z_\ell/\tau}}, \qquad j=1,\ldots,K. \] Small \(\tau\) sharpens the distribution, while large \(\tau\) flattens it. Temperature changes how strongly generation follows existing logit differences. It cannot add information to the model.
Autoregressive generation is sequential even though training across positions is parallel: token \(i+1\) must be chosen before it can become part of the prefix for token \(i+2\). This distinction explains why the same architecture can train efficiently yet generate one token at a time.
Multiple heads provide several routing patterns, the per-position multilayer perceptron transforms their outputs, and normalization and residual paths support depth. Causal masking makes each output a valid next-token prediction, and autoregressive feedback converts those predictions into a sequence.