Day 47: Feedforward Networks (FFN), SwiGLU, and Stacking Blocks

1. FFN(x) = act(x W1 + b1) W2 + b2, with x: (T, 64), W1: (64, 256), b1: (256), x W1 + b1: (T, 256), W2: (256, 64), b2: (64), output: (T, 64). The original Transformer uses ReLU; GPT-2 and nanoGPT use GELU.

2. Position-wise means the FFN is applied to each token vector x_t separately with the same weights at every position. Row t of the output depends only on row t of the input, so the FFN can be computed for all T rows as one batched matrix product. It adds no communication between tokens; attention provides that.

3. W1: 64 * 256 = 16,384. W2: 256 * 64 = 16,384. Total without biases: 32,768. This is 2 * d_model * d_ff, with d_ff = 4 * d_model.

4. b1 adds 256 and b2 adds 64, so the total is 32,768 + 320 = 33,088 parameters.

5. Three matrices of d_model x d_ff' cost 3 * 64 * d_ff'. Setting this equal to 2 * 64 * 256 gives d_ff' = (2/3) * 256 ≈ 170.7; rounding to 171 gives 3 * 64 * 171 = 32,832 parameters (bias-free, as in LLaMA-style SwiGLU), within 0.2% of 32,768. In practice d_ff' is often rounded up to a hardware-friendly multiple, which shifts the count slightly.

Common errors: counting only W1; writing the bias of W1 with size 64; assuming SwiGLU keeps d_ff = 256, which would give 49,152 parameters.

1. z = x W1 with x = [1, 2]:
z_1 = 1*1 + 2*0 = 1
z_2 = 1*0 + 2*2 = 4
z_3 = 1*(-1) + 2*1 = 1
z = [1, 4, 1].

2. h = ReLU([1, 4, 1]) = [1, 4, 1]. All entries are positive, so ReLU changes nothing here.

3. y = h W2:
y_1 = 1*1 + 4*0 + 1*(-1) = 0
y_2 = 1*0 + 4*1 + 1*1 = 5
y = [0, 5].

4. Residual: x + y = [1, 2] + [0, 5] = [1, 7]. The FFN output is an update added to the stream; it does not replace x.

5. Variant x = [2, -1]:
z = [2*1 + (-1)*0, 2*0 + (-1)*2, 2*(-1) + (-1)*1] = [2, -2, -3].
h = ReLU(z) = [2, 0, 0]: units 2 and 3 are switched off.
y = 2*[1, 0] + 0*[0, 1] + 0*[-1, 1] = [2, 0].
Only the rows of W2 for active units contribute, so the FFN acts as a lookup of the rows selected by the positive pre-activations.

Common errors: multiplying x by the rows of W1 instead of its columns; applying ReLU after W2; forgetting that 1*(-1) is negative before adding 2*1.

1. Sequence for two Pre-LN blocks:
x0 (embeddings)
block 1: a = x0 + Attn1(LN1a(x0)); x1 = a + FFN1(LN1b(a))
block 2: b = x1 + Attn2(LN2a(x1)); x2 = b + FFN2(LN2b(b))
In GPT-style models a final LayerNorm and the LM head follow x2.

2. Every point on the residual stream has shape (B, T, d_model) = (2, 4, 64): x0, LN1a(x0), Attn1 output, a, LN1b(a), FFN1 output, x1, and the same for block 2. The dtype is also unchanged (for example float32 in, float32 out).

3. The residual addition x + Sublayer(x) needs identical shapes, so each sublayer must map (B, T, d_model) back to (B, T, d_model). Because every block has the same interface, blocks can be stacked, removed or repeated without adapter projections, and one initialization scheme (for example nanoGPT's scaled init of the output projections by 1/sqrt(2L)) applies to every layer. Depth adds capacity by applying more successive updates to the same stream, not by changing the interface.

4. With 4 heads, head_dim = 64/4 = 16:
Q, K, V per head: (B, n_heads, T, head_dim) = (2, 4, 4, 16)
scores Q K^T / sqrt(16): (2, 4, 4, 4), one T x T matrix per head
attention output after merging heads: (2, 4, 64)
FFN hidden layer: (2, 4, 256), projected back to (2, 4, 64).
These are the only places where the shape leaves (2, 4, 64), and each returns to it before the residual addition.

Common errors: writing Block2(x) instead of Block2(Block1(x)); saying the scores are (2, 4, 64); putting LayerNorm after the residual addition while calling the block Pre-LN.

1. Attention: 4 * 768^2 = 4 * 589,824 = 2,359,296 parameters.

2. FFN: 2 * 768 * 3072 = 4,718,592 parameters.

3. Block weights: 2,359,296 + 4,718,592 = 7,077,888. The FFN holds 4,718,592 / 7,077,888 = 2/3 (66.7%) and attention 1/3 (33.3%), a 2:1 ratio. This follows from d_ff = 4 * d_model: 8 d^2 against 4 d^2.

4. With biases and LayerNorms:
attention: 2,359,296 + 4 * 768 = 2,362,368
FFN: 4,718,592 + 3072 + 768 = 4,722,432
two LayerNorms: 2 * (768 + 768) = 3,072
block total: 7,087,872.

5. Full model:
12 blocks: 12 * 7,087,872 = 85,054,464
token embeddings: 50257 * 768 = 38,597,376
position embeddings: 1024 * 768 = 786,432
final LayerNorm: 1,536
total: 124,439,808, the "124M" of GPT-2 small. The tied LM head reuses the token embedding matrix, so it adds nothing.

6. Matrix products cost about 2 FLOPs per weight per token: 2 * 7,077,888 = 14,155,776 FLOPs. The attention scores Q K^T and the weighted sum over V add about 2 * 2 * T * d_model = 4 * 1024 * 768 = 3,145,728 FLOPs per token at T = 1024 (about half that if the causal mask is exploited). So one block costs about 17.3 MFLOPs per token forward. The FFN share of FLOPs (9,437,184, about 54.5%) is below its 2/3 share of parameters once the score products are counted, and that gap grows with T.

Common errors: forgetting the output projection O; adding a separate LM head to a tied model; treating parameter count and FLOPs as the same ratio.