Diagnosing Character-Level Model Learning from Loss and Generated Samples

Welcome to the final module. You have already built the pieces of a small masked text diffusion model: character IDs, corruption, a bidirectional denoiser, a masked-position loss, and iterative generation. Now the question is whether those pieces learn together.

Today you’ll run a compact training loop in Colab, plot its loss, and inspect generated text at several checkpoints. The aim is not fluent language. It is to recognize evidence that the denoiser is learning the patterns of a tiny corpus—and to know what that evidence cannot tell you.

What to watch during training

In each update, the model receives a batch of clean sentences with some valid characters replaced by masks. It predicts the original characters at those masked positions. Backpropagation adjusts its weights to reduce that prediction loss; generation then tests what happens when the trained model starts with every position masked.

For a quick refresher on the PyTorch loop, watch these two short passages from Hugging Face’s “Write your training loop in PyTorch.” Its example is not a diffusion model, but the optimization steps are the same.

A loss curve answers a narrower question than a generated sample. Training loss measures performance on masked characters from the training corpus; a sample shows the behavior of an entire, iterative generation run. Both are useful, but neither alone establishes that the model generalizes to new text.

Run the mini-model in Colab

Create a new Colab notebook and run the following cells in order. They are self-contained, so you do not need to reconstruct earlier notebook state. A GPU is optional for this small run. Allow roughly 8–12 minutes to inspect and run the cells, depending on Colab’s available hardware.

The corpus deliberately has only 16 sentence patterns. Some are shorter than others, so padding remains part of the training batch, but padding is never a prediction target.

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

torch.manual_seed(7)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device.type == "cpu":
    torch.set_num_threads(2)
print("Device:", device)

subjects = ("the cat", "the dog", "a fox", "a pup")
verbs = ("runs.", "naps.", "sits.", "hops.")
phrases = [f"{subject} {verb}" for subject in subjects for verb in verbs]

PAD, MASK = 0, 1
characters = sorted(set("".join(phrases)))
stoi = {char: index + 2 for index, char in enumerate(characters)}
itos = {index: char for char, index in stoi.items()}
vocab_size = len(stoi) + 2
max_len = max(map(len, phrases))

rows = phrases * 8
data = torch.full((len(rows), max_len), PAD, dtype=torch.long)
for row, text in enumerate(rows):
    data[row, :len(text)] = torch.tensor([stoi[char] for char in text])

loader = DataLoader(data, batch_size=32, shuffle=True)
print("Sentences:", len(phrases), "| batch shape:", next(iter(loader)).shape)

The denoiser below adds token, position, and current-mask-fraction information. Its Transformer encoder can use context on both sides of a mask. The padding mask prevents padded positions from serving as attention keys.

class TinyDenoiser(nn.Module):
    def __init__(self):
        super().__init__()
        width = 48
        self.token = nn.Embedding(vocab_size, width, padding_idx=PAD)
        self.position = nn.Embedding(max_len, width)
        self.noise = nn.Linear(1, width)

        layer = nn.TransformerEncoderLayer(
            d_model=width,
            nhead=4,
            dim_feedforward=96,
            dropout=0.0,
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=1)
        self.output = nn.Linear(width, vocab_size)

    def forward(self, tokens, mask_fraction):
        _, length = tokens.shape
        positions = torch.arange(length, device=tokens.device)
        hidden = (
            self.token(tokens)
            + self.position(positions)[None, :, :]
            + self.noise(mask_fraction[:, None])[:, None, :]
        )
        hidden = self.encoder(
            hidden, src_key_padding_mask=tokens.eq(PAD)
        )
        return self.output(hidden)

model = TinyDenoiser().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.003)

For monitoring, use the irreversible confidence-based sampler you encountered earlier: at each pass it fills some of the remaining masks, prioritizing the model’s most confident character predictions. It is intentionally small; this lesson is about observing training rather than modifying the sampler.

