W07 H3: Build and Inspect a Tiny Transformer

Machine Learning in Engineering Science | NCKU + LAiMM | Fall 2026

Teaching date: Monday 19 October 2026. Guided notebook time: 40 minutes within a 60-minute module.

What you will be able to explain

- Trace embeddings, positional information, masked attention, and a feed-forward block.

- Train a tiny decoder on synthetic sequences with shifted targets.

- Test causality and permutation behavior, and document the limits of the experiment.

Start here, even if this is your first notebook

1. In Colab choose Runtime > Change runtime type > CPU. In Jupyter choose a Python 3 kernel.

2. Run cells from top to bottom with Shift+Enter. Read the Markdown before each code cell.

3. The setup installs only missing free packages. A network connection is needed for first installation; the data itself is generated locally.

4. If a variable is undefined, restart the kernel and run all earlier cells. Do not skip a setup cell.

5. Save a personal copy before changing parameters. Every core cell is complete: there are no hidden files, tokens, paid services, or unfinished code blocks.

Data: synthetic teaching data, not engineering measurements. Record your actual output; a model is not required to outperform a baseline.

Practice policy: in-class exercise and optional Homework extension, not a new weekly graded assignment.

Mechanism reference: our complete decoder block

Token IDs index learned vectors. Add a learned position vector, split projected queries/keys/values into heads, compute causal attention, and concatenate/project heads back to model width. Our pre-normalization block is

$$u=x+\operatorname{Attention}(\operatorname{LayerNorm}(x)),$$

$$y=u+\operatorname{FFN}(\operatorname{LayerNorm}(u)).$$

The feed-forward network expands width 32 to 128, applies GELU, and returns to 32. It uses shared parameters independently at each position. Attention mixes positions; the MLP transforms local features. Residual additions provide an identity route and require matching widths.

This decoder-only teaching variant differs from the original 2017 encoder-decoder with post-normalization and ReLU. Cross-entropy receives raw logits and shifted targets. Training processes known positions in parallel under a mask; generation appends unknown next tokens sequentially.

1. Set up a CPU experiment

Run independently of H1 and H2. This toy decoder is an educational implementation, not a pretrained LLM.

import sys, subprocess, importlib.util, platform, json, random, hashlib
from pathlib import Path
required = {"numpy": "numpy", "matplotlib": "matplotlib", "sklearn": "scikit-learn"}
missing = [package for module, package in required.items() if importlib.util.find_spec(module) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", *missing])
import numpy as np
import matplotlib.pyplot as plt
SEED = 2026
random.seed(SEED)
np.random.seed(SEED)
rng = np.random.default_rng(SEED)
ARTIFACTS = Path("artifacts")
ARTIFACTS.mkdir(exist_ok=True)
print("Python:", platform.python_version(), "NumPy:", np.__version__, "Seed:", SEED)

if importlib.util.find_spec("torch") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "torch"])
import torch
from torch import nn
import copy, time
torch.set_num_threads(1)
torch.manual_seed(SEED)
torch.use_deterministic_algorithms(True)
device = torch.device("cpu")
print("PyTorch:", torch.__version__, "Device:", device)

Checkpoint. CPU is sufficient. Generation is not evidence of factual knowledge.

2. Define a reproducible symbolic sequence task

We use a small repeating process with eight symbolic states. Training, validation, and test are separate generated streams with different seeds and small random perturbations. The target at each position is the next token; inputs never include that target position. No natural-language corpus is downloaded.

VOCAB, CONTEXT = 8, 16
def make_stream(seed,n=1800):
    generator=np.random.default_rng(seed)
    motif=np.array([0,1,2,3,4,5,6,7,6,5,4,3,2,1],dtype=np.int64)
    stream=np.resize(motif,n).copy()
    corrupt=generator.random(n)<0.04
    stream[corrupt]=generator.integers(0,VOCAB,corrupt.sum())
    return torch.tensor(stream,dtype=torch.long)
