The Bigram Language Model

The simplest language model worth training predicts the next character from the current character alone. Such a bigram model is like trying to write Shakespeare while seeing only the last letter typed. It cannot produce coherent text, but it introduces every piece of machinery the full Transformer will reuse: representing discrete tokens as vectors, turning scores into a probability distribution over the vocabulary, measuring how good a prediction was, sampling new text, and the training loop that improves the model. This is phase 2 of the walkthrough; its input is the batches built in Tiny Shakespeare Data Pipeline.

The code for this phase is in steps/2_the_bigram_model. Reference losses quoted below were computed from the text with tools/experiments/gpt/dataset_stats.py and tools/experiments/gpt/initial_loss.py in this site's repository; no neural model was trained for this article.

Why predict from one character

Any probability for a whole string can be written exactly as a product of next-token probabilities, each conditioned on everything before it:

$$P(s_1 s_2 \ldots s_n) = P(s_1)\, P(s_2 \mid s_1) \cdots P(s_n \mid s_1 \ldots s_{n-1}).$$

The difficulty is that there are far too many histories $s_1 \ldots s_{j-1}$ to give each its own parameters. The classical answer, set out for machine translation by Brown et al. (1990, Section 2), is to group histories into equivalence classes: in an $n$-gram model two histories count as the same if they agree in their last $n-1$ tokens, so in a bigram model only the last token matters. Brown et al. worked with words and a trigram model; the model here works with characters and a bigram model, but the principle is the same: learn the probability of what comes next from what came immediately before. A GPT is the same factorization with the equivalence classes removed: it conditions on the whole visible context.

The model is an embedding table

The bigram model is a single PyTorch nn.Embedding layer of size vocab_size × vocab_size, that is 65 × 65. An embedding layer is a lookup table: given an integer index, it returns that row of a learned matrix. Here row $i$ holds 65 scores, one for each character that might follow character $i$. Step file: train_gpt_2_3.py, which covers the walkthrough's sections on the architecture, the forward pass and the shape test.

--- train_gpt_1_12.py
+++ train_gpt_2_3.py
@@ -155 +156,26 @@
 print(f"Target batch (yb): {yb}")
+
+####################################################################################################
+# Define the Model
+
+# Define the Bigram Language Model
+class BigramLanguageModel(nn.Module):
+    def __init__(self, vocab_size):
+        super().__init__()
+        # Token embedding table (vocab_size x vocab_size)
+        self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)
+
+    def forward(self, idx, targets=None):
+        # idx and targets are both (B, T) tensors of integers
+        logits = self.token_embedding_table(idx)  # (B, T, C), where C is vocab_size
+
+        # Return the logits (predictions for the next character)
+        return logits
+
+# Instantiate the model
+model = BigramLanguageModel(vocab_size)
+
+# Test the model with a batch of input data
+logits = model(xb)
+
+print(f"Logits shape: {logits.shape}")  # Should be (batch_size, block_size, vocab_size)

Embeddings usually map tokens into a smaller space, as the next phase will do. Here the row length equals the vocabulary size because each row is read directly as the model's scores for the next character. The model has 65 × 65 = 4,225 parameters, and the class keeps the name BigramLanguageModel to the end of the walkthrough, long after it has become a Transformer.

The forward pass and its shapes

The forward pass looks up the row of every input character. With an input batch idx of shape (B, T), where B is the batch size and T the sequence length, the output has shape (B, T, C), where C, the number of "channels", is here the vocabulary size. For the batch of 4 sequences of 8 characters from phase 1, the logits have shape (4, 8, 65): for each of the 32 positions, a score for each of the 65 possible next characters.

These scores are logits: real numbers, not probabilities. If the model learns its job, the row for q will hold a large logit for u, and after To b the row for b will favour e, though the bigram model sees only the b. The same input and output shapes carry through to the final Transformer; only the computation between them changes.

Cross-entropy loss

Training needs a single number that says how wrong the predictions were. Cross-entropy turns each row of logits into probabilities with the softmax function and charges the model the negative log of the probability it gave to the character that actually came next:

$$p_k = \frac{e^{z_k}}{\sum_{j=1}^{V} e^{z_j}}, \qquad \text{loss} = -\frac{1}{N} \sum_{i=1}^{N} \ln p_{y_i},$$

