Sampling in Language Models

Imagine you could freeze a language model at the moment it receives one token of input and then follow that token step by step until the model produces a prediction. Along the way the token passes through several distinct stages, and each stage has a specific job. Learning the components of a transformer is largely a matter of naming those stages, understanding what each one does to the data, and seeing how they are arranged.

This article takes that route. We begin with the path a token travels from input to output, then zoom into the parts that repeat. Sampling is the final step in the course's story: the probabilities the model produces at the exit are what it uses to choose its outputs. The focus here is the architecture that produces those probabilities.

A transformer-based language model has an embedding module at the entry, a stack of transformer blocks in the middle, and an output layer at the exit. Two of those three stages are used only once per prediction. The transformer blocks are the part that repeats, and most of the model's internal machinery lives there.

Read the diagram from left to right. The token enters as an ID, is converted into a vector by the embedding module, is transformed repeatedly by the transformer blocks, and leaves the model through the output layer as a set of token probabilities. The vector changes in meaning at each stage, but the overall journey is one-directional, from input to output. The blocks in the middle are drawn as a repeated element because the same template is applied many times in sequence.

Before any transformer block sees the input, the model passes it through an embedding module. This module consists of two matrices: the embedding matrix and the positional embedding matrix. The embedding matrix converts tokens into embedding vectors, and the positional embedding matrix converts token positions into embedding vectors. Together they turn an ordered sequence of token IDs into vectors that carry both identity and location information.

The diagram above uses illustrative token IDs. The principle is what matters: each token ID selects one row of the embedding matrix, so the same token always starts from the same vector at the entry of the model. The positional embedding matrix works the same way for positions, which is how the model knows where a token sits in the sequence, not just which token it is.

The transformer blocks are where most of the model's internal machinery lives. A transformer architecture is built by stacking multiple transformer blocks. The exact contents of a block can vary from one model to another, but in general each block contains two sub-modules: an attention module and an MLP (multi-layer perceptron) module. Think of the block as a fixed template, and the model as many copies of that template placed one after another.

Read the concept map as a containment tree plus an order. The block contains the two sub-modules, and within the block the attention module processes the data before the MLP module does. The arrow between them marks that order. The map also shows what each sub-module is built from, which the next two subsections describe in more detail.

The source establishes that the MLP's linear layers are separated by nonlinear activation functions, and it names ReLU and GELU as examples. A useful way to see why the separation matters is to ask what would happen without it. If you stacked two linear layers with no activation between them, the result would still be one linear transformation. Feeding the output of one linear layer into another simply produces a different linear map. Stacking many linear layers therefore gains nothing in expressive power over a single one.

A nonlinear activation breaks that collapse. Because the function applied between the linear layers is not a straight-line transformation, the MLP can represent patterns that no single linear map could. This is the general reason the MLP needs nonlinear activations, and it is the reason the MLP is described as a stack of linear layers with activations between them rather than as a plain chain of matrix multiplications.

After the last transformer block has produced its output vectors, the output layer takes over. Its job is to map those vectors into token probabilities. The output layer is also called the unembedding layer or the model head. It typically consists of a single matrix, which projects each output vector onto the vocabulary so that every token receives a score. These probabilities are what the model uses to sample its outputs, a process covered elsewhere in the course.

The output layer is the mirror image of the embedding module. The embedding module turns token IDs into vectors at the entry, and the output layer turns vectors back into a distribution over tokens at the exit. Both are used once per prediction, while the blocks between them repeat.

The size of a transformer model is set by the dimensions of its building blocks. Three values matter most. The first is the model dimension, written $d_{model}$ and described in the source as the model's hidden size. It determines the sizes of the key, query, value, and output projection matrices inside each transformer block. The second is the number of transformer blocks, which the source notes is often called the number of layers. The third is the feedforward dimension, which is the width of the linear layers in the MLP module.

The source states that larger values for these dimensions produce larger models. Each of the three values grows a different part of the model: the model dimension grows the projection matrices in every block, the feedforward dimension grows the linear layers in every MLP, and the number of blocks adds more copies of the whole template. Because the blocks are repeated, a change to the per-block sizes is multiplied across every block in the stack.