Neural Language-Model Baselines and Deep-Learning Foundations

Derive softmax and cross-entropy, train baseline models, and connect the reference gradients to the reusable PyTorch path.

A Transformer is easier to debug after its objective and gradients have been reduced to inspectable cases. This chapter derives softmax cross-entropy and keeps NumPy implementations as transparent counter-examples. The reusable baseline path then moves to PyTorch tensors and modules, where the same next-token objective can be optimized, inspected, and handed to the decoder in Chapter 04.

The point of the NumPy cells is not to maintain a second training stack. They expose the mathematics and failure modes; proof_lm owns the PyTorch baselines consumed by later chapters.

Next-token likelihood

For a batch of B examples with vocabulary size V, let z_i\in\mathbb{R}^V be the logits and y_i the target index. The softmax distribution is

p_{i,v}=\frac{\exp(z_{i,v})}{\sum_{u=0}^{V-1}\exp(z_{i,u})},

and the mean cross-entropy is

\mathcal{L}(z,y)=-\frac{1}{B}\sum_{i=1}^B\log p_{i,y_i}.

Teacher forcing supplies the observed previous tokens to every training example. The model therefore learns the conditional distribution at each position while the target sequence remains known; generation later has to feed back its own sampled tokens.

Before introducing a model, make the numerical behavior of these two functions explicit.

import numpy as np


SEED = 31


def log_softmax(logits, axis=-1):
    logits = np.asarray(logits, dtype=float)
    shifted = logits - np.max(logits, axis=axis, keepdims=True)
    return shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))


def softmax(logits, axis=-1):
    return np.exp(log_softmax(logits, axis=axis))


def cross_entropy(logits, targets):
    logits = np.asarray(logits, dtype=float)
    targets = np.asarray(targets, dtype=np.int64)
    if logits.ndim != 2 or targets.shape != (logits.shape[0],):
        raise ValueError("logits must have shape (batch, vocab) and targets (batch,)")
    return -np.mean(log_softmax(logits)[np.arange(len(targets)), targets])


rng = np.random.default_rng(SEED)
logits = rng.normal(size=(4, 5))
targets = np.array([0, 3, 1, 4])
probabilities = softmax(logits)
loss = cross_entropy(logits, targets)
print("loss:", round(float(loss), 6))
print("row probability sums:", probabilities.sum(axis=1))

assert np.allclose(probabilities.sum(axis=1), 1.0)
assert np.allclose(cross_entropy(logits + 17.0, targets), loss)
assert np.all(probabilities > 0.0)
assert np.isfinite(loss)

loss: 1.636563
row probability sums: [1. 1. 1. 1.]

Subtracting the largest logit before exponentiating leaves every probability unchanged because the same factor cancels in the numerator and denominator. It prevents overflow when a later model produces a large activation. The loss is finite and the row sums equal one, so these invariants can be tested before any optimizer is involved.

Cross-entropy gradients

For one-hot target vectors e_{y_i}, differentiating the log-softmax gives

\frac{\partial\mathcal{L}}{\partial z_i} =\frac{1}{B}\left(p_i-e_{y_i}\right).

For a linear classifier z=XW+b, the chain rule then gives

\frac{\partial\mathcal{L}}{\partial W}=X^\top\frac{\partial\mathcal{L}}{\partial z},\qquad \frac{\partial\mathcal{L}}{\partial b}=\sum_i\frac{\partial\mathcal{L}}{\partial z_i}.

A finite-difference check compares this hand derivation with the slope measured by perturbing one parameter at a time. The check is small enough to inspect and strong enough to catch a missing batch factor or a wrong target index.

def linear_loss_and_grad(weights, bias, inputs, targets):
    logits = inputs @ weights + bias
    probabilities = softmax(logits)
    loss = cross_entropy(logits, targets)
    dlogits = probabilities.copy()
    dlogits[np.arange(len(targets)), targets] -= 1.0
    dlogits /= len(targets)
    return loss, inputs.T @ dlogits, dlogits.sum(axis=0)


def numerical_gradient(parameter, loss_function, epsilon=1e-5):
    gradient = np.zeros_like(parameter)
    for index in np.ndindex(parameter.shape):
        original = parameter[index]
        parameter[index] = original + epsilon
        plus = loss_function()
        parameter[index] = original - epsilon
        minus = loss_function()
        parameter[index] = original
        gradient[index] = (plus - minus) / (2.0 * epsilon)
    return gradient


