Transformer as a Language Model

The previous chapter explained the transformer block: self-attention, positional information, feed-forward networks, residual connections, and layer normalization.

To build a GPT-style model, we use that architecture in a specific way.

How do we arrange a transformer so it can predict the next token?

The answer is a decoder-only transformer. It reads a prefix, blocks access to future tokens with a causal mask, and produces one contextual representation at every position.

Imagine the model reading Jessie ordered a. The representation at the final position contains information from Jessie, ordered, and a. That final representation is the model's summary of the current context.

A language model adds one more piece: an output head that turns that representation into scores for possible next tokens. The detailed probability and loss story belongs to the Language Modelling lesson. This bridge lesson is about how the transformer is arranged so that prediction is possible.

Causal Context

During next-token prediction, a position may use only itself and earlier tokens. It must not peek at the answer. That is why decoder-only transformers use causal masking.

Input Shape

After text is converted into token IDs, each token is looked up in an embedding matrix. A sequence of $n$ tokens becomes:

Positional information is added so the model can distinguish the same tokens in different orders.

Causal Self-Attention

The transformer still computes queries, keys, and values:

But the attention matrix is masked before softmax so future positions are invisible:

Here $M$ contains large negative values wherever a token would otherwise attend to the future.

Context Window

A decoder-only transformer can only attend to tokens inside its context window. If the context window is 4,000 tokens, then token 4,001 cannot directly see token 1 unless the system provides that information again.

This is why context windows matter so much in LLM applications: they define how much text the model can condition on at once.

Output Head

Each position leaves the transformer as a contextual vector $h_t$. A language-model head maps that vector to one score per vocabulary token:

These scores are logits. The next lessons explain how raw text becomes tokens and how logits become a trained language-modelling objective.

This sketch shows the architecture only: token embeddings, positional embeddings, causal transformer block, and output head.

import torch
from torch import nn

class DecoderOnlyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, max_len):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_len, d_model)
        self.block = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=num_heads,
            dim_feedforward=4 * d_model,
            batch_first=True,
        )
        self.output_head = nn.Linear(d_model, vocab_size)

    def forward(self, token_ids):
        batch_size, seq_len = token_ids.shape
        positions = torch.arange(seq_len, device=token_ids.device)

        x = self.token_embedding(token_ids)
        x = x + self.position_embedding(positions)[None, :, :]

        causal_mask = torch.triu(
            torch.ones(seq_len, seq_len, device=token_ids.device),
            diagonal=1,
        ).bool()

        h = self.block(x, src_mask=causal_mask)
        logits = self.output_head(h)
        return logits

The result is not yet a trained language model. It is the architecture that the language-modelling objective will train.