@torch.no_grad()
def generate(length, steps=8):
    model.eval()
    tokens = torch.full((1, length), MASK, device=device)

    for step in range(steps):
        fraction = tokens.eq(MASK).float().mean(dim=1)
        probabilities = model(tokens, fraction)[..., 2:].softmax(dim=-1)
        confidence, prediction = probabilities.max(dim=-1)

        confidence = confidence.masked_fill(tokens.ne(MASK), -1.0)
        target_filled = (step + 1) * length // steps
        new_count = target_filled - tokens.ne(MASK).sum().item()

        if new_count:
            positions = confidence[0].topk(new_count).indices
            tokens[0, positions] = prediction[0, positions] + 2

    return "".join(itos[int(token)] for token in tokens[0])

print("Before training:", repr(generate(13)), repr(generate(11)))

Now train. Each batch receives freshly sampled corruption. Two rows are fully masked on every update, giving the model regular practice with the starting state used by the generator. Loss is calculated only at corrupted, non-padding positions and excludes PAD and MASK from the target character choices.

history = []

for epoch in range(1, 31):
    model.train()
    total_loss = 0.0
    total_masked = 0

    for clean in loader:
        clean = clean.to(device)
        valid = clean.ne(PAD)

        noise_level = torch.rand(clean.size(0), 1, device=device)
        corrupted_positions = (
            (torch.rand(clean.shape, device=device) < noise_level)
            & valid
        )
        corrupted_positions[:2] = valid[:2]

        corrupted = clean.masked_fill(corrupted_positions, MASK)
        fraction = (
            corrupted_positions.sum(dim=1).float()
            / valid.sum(dim=1).clamp_min(1)
        )

        optimizer.zero_grad(set_to_none=True)
        logits = model(corrupted, fraction)
        loss = F.cross_entropy(
            logits[corrupted_positions][:, 2:],
            clean[corrupted_positions] - 2,
        )
        loss.backward()
        optimizer.step()

        count = corrupted_positions.sum().item()
        total_loss += loss.item() * count
        total_masked += count

    history.append(total_loss / total_masked)

    if epoch in (1, 10, 20, 30):
        print(
            f"epoch {epoch:2d} | loss {history[-1]:.3f} | "
            f"length 13: {generate(13)!r} | "
            f"length 11: {generate(11)!r}"
        )

plt.plot(range(1, len(history) + 1), history, marker=".")
plt.xlabel("Epoch")
plt.ylabel("Cross-entropy per corrupted character")
plt.title("Training denoising loss")
plt.grid(alpha=0.3)
plt.show()

An epoch here is one pass through the repeated corpus. The plotted value is weighted by the number of masked characters in each batch, rather than averaging batch losses equally. Because the masks change every time, individual points can wobble even when the overall trend is downward.

Diagnose what you see

Read the learning-curve discussion below while your run finishes. Its examples use conventional training and validation curves, not masked diffusion; focus on the distinction between learning the training task and generalizing beyond it.

Use two observations together to interpret your run:

• If loss generally falls and later samples begin to resemble the corpus—perhaps recognizable subjects, verbs, spaces, and periods—the denoiser is learning this small training task. A sample may still be malformed: generation repeatedly feeds the model its own partially completed output, unlike a single training prediction.
• If loss stays high and samples remain mostly arbitrary characters, the run has not shown convincing learning. First check that corrupted positions exist, that targets are clean characters rather than mask IDs, and that optimizer updates execute.
• If loss falls but samples remain poor, do not call the run a failure solely from the curve. The curve tests reconstruction with varying amounts of clean context; generation begins with none. Inspect whether fully masked examples reached training and whether all masks were eventually filled.

There is no validation split in this quick experiment. You therefore cannot conclude from its downward training curve that the model generalizes, nor diagnose overfitting. Likewise, these repeated, template-based sentences do not test whether it can write broadly useful language. They make the training signal easy to see.

Takeaways

You have now run the complete miniature system: fresh character masking, bidirectional prediction, masked-position optimization, a loss plot, and periodic generation from fully masked text. Read the loss as evidence about denoising the training corpus, and the samples as a complementary check on iterative generation. Neither is a substitute for evaluation on held-out text.

Next, you’ll adapt the sampler to denoise one block of positions at a time, then freeze that block before moving to the next.