check_rng = np.random.default_rng(SEED + 1)
inputs = check_rng.normal(size=(3, 2))
weights = check_rng.normal(scale=0.2, size=(2, 4))
bias = check_rng.normal(scale=0.2, size=4)
targets = np.array([1, 0, 3])
loss, analytic_weights, analytic_bias = linear_loss_and_grad(weights, bias, inputs, targets)
numeric_weights = numerical_gradient(
    weights, lambda: linear_loss_and_grad(weights, bias, inputs, targets)[0]
)
numeric_bias = numerical_gradient(
    bias, lambda: linear_loss_and_grad(weights, bias, inputs, targets)[0]
)
print("linear loss:", round(float(loss), 6))
print("maximum gradient error:", max(
    np.max(np.abs(analytic_weights - numeric_weights)),
    np.max(np.abs(analytic_bias - numeric_bias)),
))
np.testing.assert_allclose(analytic_weights, numeric_weights, atol=1e-8, rtol=1e-6)
np.testing.assert_allclose(analytic_bias, numeric_bias, atol=1e-8, rtol=1e-6)

linear loss: 1.343788
maximum gradient error: 1.7439993893475503e-11

The maximum error is at numerical-difference scale, which validates the complete softmax-to-linear chain for this tiny model. The division by batch size is visible in the agreement; omitting it would produce a gradient exactly three times too large here. This is the manual reference against which the multilayer implementation can be checked conceptually.

Bigram baselines

A bigram model assigns one categorical distribution to each preceding token. Let C_{a,b} count transitions from token a to token b. The maximum-likelihood row is

\hat p(b\mid a)=\frac{C_{a,b}}{\sum_u C_{a,u}},

and the minimum achievable average training loss is the empirical conditional entropy

H_{\mathrm{emp}}=-\frac{1}{N}\sum_{a,b}C_{a,b}\log\hat p(b\mid a).

Train a table of logits with gradient descent and compare its loss with this independently calculated number. The comparison tests the data indexing, softmax, gradient, and update rule together.

def bigram_loss_and_grad(table, previous, targets):
    logits = table[previous]
    probabilities = softmax(logits)
    loss = cross_entropy(logits, targets)
    gradient = np.zeros_like(table)
    np.add.at(gradient, previous, probabilities)
    np.add.at(gradient, (previous, targets), -1.0)
    gradient /= len(targets)
    return loss, gradient


stream = np.array([
    0, 1, 0, 2, 0, 1, 0, 3, 1, 0, 1, 2, 1, 0, 2, 0, 3,
] * 18, dtype=np.int64)
vocab_size = 4
previous, next_token = stream[:-1], stream[1:]
counts = np.zeros((vocab_size, vocab_size), dtype=np.int64)
np.add.at(counts, (previous, next_token), 1)
empirical_probabilities = np.zeros_like(counts, dtype=float)
row_totals = counts.sum(axis=1)
nonempty_rows = row_totals > 0
empirical_probabilities[nonempty_rows] = (
    counts[nonempty_rows] / row_totals[nonempty_rows, None]
)
observed = counts > 0
empirical_entropy = -np.sum(
    counts[observed] * np.log(empirical_probabilities[observed])
) / len(next_token)

table = np.random.default_rng(SEED + 2).normal(scale=0.1, size=(vocab_size, vocab_size))
losses = []
for step in range(1800):
    loss, gradient = bigram_loss_and_grad(table, previous, next_token)
    table -= 1.4 * gradient
    if step % 100 == 0 or step == 1799:
        losses.append(loss)

final_loss = bigram_loss_and_grad(table, previous, next_token)[0]
print("initial recorded loss:", round(float(losses[0]), 4))
print("final loss:", round(float(final_loss), 6))
print("empirical bigram entropy:", round(float(empirical_entropy), 6))
assert final_loss < losses[0]
assert final_loss - empirical_entropy < 0.02

initial recorded loss: 1.4161
final loss: 0.787097
empirical bigram entropy: 0.785596

The trained table approaches the entropy of the observed transition distribution because each row can represent its own conditional probabilities. It cannot improve on that entropy on the same empirical distribution: cross-entropy equals entropy plus a nonnegative KL divergence, and the divergence is zero at the maximum-likelihood rows. The remaining gap is finite-step optimization, not a missing hidden state.

Fixed-context baselines

The bigram table sees one previous token. A fixed-context MLP sees C previous tokens, looks up an embedding for each, concatenates the embeddings, and maps them through a nonlinear hidden layer:

X\mapsto \tanh(\operatorname{concat}(E[x_{t-C}],\ldots,E[x_{t-1}])W_1+b_1)W_2+b_2.

The context length is a data-shape contract. A model trained with C=3 cannot silently receive a two-token input without padding or a defined shorter-context rule. The next cell keeps the forward cache explicit because the backward pass needs each intermediate array.

def make_context_dataset(tokens, context_length):
    tokens = np.asarray(tokens, dtype=np.int64)
    if len(tokens) <= context_length:
        return np.empty((0, context_length), dtype=np.int64), np.empty(0, dtype=np.int64)
    inputs = np.stack([
        tokens[start:start + context_length]
        for start in range(len(tokens) - context_length)
    ])
    targets = tokens[context_length:]
    return inputs, targets


def init_mlp(vocab_size, context_length, embedding_dim=5, hidden_dim=24, seed=0):
    init_rng = np.random.default_rng(seed)
    scale = 1.0 / np.sqrt(embedding_dim * context_length)
    return {
        "embedding": init_rng.normal(scale=0.4, size=(vocab_size, embedding_dim)),
        "w1": init_rng.normal(scale=scale, size=(embedding_dim * context_length, hidden_dim)),
        "b1": np.zeros(hidden_dim),
        "w2": init_rng.normal(scale=1.0 / np.sqrt(hidden_dim), size=(hidden_dim, vocab_size)),
        "b2": np.zeros(vocab_size),
    }


def mlp_forward(parameters, inputs):
    embedding = parameters["embedding"][inputs]
    flat_embedding = embedding.reshape(len(inputs), -1)
    hidden_pre = flat_embedding @ parameters["w1"] + parameters["b1"]
    hidden = np.tanh(hidden_pre)
    logits = hidden @ parameters["w2"] + parameters["b2"]
    cache = (inputs, embedding, flat_embedding, hidden, logits)
    return logits, cache


def mlp_loss_and_grads(parameters, inputs, targets):
    logits, cache = mlp_forward(parameters, inputs)
    loss = cross_entropy(logits, targets)
    probabilities = softmax(logits)
    dlogits = probabilities.copy()
    dlogits[np.arange(len(targets)), targets] -= 1.0
    dlogits /= len(targets)
    _, embedding, flat_embedding, hidden, _ = cache
    dw2 = hidden.T @ dlogits
    db2 = dlogits.sum(axis=0)
    dhidden = dlogits @ parameters["w2"].T
    dhidden_pre = dhidden * (1.0 - hidden ** 2)
    dw1 = flat_embedding.T @ dhidden_pre
    db1 = dhidden_pre.sum(axis=0)
    dflat_embedding = dhidden_pre @ parameters["w1"].T
    dembedding = dflat_embedding.reshape(embedding.shape)
    d_embedding_table = np.zeros_like(parameters["embedding"])
    for position in range(inputs.shape[1]):
        # Accumulate repeated token occurrences instead of overwriting them.
        np.add.at(d_embedding_table, inputs[:, position], dembedding[:, position])
    gradients = {
        "embedding": d_embedding_table,
        "w1": dw1,
        "b1": db1,
        "w2": dw2,
        "b2": db2,
    }
    return loss, gradients


def parameter_count(parameters):
    return sum(value.size for value in parameters.values())


text = "abacabadabacabae" * 16
alphabet = sorted(set(text))
stoi = {symbol: index for index, symbol in enumerate(alphabet)}
tokens = np.asarray([stoi[symbol] for symbol in text], dtype=np.int64)
context_length = 3
inputs, targets = make_context_dataset(tokens, context_length)
split = int(0.75 * len(inputs))
train_inputs, valid_inputs = inputs[:split], inputs[split:]
train_targets, valid_targets = targets[:split], targets[split:]
parameters = init_mlp(len(alphabet), context_length, seed=SEED + 3)
print("dataset shapes:", train_inputs.shape, valid_inputs.shape)
print("vocabulary:", alphabet, "parameters:", parameter_count(parameters))
assert train_inputs.shape[1] == context_length
assert len(train_inputs) == len(train_targets)

dataset shapes: (189, 3) (64, 3)
vocabulary: ['a', 'b', 'c', 'd', 'e'] parameters: 534

The embedding gradient uses np.add.at because the same token can occur several times in one batch. Ordinary indexed assignment would retain only the last contribution and silently change the derivative. The hidden layer gives the model a nonlinear way to combine positions, while the fixed concatenation still makes the receptive field exactly three tokens.