where $z$ is a row of $V = 65$ logits, $y_i$ is the true next character at position $i$, and the average runs over all $N = B \times T$ positions of the batch. The loss is a measure of surprise: a confident correct prediction ($p_y$ near 1) costs almost nothing, and a confident wrong one costs a great deal. Step file: train_gpt_2_4.py.

--- train_gpt_2_3.py
+++ train_gpt_2_4.py
@@ -8,2 +8,4 @@
 import torch.nn as nn
+
+import torch.nn.functional as F

@@ -181 +183,23 @@
 print(f"Logits shape: {logits.shape}")  # Should be (batch_size, block_size, vocab_size)
+
+####################################################################################################
+# Evaluate the Loss Function
+
+# Define the loss function (cross-entropy)
+def compute_loss(logits, targets):
+    # Reshape logits to be (B * T, C)
+    B, T, C = logits.shape
+    logits = logits.view(B * T, C)
+
+    # Reshape targets to be (B * T)
+    targets = targets.view(B * T)
+
+    # Compute cross-entropy loss
+    loss = F.cross_entropy(logits, targets)
+    return loss
+
+# Test the model and compute the loss
+logits = model(xb)
+loss = compute_loss(logits, yb)
+
+print(f"Loss: {loss.item()}")

PyTorch's F.cross_entropy takes logits of shape (N, C) and targets of shape (N), so compute_loss flattens the batch and time dimensions together with view. Each of the 32 positions is treated as an independent prediction.

Three reference values make the loss interpretable. A model that spreads probability evenly over the 65 characters has loss $-\ln(1/65) = \ln 65 \approx 4.174$. In the run described in the walkthrough, the untrained model's loss on its first batch was about 4.87, higher than that. The reason is the initialization rather than anything about Shakespeare: nn.Embedding starts its weights as independent standard normal values, so each row of logits is 65 random numbers, and a softmax over them is confidently arbitrary, not uniform. The expected loss of such logits is close to $\ln 65 + \tfrac{1}{2} \approx 4.67$ (a simulation of a million rows gives 4.66); over 30 freshly drawn tables scored on the same 32 positions, the median loss was 4.697 (interquartile range 4.571 to 4.802, extremes 4.368 and 5.077), so 4.87 is within the ordinary spread of starting values. Exponentiating a loss gives the perplexity, the number of equally likely choices the model is effectively hesitating between, which is used to evaluate topic models in Evaluating Topic Models: the uniform model has perplexity 65.

Generating text by sampling

Generation runs the model on the text so far, keeps only the logits at the last position, turns them into probabilities with softmax, draws one character from that distribution with torch.multinomial, appends it, and repeats. Step file: train_gpt_2_5.py.

--- train_gpt_2_4.py
+++ train_gpt_2_5.py
@@ -166,3 +166,2 @@
         super().__init__()
-        # Token embedding table (vocab_size x vocab_size)
         self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)
@@ -170,7 +169,25 @@
     def forward(self, idx, targets=None):
-        # idx and targets are both (B, T) tensors of integers
-        logits = self.token_embedding_table(idx)  # (B, T, C), where C is vocab_size
-
-        # Return the logits (predictions for the next character)
+        logits = self.token_embedding_table(idx)  # (B, T, C)
         return logits
+
+    def generate(self, idx, max_new_tokens):
+        # idx is (B, T), where B is batch size and T is current sequence length
+        for _ in range(max_new_tokens):
+            # Get predictions for the current sequence
+            logits = self(idx)  # (B, T, C)
+
+            # Focus only on the last time step to predict the next token
+            logits = logits[:, -1, :]  # (B, C)
+
+            # Convert logits to probabilities using softmax
+            probs = F.softmax(logits, dim=-1)  # (B, C)
+
+            # Sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)
+
+            # Concatenate the new prediction to the sequence
+            idx = torch.cat([idx, idx_next], dim=1)  # (B, T+1)
+
+        return idx
+

@@ -205 +222,15 @@
 print(f"Loss: {loss.item()}")
+
+####################################################################################################
+# Generate Text
+
+# Start with a batch of size 1x1 containing the index for '\n'
+idx = torch.zeros((1, 1), dtype=torch.long)
+
+# Generate 100 tokens from the Bigram model
+generated_sequence = model.generate(idx, max_new_tokens=100)
+
+# Convert the generated sequence back to text
+generated_text = decode(generated_sequence[0].tolist())
+
+print(generated_text)

