Lesson 6

Bigram language models and probabilistic baselines

Why start with a bigram

A bigram model makes a strong simplifying assumption: the next token depends only on the current token. That is not enough for rich language, but it contains the complete learning pipeline in a small form. We can prepare pairs, produce logits, calculate cross-entropy loss, run backpropagation, update weights, and sample predictions. Later models improve the context without changing this basic training idea.

The model stores one row of learned scores for every possible input token. If the vocabulary has size V, its table has shape (V, V). Looking up token h selects the row for h. The row contains one score for every possible next token.

Terms and shapes

Input IDs may have shape (B, T). Embedding lookup with nn.Embedding(V, V) returns logits of shape (B, T, V). Cross-entropy treats the last axis as classes after we flatten the first two axes.

Worked example: learn repeating text

Calculate before training: suppose token a is followed by b three times and by c once. The count-based probabilities are P(b|a)=3/4 and P(c|a)=1/4. A trained bigram row should prefer b, although its exact probabilities depend on optimization.

The tiny training text repeats hi . Its common transitions are h → i, i → space, and space → h. We use every adjacent pair. self.table(ids) is lookup syntax: each ID selects one trainable row.

import torch
from torch import nn

torch.manual_seed(7)
text = "hi " * 20
characters = ["h", "i", " "]
to_id = {c: i for i, c in enumerate(characters)}
ids = torch.tensor([to_id[c] for c in text], dtype=torch.long)
x, y = ids[:-1], ids[1:]

class BigramModel(nn.Module):
    def __init__(self, vocabulary_size):
        super().__init__()
        self.table = nn.Embedding(vocabulary_size, vocabulary_size)

    def forward(self, token_ids):
        return self.table(token_ids)

model = BigramModel(len(characters))
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

with torch.no_grad():
    first_loss = nn.functional.cross_entropy(model(x), y).item()

for _ in range(200):
    optimizer.zero_grad()
    logits = model(x)                    # shape: (tokens, V)
    loss = nn.functional.cross_entropy(logits, y)
    loss.backward()
    optimizer.step()

with torch.no_grad():
    final_logits = model(x)
    final_loss = nn.functional.cross_entropy(final_logits, y).item()
    next_after_h = final_logits[0].argmax().item()

print(final_logits.shape)
print(final_loss < first_loss)
print(characters[next_after_h])

# Expected:
# torch.Size([59, 3])
# True
# i

Cross-entropy accepts raw logits, so we did not call softmax before the loss. Internally, the loss compares all three scores with the correct target ID. Repeated updates raise the useful transition scores. The final loss is not zero because the last incomplete boundary and limited data can still create uncertainty, but it should be lower than the initial loss.

Generate one character at a time

A common wrong interpretation is that a bigram remembers the whole generated prefix. It only uses the current token to choose the next row of logits. Two different prefixes ending in the same token therefore receive the same next-token distribution.

Generation feeds the latest chosen token back into the same model. Greedy decoding takes the largest probability. Sampling with torch.multinomial can produce variety, but its output depends on the random seed.

current = torch.tensor([to_id["h"]])
generated = [current.item()]

with torch.no_grad():
    for _ in range(5):
        logits = model(current)
        current = logits.argmax(dim=-1)
        generated.append(current.item())

result = "".join(characters[i] for i in generated)
print(repr(result))  # Expected: 'hi hi '

This result only uses the current token. If two different sentences end in h, the model gives them the same next-token distribution. That memory limit motivates embeddings of longer contexts, recurrent networks, and attention.

Try it

Add a new token ! and train on hi! hi! . What shape must the bigram table have, and what should follow i most often?

Recap

A bigram model learns a V × V score table. Each input ID selects one row of next-token logits. Cross-entropy rewards the correct next-token class, and an optimizer changes the table. Generation repeatedly predicts and feeds back one token. The model is useful for learning the pipeline, but one-token context is a strict limit.

References: PyTorch Embedding and cross-entropy documentation.