def copy_parameters(parameters):
    return {name: value.copy() for name, value in parameters.items()}


def evaluate_mlp(parameters, inputs, targets):
    return mlp_loss_and_grads(parameters, inputs, targets)[0]


def train_mlp(parameters, inputs, targets, steps, learning_rate):
    losses = []
    for _ in range(steps):
        loss, gradients = mlp_loss_and_grads(parameters, inputs, targets)
        for name in parameters:
            parameters[name] -= learning_rate * gradients[name]
        losses.append(loss)
    return np.asarray(losses)


tiny_parameters = init_mlp(len(alphabet), context_length, hidden_dim=32, seed=SEED + 4)
tiny_losses = train_mlp(
    tiny_parameters, train_inputs[:4], train_targets[:4], steps=400, learning_rate=0.18
)
main_parameters = init_mlp(len(alphabet), context_length, seed=SEED + 5)
main_losses = train_mlp(
    main_parameters, train_inputs, train_targets, steps=260, learning_rate=0.12
)
train_loss = evaluate_mlp(main_parameters, train_inputs, train_targets)
valid_loss = evaluate_mlp(main_parameters, valid_inputs, valid_targets)
print("tiny-batch loss:", round(float(tiny_losses[0]), 4), "->", round(float(tiny_losses[-1]), 6))
print("main train loss:", round(float(main_losses[0]), 4), "->", round(float(train_loss), 6))
print("main validation loss:", round(float(valid_loss), 6))
assert tiny_losses[-1] < 0.08
assert main_losses[-1] < main_losses[0]
assert np.isfinite(valid_loss)

tiny-batch loss: 1.8444 -> 0.001827
main train loss: 1.5849 -> 0.273499
main validation loss: 0.269394

The tiny-batch run is a debugging test, not a generalization result. A model with enough parameters should drive a handful of fixed examples close to zero; failure points to an indexing, gradient, or update error before hyperparameter comparisons begin. The main run separates the decreasing training loss from validation loss. Their gap is a capacity-and-data diagnostic, not proof that the larger model has learned a useful rule.

Sampling and generation

Training selects the parameters that assign probability to observed targets. Generation selects a token from the resulting distribution. Temperature \tau rescales logits as z/\tau; lower values concentrate probability around the largest logits. Top-k sampling sets all but the k largest logits to zero probability before drawing. Greedy decoding is the limiting case k=1.

Use the same prompt and a named random generator for each comparison so the decoding policy, rather than hidden global state, determines the difference.

def sample_next(logits, random_generator, temperature=1.0, top_k=None):
    logits = np.asarray(logits, dtype=float).copy()
    if temperature <= 0.0:
        raise ValueError("temperature must be positive")
    if top_k is not None:
        if top_k < 1:
            raise ValueError("top_k must be positive")
        keep = min(int(top_k), len(logits))
        candidate_indices = np.argpartition(logits, -keep)[-keep:]
        filtered = np.full_like(logits, -np.inf)
        filtered[candidate_indices] = logits[candidate_indices]
        logits = filtered
    probabilities = softmax(logits / temperature)
    return int(random_generator.choice(len(logits), p=probabilities))


def generate(parameters, prompt, steps, random_generator, temperature=1.0, top_k=None):
    generated = list(prompt)
    for _ in range(steps):
        context = np.asarray(generated[-context_length:], dtype=np.int64)[None, :]
        logits = mlp_forward(parameters, context)[0][0]
        generated.append(sample_next(logits, random_generator, temperature, top_k))
    return "".join(alphabet[index] for index in generated)


prompt = [stoi[symbol] for symbol in "aba"]
print("greedy:     ", generate(main_parameters, prompt, 18, np.random.default_rng(SEED + 6), top_k=1))
print("temperature:", generate(main_parameters, prompt, 18, np.random.default_rng(SEED + 7), temperature=1.2))
print("top-k=2:    ", generate(main_parameters, prompt, 18, np.random.default_rng(SEED + 8), temperature=0.9, top_k=2))
assert sample_next(np.array([0.1, 0.8, 0.2]), np.random.default_rng(SEED), top_k=1) == 1
assert len(generate(main_parameters, prompt, 5, np.random.default_rng(SEED))) == len(prompt) + 5