Generation starts from a (1, 1) tensor holding index 0, which is the newline character (the first character in sorted order), a natural "start of a speech" in this text. Sampling, rather than always taking the most probable character, lets the model produce varied continuations of the same start, the behaviour seen in the two poems on the collection page. Before training, the output is random characters; the same function later shows what the model has learned.

The bigram model computes logits for every position of the growing sequence and then discards all but the last, which is wasteful for a model that only looks at one character. The code is written this way so that it keeps working unchanged once the model looks at more context.

Training the model with Adam

The training loop repeats four operations: sample a batch, compute the loss, compute the gradient of the loss with respect to every parameter by backpropagation (loss.backward()), and let the optimizer adjust the parameters (optimizer.step()), clearing the old gradients first with optimizer.zero_grad(). Step file: train_gpt_2_6.py.

--- train_gpt_2_5.py
+++ train_gpt_2_6.py
@@ -223,2 +223,38 @@

+
+####################################################################################################
+# Training the Model
+
+# Set up the Adam optimizer with a higher learning rate
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+# Set the new batch size
+batch_size = 32
+
+# Training loop
+def train_model(steps):
+    for step in range(steps):
+        # Sample a batch of training data
+        xb, yb = get_batch('train')
+
+        # Step 4: Compute the logits and loss
+        logits = model(xb)
+        loss = compute_loss(logits, yb)
+
+        # Step 5: Zero out gradients from the previous step
+        optimizer.zero_grad()
+
+        # Step 6: Compute gradients (backpropagation)
+        loss.backward()
+
+        # Step 7: Update the model's parameters using the optimizer
+        optimizer.step()
+
+        # Print the loss every 100 steps
+        if step % 100 == 0:
+            print(f"Step {step}, Loss: {loss.item()}")
+
+# Train the model for 10,000 steps
+train_model(10000)
+
 ####################################################################################################

The optimizer is Adam (Kingma and Ba, 2015), which gives every parameter its own effective step size, computed from running averages of that parameter's gradients and squared gradients: the update divides the first by the square root of the second, so parameters whose gradients are consistently large do not take proportionally larger steps. The learning rate is 1e-3, and the loop runs 10,000 steps, printing the loss of the current batch every 100.

One detail of this step file differs from the walkthrough's text. The walkthrough calls get_batch('train', block_size=8, batch_size=32). The step file instead sets a global batch_size = 32 and calls get_batch('train'), but the get_batch defined in phase 1 has its own default argument batch_size=4, which takes precedence, so this step actually trains on batches of 4. The next step file fixes this by having get_batch read the global settings.

What a bigram table can learn

Because the bigram model is a free table with one row per character, the best it can possibly do on the training text is known in advance. For each previous character $a$, the cross-entropy over the positions following $a$ is minimized by setting the row's softmax probabilities equal to the observed frequencies of the following characters, $\hat{P}(b \mid a) = \text{count}(ab) / \text{count}(a)$. (This is the familiar fact that the maximum-likelihood estimate of a categorical distribution is its relative frequencies; the softmax can represent any distribution with nonzero entries, and approaches the zero ones as logits go to minus infinity.) Training can therefore only approach the count-based model, and the counts show what it holds:

• after q (563 times in the training text), the next character is always u;
• after T, it is h 46.2 percent of the time, then o (13.8 percent), E (7.6) and : (6.4), the last two from speaker names written in capitals, such as GLOUCESTER: and JULIET:;
• after a space, the most common next characters are t (14.1 percent), a (7.9), s (7.2) and h (7.1);
• after a newline, another newline (17.7 percent), which separates speeches.

Of the 4,225 possible character pairs, 1,380 occur in the training text. The table below places the bigram model between the trivial baselines and the final Transformer. Losses are average cross-entropy per character in nats, rounded to two decimals; the counts-based validation figure uses add-one smoothing so that the 187 places in the validation text where a character pair occurs that never occurs in training (23 distinct pairs) do not get probability zero.

No amount of training will take a bigram model much below about 2.45 on this text; going further requires looking at more than one character, which is the subject of the next phase.

