Lesson 7

Fixed-window neural language models

Introduction

Two sequences containing three token IDs each become embeddings with shape two by three by four. Flattening each context gives 12 features, which an MLP converts into five next-token logits.

Learning goal: Build a fixed-window neural language model and trace embeddings, flattened contexts, hidden features, and output logits.
Before you start: Token batches, embedding lookup, linear layers, activation functions, logits, and cross-entropy.

Lesson plan

Explain why raw token IDs need learned vectors before combining context positions.
Trace every tensor shape through lookup, flattening, hidden activation, and output scores.
Run one complete learning step and test a changed context or vocabulary size.

Why move beyond token IDs?

A token ID is only a label. ID 8 is not eight times ID 1, and nearby IDs do not have to represent similar tokens. An embedding table converts each ID into a trainable vector of floating-point features. During training, tokens that are useful in similar contexts may develop useful relationships in this feature space.

A bigram uses only one previous token. We can instead collect a fixed number of earlier tokens, embed each one, join their vectors, and pass the result through a multilayer perceptron, or MLP. An MLP is a sequence of linear layers with a nonlinear activation between them. The nonlinearity lets the network represent relationships more flexible than one large linear formula.

Terms and shapes

EmbeddingA learned vector used to represent one token.

Embedding dimensionThe number of features in each token vector, written C.

Context windowThe fixed number of earlier tokens, written T.

MLPA feed-forward network made from linear layers and nonlinear activations.

Hidden layerAn internal representation between input and output.

ActivationA nonlinear function such as Tanh or ReLU.

Input IDs have shape (B, T). nn.Embedding(V, C) changes them to shape (B, T, C). Flattening the final two axes gives (B, T × C). The output layer produces (B, V) logits, one score for each possible next token.

Worked example: three tokens predict one

Trace the shapes first: with batch size 4, context length 3, and embedding width 5, lookup produces (4, 3, 5). Flattening the last two axes should produce (4, 15). Predict the final logit shape from the vocabulary size before running the code.

The following CPU model receives three earlier token IDs. The class syntax follows Lesson 3: self is the current model object, __init__ stores layers, and forward defines the calculation. The code only checks the data flow; training would use many context-target examples.

import torch
from torch import nn

class ContextMLP(nn.Module):
    def __init__(self, vocabulary_size, context_size, embedding_size):
        super().__init__()
        self.context_size = context_size
        self.embedding_size = embedding_size
        self.embedding = nn.Embedding(vocabulary_size, embedding_size)
        self.hidden = nn.Linear(context_size * embedding_size, 12)
        self.output = nn.Linear(12, vocabulary_size)

    def forward(self, token_ids):
        vectors = self.embedding(token_ids)
        flat = vectors.reshape(token_ids.shape[0], -1)
        hidden = torch.tanh(self.hidden(flat))
        return self.output(hidden)

torch.manual_seed(4)
model = ContextMLP(vocabulary_size=5, context_size=3, embedding_size=4)
x = torch.tensor([[0, 1, 2], [2, 3, 4]], dtype=torch.long)
vectors = model.embedding(x)
logits = model(x)

print(tuple(x.shape))
print(tuple(vectors.shape))
print(tuple(logits.shape))

# Expected:
# (2, 3)
# (2, 3, 4)
# (2, 5)

There are two examples, three positions per example, and four features per position. Each example therefore reaches the first linear layer as twelve values. The hidden layer returns twelve new features; this happens to use the same size, but it does not have to. The output returns five logits. Cross-entropy can compare those logits with a target tensor of shape (2,).

One complete learning step

Token IDs only select rows; gradients do not update the integer IDs. They update the selected embedding rows and the MLP weights. If every embedding row changes after one batch, check whether weight decay or another global update rule is active.

targets = torch.tensor([3, 0], dtype=torch.long)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

optimizer.zero_grad()
logits = model(x)
loss = nn.functional.cross_entropy(logits, targets)
loss.backward()
optimizer.step()

print(loss.ndim)  # Expected: 0 (a scalar)

The gradient flows backward through both linear layers and into only the embedding rows selected by the batch. An embedding is not a fixed dictionary definition; it is part of the learned model. A larger embedding dimension gives more representational capacity but also adds parameters and can overfit small data.

Common pitfalls

Embedding inputs must be integer IDs, normally torch.long.
Do not include the target token inside its own context; that leaks the answer.
Keep the axis order clear before flattening. Accidentally mixing the batch axis combines different examples.
The first linear layer must accept exactly context_size × embedding_size values.
A fixed context window forgets anything before the window.

Try it

Suppose B = 8, T = 4, C = 6, and V = 20. Give the shape after embedding, after flattening, and after the output layer. What in_features should the first linear layer use?

Reveal the worked answer

Embedding produces (8, 4, 6). Flatten only the position and feature axes to get (8, 24). The first linear layer therefore needs in_features=24. The final logits have shape (8, 20), one row per example and one column per vocabulary token.

x = torch.zeros(8, 4, dtype=torch.long)
embedding = nn.Embedding(20, 6)
flat = embedding(x).reshape(8, -1)
print(tuple(flat.shape))  # Expected: (8, 24)

Recap

An embedding maps a categorical ID to a learned vector. A context MLP joins vectors from several earlier positions, processes them with a nonlinear hidden layer, and returns next-token logits. Track the transformation (B,T) → (B,T,C) → (B,T×C) → (B,V). This model sees more context than a bigram, but its context size remains fixed.

References: PyTorch Embedding and PyTorch Linear.

← Lesson 6: Bigram Lesson 8: RNN →