16 | The Complete Model: Follow the Character “欢” Through Every Step

Build a Small Model, Understand Large Models · Lesson 16

In Lesson 09, looking only at the final character “欢” could not distinguish whether the preceding person was Xiaoding or Xiaozhou. Now, when the complete model reads “早上，小丁喜欢” (“In the morning, Xiaoding likes”), it assigns the next character “红” (“red”) probability 0.999765. How did the added components change the prediction?

For this lesson, we borrow parameters from a completed pretraining run to observe how the components work together. The next lesson will examine how those parameters were trained. As when debugging a program, we hold the input and parameters fixed and stop at several points: are the input IDs correct, when does the row for “欢” change, and which step finally produces candidate-character scores?

Start with These 14 IDs

Our fixed text is:

早上，小丁喜欢红色的苹果。

It means “In the morning, Xiaoding likes red apples.” The series encodes text character by character. Adding BOS before the 13 characters gives 14 IDs:

[1, 32, 7, 76, 27, 6, 22, 44, 55, 57, 50, 58, 39, 5]

Entry 0 is BOS, entry 1 is “早”, entry 7 is “欢”, and entry 13 is “。”. Lock onto position 7 and use that same index at every later checkpoint, so we do not accidentally switch to another character.

The complete computation follows this route:

IDs [1, 14] → embedding lookup [1, 14, 64] → two model blocks → final RMSNorm → candidate scores [1, 14, 79]

The first number, 1, is the number of samples; 14 is the number of positions. The internal representation width is 64, and the candidate vocabulary size is 79. They answer different questions: “How many numbers represent each position?” and “How many candidates need scores?” They are not interchangeable.

The number of positions remains 14 throughout. What changes is the content of each row: first one ID, then 64 representation values, and finally 79 candidate scores. Position 7 can use the prefix from BOS through “欢”, but cannot read the following “红色的苹果。” (“red apples”). Holding the whole sentence in one tensor does not remove this restriction.

The 3D model inspector arranges these checkpoints in layers. Select a position and coordinate, then compare stages to inspect residual addition, final normalization, and how the 64 entries contribute to a candidate's score. It reads the same verified pretrained record; changing views does not train the model. The tables and code below also provide the complete verification path.

Stay with the Row for “欢”

The first stop is the lookup. The ID for “欢”, 44, selects row 44 of the Embedding table. These 64 values depend only on the character and the table's parameters; they have not read “小丁” (“Xiaoding”). Stopping here lets us compare them with the representations after contextual processing.

The first block gives this row a path for reading the prefix, then processes the aggregated coordinates with the feed-forward network. Two residual connections preserve direct paths for the incoming representations. The second block receives that result and processes it again with another set of parameters. Both block outputs still have shape [1, 14, 64], so a shape check alone cannot reveal that their values have changed. We will inspect actual slices shortly.

The 64 internal coordinates do not map one-to-one to characters. Scoring the 79 vocabulary candidates requires a final projection.

After the second block, another RMSNorm is applied. A bias-free linear layer then projects each 64-entry vector to 79 scores. For candidate j, this is:

logit_j=h₁W_j1+h₂W_j2+…+h₆₄W_j64

Here, h is the final normalized vector, and W_j is the output matrix row for candidate j. This is exactly the row-by-row weighted sum from Lesson 2, now with 64 input values instead of two.

This is where [1, 14, 64] becomes [1, 14, 79]. The returned logits are scores that have not been normalized into probabilities. Training sends them to cross-entropy; inspecting candidates or selecting a character uses Softmax. The probability of “红” at the start of the lesson comes from this entire row of scores.

Connect the Checkpoints in One forward Method

First write the route in code and create an untrained random_model, checking only its structure, shapes, and parameter count. The pretrain_trace read later is the record of the pretrained model from the opening; these two models have different parameters.

The variable x is reassigned along the route. Each step must produce the shape the next step expects:

import torch
from torch import nn
from tiny_lm.model import ModelConfig
from tiny_lm.layers import RMSNorm, TransformerBlock
from tiny_lm.tokenizer import CharTokenizer

