Train it. Test what it learned.

We have assembled the pieces. Now the model must learn a task. Given a source such as [4, 7, 9], we want it to generate [4, 7, 9, EOS]. We choose copying because the target is unambiguous, the data can be generated locally, and every stage of an encoder–decoder transformer is still exercised.

Download and unzip the complete code bundle . It contains the NumPy forward calculations, the PyTorch model, this training script, numerical checks, and a recorded example run. No dataset download or GPU is required.

1. Define the task before the model

Our vocabulary has twelve IDs. PAD=0, BOS=1, and EOS=2 are reserved; ID 3 is unused in this synthetic task. The actual source symbols are integers 4 through 11. A source has one, two, or three symbols.

We enumerate every such source: \(8+8^2+8^3=584\) distinct sequences. A seeded permutation assigns 467 to training and 117 to evaluation. The sets are disjoint. Evaluation therefore checks unseen arrangements within this small bounded domain, rather than simply reusing training examples.

The source is padded to length three. The decoder input and target are padded to length four. We append EOS to the target, prepend BOS to the decoder input, and align the loss one step ahead.

def dataset(seed: int = 7) -> tuple[list[list[int]], list[list[int]]]:
    sequences = [list(items) for length in [1, 2, 3]
                 for items in itertools.product(range(4, 12), repeat=length)]
    generator = torch.Generator().manual_seed(seed)
    order = torch.randperm(len(sequences), generator=generator).tolist()
    cut = int(0.8 * len(order))
    return ([sequences[i] for i in order[:cut]],
            [sequences[i] for i in order[cut:]])


def batch(sequences: list[list[int]]) -> tuple[Tensor, Tensor, Tensor]:
    source = torch.zeros(len(sequences), 3, dtype=torch.long)
    decoder_input = torch.zeros(len(sequences), 4, dtype=torch.long)
    target = torch.zeros(len(sequences), 4, dtype=torch.long)
    for i, tokens in enumerate(sequences):
        source[i, :len(tokens)] = torch.tensor(tokens)
        decoder_input[i, :len(tokens) + 1] = torch.tensor([1] + tokens)
        target[i, :len(tokens) + 1] = torch.tensor(tokens + [2])
    return source, decoder_input, target

For source [4, 7, 9], the decoder input is [BOS, 4, 7, 9] and the target is [4, 7, 9, EOS]. For source [6], the decoder input is [BOS, 6, PAD, PAD] and the target is [6, EOS, PAD, PAD].

2. Assemble the complete PyTorch model

The source and target share an embedding table in this tiny task because they use the same symbols. That is a design choice, not a requirement of encoder–decoder attention. The model adds sinusoidal positions, applies one encoder block and one decoder block, and maps final target states to twelve logits.

class TinyTransformer(nn.Module):
    def __init__(self, vocabulary: int = 12, width: int = 32, heads: int = 4, inner: int = 64) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocabulary, width, padding_idx=0)
        self.encoder = EncoderBlock(width, heads, inner)
        self.decoder = DecoderBlock(width, heads, inner)
        self.encoder_norm, self.decoder_norm = nn.LayerNorm(width), nn.LayerNorm(width)
        self.readout = nn.Linear(width, vocabulary)

    def embed(self, ids: Tensor) -> Tensor:
        x = self.embedding(ids)
        return x + positions(ids.shape[1], x.shape[-1], device=x.device, dtype=x.dtype)

    def encode(self, source: Tensor) -> tuple[Tensor, Tensor]:
        visible = (source != 0)[:, None, None, :]
        memory = self.encoder_norm(self.encoder(self.embed(source), visible))
        return memory, visible

    def decode(self, target_prefix: Tensor, memory: Tensor, source_visible: Tensor) -> Tensor:
        _, target_visible = attention_masks(target_prefix, target_prefix)
        y = self.decoder(self.embed(target_prefix), memory, target_visible, source_visible)
        return self.readout(self.decoder_norm(y))

    def forward(self, source: Tensor, target_prefix: Tensor) -> Tensor:
        memory, visible = self.encode(source)
        return self.decode(target_prefix, memory, visible)

The full file contains the previously introduced EncoderBlock, DecoderBlock, attention_masks, and positions definitions. It does not wrap nn.Transformer; the information flow remains visible.