train_stream=make_stream(SEED)
val_stream=make_stream(SEED+1,600)
test_stream=make_stream(SEED+2,600)
def windows(stream):
    # Nonoverlapping blocks here keep sample accounting easy to inspect.
    starts=range(0,len(stream)-CONTEXT,CONTEXT)
    return (torch.stack([stream[i:i+CONTEXT] for i in starts]),
            torch.stack([stream[i+1:i+CONTEXT+1] for i in starts]))
train_x,train_y=windows(train_stream)
val_x,val_y=windows(val_stream)
test_x,test_y=windows(test_stream)
print("Example input:",train_x[0].tolist())
print("Shifted target:",train_y[0].tolist())
assert torch.equal(train_x[0,1:],train_y[0,:-1])

Checkpoint. The final target token lies just beyond the input window. Do not report performance as generalization to a different physical process: all streams share the same motif.

3. Write attention with explicit axes and masks

Representations have shape (B,T,D). Split D into H heads of width d=D/H. Attention scores have shape (B,H,T,T), with query positions on rows and key positions on columns. Store weights for inspection, but detach them so diagnostic storage does not retain training graphs.

class CausalAttention(nn.Module):
    def __init__(self,d_model=32,heads=4):
        super().__init__()
        assert d_model%heads==0
        self.heads,self.d_head=heads,d_model//heads
        self.qkv=nn.Linear(d_model,3*d_model)
        self.proj=nn.Linear(d_model,d_model)
        self.last_weights=None
    def forward(self,x):
        B,T,D=x.shape
        q,k,v=self.qkv(x).chunk(3,dim=-1)
        def split(z): return z.reshape(B,T,self.heads,self.d_head).transpose(1,2)
        q,k,v=map(split,(q,k,v))
        scores=q@k.transpose(-2,-1)/(self.d_head**0.5)
        allowed=torch.ones(T,T,dtype=torch.bool,device=x.device).tril()
        weights=scores.masked_fill(~allowed,float("-inf")).softmax(dim=-1)
        self.last_weights=weights.detach()
        mixed=(weights@v).transpose(1,2).contiguous().reshape(B,T,D)
        return self.proj(mixed)

Checkpoint. The softmax dimension is the last dimension: source keys. Upper-triangular weights must be exactly zero.

4. Complete the block and output head

This is a pre-normalization decoder-only block, a common educational variant. The 2017 paper used a full encoder-decoder with post-normalization. Residual paths preserve an identity route; the feed-forward network transforms each position separately with shared weights.

class Block(nn.Module):
    def __init__(self,d=32,heads=4):
        super().__init__()
        self.ln1,self.ln2=nn.LayerNorm(d),nn.LayerNorm(d)
        self.attn=CausalAttention(d,heads)
        self.ff=nn.Sequential(nn.Linear(d,4*d),nn.GELU(),nn.Linear(4*d,d))
    def forward(self,x):
        x=x+self.attn(self.ln1(x))
        return x+self.ff(self.ln2(x))

class TinyTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.tokens=nn.Embedding(VOCAB,32)
        self.positions=nn.Embedding(CONTEXT,32)
        self.block=Block()
        self.norm=nn.LayerNorm(32)
        self.head=nn.Linear(32,VOCAB)
    def forward(self,ids,use_positions=True):
        B,T=ids.shape
        assert T<=CONTEXT
        x=self.tokens(ids)
        if use_positions: x=x+self.positions(torch.arange(T,device=ids.device))[None,:,:]
        return self.head(self.norm(self.block(x)))

torch.manual_seed(SEED)
model=TinyTransformer().to(device)
print("Parameters:",sum(p.numel() for p in model.parameters()))
assert model(train_x[:2]).shape==(2,CONTEXT,VOCAB)

Checkpoint. The model returns logits (B,T,V), not token IDs or probabilities. The vocabulary size is eight, not the embedding width 32.

