Coding Hubs School of AI – Best AI & Full Stack Courses in Delhi NCR | 100% Placement
Limited Offer: Get 50% OFF on AI & Full Stack Courses
📞 Call Now: +91 8448811540
Back to Deep Learning Notes
Topic #442

Feed-Forward Network (Transformer)

Every Transformer encoder and decoder layer contains a small position-wise feed-forward network — a simple 2-layer MLP, applied identically and independently to every token's representation, right after the attention sublayer.

Formula

\[ \text{FFN}(\mathbf{x}) = \phi(\mathbf{x}\mathbf{W}_1+\mathbf{b}_1)\mathbf{W}_2+\mathbf{b}_2 \]

\(\phi\) is a non-linear activation — the original paper used ReLU; many modern Transformer implementations use GELU (see GELU). \(\mathbf{W}_1\) typically expands the dimensionality significantly (e.g. from \(d_{\text{model}}=512\) to an inner dimension of \(2048\)), and \(\mathbf{W}_2\) projects it back down to \(d_{\text{model}}\).

"Position-Wise" — The Key Detail

This feed-forward network is applied independently to each position in the sequence — the exact same two weight matrices \(\mathbf{W}_1, \mathbf{W}_2\) are used for every single token, with no mixing of information across positions inside this sublayer (that cross-position mixing already happened in the preceding attention sublayer). It's functionally identical to running the same small MLP separately on each token's vector.

Why This Sublayer Exists at All

Self-attention is fundamentally a linear operation applied to the values (a weighted sum) — it can mix information across positions, but it doesn't add much non-linear transformation capacity on its own. The feed-forward network provides exactly that: a genuine non-linear transformation, applied per-position, giving the model additional representational capacity to process what attention has gathered — directly recalling the argument from Linear Transformations about why non-linearities are essential between linear operations.

Numerical Example — The Dimension Expansion

A common configuration: \(d_{\text{model}}=512\), inner dimension \(d_{ff}=2048\) — a 4x expansion. For a single token's 512-dim vector, \(\mathbf{W}_1\) projects it up to 2048 dimensions, ReLU/GELU zeroes out or reshapes negative values, then \(\mathbf{W}_2\) projects back down to 512 — the token's representation leaves this sublayer the same shape it entered, but has passed through a much higher-dimensional intermediate space where the non-linear transformation actually happens.

Code

import torch
import torch.nn as nn

class PositionWiseFFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.activation = nn.GELU()

    def forward(self, x):   # x: (batch, seq_len, d_model)
        return self.linear2(self.activation(self.linear1(x)))
        # applied identically to EVERY position -- no cross-position mixing here

ffn = PositionWiseFFN(d_model=512, d_ff=2048)
x = torch.randn(1, 10, 512)   # 10 tokens
output = ffn(x)
print(output.shape)   # (1, 10, 512) -- same shape; same FFN applied independently to each of the 10 tokens

Common Mistakes

  • Assuming the feed-forward network mixes information across different token positions — it explicitly does not; all cross-position information mixing in a Transformer layer happens exclusively in the attention sublayer.
  • Underestimating how much of a Transformer's total parameter count lives in these feed-forward layers — in many architectures, the FFN sublayers (due to their dimension expansion) account for a majority of the model's total parameters, more than the attention mechanism itself.

Interview Relevance

Q: "Why does a Transformer layer need a feed-forward sublayer if self-attention already processes the sequence?" Self-attention's core operation — a weighted sum of value vectors — is largely a linear combination; it mixes information across positions but doesn't add much genuine non-linear transformation capacity on its own. The position-wise feed-forward network adds that missing non-linear processing, applied independently to each token's already-attention-enriched representation, following the same "linear needs non-linearity between layers" principle from Linear Transformations.

Practice Question

If \(d_{\text{model}}=256\) and \(d_{ff}=1024\), what are the shapes of \(\mathbf{W}_1\) and \(\mathbf{W}_2\) in the feed-forward sublayer?

Related DL Notes

Want to go beyond the notes?

Join Coding Hubs School of AI's Deep Learning course — live mentorship, real projects, and 100% placement support.

Enroll Now — Free Demo Available
💬 Talk to Advisor
1
WhatsApp

Latest from Our Blog

Insights on AI, Data Science, Full Stack & Career

View All Articles →