A Short Course for Modern Large Language ModelsBlock walkthrough

Lesson 23 of 23 · Capstone

Walk one block without losing a dimension.

Your win: narrate every operation in a pre-norm causal Transformer block with RoPE, including its shape and purpose.

20 minutesNeeds: Lessons 1–22Outcome: teach the complete algorithm
Final mission: start with \(X^{(\ell)}\), finish with \(X^{(\ell+1)}\), and explain exactly where position, comparison, probability, content, and residual updates enter.

Step through the neural network block

Input → operation → output

The complete pre-norm block

  1. 1Normalize: \(\bar X=\operatorname{Norm}(X^{(\ell)})\)\(B\times T\times d_{\text{model}}\)
  2. 2Project: \(Q=\bar XW_Q,\ K=\bar XW_K,\ V=\bar XW_V\)then split into \(h\) heads
  3. 3Encode position: \(Q'=\operatorname{RoPE}(Q),\ K'=\operatorname{RoPE}(K)\)\(V\) unchanged
  4. 4Compare and mask: \(A=Q'K'^{\mathsf T}/\sqrt{d_h}+M_{\text{causal}}\)\(B\times h\times T\times T\)
  5. 5Blend values: \(O_r=\operatorname{softmax}(A_r)V_r\)each head \(T\times d_h\)
  6. 6Join heads: \(O_{\text{attn}}=\operatorname{Concat}(O_1,\ldots,O_h)W_O\)\(B\times T\times d_{\text{model}}\)
  7. 7Add: \(Y=X^{(\ell)}+O_{\text{attn}}\)first residual
  8. 8Transform and add: \(X^{(\ell+1)}=Y+\operatorname{FFN}(\operatorname{Norm}(Y))\)second residual
Complete decoder-only Transformer block from hidden states through attention, residual paths, normalization, and SwiGLU
The whole block in one view. Follow the numbered path top to bottom, then use the worked example below to substitute actual matrices at the attention stages.

Toy walkthrough: \(B=1,T=2,d_{\text{model}}=2,h=1\)

Use RMSNorm with unit scale, identity projection matrices, positions \(m=0,1\), RoPE angle step \(\theta=\pi/2\), and

\[ X^{(\ell)}=\begin{bmatrix}1&1\\1&-1\end{bmatrix}. \]

Each row has RMS \(1\), so \(\bar X=X^{(\ell)}\). Identity projections give \(Q=K=V=\bar X\).

At position 0, \((1,1)\) stays \((1,1)\). At position 1,

\[ R_{\pi/2}\begin{bmatrix}1\\-1\end{bmatrix} =\begin{bmatrix}0&-1\\1&0\end{bmatrix} \begin{bmatrix}1\\-1\end{bmatrix} =\begin{bmatrix}1\\1\end{bmatrix}. \]

Thus \(Q'=K'=\begin{bmatrix}1&1\\1&1\end{bmatrix}\). The scaled score and masked score matrices are

\[ S=\frac{Q'K'^{\mathsf T}}{\sqrt2} =\begin{bmatrix}\sqrt2&\sqrt2\\\sqrt2&\sqrt2\end{bmatrix}, \qquad A=\begin{bmatrix}\sqrt2&-\infty\\\sqrt2&\sqrt2\end{bmatrix}. \]

Row-wise softmax gives \(P=\begin{bmatrix}1&0\\0.5&0.5\end{bmatrix}\). Therefore

\[ O=PV =\begin{bmatrix}1&0\\0.5&0.5\end{bmatrix} \begin{bmatrix}1&1\\1&-1\end{bmatrix} =\begin{bmatrix}1&1\\1&0\end{bmatrix}. \] \[ Y=X^{(\ell)}+O =\begin{bmatrix}2&2\\2&-1\end{bmatrix}. \]

After the second normalization, suppose the learned SwiGLU produces the toy update \(F=\begin{bmatrix}0.1&-0.2\\0.3&0.1\end{bmatrix}\). The block output is

\[ X^{(\ell+1)}=Y+F =\begin{bmatrix}2.1&1.8\\2.3&-0.9\end{bmatrix}. \]

The chosen \(F\) stands in for the SwiGLU calculation from Lesson 22; every attention-stage number above is calculated explicitly.

Teach-back checkpoints

RoPE's exact job

Rotate Q and K pairs by position-dependent phases so \(q_m'^{\mathsf T}k_n'=q_m^{\mathsf T}R_{n-m}k_n\). It does not rotate V.

Attention's exact job

Use masked query-key comparisons to build probability weights, then blend value vectors from allowed positions.

Code checkpoint · assemble the TransformerModel

The block pieces from Lessons 15–22 now form a decoder-only language model. Every block preserves \([B,T,D]\); the language-model head converts the final features into one logit per vocabulary item.

Show the final assembly
class TransformerModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            TransformerBlock(config) for _ in range(config.num_layers)
        )
        self.final_norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight

    def forward(self, input_ids):
        x = self.token_embedding(input_ids)  # [B, T] -> [B, T, D]
        for layer in self.layers:
            x = layer(x)                    # [B, T, D] -> [B, T, D]
        return self.lm_head(self.final_norm(x))  # [B, T, vocab]

Open the annotated complete implementation → or download transformer_model.py.

Trace it: Starting from input_ids, name the final axis at each of the three comments.

Retrieval check

Where does RoPE occur in this block?

Capstone practice

  1. Without looking above, write the eight stages of the block in order.
  2. For \(B=2,T=5,d_{\text{model}}=12,h=3\), make a shape ledger for split Q, one head's score matrix, concatenated output, and final block output.
  3. Repeat the toy multiplication \(PV\) and residual addition \(X^{(\ell)}+O\).
  4. Explain why \(R_m^{\mathsf T}R_n=R_{n-m}\) matters specifically inside Step 4.
  5. Name the two places where information is mixed across features and the place where it is mixed across tokens.
  6. Teach the block aloud in under two minutes without saying “it just learns it.” Name the input and output of every stage.
Check solutions
  1. Normalize; project Q/K/V; apply RoPE to Q/K; score and causally mask; softmax and blend V; concatenate heads and apply \(W_O\); first residual; normalize, SwiGLU, and second residual.
  2. Split Q: \(2\times3\times5\times4\); one head score per batch: \(5\times5\); concatenated output: \(2\times5\times12\); final output: \(2\times5\times12\).
  3. \(O=\begin{bmatrix}1&1\\1&0\end{bmatrix}\); \(Y=\begin{bmatrix}2&2\\2&-1\end{bmatrix}\).
  4. It turns the rotated query-key dot product into a comparison that depends on relative displacement \(n-m\).
  5. Projection/output matrices and the FFN mix features; attention's \(PV\) mixes information across allowed token positions.
  6. Use the reference sheet to check omissions after speaking, not during the first attempt.

Primary sources: Vaswani et al.; Su et al., RoFormer; Touvron et al., LLaMA.

Record your two-minute explanation and send its transcript plus the capstone answers to the teaching agent.