MLP language model

Beyond counts: learning distributed representations

A bigram model stores one probability vector per vocabulary token — if you have 8,000 tokens, you have 8,000 rows and no sharing between similar contexts. "The cat" and "the dog" do not inform each other because only the immediate predecessor matters and each row is independent.

An MLP language model fixes the first problem by using a fixed context window of the last k tokens and fixes the second by embedding tokens into a shared continuous space. Similar tokens and similar contexts produce similar hidden states; the network generalizes to combinations never seen in training.

This is the same architectural family you used for classification in the foundations module — embeddings, linear layers, nonlinearity, softmax — rearranged for sequential prediction. You are one step away from attention: the MLP LM uses a fixed window; the transformer learns which past tokens to look at.

Callout — same loss, richer function: The training objective is still next-token cross-entropy. You are upgrading the function class from a lookup table to a neural net.

Architecture: context window → logits

Given token IDs [t_{-k+1}, ..., t_{-1}], predict distribution over t_0:

inputs:  (batch, context_len)           # integer token IDs
embed:   (batch, context_len, embed_dim) # Embedding lookup
flatten: (batch, context_len * embed_dim)
mlp:     (batch, hidden) → (batch, vocab_size)
softmax: probabilities over vocab
loss:    cross_entropy(logits, target_token)

For a context length of 8 and embed_dim of 64, the first linear layer sees 512 inputs. Each position's embedding captures semantic hints; the MLP learns which combinations predict the next token.

Character vs token MLP

Character models see more positions but a smaller vocab (~256 bytes). Token models match your BPE pipeline from lesson 3.1. For this project, either works — pick one and stay consistent through lesson 3.5.

Training data: sliding window

From a token sequence [x_1, x_2, ..., x_N], create training examples:

context = [x_1, ..., x_k]     → target = x_{k+1}
context = [x_2, ..., x_{k+1}] → target = x_{k+2}
...

Every contiguous window is one example. A 10k-token corpus with context 8 yields ~10k training rows — plenty for a tiny MLP on a small vocab.

Split train/val by sequence or by time block, not by random windows from the same document (that leaks adjacent context across splits). For a single file like tiny Shakespeare, a simple 90/10 split on windows is acceptable at this scale.

Training loop (same as lesson 2.5):

for context, target in train_loader:
    logits = model(context)           # (batch, vocab)
    loss = F.cross_entropy(logits, target)
    loss.backward()
    optimizer.step()

Log train and val loss per epoch. Val loss should track train loss early; divergence means overfitting — your MLP may be large relative to corpus size.

Sampling from the MLP LM

Generation mirrors the bigram loop, but context is a tensor of the last k tokens:

@torch.no_grad()
def generate(model, start_ids, max_new_tokens):
    context = list(start_ids)
    for _ in range(max_new_tokens):
        x = torch.tensor([context[-k:]], dtype=torch.long)
        logits = model(x)[0]
        probs = F.softmax(logits / temperature, dim=-1)
        next_id = torch.multinomial(probs, 1).item()
        context.append(next_id)
    return context

Use model.eval() and torch.no_grad() — no gradients during inference.

Qualitative expectations on a small corpus after modest training:

• Local spelling and short phrases resemble the training text.
• Long-range coherence is weak (fixed window cannot remember beginnings).
• Repeated n-grams may appear (model finds low-loss shortcuts).

Compare samples side-by-side with your bigram model from lesson 3.2. The MLP should show slightly better phrase structure within the context window even if both models ramble globally.

Capacity, overfitting, and hyperparameters

Knobs that matter on CPU-scale training:

Start small: embed_dim=64, hidden=128, context=8. Scale up only if underfitting.

Callout — qualitative beats benchmark here: Your tiny MLP will not beat GPT. Success is: val loss decreases, samples look more corpus-like than bigram output, and you can explain the architecture in an interview.

Why this model still hits a wall

Fixed context is the bottleneck. No matter how wide the MLP, token at position 100 cannot directly attend to token at position 1 if context_len is 8. Convolutional and recurrent architectures partially address this; self-attention (lesson 3.4) addresses it directly and scales better on parallel hardware.

The MLP LM is pedagogically valuable because you train it with the exact same loss and sampling code you will reuse in the transformer — only the forward method changes.

Engineering problem (staff framing)

A fixed-window MLP LM teaches capacity limits before attention: context is a flattened bag of last-k embeddings.

Diagram — Fixed-window MLP LM

flowchart TD
  T1[tok t-k] --> Cat[Concat embeddings]
  T2[tok t-1] --> Cat
  Cat --> MLP[MLP]
  MLP --> Logits[Vocab logits]

Precise definitions & mental model

Context window as hyperparameter; embedding+MLP; no content-based routing across positions.

Tradeoffs — when to use what

Simple & parallel vs cannot learn long-range selectively.

Failure modes (interview + on-call)

k too small → no syntax; k too large → param blowup in concat MLP.

Production & OSS practices

Pedagogical baseline; production uses transformers — keep MLP as unit test of training infra.

Deep dive (FAANG / OSS bar)

Push «mlp-language-model» past tutorial depth: write the interface contract (inputs/outputs/invariants), list three measurable metrics, and name two degrade modes if the happy path fails. Add a short threat note: what an attacker or noisy tool result could do, and which layer catches it (schema, policy, HITL, or eval gate).

flowchart LR
  Contract[Interface contract] --> Metrics
  Metrics --> Degrade[Degrade modes]
  Degrade --> Threat[Threat + control]

Micro-project: Train + sample

In m3/mlp_lm/:

1. Implement MLPLM with configurable context_len, embed_dim, hidden_dim, vocab_size.
2. Build sliding-window dataset from your tokenized corpus (reuse tokenizer artifacts).
3. Train with PyTorch; save checkpoints/best.pt and history.json.
4. Implement sampling with temperature=1.0; generate ≥200 tokens from 3 prompts.
5. Write COMPARE.md: paste one bigram sample and one MLP sample for the same start context; note differences in local coherence, repetition, and val perplexity if computed.

Commit samples.txt and training script. Document hyperparameters and hardware (CPU ok).

Checklist

• MLP LM trains with decreasing val loss
• Samples generated from saved checkpoint
• COMPARE.md with bigram vs MLP qualitative analysis
• Architecture hyperparameters documented in README