class TinyLM(nn.Module):
    """将整数编号 [B,T] 转为原始候选分数 [B,T,V]，不在前向内选字。"""
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        # ModuleList 注册各块的独立参数，模型的优化器才能找到它们。
        self.blocks = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layers)])
        self.final_norm = RMSNorm(config.d_model, config.norm_eps)
        self.output = nn.Linear(config.d_model, config.vocab_size,
                                bias=False)
        self.apply(self._initialize)

    @staticmethod
    def _initialize(module):
        # 随机填写 Linear/Embedding；RMSNorm 的缩放参数保留初值 1。
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0., std=.02)

    def forward(self, ids):
        # 这里只接受一批编号，不把标签或损失计算混进模型前向。
        if ids.ndim != 2 or not 0 < ids.shape[1] <= self.config.max_seq_len:
            raise ValueError("需要 [B,T]，且 1 <= T <= 32")
        x = self.embedding(ids)  # [B,T] → [B,T,D]，按编号读取参数表行。
        # 各块改写位置表示，保留 [B,T,D] 的轴与宽度。
        for block in self.blocks:
            x = block(x)
        # 最后归一后，将 D 项内部坐标映射成 V 个候选分数；此处没有 Softmax。
        return self.output(self.final_norm(x))

torch.manual_seed(42)
tokenizer = CharTokenizer.for_course()
# 这是一份新建的随机模型，尚未训练；不要与下段的预训练记录混为一份参数。
random_model = TinyLM(ModelConfig(vocab_size=tokenizer.vocab_size))
text = "早上，小丁喜欢红色的苹果。"
ids = torch.tensor([tokenizer.encode(text, add_bos=True)])
logits = random_model(ids)
print(ids.shape)     # torch.Size([1, 14])
print(logits.shape)  # torch.Size([1, 14, 79])
print(sum(p.numel() for p in random_model.parameters()))  # 92352

ModuleList registers both blocks as model submodules. A plain Python list can be iterated over but does not integrate them into the model's parameter management in the same way. apply visits modules and fills the linear and embedding weights with random values of mean 0 and standard deviation 0.02. The RMSNorm scale parameters still start at 1.

Stop at Each Checkpoint in the Saved Record

Read the pretrained model's forward record for this same sentence and select position 7. We display only the first four coordinates at each checkpoint; the other 60 still participate in the actual computation:

Change Only the Prefix: What Happens to the Same “欢”?

Changed values alone do not establish that the prefix had an effect. Using the same pretrained weights, compare 早上，小丁喜欢 and 早上，小周喜欢—“In the morning, Xiaoding likes” and “In the morning, Xiaozhou likes.” Supply only the prefixes, without their following answers.

Both end in “欢”, with ID 44, so their initial 64-entry lookup vectors are identical. After two blocks, the preceding “丁” and “周” have paths that can affect this row, and the outputs differ.

The two grids share a color scale: teal is positive, purple is negative, and deeper color means a larger absolute value. Each 8×8 grid is still one 64-dimensional vector, not 64 characters. A cell's color alone does not justify calling it a “fruit feature.”

With the “小丁” prefix, “红” has probability about 99.9765%. With “小周”, “黄” (“yellow”) has probability about 99.9736%. The character baseline in Lesson 09 looked up only the current “欢” and gave both people the same distribution. The prefix now truly changes the next-character prediction. No parameters were updated between the two runs; the input and the representations computed along the way changed.

from tiny_lm.checkpoint import load_checkpoint
from tiny_lm.trace import parameter_digest

# 加载同一份预训练权重；本段没有标签、反向或优化器更新。
context_model, context_tokenizer, _ = load_checkpoint(
    "artifacts/first_run/pretrain.pt")
before_digest = parameter_digest(context_model)
embeddings, representations = [], []
for prefix in ["早上，小丁喜欢", "早上，小周喜欢"]:
    # 两个前缀都以“欢”结尾，只在前面的人名处差一个字。
    prefix_ids = torch.tensor([context_tokenizer.encode(prefix, add_bos=True)])
    with torch.no_grad():
        scores, detail = context_model(prefix_ids, return_trace=True)
    embeddings.append(detail["embedding"][0, -1])
    representations.append(detail["blocks"][-1]["block_output"][0, -1])
    # 对最后位置的完整79项归一化，再查看“红”和“黄”，不单独重分这两项。
    probabilities = scores[0, -1].softmax(-1)
    print(prefix, {c: float(probabilities[context_tokenizer.stoi[c]])
                   for c in ["红", "黄"]})

assert torch.equal(embeddings[0], embeddings[1])  # 当前字符相同，查表结果完全相同。
assert not torch.allclose(representations[0], representations[1])
assert parameter_digest(context_model) == before_digest  # 比较期间参数没有改变。

Inspect the 64-Term Weighted Sum for “红”

At this position, the output head gives “红” a logit of about 13.768686 and “紫” (“purple”) about 0.992895. Softmax over all 79 entries gives “红” probability 0.999765. This is a measured prediction from a trained model, and it happens to match the actual next character.