5. Establish a training-only bigram baseline

A bigram baseline estimates P(next|current) from training transitions. Add-one smoothing avoids zero probabilities. Our motif contains ambiguous next states, so longer context may help. Compare NLL in nats/token on identical test targets.

counts=np.ones((VOCAB,VOCAB),dtype=np.float64)
for a,b in zip(train_stream[:-1],train_stream[1:]): counts[int(a),int(b)]+=1
bigram=counts/counts.sum(axis=1,keepdims=True)
def bigram_nll(x,y): return float(-np.log(bigram[x.numpy(),y.numpy()]).mean())
print("Uniform baseline NLL:",np.log(VOCAB))
print("Bigram validation NLL:",bigram_nll(val_x,val_y))

Checkpoint. A uniform eight-token model has NLL log(8), about 2.079. This is a mathematical reference, not a measured model score.

6. Train with shifted-target cross-entropy

All positions train in parallel, but the causal mask prevents future access. Flatten (B,T,V) into (B*T,V) for cross-entropy. Select the best state using validation NLL and inspect test only after selection.

criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.AdamW(model.parameters(),lr=0.003,weight_decay=0.01)
best,best_state=float("inf"),None
curve=[]
start=time.perf_counter()
for epoch in range(45):
    model.train()
    order=torch.randperm(len(train_x))
    losses=[]
    for batch in order.split(16):
        optimizer.zero_grad()
        logits=model(train_x[batch])
        loss=criterion(logits.reshape(-1,VOCAB),train_y[batch].reshape(-1))
        loss.backward(); nn.utils.clip_grad_norm_(model.parameters(),1.0); optimizer.step()
        losses.append(loss.item())
    model.eval()
    with torch.no_grad():
        val_loss=criterion(model(val_x).reshape(-1,VOCAB),val_y.reshape(-1)).item()
    curve.append((float(np.mean(losses)),val_loss))
    if val_loss<best: best,best_state=val_loss,copy.deepcopy(model.state_dict())
model.load_state_dict(best_state)
print("Training seconds:",round(time.perf_counter()-start,2))
plt.plot(curve); plt.legend(["train NLL","validation NLL"])
plt.xlabel("Epoch"); plt.ylabel("Nats/token"); plt.show()

Checkpoint. The validation curve determines the retained checkpoint. If you change architecture, do not tune against test NLL.

7. Evaluate and inspect a head

NLL measures probability assigned to true next tokens. Accuracy measures the fraction of argmax choices that are correct. Inspect a real trained attention matrix and state what its axes mean.

model.eval()
with torch.no_grad():
    logits=model(test_x)
    test_nll=criterion(logits.reshape(-1,VOCAB),test_y.reshape(-1)).item()
    accuracy=(logits.argmax(-1)==test_y).float().mean().item()
print({"bigram_test_nll":bigram_nll(test_x,test_y),"transformer_test_nll":test_nll,
       "transformer_test_accuracy":accuracy})
head=model.block.attn.last_weights[0,0].numpy()
assert np.all(head[np.triu_indices(CONTEXT,1)]==0)
plt.figure(figsize=(6,5)); plt.imshow(head,cmap="YlGnBu",vmin=0,vmax=1)
plt.xlabel("Key/source position"); plt.ylabel("Query/destination position")
plt.title("Actual trained weights: test example 0, head 0"); plt.colorbar(); plt.show()

Checkpoint. The heatmap is an actual run output. The zero upper triangle is a mask property, not a relationship the model discovered.

8. Perform a future-token intervention

Change only tokens at positions >=8. Under causal attention, logits at positions 0 through 7 must remain unchanged because the changed tokens lie in their future. This checks architecture behavior even if model accuracy is poor.

original=test_x[:1].clone()
changed=original.clone(); changed[:,8:]=(changed[:,8:]+3)%VOCAB
with torch.no_grad():
    old_logits=model(original)
    new_logits=model(changed)
