A Decoder Block from Input to Logits

In Causal Masks and Attention Heads, you learned how causal masking prevents future-token leakage and how Multi-Query and Grouped-Query Attention reduce KV cache memory while preserving quality. You now have all the attention machinery in place. This lesson completes Part III by assembling the full decoder block: residuals, normalization, feed-forward networks, stacking multiple blocks, and the final projection that turns hidden states into next-token logits.

---

From Attention to a Complete Decoder Block

We have covered the attention sublayer in depth. A transformer decoder block contains two sublayers per layer:

1. Causal self-attention — mixes information across token positions
2. Feed-forward network (FFN) — processes each token position independently

Both sublayers are wrapped with residual connections and normalization. The modern standard (used by Llama, Mistral, Gemma, and many others) is Pre-LN: normalization happens before each sublayer, not after.

The Anatomy of a Modern Decoder Block (Pre-LN)

Original Post-LN (Vaswani et al., 2017):

x = x + Sublayer(x)      # Residual
x = LayerNorm(x)         # Normalize AFTER addition

Modern Pre-LN (Llama, Mistral, etc.):

x_norm = RMSNorm(x)      # Normalize BEFORE sublayer
sub_out = Sublayer(x_norm)
x = x + sub_out          # Residual AFTER sublayer

Pre-LN stabilizes gradient flow, enabling training of 80+ layer models without divergence. The residual connection becomes an uninterrupted highway for gradients to flow backward.

For layer l, the complete pipeline is:

x_norm1  = RMSNorm_1(x)
attn_out = CausalAttention(x_norm1)
x_mid    = x + attn_out           # Residual connection 1

x_norm2  = RMSNorm_2(x_mid)
ffn_out  = SwiGLU(x_norm2)
x_out    = x_mid + ffn_out        # Residual connection 2

Key observations:

• Two RMSNorm operations per layer (one before attention, one before FFN)
• Two residual additions per layer
• The input x is never normalized in-place — normalization produces a new tensor
• x_mid carries the accumulated representation forward

---

Residual Connections: The Feature Highway

What Is a Residual Connection?

A residual (or skip) connection adds the input of a sublayer directly to its output:

x_out = x_in + f(x_in)

Physical Intuition

Think of x_in as the accumulated belief state or context representation at this layer. The sublayer f(x) computes only a refinement (a delta) to add to that state. The model does not replace the representation — it adjusts it.

Why Residuals Preserve Gradients

During backpropagation, the gradient of the output with respect to the input is:

∂x_out / ∂x_in = I + ∂f / ∂x_in

The identity matrix I means gradients can flow straight through the residual path, unchanged. Even if ∂f / ∂x_in becomes small (vanishing), the I term ensures the gradient never disappears completely. This is why very deep transformers can train stably with Pre-LN.

---

Feed-Forward Networks: From ReLU to SwiGLU

Classical FFN (GPT-1/2 Style)

FFN(x) = ReLU(x @ W_1) @ W_2

• W_1 ∈ ℝ^(d_model × d_ff) — up-projection
• W_2 ∈ ℝ^(d_ff × d_model) — down-projection
• Typically d_ff = 4 × d_model

Swish / SiLU Activation

Swish (also called SiLU) is a smooth, non-monotonic activation:

Swish(z) = z · σ(z) = z / (1 + e^(-z))

Properties:

• Smooth everywhere (no sharp corner like ReLU at 0)
• Allows tiny negative gradients (not exactly zero for negative inputs)
• Self-gating: the value is modulated by its own sigmoid

SwiGLU (Shazeer, 2020): Gated Linear Unit with Swish

SwiGLU uses three weight matrices instead of two:

• Gate projection: W_gate ∈ ℝ^(d_model × d_ff)
• Up projection: W_up ∈ ℝ^(d_model × d_ff)
• Down projection: W_down ∈ ℝ^(d_ff × d_model)

Formula:

SwiGLU(x) = ( Swish(x @ W_gate) ⊙ (x @ W_up) ) @ W_down

Where ⊙ is element-wise (Hadamard) multiplication.