Our teaching model uses D = 32, four attention heads, an FFN inner width of 64, pre-norm residual blocks, and final stack normalization. It omits dropout and the original paper’s embedding scaling and training recipe. This is a small working architecture, not a checkpoint-compatible reimplementation of the 2017 base model.

For comparison, demo_forward() in numpy_core.py executes the same sequence of architectural operations with untrained arrays. NumPy in this course does not implement automatic differentiation. We use PyTorch to optimize the model rather than hide a separate hand-written backpropagation engine inside the explanation.

3. Turn scores into a learning objective

For each non-padding target position, the model produces one logit per vocabulary symbol. Cross-entropy penalizes assigning low probability to the correct symbol. If z contains the logits and c is the correct ID, the single-position loss is:

Equivalently, this is the negative log of the softmax probability assigned to c. A confidently correct prediction has low loss; a confidently wrong one has high loss. We average over valid target positions, including EOS, while ignoring PAD.

def training_loss(model: TinyTransformer, source: Tensor, decoder_input: Tensor, target: Tensor) -> Tensor:
    logits = model(source, decoder_input)  # (B, T, vocabulary)
    return F.cross_entropy(
        logits.reshape(-1, logits.shape[-1]),
        target.reshape(-1),
        ignore_index=0,                   # Padding is not a prediction target.
    )                                    # Pass raw logits; do not softmax first.

PyTorch’s cross_entropy / CrossEntropyLoss ↗ expects raw logits for this target-ID usage. It performs the log-softmax calculation internally. Applying softmax first would feed probabilities into an interface expecting scores and change the objective.

Our flattening turns (B, T, M) into (B*T, M) and (B, T) into (B*T,). It combines batch and target-position axes while preserving the vocabulary axis. The ignored padding ID removes unused target positions from the objective.

4. What a training step changes

The forward pass computes logits and loss from the current parameters. loss.backward() applies the chain rule to compute gradients: how a small parameter change would change the loss locally. The optimizer then updates the parameters. This includes embeddings, attention projections, normalization gains and biases, FFNs, and the vocabulary head.

Gradients are not stored attention maps. They are derivatives associated with parameters and intermediate computations. Attention weights are forward-pass activations used while those derivatives are computed.

The training loop clears accumulated gradients, computes a new loss, backpropagates, clips the total gradient norm, and applies Adam. Clipping limits unusually large update directions; it does not solve an incorrect objective or leaking mask.

def train(steps: int = 800, seed: int = 7) -> tuple[TinyTransformer, TrainingResult]:
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    train_sequences, held_out = dataset(seed)
    model = TinyTransformer()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    generator = torch.Generator().manual_seed(seed + 1)
    initial_source, initial_input, initial_target = batch(train_sequences[:32])
    initial_loss = training_loss(model, initial_source, initial_input, initial_target).item()
    start = time.perf_counter()
    for _ in range(steps):
        indices = torch.randint(len(train_sequences), (32,), generator=generator)
        source, decoder_input, target = batch([train_sequences[i] for i in indices.tolist()])
        optimizer.zero_grad(set_to_none=True)
        loss = training_loss(model, source, decoder_input, target)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

    source, decoder_input, target = batch(held_out)
    model.eval()
    with torch.no_grad():
        held_out_loss = training_loss(model, source, decoder_input, target).item()
        generated = generate(model, source)
    generated = F.pad(generated, (0, target.shape[1] - generated.shape[1]))
    exact_match = (generated == target).all(dim=1).float().mean().item()
    result: TrainingResult = {
        "seed": seed, "steps": steps, "train_sequences": len(train_sequences),
        "held_out_sequences": len(held_out), "initial_train_loss": round(initial_loss, 6),
        "held_out_loss": round(held_out_loss, 6), "greedy_exact_match": round(exact_match, 6),
        "seconds": round(time.perf_counter() - start, 2),
        "numpy_torch_note": "NumPy demonstrates forward passes; PyTorch trains the model.",
        "examples": [
            {"source": source[i].tolist(), "expected": target[i].tolist(),
             "generated": generated[i].tolist()} for i in range(3)
        ],
    }
    return model, result

The fixed seed controls initialization, data splitting, and batch sampling in this run. Exact floating-point behavior and runtime can still vary across versions and hardware. The point is to make the experiment inspectable and repeatable, not to promise bit-for-bit equivalence on every device.

5. Generate without seeing the target

