Full LLM

A full language model turns a prompt into a probability distribution for the next token.

We have already encountered almost all of the components required to make an LLM. An LLM consists of tokenisation and vector embedding, followed by a series of transformer blocks. We then use a final LayerNorm and project the final resulting vector into an array of logits, with one logit for each token in the vocabulary:

Full language model pipeline A prompt is tokenised and embedded. Its vectors pass through a stack of transformer blocks, a final LayerNorm, and a vocabulary projection. Softmax turns the resulting logits into a probability for every possible next token. Prompt “The cat sat” Tokenisation text → token IDs Token + position embeddings one vector for every prompt token Transformer stack Transformer block 1 attention + feed-forward Transformer block 2 attention + feed-forward … Transformer block N Final LayerNorm Vocabulary projection one logit for every token in the vocabulary Softmax → next-token probabilities the 56% on 27% … 17%

The code now looks like this:

import torch
import torch.nn as nn


class LanguageModel(nn.Module):

    def __init__(
        self,
        vocab_size,
        max_sequence_length,
        model_dim,
        head_dim,
        num_heads,
        hidden_dim,
        num_layers,
    ):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, model_dim)
        self.position_embedding = nn.Embedding(max_sequence_length, model_dim)
        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(model_dim, head_dim, num_heads, hidden_dim)
            for _ in range(num_layers)
        ])
        self.final_norm = nn.LayerNorm(model_dim)
        self.vocabulary_projection = nn.Linear(model_dim, vocab_size)

    def forward(self, token_ids):
        positions = torch.arange(
            token_ids.size(0),
            device=token_ids.device,
        )
        x = self.token_embedding(token_ids)
        x = x + self.position_embedding(positions)

        for block in self.transformer_blocks:
            x = block(x)

        x = self.final_norm(x)
        last_token_vector = x[-1]
        return self.vocabulary_projection(last_token_vector)

The result of this is an array with the same size as the number of tokens in the vocabulary of our LLM. We then use softmax to interpret this as a probability distribution over the next token to be generated:

next_token_logits = model(token_ids)
next_token_probabilities = torch.softmax(next_token_logits, dim=-1)

One way of obtaining the next token at this point is to just use the most likely token:

next_token_id = torch.argmax(next_token_probabilities).item()

We can also make our LLM more variable by instead randomly sampling the probability distribution of tokens.

next_token_id = torch.multinomial(
    next_token_probabilities,
    num_samples=1,
).item()