Moving the computation to a GPU

A GPU executes the matrix operations of a neural network in parallel, and all tensors in a computation must live on the same device. The step that adds GPU support (step file train_gpt_2_9.py, which implements the walkthrough's sections 2.7 to 2.9 together) chooses the device once, moves the model there, and moves every new batch there as it is created:

####################################################################################################
# Optimizing for GPU

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)

def get_batch(split):
    data = train_data if split == 'train' else val_data

    ix = [random.randint(0, len(data) - block_size - 1) for _ in range(batch_size)]
    x_batch = torch.stack([data[i:i + block_size] for i in ix]).to(device)
    y_batch = torch.stack([data[i + 1:i + block_size + 1] for i in ix]).to(device)

    return x_batch, y_batch

On a computer without a CUDA GPU, device is 'cpu' and everything still runs. The step file also gathers the hyperparameters at the top of the program and seeds both random number generators:

# Hyperparameters
batch_size = 32
block_size = 8
learning_rate = 1e-3
eval_interval = 100
eval_iters = 100
num_train_steps = 10000

# Setting random seeds for reproducibility
random.seed(42)
torch.manual_seed(42)

One gap remains in this and the following step files: the generation code at the end creates its starting tensor with torch.zeros((1, 1), dtype=torch.long) on the CPU. On a GPU machine the model's weights are on the GPU, and PyTorch refuses to mix devices, so that last call fails; adding device=device fixes it. The finished program's generate_text function creates its prompts with .to(device).

Estimating the loss over many batches

The loss of a single batch is noisy: some batches are harder than others. To judge progress, estimate_loss averages the loss over 100 batches from the training set and 100 from the validation set:

# Estimate Loss

def estimate_loss():
    model.eval()
    losses = {'train': [], 'val': []}

    with torch.no_grad():
        for split in ['train', 'val']:
            for _ in range(eval_iters):
                xb, yb = get_batch(split)
                logits = model(xb)
                loss = compute_loss(logits, yb)
                losses[split].append(loss.item())

    model.train()
    return {split: sum(losses[split]) / eval_iters for split in losses}

Two calls make the estimate cheaper and cleaner. torch.no_grad() stops PyTorch from recording the operations needed for backpropagation, which saves memory and time when no update follows. model.eval() switches layers that behave differently in training, such as dropout (added in phase 5), to their evaluation behaviour; model.train() switches them back. The bigram model has no such layers yet, but the habit is in place before it matters.

The two averages answer different questions. The training loss says how well the model fits the text it learns from; the validation loss says how well that carries over to text it has never seen. A widening gap between them is the signature of overfitting, which becomes a real risk only once the model is large enough to memorize, in Scaling the Shakespeare GPT.

The training loop with periodic evaluation

The loop itself is unchanged in substance: sample, compute the loss, backpropagate, update. What is new is that every eval_interval steps it pauses to print the averaged training and validation losses, rather than the noisy loss of one batch:

# Training Loop

optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

def train_model():
    for step in range(num_train_steps):
        xb, yb = get_batch('train')
        logits = model(xb)
        loss = compute_loss(logits, yb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if step % eval_interval == 0:
            losses = estimate_loss()
            print(f"Step {step}, Train Loss: {losses['train']}, Val Loss: {losses['val']}")

train_model()

This loop, with the hyperparameters at the top of the file, is the template the rest of the walkthrough keeps. The model inside it changes in Building Self-Attention Step by Step.

References

• Brown, P. F., Cocke, J., Della Pietra, S. A., Della Pietra, V. J., Jelinek, F., Lafferty, J. D., Mercer, R. L., and Roossin, P. S. (1990). A statistical approach to machine translation. Computational Linguistics, 16(2), 79–85.
• Karpathy, A. (2023). Let's build GPT: from scratch, in code, spelled out [Video lecture]. https://www.youtube.com/watch?v=kCc8FmEb1nY. Code: https://github.com/karpathy/ng-video-lecture
• Kingma, D. P., and Ba, J. (2015). Adam: A method for stochastic optimization. In Proceedings of the 3rd International Conference on Learning Representations (ICLR 2015). arXiv:1412.6980.
• Code for this article: gpt_text_model/steps/2_the_bigram_model in github.com/mdb42/ml-demos.