Teacher-forced evaluation supplies the correct previous symbols. That is useful for measuring loss, but it does not test what happens after the model makes its own mistake. We also run greedy autoregressive generation.

Start every target with BOS. Encode the source once. Decode the current prefix, choose the highest-scoring next token, append it, and repeat until EOS or the length limit. Finished batch items receive padding on subsequent steps.

@torch.no_grad()
def generate(model: TinyTransformer, source: Tensor, max_new_tokens: int = 4, bos: int = 1, eos: int = 2, pad: int = 0) -> Tensor:
    was_training = model.training
    model.eval()
    try:
        memory, visible = model.encode(source)
        prefix = torch.full((source.shape[0], 1), bos, dtype=torch.long, device=source.device)
        finished = torch.zeros(source.shape[0], dtype=torch.bool, device=source.device)
        for _ in range(max_new_tokens):
            logits = model.decode(prefix, memory, visible)[:, -1, :]
            token = logits.argmax(dim=-1)  # Greedy decoding for this experiment.
            token = torch.where(finished, pad, token)
            prefix = torch.cat([prefix, token[:, None]], dim=1)
            finished |= token == eos
            if finished.all():
                break
        return prefix[:, 1:]
    finally:
        model.train(was_training)

This simple decoder recomputes the target prefix at every step. It does not implement a KV cache. Source memory is still reused. That keeps inference easy to compare with the full teacher-forced forward pass; a cache is a later optimization whose correctness must preserve the same causal computation.

The max_new_tokens limit includes the generated EOS token. It is a stopping limit, not a guarantee that the model will produce a correct end marker.

6. A measured run

Run the experiment from the directory containing the downloaded files:

python verify.py
python train_copy.py --steps 800 --seed 7

A local CPU run with Python 3.12, NumPy 2.5.2, and PyTorch 2.14.0 produced the following results. The recorded JSON output contains the original measurements and sample predictions.

One held-out example was [11, 10, 11]. The generated result was [11, 10, 11, 2], where 2 is EOS. The model copied a repeated symbol and terminated correctly. The exact-match metric includes all target positions and correct termination; it is stricter than counting individual correct tokens.

This result demonstrates successful learning within a small synthetic domain. It does not demonstrate natural-language understanding, translation ability, or generalization to longer sequences. One evaluation sequence failed exact match in this run. Small average loss does not imply every generated sequence is correct.

7. Test mechanisms as well as metrics

The supplied verify.py checks numerical agreement between NumPy and PyTorch, normalized attention rows, exactly zero masked weights, and invariance to changing masked values. It checks that gradients are finite and that all-masked rows are rejected.

At the full-model level, it checks three especially useful properties:

• Future-token invariance: changing a later target input does not change earlier logits.
• Source-padding invariance: appending padding to the source does not change valid target predictions.
• Prefix equivalence: decoding each available prefix agrees with the corresponding position of the full causally masked pass.

These tests tell us whether the intended information flow is implemented. They complement the learning result rather than replacing it. A shape test alone would miss several errors these properties catch.

8. Experiments that build on the result

First, train for fewer steps and compare teacher-forced loss with greedy exact match. Notice that the two metrics answer different questions. Then inspect the failed generated sequences: did the model choose the wrong symbol, repeat one, or stop at the wrong time?

Next, try removing the causal mask only as a diagnostic experiment, then rerun the future-token invariance check. An invalid model may achieve a deceptively good training loss by seeing later target inputs. The mechanism check should catch the problem even before you trust its score.

Finally, change the task to reversing the source sequence. Update the target construction, keep the source split disjoint, and evaluate generation again. This changes the required source-to-target alignment while retaining all the components we have built.

Where the pieces now fit

Tokens become IDs; IDs select embeddings; positional signals label where those embeddings occur. Queries and keys determine weights over visible sources, and values supply the information being mixed. Heads compute parallel projected views. Residual updates preserve a path through the block, normalization operates within tokens, and FFNs transform contextual features.

An encoder contextualizes the source. A causal decoder builds the target, using cross-attention when it reads a separate source memory. A vocabulary head supplies logits, the loss supplies a training signal, and autoregressive decoding tests what the model can produce without the answer in front of it.

You can now read the architecture as a sequence of operations rather than a collection of labels. For the next layer of questions, explore the KV cache , linear attention , or DETR’s object-query decoder .