In article two, we trained a language model using softmax, loss, and gradient descent. It learned how to begin and end a sequence. But every prediction still saw only one previous token.
That limitation shows up after l. In our examples, it can be followed by l, o, or d. One row of weights has to cover all three situations. This time, we give a prediction two tokens of context. The contexts el, ll, and rl can now select different rows.
The change is small enough to follow in plain Python. We then compare the models using the same training settings and a reproducible set of generated outputs.
1. Give each example two input tokens
We preserve the one-token model in bigram_checkpoint.py and begin the new experiment in context_two.py. The three training strings are unchanged: hello world, hello, and world. The vocabulary still contains eight characters plus <START> and <END>.
The difference is the shape of an input. Instead of a single token ID, it is an ordered pair of IDs, represented by a tuple. We place two start tokens before each example so even the first prediction has a complete context.
inputs = []
targets = []
for example in training_texts:
tokens = [start_id, start_id] + encode(example) + [end_id]
for position in range(2, len(tokens)):
context = (tokens[position - 2], tokens[position - 1])
inputs.append(context)
targets.append(tokens[position])At each position, the two preceding tokens become the input, and the token at that position becomes the target. The first five examples are:
('<START>', '<START>') → 'h'
('<START>', 'h') → 'e'
('h', 'e') → 'l'
('e', 'l') → 'l'
('l', 'l') → 'o'We still have 24 training pairs. We predict the same sequence of targets as before, including the end of each example. We have changed the information available for each prediction, not added extra targets or mixed examples together.
A model that predicts one token from the previous two is often called a trigram model: two context tokens plus the predicted token make a group of three.
2. Give each context a row of weights
The one-token model selected a row with weights[current_id]. The new model uses a dictionary keyed by a two-token tuple:
weights[(first_id, second_id)]We allocate a row for every possible pair of tokens. Each row has one adjustable score for each possible next token.
weights = {}
for first_id in range(len(vocab)):
for second_id in range(len(vocab)):
context = (first_id, second_id)
row = []
for next_id in range(len(vocab)):
row.append(0.0)
weights[context] = rowThere are ten token choices in each context position, so we have 10 × 10 = 100 context rows. Each row contains ten scores: 1,000 weights in total, compared with 100 in the previous model.
Only 13 of those 100 contexts actually occur as inputs in our training set. Their rows receive updates. The other 87 rows stay at their initial values. Allocating every possible context is simple, but much of this table is unused.
3. Zero scores are an equal starting point
All the scores initially read 0.0. That does not mean every token has zero probability. Softmax turns ten equal scores into ten equal probabilities, each 0.1.
We reuse the same softmax function and negative-log loss. Only the way we look up a row changes:
total_loss = 0
for context, target_id in zip(inputs, targets):
predicted = softmax(weights[context])
probability = predicted[target_id]
total_loss += -math.log(probability)
average_loss = total_loss / len(targets)
print("Untrained model loss:", round(average_loss, 3))The untrained average loss is 2.303 for both models, because both initially give every correct next token probability 1/10. More context creates the capacity to distinguish situations. Training makes the scores use that capacity.
4. Train the row for the whole context
The update rule is familiar:
gradient = predicted probability − target
new weight = old weight − learning rate × gradientThe target is 1 for the actual next token and 0 for the others. We subtract a small multiple of the gradient from each score in the selected context row.
learning_rate = 0.1
for epoch in range(100):
for context, target_id in zip(inputs, targets):
predicted = softmax(weights[context])
for candidate_id in range(len(vocab)):
target = 0.0
if candidate_id == target_id:
target = 1.0
gradient = predicted[candidate_id] - target
weights[context][candidate_id] -= learning_rate * gradient
total_loss = 0
for context, target_id in zip(inputs, targets):
predicted = softmax(weights[context])
total_loss += -math.log(predicted[target_id])
average_loss = total_loss / len(targets)
if (epoch + 1) % 10 == 0:
print(
"Epoch:", epoch + 1,
"Loss:", round(average_loss, 3),
)The examples ll → o and rl → d shared the same l row in the one-token model. They now update separate rows, so each can become a confident prediction.
Some uncertainty remains. The context lo is followed by a space in hello world and by <END> in hello. The start context also has two possible targets: h and w. Two tokens of context cannot resolve every ambiguity in this dataset.
5. Move a two-token window through the output
Generation begins with the context (<START>, <START>). After sampling a token, we keep the second token from the old context and add the new one:
context = (context[1], next_id)For example, after the context ('h', 'e') generates l, the next context is ('e', 'l'). We have a moving window of two tokens, not a memory of the entire string.
def generate(max_new_tokens):
context = (start_id, start_id)
generated_ids = []
for _ in range(max_new_tokens):
predicted = softmax(weights[context])
predicted[start_id] = 0.0
next_id = random.choices(
population=range(len(vocab)),
weights=predicted,
k=1,
)[0]
if next_id == end_id:
break
generated_ids.append(next_id)
# Keep the newest two tokens.
context = (context[1], next_id)
return decode(generated_ids)As before, <START> is excluded from sampling, <END> stops generation, and the 30-token limit bounds the output. Softmax still assigns some probability to unlikely continuations. If sampling reaches a context absent from training, that row gives a uniform distribution before we exclude <START>.
This helps explain why the output can still wander even when common fragments become more recognizable.
6. Make the comparison fair
Looking at a few attractive outputs is not enough. We use a separate comparison script, written by Codex, to train and sample both models under matched conditions.
| Setting | Both models |
|---|---|
| Training strings | hello world, hello, world |
| Vocabulary | 10 tokens |
| Targets per epoch | The same 24 targets, in the same order |
| Initial scores | All zero |
| Training | 100 epochs, learning rate 0.1 |
| Updates per model | 2,400 |
| Generation | Exclude START; stop at END or 30 tokens |
| Sampling | One output for each seed from 0 to 999 |
The one-token lesson file includes an extra single-example update before its epoch loop. The comparison loads the lesson files, discards their demonstration output, and resets all weights before training either model. That removes the extra update from the measurement without editing the original lesson.
The script uses each lesson’s softmax and generation functions, with the same gradient rule for training. It checks that the vocabulary and targets match, and that each two-token input ends with the corresponding one-token input. The full results include source-file hashes, the loss after every epoch, and the sampling counts.
7. What changed in the measurements
| Measurement | One-token context | Two-token context |
|---|---|---|
| Allocated scores | 100 | 1,000 |
| Observed / allocated contexts | 9 / 10 | 13 / 100 |
| Initial loss | 2.303 | 2.303 |
| Epoch 10 loss | 1.281 | 1.134 |
| Epoch 20 loss | 0.911 | 0.650 |
| Epoch 50 loss | 0.657 | 0.303 |
| Epoch 100 loss | 0.585 | 0.208 |
| Counting reference loss | 0.528 | 0.137 |
The extra context lets the model fit the observed transitions with less uncertainty. The counting references calculate probabilities directly from observed next-token frequencies, separately for each context size. They show how much ambiguity remains in this particular training set. Our finite training run approaches those references.
We score the actual next token in the training sequence, using its actual preceding context. We do not generate a string and then score it against a preferred answer. Loss is averaged over the full softmax probabilities, including the start-token output slot; the start-token exclusion happens only during generation.
8. Compare the generated text too
Here are the first five seeds from the experiment, in order. The examples are not selected for attractive results. Quotes make spaces and empty strings visible.
| Seed | One-token output | Two-token output |
|---|---|---|
| 0 | 'worlllorllorlorlo' | 'world' |
| 1 | 'held' | 'hellodow wor lre' |
| 2 | 'wo rlorlllld' | 'world' |
| 3 | 'hellldeld' | 'hello wor' |
| 4 | 'held' | 'hello world' |
The longer-context model produces complete training strings more often, but it still has odd outputs. Using the same seed makes the experiment reproducible; it does not force the models to make the same token choices or stop at the same point.
Across all 1,000 outputs per model, the one-token model exactly reproduced one of the three training strings 65 times (6.5%). The two-token model did so 611 times (61.1%). Each produced 13 empty strings; six one-token outputs and four two-token outputs reached the 30-character limit.
Exact training-string matches are a narrow description of this toy experiment, not an accuracy score for language. An unfamiliar output is not automatically bad, and reproducing training text is not proof of understanding. These counts give us a wider view of the sampling behavior than five examples alone.
9. More context has a cost
With a vocabulary of size V and context length C, this full-table design allocates V^(C + 1) scores. At ten tokens, that is 100 scores for one-token context, 1,000 for two-token context, and 10,000 for three-token context.
Longer contexts also occur less often. Even this tiny two-token model has 87 allocated rows with no training examples. Adding more context does not teach those rows how to behave.
Our result is encouraging and specific: two-token context fits these three training strings better under the settings we measured. We have not held out a separate evaluation set, tested a larger corpus, or established how either model performs on unseen text.
The next useful questions are how to evaluate on text kept out of training, and how to share what the model learns across different contexts. Embeddings and neural networks will give us ways to explore that without allocating an independent row for every possible history.
Full lesson code
The model snapshots preserve Mark’s code as written. The comparison script is separately credited to Codex. Download all three Python files into the same directory to repeat the experiment; they use only Python’s standard library.
Mark’s two-token model
# Part three: predict the next token from two previous tokens.
import math
import random
training_texts = [
"hello world",
"hello",
"world",
]
vocab = sorted(set("".join(training_texts)))
vocab.append("<START>")
start_id = len(vocab) - 1
vocab.append("<END>")
end_id = len(vocab) - 1
char_to_id = {}
for token_id, character in enumerate(vocab):
char_to_id[character] = token_id
def encode(text):
token_ids = []
for character in text:
token_ids.append(char_to_id[character])
return token_ids
def decode(token_ids):
characters = []
for token_id in token_ids:
characters.append(vocab[token_id])
return "".join(characters)
# Next: build training examples with two tokens of context.
inputs = []
targets = []
for example in training_texts:
tokens = [start_id, start_id] + encode(example) + [end_id]
for position in range(2, len(tokens)):
context = (tokens[position - 2], tokens[position - 1])
inputs.append(context)
targets.append(tokens[position])
for context, target in zip(inputs, targets):
first_id, second_id = context
print(
repr(vocab[first_id]),
repr(vocab[second_id]),
"->",
repr(vocab[target]),
)
weights = {}
for first_id in range(len(vocab)):
for second_id in range(len(vocab)):
context = (first_id, second_id)
row = []
for next_id in range(len(vocab)):
row.append(0.0)
weights[context] = row
print("Context rows:", len(weights))
print("Total weights:", len(weights) * len(vocab))
he_context = (char_to_id["h"], char_to_id["e"])
print("Scores after 'he':", weights[he_context])
def softmax(scores):
largest = max(scores)
exponentials = []
for score in scores:
exponentials.append(math.exp(score - largest))
total = sum(exponentials)
return [value / total for value in exponentials]
total_loss = 0
for context, target_id in zip(inputs, targets):
predicted = softmax(weights[context])
probability = predicted[target_id]
total_loss += -math.log(probability)
average_loss = total_loss / len(targets)
print("Untrained model loss:", round(average_loss, 3))
learning_rate = 0.1
for epoch in range(100):
for context, target_id in zip(inputs, targets):
predicted = softmax(weights[context])
for candidate_id in range(len(vocab)):
target = 0.0
if candidate_id == target_id:
target = 1.0
gradient = predicted[candidate_id] - target
weights[context][candidate_id] -= learning_rate * gradient
total_loss = 0
for context, target_id in zip(inputs, targets):
predicted = softmax(weights[context])
total_loss += -math.log(predicted[target_id])
average_loss = total_loss / len(targets)
if (epoch + 1) % 10 == 0:
print(
"Epoch:", epoch + 1,
"Loss:", round(average_loss, 3),
)
def generate(max_new_tokens):
context = (start_id, start_id)
generated_ids = []
for _ in range(max_new_tokens):
predicted = softmax(weights[context])
predicted[start_id] = 0.0
next_id = random.choices(
population=range(len(vocab)),
weights=predicted,
k=1,
)[0]
if next_id == end_id:
break
generated_ids.append(next_id)
# Keep the newest two tokens.
context = (context[1], next_id)
return decode(generated_ids)
for _ in range(5):
print(repr(generate(30)))
Mark’s one-token checkpoint
import math
import random
training_texts = [
"hello world",
"hello",
"world",
]
vocab = sorted(set("".join(training_texts)))
vocab.append("<START>")
start_id = len(vocab) - 1
vocab.append("<END>")
end_id = len(vocab) - 1
print(vocab)
char_to_id = {}
for token_id, character in enumerate(vocab):
char_to_id[character] = token_id
print(char_to_id)
def encode(text):
token_ids = []
for character in text:
token_ids.append(char_to_id[character])
return token_ids
print(encode("hello"))
def decode(token_ids):
characters = []
for token_id in token_ids:
characters.append(vocab[token_id])
return "".join(characters)
print(decode([3, 2, 4, 4, 5]))
original = "hello world"
reconstructed = decode(encode(original))
assert reconstructed == original
print("Round trip passed!")
inputs = []
targets = []
for example in training_texts:
tokens = [start_id] + encode(example) + [end_id]
inputs.extend(tokens[:-1])
targets.extend(tokens[1:])
for current_id, next_id in zip(inputs, targets):
current_character = vocab[current_id]
next_character = vocab[next_id]
print(repr(current_character), "->", repr(next_character))
counts = {}
for current_id, next_id in zip(inputs, targets):
if current_id not in counts:
counts[current_id] = {}
following = counts[current_id]
if next_id not in following:
following[next_id] = 0
following[next_id] += 1
l_id = char_to_id["l"]
for next_id, count in counts[l_id].items():
print(repr(vocab[next_id]), count)
probabilities = {}
for current_id, following in counts.items():
total = sum(following.values())
probabilities[current_id] = {}
for next_id, count in following.items():
probabilities[current_id][next_id] = count / total
for next_id, probability in probabilities[l_id].items():
print(repr(vocab[next_id]), probability)
def generate(start_character, max_new_tokens):
current_id = char_to_id[start_character]
genereated_ids = [current_id]
for _ in range(max_new_tokens):
if current_id not in probabilities:
break
following = probabilities[current_id]
next_id = random.choices(
population=list(following.keys()),
weights=list(following.values()),
k=1,
)[0]
if next_id == end_id:
break
genereated_ids.append(next_id)
current_id = next_id
return decode(genereated_ids)
print(generate("h", 30))
total_loss = 0
for current_id, next_id in zip(inputs, targets):
probability = probabilities[current_id][next_id]
loss = -math.log(probability)
total_loss += loss
print(
repr(vocab[current_id]),
"->",
repr(vocab[next_id]),
"loss:",
round(loss, 3)
)
average_loss = total_loss / len(targets)
print("Average loss:", round(average_loss, 3))
def softmax(scores):
largest = max(scores)
exponentials = []
for score in scores:
exponentials.append(math.exp(score - largest))
total = sum(exponentials)
return [value / total for value in exponentials]
print(softmax([0, 0, 0]))
print(softmax([2, 0, 0]))
weights = []
for current_id in range(len(vocab)):
row = []
for next_id in range(len(vocab)):
row.append(0.0)
weights.append(row)
total_loss = 0
for current_id, next_id in zip(inputs, targets):
scores = weights[current_id]
predicted_probabilities = softmax(scores)
probability = predicted_probabilities[next_id]
total_loss += -math.log(probability)
average_loss = total_loss / len(targets)
print("Untrained model loss:", round(average_loss, 3))
learning_rate = 0.1
current_id = inputs[0] # "<START>"
next_id = targets[0] # "h"
predicted = softmax(weights[current_id])
print("Before:", -math.log(predicted[next_id]))
for candidate_id in range(len(vocab)):
target = 0.0
if candidate_id == next_id:
target = 1.0
gradient = predicted[candidate_id] - target
weights[current_id][candidate_id] -= learning_rate * gradient
updated = softmax(weights[current_id])
print("After: ", -math.log(updated[next_id]))
for epoch in range(100):
for current_id, next_id in zip(inputs, targets):
predicted = softmax(weights[current_id])
for candidate_id in range(len(vocab)):
target = 0.0
if candidate_id == next_id:
target = 1.0
gradient = predicted[candidate_id] - target
weights[current_id][candidate_id] -= learning_rate * gradient
# Measure loss after this epoch's updates.
total_loss = 0
for current_id, next_id in zip(inputs, targets):
predicted = softmax(weights[current_id])
total_loss += -math.log(predicted[next_id])
average_loss = total_loss / len(targets)
if (epoch + 1) % 10 == 0:
print(
"Epoch:", epoch + 1,
"Loss:", round(average_loss, 3)
)
def generate_learned(max_new_tokens):
current_id = start_id
generated_ids = []
for _ in range(max_new_tokens):
predicted = softmax(weights[current_id])
# START is an input marker, never generated text.
predicted[start_id] = 0.0
next_id = random.choices(
population=range(len(vocab)),
weights=predicted,
k=1,
)[0]
if next_id == end_id:
break
generated_ids.append(next_id)
current_id = next_id
return decode(generated_ids)
print("Counting model:", generate("h", 30))
for _ in range(5):
print("Learned model: ", repr(generate_learned(30)))
d_id = char_to_id["d"]
predicted = softmax(weights[d_id])
print("After 'd':")
for token_id, probability in enumerate(predicted):
print(
repr(vocab[token_id]),
f"{probability:.1%}",
)
Codex comparison script
"""Codex comparison harness: identical training settings, unchanged lesson files.
Run: python3 compare_models.py
Save the measurements: python3 compare_models.py --json results.json
Uses Python's standard library only.
"""
import argparse
import contextlib
import hashlib
import io
import json
import math
from pathlib import Path
import random
import runpy
ROOT = Path(__file__).resolve().parent
EPOCHS = 100
LEARNING_RATE = 0.1
SAMPLE_COUNT = 1000
MAX_TOKENS = 30
def load_lesson(filename, generator_name):
path = ROOT / filename
with contextlib.redirect_stdout(io.StringIO()):
lesson = runpy.run_path(str(path))
# The lessons run their demonstrations on load. Reset every weight in place
# so both comparisons start from zero, including the bigram's extra update.
weights = lesson['weights']
rows = weights.values() if isinstance(weights, dict) else weights
for row in rows:
for token_id in range(len(row)):
row[token_id] = 0.0
lesson['comparison_generator'] = lesson[generator_name]
lesson['source_sha256'] = hashlib.sha256(path.read_bytes()).hexdigest()
return lesson
def average_loss(lesson):
total = 0.0
for context, target_id in zip(lesson['inputs'], lesson['targets']):
probabilities = lesson['softmax'](lesson['weights'][context])
total += -math.log(probabilities[target_id])
return total / len(lesson['targets'])
def train(lesson):
history = [{'epoch': 0, 'loss': average_loss(lesson)}]
for epoch in range(1, EPOCHS + 1):
for context, target_id in zip(lesson['inputs'], lesson['targets']):
row = lesson['weights'][context]
predicted = lesson['softmax'](row)
for candidate_id in range(len(lesson['vocab'])):
target = 1.0 if candidate_id == target_id else 0.0
row[candidate_id] -= LEARNING_RATE * (predicted[candidate_id] - target)
history.append({'epoch': epoch, 'loss': average_loss(lesson)})
return history
def counting_reference(lesson):
counts = {}
for context, target_id in zip(lesson['inputs'], lesson['targets']):
following = counts.setdefault(context, {})
following[target_id] = following.get(target_id, 0) + 1
total_loss = 0.0
for context, target_id in zip(lesson['inputs'], lesson['targets']):
following = counts[context]
total_loss -= math.log(following[target_id] / sum(following.values()))
return total_loss / len(lesson['targets'])
def summarize(lesson):
history = train(lesson)
rows = lesson['weights']
parameter_count = sum(len(row) for row in (rows.values() if isinstance(rows, dict) else rows))
samples = []
matches = empty = at_limit = 0
# Reset the seed for each output in each model: reproducible samples without
# selecting only the attractive outputs. Different models may draw different
# numbers of tokens; identical seeds do not imply identical token choices.
saved_random_state = random.getstate()
try:
for seed in range(SAMPLE_COUNT):
random.seed(seed)
output = lesson['comparison_generator'](MAX_TOKENS)
if seed < 5:
samples.append({'seed': seed, 'text': output})
matches += output in lesson['training_texts']
empty += output == ''
at_limit += len(output) == MAX_TOKENS
finally:
random.setstate(saved_random_state)
return {
'source_sha256': lesson['source_sha256'],
'allocated_contexts': len(rows),
'observed_contexts': len(set(lesson['inputs'])),
'parameters': parameter_count,
'training_pairs': len(lesson['targets']),
'history': history,
'counting_reference_loss': counting_reference(lesson),
'samples': samples,
'sample_count': SAMPLE_COUNT,
'training_string_matches': matches,
'empty_outputs': empty,
'outputs_at_character_limit': at_limit,
}
def compare():
one = load_lesson('bigram_checkpoint.py', 'generate_learned')
two = load_lesson('context_two.py', 'generate')
assert one['training_texts'] == two['training_texts']
assert one['vocab'] == two['vocab']
assert one['targets'] == two['targets']
assert one['inputs'] == [context[-1] for context in two['inputs']]
result = {
'training_texts': one['training_texts'],
'vocab': one['vocab'],
'epochs': EPOCHS,
'learning_rate': LEARNING_RATE,
'updates_per_model': EPOCHS * len(one['targets']),
'max_generated_tokens': MAX_TOKENS,
'sample_seeds': [0, SAMPLE_COUNT - 1],
'loss_scope': 'Full softmax, average over the training pairs; not held-out evaluation.',
'generation_rule': 'Exclude START from sampling, stop at END or the 30-token cap.',
'one_token': summarize(one),
'two_tokens': summarize(two),
}
return result
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--json', type=Path, help='Write the full measurements to this file.')
args = parser.parse_args()
result = compare()
print('Same three examples, 24 pairs, zero initialization, 100 epochs, learning rate 0.1.')
print('All losses below are training losses, not held-out scores.\n')
for key in ('one_token', 'two_tokens'):
item = result[key]
print(key)
print(' Weights:', item['parameters'])
print(' Loss:', round(item['history'][0]['loss'], 6), '->', round(item['history'][-1]['loss'], 6))
print(' Counting reference:', round(item['counting_reference_loss'], 6))
print(' Training-string matches:', item['training_string_matches'], '/', SAMPLE_COUNT)
print(' Empty / at length limit:', item['empty_outputs'], '/', item['outputs_at_character_limit'])
for sample in item['samples']:
print(' Seed', sample['seed'], ':', repr(sample['text']))
if args.json:
args.json.write_text(json.dumps(result, indent=2) + '\n')
print('\nSaved measurements to', args.json)
if __name__ == '__main__':
main()
Download recorded comparison results
python3 compare_models.py --json comparison-results.jsonThe saved results include all 101 loss checkpoints, from initialization through epoch 100. Sample strings were recorded on this machine’s Python 3.14 runtime; different Python implementations or versions may produce different seeded samples.
Try both models in your browser
Train both models on the same examples and compare their losses and generated text. The defaults match the article’s training settings. Each run starts the scores at zero; the seed controls generation. The demo generates five samples per model, rather than the full 1,000-sample experiment above.
Ready. Python loads when you run the experiment.
Training happens in a separate browser worker on your device. Your examples are never uploaded. Stop cancels the worker; each training run has a 15-second limit after Python loads.
One token of context
Average training loss
Two tokens of context
Average training loss
These losses measure the training examples. Sample text is random; neither lower training loss nor a familiar string proves better performance on unseen text.
Inspect loss during training
| Epoch | One-token loss | Two-token loss |
|---|
Python runs on your device through Pyodide. Your training text stays in this tab. The first run downloads about 13 MiB of runtime files before browser compression; later runs reuse the loaded runtime. No NumPy or TensorFlow is used.
Mark’s reflection
Wow. Very impressed with how simply adding two characters to the training data can massively improve the results. Is this where scaling begins?