prefix_error=(old_logits[:,:8]-new_logits[:,:8]).abs().max().item()
print("Maximum earlier-logit change:",prefix_error)
assert torch.allclose(old_logits[:,:8],new_logits[:,:8],atol=1e-6)

Checkpoint. A passing test supports the no-future-access property for this implementation and example. It does not validate the dataset or guarantee every possible deployment behavior.

9. Show why unmasked attention needs position information

For unmasked self-attention without positions, permuting token rows permutes output rows in the same way. We test the isolated mathematical operation. A causal mask itself introduces an order constraint, so do not apply this permutation claim unchanged to a masked decoder.

z=torch.tensor([[1.,0.],[0.,1.],[1.,1.]])
p=torch.tensor([2,0,1])
def unmasked_attention(z): return (z@z.T/(z.shape[-1]**0.5)).softmax(-1)@z
assert torch.allclose(unmasked_attention(z[p]),unmasked_attention(z)[p],atol=1e-6)
print("Unmasked attention is permutation-equivariant without positions.")
with torch.no_grad():
    no_position_nll=criterion(model(test_x,use_positions=False).reshape(-1,VOCAB),test_y.reshape(-1)).item()
print("Remove learned positions at inference NLL:",no_position_nll)

Checkpoint. Removing positions from a trained model is an inference intervention, not a fair retraining comparison. It may worsen, improve, or barely change this toy task; report the observed result.

10. Generate sequentially from a prefix

Training processes many known target positions together. Autoregressive generation still produces one new token at a time and feeds it back. Cropping retains only CONTEXT tokens; learned absolute positions restart in each cropped window in this teaching implementation.

generated=test_x[:1,:4].clone()
for _ in range(20):
    with torch.no_grad():
        next_logits=model(generated[:,-CONTEXT:])[:,-1,:]
        next_id=next_logits.argmax(-1,keepdim=True)
    generated=torch.cat([generated,next_id],dim=1)
print("Generated symbolic sequence:",generated[0].tolist())
record={"seed":SEED,"vocab":VOCAB,"context":CONTEXT,"torch":torch.__version__,
        "architecture":"one-block pre-norm decoder; 4 heads; width 32; learned positions",
        "training_data_sha256":hashlib.sha256(train_stream.numpy().tobytes()).hexdigest(),
        "validation_nll":best,"test_nll":test_nll,"test_accuracy":accuracy,
        "bigram_test_nll":bigram_nll(test_x,test_y),"causal_prefix_max_error":prefix_error,
        "position_removal_nll":no_position_nll,"generated_tokens":generated[0].tolist()}
(ARTIFACTS/"w07-h3-transformer.json").write_text(json.dumps(record,indent=2))
torch.save(model.state_dict(),ARTIFACTS/"w07-h3-transformer.pt")

Checkpoint. These symbols demonstrate sequence modeling only. A token generator does not automatically have tools, memory management, permissions, or verification required by an agent.

11. Independent experiment and exit ticket

Choose one intervention: change heads from 4 to 2 in Block; change CONTEXT from 16 to 8; or change the corruption rate from 0.04 to 0.10. Make a prediction, rerun from data/model creation, and compare validation NLL before observing test results.

Explain back:

1. Where are the learned parameters in the block? Which weights are computed dynamically for each input?

2. Why does parallel training not imply parallel generation of unknown next tokens?

3. What does O(T^2) refer to in dense attention? At T=16, scores per head are 256; doubling T produces 1024.

4. Why is an attention heatmap not a causal explanation? Name one controlled intervention you performed.

5. Why is this small decoder not an LLM or a complete agent?

Stop when: you have baseline and Transformer NLL, a trained head plot, a passing causal test, and an evidence JSON.

Sources

- Vaswani et al., Attention Is All You Need

- Dive into Deep Learning: Transformer architecture

- Karpathy, ng-video-lecture