Make the Frozen Decoder Trainable in PyTorch

Map every audited decoder operation to a registered PyTorch module while preserving pre-norm order, causal masking, true weight tying, tensor shapes, and the 368-parameter ledger.

Chapter 6 represented each operation directly. PyTorch packages the same operations into modules and records the computation needed for gradients. The architecture remains the specification; the framework is an implementation.

Map Every Module to Known Mathematics

PyTorch object	Chapter 6 operation	Learned entries
`nn.Embedding(8, 4)`	token row lookup $E[I]$	32
position `nn.Parameter(4, 4)`	learned $P[:T]$	16
`nn.MultiheadAttention(4, 2, bias=False)`	$W_Q,W_K,W_V,W_O$	64 per block
two `nn.LayerNorm(4)`	two pre-norm operations	16 per block
`Linear(4,8)`, ReLU, `Linear(8,4)`	biased position-wise MLP	76 per block
final `nn.LayerNorm(4)`	$N_f$	8
tied `nn.Linear(4,8,bias=False)`	$E^\top$ readout	0 additional

Two blocks contribute $2(64+16+76)=312$ . The total is

32+16+312+8=368.

Register One Shared Readout Parameter

Creating a second linear layer and copying embedding values into it does not tie the weights. Assign the same parameter object:

Then test identity and the complete ledger:

The shared table participates in two parts of the forward pass. During backpropagation, its gradient accumulates contributions from input lookup and vocabulary readout before the optimizer updates the one stored parameter.

Preserve Pre-Norm Order and the Causal Mask

One block retains the audited update order:

The three repeated normalized arguments mean queries, keys, and values come from the same residual stream. The mask must contain $-\infty$ strictly above the diagonal so those entries receive zero probability after softmax.

Verify Interfaces before Training

For a batch inputs with shape (B,T), assert:

Repeat Chapter 6's causal intervention: change a later input token and confirm that earlier logits remain equal. A framework model that returns the right shape can still contain future leakage.

Q1. Find a false weight tie

An implementation copies the embedding weights into the readout once during initialization:

readout.weight.data.copy_(embedding.weight.data)

It then reports 400 learned scalar parameters. What is wrong?

Choose the parameter error

The values were copied into an independent readout parameter instead of sharing the embedding parameter object.

The copied matrix has the wrong shape for vocabulary logits.

Copying values makes both parameters permanently receive identical gradients.

Select one choice, then check.

Hint

The tied model has 368 entries, not 400.

Solution

The readout owns a separate

8\times4

parameter. Assign the same parameter object instead of copying its current values.

Framework Reference

The complete experiment freezes one tested PyTorch version for regression values, but these links should be checked when running under a newer release.