Teaching my language model to learn

1. Give each example a beginning and an ending

The first model started with a character we supplied, such as "h". It stopped when it reached a character with no recorded continuation. We can represent beginnings and endings explicitly instead.

Our final training set contains three short examples. We collect their characters into a shared vocabulary and append two special tokens. <START> and <END> each get one ID; they are not encoded as individual angle brackets and letters.

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

The original eight character IDs stay the same. <START> is ID 8, and <END> is ID 9. Our existing encode and decode functions still handle ordinary text.

We surround each encoded example with the two markers, then build input/target pairs within that example:

inputs = []
targets = []

for example in training_texts:
    tokens = [start_id] + encode(example) + [end_id]

    inputs.extend(tokens[:-1])
    targets.extend(tokens[1:])

extend adds the items from one list to another. Building the pairs separately prevents an accidental transition from the end of one example into the next example.

The first training pair is now <START> → h. Two examples begin with h, and one begins with w. There are also examples of d → <END> and o → <END>. The markers give the model something concrete to learn about boundaries.

2. Give a prediction a score

To improve predictions, we need to measure them. Our loss is the negative logarithm of the probability assigned to the actual next token:

loss = -math.log(probability)

This uses math from Python’s standard library. A confident correct prediction receives a small penalty. Giving the correct token very little probability produces a large penalty.

For probabilities between zero and one, the logarithm is negative. The minus sign makes the penalty positive. Logarithms also turn multiplication into addition: multiplying the probabilities along a sequence corresponds to adding its log probabilities. Here we average the negative log probabilities over our training pairs.

We first score the counting model from part one. Each pair is looked up in its existing probability table:

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))

On these three examples, the counting model’s average loss is approximately 0.528. This is a useful training-data reference. It is not evidence about performance on unseen text; we are scoring the same examples used to collect the counts.

3. Turn adjustable scores into probabilities

Our new model starts with scores that we can change. A score can be any real number, but sampling needs nonnegative weights and our loss needs probabilities. Softmax converts the scores into positive probabilities whose sum is one, apart from floating-point rounding.

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]

math.exp raises the number e to a power. We exponentiate every score, add the results, and divide each result by that total. Subtracting the largest score first keeps the exponentials from overflowing without changing the resulting probabilities.

softmax([0, 0, 0]) → approximately [0.333, 0.333, 0.333]
softmax([2, 0, 0]) → approximately [0.787, 0.107, 0.107]

Equal scores give equal probabilities. Raising one score increases its share of the total. Training will make those adjustments using the examples.

4. Make a table of weights

We give every current token its own row of scores. Each column represents a possible next token. With ten tokens, this is a 10 × 10 table: 100 adjustable numbers.

weights = []

for current_id in range(len(vocab)):
    row = []

    for next_id in range(len(vocab)):
        row.append(0.0)

    weights.append(row)

Every score starts at zero, so each row initially assigns a probability of 1/10 to every token. The untrained average loss is therefore -math.log(0.1), or about 2.303.

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))

Each row is a separate list. When we update the row for <START>, we change the model’s predictions at the beginning of a sequence. Predictions after h live in a different row.

5. Make one learning update

Our first example is <START> → h. For this example, the target is 1 for h and 0 for every other possible next token. With softmax and this negative-log loss, the derivative with respect to each score has a compact form:

gradient = predicted probability − target
new weight = old weight − learning rate × gradient

The gradient tells us how the loss changes with that score. We move a small distance in the opposite direction. The learning rate controls that distance.

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]))

Initially, h has probability 0.1. Its gradient is 0.1 − 1 = −0.9; subtracting 0.1 times that gradient raises its score by 0.09. Every other score in the row falls by 0.01.

The loss for this one example drops from about 2.303 to 2.213. We have written a gradient descent update ourselves. No automatic-differentiation library is calculating the derivative for us.

6. Repeat across the examples

One complete pass through the training pairs is an epoch. We repeat the update for each pair, then measure the average loss using the weights at the end of that epoch.

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)
        )

These are selected results from running Mark’s current file. The single-example update above happens before the epoch loop.

The learned weights are approaching the probabilities estimated by counting. We do not expect zero loss: after l, for example, the training data contains several different next characters. A model that sees only l has to divide its probability among them.

We update after each pair, so this is an example-by-example training loop. A loss decrease is not guaranteed after every individual update, especially when examples compete for the same row. What we observe here is a downward trend across epochs.

7. Let the learned weights write

Generation starts at <START> with an empty output list. We look up the current row, apply softmax, sample a token, and repeat. If the sampled token is <END>, we stop before adding it to the output.

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)

We set the sampling weight of <START> to zero because it is an input marker, not something to write. random.choices uses the relative weights of the remaining tokens, so it handles that adjustment without another division. The loss above was measured using the full softmax distribution; this exclusion is a generation-time rule.

for _ in range(5):
    print("Learned model:", repr(generate_learned(30)))

There is no starting "h" argument now. The first character comes from the learned start row. Endings are learned too: in this run, the softmax probability of <END> after d reaches about 94.6%. That is a strong preference, not a guarantee.

The 30-step limit remains a backup. Generated strings vary between runs; they may repeat letters, combine fragments, or occasionally finish immediately. repr makes spaces and an empty result visible.

8. Where the model still gets lost

After l, the model cannot tell whether it has reached that letter in hello or world. Both situations select the same row of weights. Training improves that row’s probabilities, but the row still has only one token of context.

This is the next question to explore: what changes if a prediction can see two previous characters? That would let he and rl select different predictions. For now, the milestone is complete: we can measure loss, calculate gradients, train weights, and generate text with learned beginnings and endings.

Full lesson code

This is a snapshot of Mark’s tokenizer.py at this milestone, including the earlier counting model, diagnostic prints, and incremental experiments. The snippets above are excerpts from that file. The snapshot preserves his code as written.

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%}",
    )

Download tokenizer.py

Run the downloaded file locally with python3 tokenizer.py. It uses only math and random from the standard library. The tokenizer knows only this vocabulary; adding a new character requires rebuilding the vocabulary and training the weights again.

Mark’s reflection

By Mark Agate

This work made me realize how truly simple it is to implement gradient descent. What was really interesting is that the Counting-model sometimes outperformed the Learned-model at this stage. It will be fun to see how the performance improves as we continue adding features.