Causal Decoder Blocks Update Every Prefix Position

Two distinct pre-norm blocks preserve the model-width residual stream while applying causal attention at every depth. Track axes, masks, head and MLP widths, and the 156-parameter block count.

The position-aware records enter two decoder blocks. Each block uses the pre-norm equations established earlier:

U^{(\ell)}=X^{(\ell)}+ \operatorname{MHA}^{(\ell)}(N_1^{(\ell)}(X^{(\ell)})),

X^{(\ell+1)}=U^{(\ell)}+ \operatorname{MLP}^{(\ell)}(N_2^{(\ell)}(U^{(\ell)})).

For this model, every stream tensor has shape $(B,T,4)$ . The two attention heads each use width 2, and the MLP expands each token record from width 4 to 8 and contracts it to 4.

Causality Applies in Every Block

For $T=4$ , each head applies the inclusive lower-triangular visibility matrix

\begin{bmatrix} 1&0&0&0\\ 1&1&0&0\\ 1&1&1&0\\ 1&1&1&1 \end{bmatrix}.

Applying the mask only in the first block is insufficient. A later unmasked block can create a future-to-past route even if earlier representations were causal.

Depth Changes Values, Not the Interface

Block 1 and Block 2 have the same input-output shape but different parameters. Write $X^{(0)}$ , $X^{(1)}$ , and $X^{(2)}$ rather than calling all three tensors “the embedding.” The token and position tables create $X^{(0)}$ ; subsequent records are contextual residual-stream states.

Count One Block from Its Declared Conventions

Attention has four $4\times4$ matrices and no biases: 64 entries. The biased MLP has $4\times8+8+8\times4+4=76$ . Two LayerNorms each have gain and bias of length 4: 16. One block therefore has

64+76+16=156

learned entries, and two blocks have 312.

Q1. Track a decoder stack shape

For $B=3$ and $T=4$ , what is the residual-stream shape after the second block?

Your answer

Answer it first, then check.

Hint

Use

(B,T,d_{model})

Solution

The shape is

(3,4,4)