Why three matrices? The gate and up projections are independent. The gate acts as a learned, data-dependent mask on the up projection. This gives the FFN more expressive capacity per parameter.

Dimension Choice in Llama

In Llama models, d_ff ≈ (8/3) × d_model (e.g., 11,008 for 7B where d_model = 4096).

This keeps the parameter count equal to a classical 4×d_model FFN:

• Classical: 2 × d_model × (4 × d_model) = 8 × d_model²
• SwiGLU: 3 × d_model × ((8/3) × d_model) = 8 × d_model²

---

Hands-On Exercise 1: SwiGLU Hand Calculation on Scalar Features

Problem: Trace SwiGLU with scalar values.

• Input x = 2.0
• Weights: W_gate = 0.5, W_up = 1.5, W_down = 2.0

Step-by-step Solution:

Step 1: Gate projection

x @ W_gate = 2.0 × 0.5 = 1.0

Step 2: Swish activation on gate

Swish(1.0) = 1.0 × σ(1.0)
           = 1.0 / (1 + e^(-1.0))
           = 1.0 / (1 + 0.367879...)
           = 1.0 / 1.367879...
           ≈ 0.731058...

Step 3: Up projection

x @ W_up = 2.0 × 1.5 = 3.0

Step 4: Element-wise multiplication (gating)

gate_value × up_value = 0.731058... × 3.0 ≈ 2.193176...

Step 5: Down projection

result = 2.193176... × 2.0 ≈ 4.38635...

Final Answer: ≈ 4.386

Verification: You can check this by computing Swish(1.0) * 3.0 * 2.0 in any calculator.

---

From Block Output to Next-Token Logits

After L decoder blocks, the final hidden states x_final ∈ ℝ^[B, T, d_model] pass through two final operations:

1. Final RMSNorm

x_norm = RMSNorm_final(x_final)

This normalizes the accumulated representation before the output projection. It ensures the logits are not dominated by scale differences across positions or layers.

2. Unembedding Projection

logits = x_norm @ W_unembed

• W_unembed ∈ ℝ^(d_model × V)
• Output shape: [B, T, V] — logits over vocabulary for every position

During next-token prediction, we only need the last position’s logits: logits[:, -1, :] ∈ ℝ^[B, V].

---

Weight Tying: Sharing Input and Output Embeddings

The Idea (Press & Wolf, 2017)

The input embedding matrix W_embed ∈ ℝ^(V × d_model) maps token IDs to vectors. The unembedding matrix W_unembed ∈ ℝ^(d_model × V) maps vectors back to logits.

Weight tying reuses the transpose of the input embedding matrix as the unembedding matrix:

W_unembed = W_embed^T

Parameter Savings

• Untied: V × d_model (embed) + d_model × V (unembed) = 2 × V × d_model
• Tied: V × d_model (shared) = V × d_model

Savings: V × d_model parameters.

For Llama 3 8B: V = 128,256, d_model = 4,096

• Untied unembedding: 128,256 × 4,096 ≈ 525M parameters
• Tied: saves all 525M (≈6.5% of total model params)

Who Uses What?

Model Family	Weight Tying?
GPT-2	Yes
Gemma	Yes
Llama 2/3	No (untied)
Mistral	No (untied)
Qwen	No (untied)

Llama and Mistral use untied embeddings, likely because the input and output roles differ enough that separate optimization helps.

---

Hands-On Exercise 2: Residual Stream Addition and Preservation

Problem: Trace residual additions through one decoder block with 2D vectors.

• x_in = [2.0, -1.0]
• Attention output a = [0.4, 0.6]
• FFN output f = [-0.4, 1.4]

Step-by-step Solution:

Step 1: First residual (after attention)

x_mid = x_in + a
      = [2.0 + 0.4, -1.0 + 0.6]
      = [2.4, -0.4]

Step 2: Second residual (after FFN)

x_out = x_mid + f
      = [2.4 + (-0.4), -0.4 + 1.4]
      = [2.0, 1.0]

Step 3: Total change from input to output

Δ = x_out - x_in
  = [2.0 - 2.0, 1.0 - (-1.0)]
  = [0.0, 2.0]