print("\ncontext-length experiment")
context_results = {}
for candidate_context in (1, 2, 4):
    candidate_inputs, candidate_targets = make_context_dataset(tokens, candidate_context)
    candidate_split = int(0.75 * len(candidate_inputs))
    candidate_parameters = init_mlp(
        len(alphabet), candidate_context, hidden_dim=20, seed=SEED + 10 + candidate_context
    )
    candidate_losses = train_mlp(
        candidate_parameters,
        candidate_inputs[:candidate_split],
        candidate_targets[:candidate_split],
        steps=120,
        learning_rate=0.12,
    )
    candidate_valid_loss = evaluate_mlp(
        candidate_parameters,
        candidate_inputs[candidate_split:],
        candidate_targets[candidate_split:],
    )
    context_results[candidate_context] = float(candidate_valid_loss)
    print(f"context={candidate_context}: validation loss={candidate_valid_loss:.4f}")
assert set(context_results) == {1, 2, 4}

greedy:      abacabacabacabacabaca
temperature: abadabaeabacabadabaca
top-k=2:     abacabacabadabacabaca

context-length experiment
context=1: validation loss=0.6167
context=2: validation loss=0.2981
context=4: validation loss=0.1557

PyTorch baseline handoff

The NumPy cells make the objective and gradient visible. The cumulative implementation is now a PyTorch module with the same next-token interface. Run the bigram and fixed-context MLP on a small periodic stream, then compare their decreasing losses before the Transformer adds attention and a larger receptive field.

import torch

from proof_lm.baselines import (
    TorchBigramLM,
    TorchContextMLP,
    make_context_dataset,
    train_classifier,
)

torch.manual_seed(SEED)
torch_tokens = [0, 1, 0, 1, 0, 1] * 18
context_inputs, context_targets = make_context_dataset(torch_tokens, context_length=3)
bigram_inputs = context_inputs[:, -1]
bigram = TorchBigramLM(vocab_size=2)
bigram_losses = train_classifier(
    bigram, bigram_inputs, context_targets, steps=24, learning_rate=0.1
)
torch_mlp = TorchContextMLP(
    vocab_size=2, context_length=3, embedding_dim=8, hidden_dim=24
)
mlp_losses = train_classifier(
    torch_mlp, context_inputs, context_targets, steps=24, learning_rate=0.05
)
print("PyTorch bigram loss:", round(bigram_losses[0], 4), "->", round(bigram_losses[-1], 6))
print("PyTorch MLP loss:", round(mlp_losses[0], 4), "->", round(mlp_losses[-1], 6))
assert bigram_losses[-1] < bigram_losses[0]
assert mlp_losses[-1] < mlp_losses[0]

PyTorch bigram loss: 0.6931 -> 0.031538
PyTorch MLP loss: 0.5981 -> 1e-06

The three samples use identical learned logits but different support and temperature rules, so their continuations need not agree. A random generator makes stochastic sampling reproducible without making it deterministic across policies. The context experiment reports validation loss for a fixed update budget; increasing context changes both the information available and the number of parameters, so it should be read as a controlled baseline rather than an isolated claim about context length.

Summary

• Stable softmax and cross-entropy turn next-token likelihood into testable NumPy functions.
• The hand-derived linear gradient matches finite differences, including the batch normalization factor.
• A bigram table trained by gradient descent approaches the empirical conditional entropy, which is an independent target for the training loop.
• A fixed-context MLP adds embeddings and nonlinear capacity; deliberate tiny-batch overfitting detects implementation bugs.
• Temperature and top-k modify the sampling distribution after training, while validation loss measures a separate held-out property.

Chapter 04 replaces the fixed local context with causal self-attention and checks the shape and gradient contracts of a decoder block.

Exercises

Use the exercises to test the chapter’s invariants and connect the derivations to the reusable implementation. Solutions are hidden in the notebook source and are available through the course tooling when needed.

[P3.1] Softmax gradient

Gradient derivation. Starting from mean softmax cross-entropy, derive the derivative with respect to one logit vector. Explain why the batch-size factor must appear in both the weight and bias gradients.

[P3.2] Sampling policy

Sampling policy. Given logits [2.0, 1.0, 0.0], describe the support of greedy decoding and top-k sampling with k=2. Explain how increasing temperature changes the relative probabilities without changing their ordering.

[P3.3]

PyTorch baseline contract. Explain why the bigram baseline should be evaluated with the previous-token tensor context_inputs[:, -1], while the fixed-context MLP receives the full (batch, context_length) tensor. State one shape assertion that would catch swapping those inputs.

def baseline_input_shapes(context_inputs, bigram_inputs, context_length):
    # Return a pair of booleans for the bigram and MLP contracts.
    pass