We can also break down 13.768686. Each bar below is one term hᵢ×W红,ᵢ: positive contributions appear above zero, and negative contributions below. All 64 terms must be summed to obtain the score for “红”. Each of the other 78 candidates has its own parameter row and requires the same computation.

No character has been selected or appended yet. The generation loop in Lesson 18 masks PAD, BOS, and UNK before renormalizing, so its probability for the same candidate differs slightly. Do not combine the last decimal places from these two contexts as though they came from one distribution.

Read the original record directly with the following code, without retraining:

import json
from pathlib import Path

# JSON 保存已训练模型的一次前向结果；读取它不会给 random_model 加载权重。
pretrain_record = json.loads(Path(
    "artifacts/first_run/forward_trace.json").read_text(encoding="utf-8"))
pretrain_trace = pretrain_record["trace"]
assert pretrain_record["stage"] == "pretrain"
position = 7
# 取样本 0 的位置 7（欢），只显示前四个坐标，其余坐标仍参与原计算。
print(pretrain_trace["embedding"][0][position][:4])
print(pretrain_trace["blocks"][0]["block_output"][0][position][:4])
scores = torch.tensor(pretrain_trace["logits"][0][position])
# red 索引的是词表候选轴，不是 64 维表示中的某个坐标。
red = tokenizer.stoi["红"]
print(float(scores[red]), float(scores.softmax(-1)[red]))

When debugging, if the output still has only 64 entries, check whether the vocabulary projection was omitted. If decoded IDs produce the wrong characters, check the vocabulary paired with the model. The checkpoint records above help locate the step where something went wrong.

In the model inspection demo, you can select another input position and inspect its embedding lookup, each block's attention and two residual additions, and the final candidate scores. Clicking a coordinate shows its value. The page uses the saved forward record from this lesson; it does not retrain a model in the browser.

Of All These Numbers, Which Ones Does Training Change?

The vectors in the table are regenerated from the input. What the model stores long-term and modifies during training are its parameter groups. Counting them also helps detect an unregistered layer or an accidentally shared matrix. Let V be the vocabulary size, D the internal width, L the number of blocks, and F the intermediate feed-forward width. The input and output matrices are not shared, and all linear layers are bias-free.

The total is therefore:

2VD+L(4D²+3DF+2D)+D=92,352

The “per block” rows describe one block's components; in the final total, count “two complete blocks” only once. The 2 heads divide width 64 into width 32 per head, so do not multiply 4D² by the head count again. Fixed RoPE frequencies, input IDs, intermediate results, and the causal mask are not trainable parameters.

Check That the Assembled Network Can Propagate Gradients

Return to the random_model we built earlier. Supply targets shifted by one position, compute the loss, and backpropagate to check that every trainable parameter is connected to the loss. Reading a JSON record does not alter this model's parameters. The following checks connectivity only; no training has occurred yet:

from tiny_lm.losses import token_loss

# 输入为 BOS+原文；标签向后错一格，成为原文+EOS，两者同为 [1,14]。
targets = torch.tensor([tokenizer.encode(text) + [2]])  # 末尾是 EOS
# 回到前面新建的随机模型，只检查它的全部参数是否接到了损失。
random_model.zero_grad(set_to_none=True)
loss = token_loss(random_model(ids), targets)
loss.backward()
# backward 只填写梯度；这里没有优化器 step，随机参数仍未更新。
assert all(p.grad is not None and torch.isfinite(p.grad).all()
           for p in random_model.parameters())
print(float(loss.detach()))

This confirms that the parameters connect to the loss and that all gradients in this pass are finite. Unused vocabulary rows may have zero gradients. Whether the network can learn well still requires actual training.

Two Questions About Position and Structure

1. What should the scores at the BOS position predict? What should the final valid position, containing the period, predict?
2. If only the internal width changes from 64 to 128, keeping two heads, how many parameters are in one attention block's four projection matrices? Should you multiply by the head count again?

Before training, give your model two ID sequences of length 5. Write down the expected output shape, then run the forward pass. Change only the character ID at position index 3 in one sequence, and verify that positions 0–2 of that sequence and every output in the other sample remain unchanged. Finally, compute a loss, backpropagate, and check that the parameters receive finite gradients. If a check fails, investigate axes, visibility, or parameter registration before increasing the training steps.

Run this from the code companion code directory:

# 在配套代码目录运行本讲实验，核对正文中的结果。
python -m lessons.l16

Code and run records: Lesson runner · Complete model source · Actual forward record.

Source navigation: model.py · TinyLM.