Step 4: Verify decomposition

a + f = [0.4 + (-0.4), 0.6 + 1.4]
      = [0.0, 2.0]
      = Δ ✓

Interpretation: The residual stream cleanly accumulates each sublayer’s contribution. The first dimension was unchanged by the attention + FFN combined (net zero delta); the second dimension increased by 2.0. This additive property is why we can interpret individual layer contributions.

---

Hands-On Exercise 3: Full Architecture Annotation (Part III Exit Task)

Problem: Annotate the complete forward pass of a tiny 2-layer decoder model.

Model Specs:

• Batch size B = 1
• Sequence length T = 3
• Model dimension d_model = 4
• FFN hidden dimension d_ff = 8
• Vocabulary size V = 6
• Layers L = 2

Task: List the exact tensor shape at every stage.

Solution:

Stage	Operation	Tensor Shape	Notes
1	Token IDs (input)	[1, 3]	Batch, Sequence
2	Embedding lookup (W_embed ∈ [6, 4])	[1, 3, 4]	B, T, d_model
3	Layer 1 Pre-RMSNorm (attention)	[1, 3, 4]	Same shape
4	Layer 1 Attention output	[1, 3, 4]	Causal self-attention
5	Layer 1 Residual 1	[1, 3, 4]	x + attn_out
6	Layer 1 Pre-RMSNorm (FFN)	[1, 3, 4]
7	Layer 1 SwiGLU output	[1, 3, 4]	d_ff=8 internal, projected back
8	Layer 1 Residual 2	[1, 3, 4]	x_mid + ffn_out
9	Layer 2 Residual 2 (output)	[1, 3, 4]	Same shapes repeated for Layer 2
10	Final RMSNorm	[1, 3, 4]
11	Unembedding projection (W_u ∈ [4, 6])	[1, 3, 6]	Logits: B, T, V

Key Observations:

• Shapes B, T, d_model are invariant through all decoder blocks
• Only the final projection changes the last dimension to V
• During generation, we index logits[:, -1, :] → [1, 6] for next-token prediction

---

Summary Reference Table

Component	Formula / Operation	Input Shape	Output Shape	Design Reason
Input Embedding	IDs @ W_embed	[B, T]	[B, T, d]	Map tokens to vectors
Pre-RMSNorm (Attn)	RMSNorm(x)	[B, T, d]	[B, T, d]	Stabilize attention input
Causal Attention	softmax((QK^T)/√d) @ V	[B, T, d]	[B, T, d]	Mix token positions
Residual 1	x + attn_out	[B, T, d]	[B, T, d]	Gradient highway
Pre-RMSNorm (FFN)	RMSNorm(x_mid)	[B, T, d]	[B, T, d]	Stabilize FFN input
SwiGLU	(Swish(xW_g) ⊙ xW_u) @ W_d	[B, T, d]	[B, T, d]	Gated FFN, expressive
Residual 2	x_mid + ffn_out	[B, T, d]	[B, T, d]	Accumulate refinement
Final RMSNorm	RMSNorm(x_final)	[B, T, d]	[B, T, d]	Normalize before logits
Unembedding	x_norm @ W_unembed	[B, T, d]	[B, T, V]	Vocabulary logits
Weight Tying	W_unembed = W_embed^T	—	—	Save V×d params

---

Primary Source Citations

• Shazeer (2020), GLU Variants Improve Transformer — arXiv:2002.05202
• Press & Wolf (2017), Using the Output Embedding to Improve Language Models — arXiv:1608.05859
• Touvron et al. (2023), Llama 2: Open Foundation and Fine-Tuned Chat Models — arXiv:2307.09288
• Vaswani et al. (2017), Attention Is All You Need — arXiv:1706.03762 (Post-LN baseline)

---

What Comes Next

• Next lesson: Prefill: Process the Prompt and Build the Cache — Part IV begins. Trace prompt positions through causal attention, show which calculations run together, and build the initial KV cache.
• Previous lesson: Causal Masks and Attention Heads
• Track overview: Start Here: LLM Foundations