Feed Forward Network — Module 10.1

How This Lesson Fits the Module & Engineering Practice

Attention mixes information across positions. The position-wise feed-forward network (FFN) transforms each token independently with a small MLP—usually expand → activate → project back. Together with multi-head attention, residuals, and LayerNorm, the FFN completes a transformer block.

This lecture closes Module 10.1. Module 10.2 starts with the history of transformers and then assembles the full architecture from these parts.

Learning Objectives

By the end of this lesson, students should be able to:

Define the position-wise FFN and contrast it with attention.
Explain the expansion ratio (often \(4\times\)) and why it is used.
Compare ReLU and GELU activations in FFN practice.
Implement an FFN with residual + LayerNorm in PyTorch.
Place the FFN inside an encoder/decoder block recipe.
Recap how Module 10.1 pieces form a transformer layer.

Definition

The transformer feed-forward network is a two-layer MLP applied identically and independently to every position: \(\mathrm{FFN}(x) = W_2\, \sigma(W_1 x + b_1) + b_2\), where \(\sigma\) is typically ReLU or GELU and the hidden width is larger than \(d_model\) (commonly \(4\, d_model\)).

Position-Wise: What That Means

“Position-wise” means the same weights \(W_1, W_2\) are shared across tokens, but there is no mixing between positions inside the FFN. If attention gathered context, the FFN is a per-token nonlinear feature transform—like a shared MLP over the sequence axis.

Sublayer	Mixes positions?	Primary role
Self-attention / MHA	Yes	Route and aggregate context
Cross-attention	Across sequences	Retrieve from another stream
Position-wise FFN	No	Nonlinear channel transform

Expansion Ratio and Activations

Width / expansion

Hidden size \(d_{ff} \approx 4\, d_model\).
Expand → nonlinear → compress.
Most transformer parameters often live here.

Activation

Original paper: ReLU.
Many modern LMs: GELU (or variants).
Same structure; \(\sigma\) changes curvature.

Activation	Notes in FFN context
ReLU	Simple, sparse; used in “Attention Is All You Need”
GELU	Smooth; common in BERT/GPT-style stacks

Full Block Recipe (Module Capstone View)

1. Attend

MHA (+ mask / cross as needed).

2. Residual + LN

Stabilize the attention path.

3. FFN

Position-wise expand–activate–project.

4. Residual + LN

Stabilize the FFN path.

Also recall positional encoding at the stack input, and the encoder/decoder roles from earlier lectures.

PyTorch: Position-Wise FFN

import torch
from torch import nn

class PositionwiseFFN(nn.Module):
    def __init__(self, d_model: int, d_ff: int | None = None, dropout: float = 0.1):
        super().__init__()
        d_ff = d_ff or 4 * d_model  # expansion ratio 4x
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),               # or nn.ReLU()
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        # x: (B, T, d_model) — same MLP at every position
        return self.net(x)

class TransformerEncoderBlock(nn.Module):
    """Pre-norm block: attention + FFN with residuals."""
    def __init__(self, d_model: int, num_heads: int, d_ff: int | None = None):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.mha = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = PositionwiseFFN(d_model, d_ff)

    def forward(self, x):
        h = self.ln1(x)
        attn, _ = self.mha(h, h, h)
        x = x + attn
        x = x + self.ffn(self.ln2(x))
        return x

x = torch.randn(2, 16, 64)
y = TransformerEncoderBlock(64, 8)(x)
print(y.shape)  # torch.Size([2, 16, 64])

Strengths

Cheap parallel per-token MLP.
Adds capacity beyond attention mixing.
Simple to implement and scale.

Tradeoffs

Often dominates parameter count.
No cross-token mixing by itself.
Width choices affect memory heavily.

Common Mistake

Thinking the FFN is a single shared vector across time (like a global dense layer on the flattened sequence). It is applied per position with shared weights—shapes stay \((B, T, d_model)\), not \((B, T\cdot d)\).

Misconception

“Attention does all the work; the FFN is optional decoration.” Empirically FFNs hold a large share of model capacity. Attention routes information; the FFN processes it nonlinearly at each site.

Module 10.1 Map (Sibling Cross-Links)

Build the mental stack: Query / Key / Value → Attention & Scaled Dot-Product → Self / Cross → Multi-Head → Positional Encoding → Residual + LayerNorm + FFN. Encoder/decoder pages: Encoder, Decoder.

Knowledge Check

Short Answer: What does “position-wise” mean for the FFN? Answer: The same MLP is applied independently to each token position with shared weights.
True/False: The FFN mixes information between different token positions. Answer: False.
Multiple Choice: A common expansion ratio \(d_{ff}/d_model\) is: (a) 4, (b) 1/4, (c) 100. Answer: (a).
Short Answer: Name two activations used in transformer FFNs. Answer: ReLU and GELU.
True/False: Attention aggregates across positions; the FFN transforms channels per position. Answer: True.
Multiple Choice: After FFN in a standard block you typically apply: (a) residual (+ norm), (b) stemming, (c) max-pool only. Answer: (a).
Short Answer: Why can FFNs dominate parameter count? Answer: The large intermediate width (\(d_{ff}\)) creates two wide linear layers shared across positions.
True/False: Original “Attention Is All You Need” used ReLU in the FFN. Answer: True.
Multiple Choice: Next module lecture after this page: (a) History of Transformers, (b) Lemmatization, (c) VGG16. Answer: (a).
Short Answer: List the four ingredients of a typical transformer block covered in 10.1. Answer: Multi-head attention, residual connections, layer normalization, and the position-wise FFN (plus PE at the input).

Key Takeaways

FFN = shared expand–activate–project MLP at each position.
Typical width \(d_{ff} = 4\, d_model\); activations ReLU or GELU.
Complements attention: mixing vs per-token nonlinearity.
Wrapped with residual + LayerNorm like the attention sublayer.
Module 10.1 complete → 10.2 History of Transformers.

Trainer’s Guide

Hands-on idea: Count parameters in MHA vs FFN for \(d_model=512\), \(h=8\), \(d_{ff}=2048\) and discuss where capacity lives.

Discussion prompt: If you had to shrink a model, would you cut heads, \(d_model\), or \(d_{ff}\) first—and why?

Recap: The position-wise FFN finishes the transformer block with a wide per-token MLP. Module 10.1 ends here—continue to History of Transformers.