Position-Wise Feed-Forward Networks

The position-wise feed-forward network is a sublayer that maps each token's $d$-dimensional vector through a two-layer multilayer perceptron (MLP), with a non-linear activation between the layers, applied identically at every sequence position. Formally, for a token representation $x \in \mathbb{R}^d$,

where $W_1 \in \mathbb{R}^{d_{ff}\times d}$, $W_2 \in \mathbb{R}^{d \times d_{ff}}$, and $\sigma$ is a non-linear activation — originally ReLU, later GELU or Swish/SwiGLU variants in modern models. The inner dimension $d_{ff}$ is usually $4d$ (e.g. $d=512, d_{ff}=2048$ in the original Transformer; $d=768, d_{ff}=3072$ in BERT-base). The crucial qualifier is "position-wise": the same weight matrices $W_1, W_2$ are reused at every position $1, \dots, n$ in the sequence, but there is no interaction between positions inside this sublayer. If the input to the block is a matrix $X \in \mathbb{R}^{n \times d}$ (one row per token), the FFN is applied row-by-row:

This is mathematically equivalent to two back-to-back $1{\times}1$ convolutions across the sequence dimension — a framing used in the original "Attention Is All You Need" paper, which is exactly why it is called position-wise rather than sequence-wise: it behaves like a kernel-size-1 convolution, sliding independently over each position with shared weights. Contrast this with the Multi-Head Attention sublayer that precedes it in each block: attention computes a weighted combination of other tokens' values, an operation that is linear in the value vectors and whose only source of adaptivity is the attention weights themselves. Stacking two linear operations (attention's value-mixing, followed by nothing) would collapse into an expressively limited linear map. The FFN's non-linearity $\sigma$ is what allows the network to compute genuinely new features — logical AND/OR-like combinations, thresholding, sparse activation patterns — that a purely linear attention mechanism cannot represent. This is the same reason a plain linear regression stack collapses to a single linear layer without non-linear activations between them; the FFN is where the Transformer earns its non-linear expressive power.

Each Transformer block (encoder or decoder) applies attention, then the FFN, each wrapped in a residual connection and a normalization step. The full block, following the standard post-norm formulation, is: ``` z = LayerNorm(x + MultiHeadAttention(x)) y = LayerNorm(z + FFN(z)) ``` (Modern implementations often use pre-norm, applying LayerNorm before the sublayer instead — a design detail orthogonal to the FFN's internal computation.) The FFN sublayer itself, expanded into pseudocode operating on a batch of sequences: ``` function FeedForward(X, W1, b1, W2, b2):

# X: tensor of shape (batch, seq_len, d_model)

# W1: (d_model, d_ff), b1: (d_ff,)

# W2: (d_ff, d_model), b2: (d_model,) H = X @ W1 + b1 # shape (batch, seq_len, d_ff) H = activation(H) # elementwise, e.g. ReLU(x) = max(0, x) Y = H @ W2 + b2 # shape (batch, seq_len, d_model) return Y # same shape as X function TransformerBlock(X, params): A = MultiHeadAttention(X, X, X, params.attn) # self-attention Z = LayerNorm(X + A, params.ln1) F = FeedForward(Z, params.W1, params.b1, params.W2, params.b2) Y = LayerNorm(Z + F, params.ln2) return Y ``` Walking through the shapes makes the "position-wise" property concrete. The input `X` has shape `(batch, seq_len, d_model)`. The matrix multiply `X @ W1` broadcasts `W1` across the `batch` and `seq_len` axes — every one of the `batch seq_len` token vectors is multiplied by the same `W1`, independently. There is no operation here that reads from position $i$ and position $j \neq i$ simultaneously; that cross-position mixing already happened in the attention sublayer, and it will happen again in the next block's attention sublayer, but not inside this FFN. Each token's forward pass through the FFN is: 1. <strong>Expand</strong>: project from $d_{model}$ up to $d_{ff}$ (typically $4\times$ wider), giving the network a larger "workspace" to compute intermediate features. 2. <strong>Activate</strong>: apply a pointwise non-linearity (ReLU, GELU, or a gated variant like SwiGLU) to each of the $d_{ff}$ hidden units. 3. <strong>Contract</strong>: project back down from $d_{ff}$ to $d_{model}$, so the output can be added back into the residual stream. Because step 1 and step 3 are ordinary matrix multiplications, and step 2 is elementwise, the entire computation for all $n$ positions can be done as a single* batched matrix multiply of shape $(n, d) \times (d, d_{ff})$, rather than $n$ separate small operations — this is what makes the FFN extremely GPU-friendly.

Time complexity. For a single token vector, the FFN performs a $d \to d_{ff}$ matrix multiply ($O(d \cdot d_{ff})$ multiply-adds) followed by a $d_{ff} \to d$ matrix multiply (another $O(d \cdot d_{ff})$). For a sequence of length $n$, since the same computation repeats at every position independently, the total cost is

With $d_{ff} = 4d$, this is $O(4nd^2) = O(nd^2)$ — linear in sequence length, in sharp contrast to the self-attention sublayer's $O(n^2 d)$ cost, which grows quadratically because attention must compute pairwise interactions between all $n^2$ token pairs. This asymmetry matters a great deal in practice: for short-to-moderate sequences the FFN can dominate total FLOPs (since $d \gg 1$ and $n$ is small), but for very long sequences attention's $n^2$ term eventually dominates. Understanding which sublayer is the bottleneck is essential when deciding whether to optimize attention (e.g. with sparse or linear-attention variants) or the FFN (e.g. with mixture-of-experts routing).

Space complexity. The parameters are $W_1 \in \mathbb{R}^{d \times d_{ff}}$ and $W_2 \in \mathbb{R}^{d_{ff}\times d}$, giving $2 \cdot d \cdot d_{ff}$ weights (plus $d_{ff} + d$ biases, negligible in comparison). With $d_{ff} = 4d$, this is $8d^2$ parameters per FFN sublayer, versus roughly $4d^2$ for the multi-head attention sublayer's four projection matrices ($W_Q, W_K, W_V, W_O$, each $d \times d$). This is precisely why the FFN sublayers account for about two-thirds of a Transformer's total parameter count, a fact often surprising to newcomers who assume attention "is" the Transformer. The activation memory during the forward pass requires storing the intermediate hidden tensor of shape $(n, d_{ff})$, which is $4\times$ larger than the token representations themselves — a real consideration for memory-constrained training, motivating techniques like activation checkpointing.

Correctness / why this design works. The argument for the FFN's necessity is an expressivity argument, not a convergence proof, and it rests on two facts. Fact 1 (attention alone is a linear map in its values). Self-attention output for token $i$ is $\sum_j \alpha_{ij} (X_j W_V)$, a convex combination (in $\alpha$) of linearly-projected value vectors. For fixed attention weights $\alpha$, this is an affine function of the token embeddings. Composing this with the residual add and LayerNorm (which is a nonlinear but position-wise renormalization, not a feature-mixing nonlinearity) does not introduce the kind of nonlinear feature interaction that a hidden-layer MLP provides. Fact 2 (a composition of purely linear sublayers cannot exceed linear expressivity). If every sublayer in a stack were linear, the entire depth-$L$ network would collapse to a single linear transformation (a composition of linear maps is linear), regardless of $L$ — exactly the classical argument for why deep linear networks gain no expressive power over depth. The universal approximation theorem for two-layer MLPs (Cybenko, 1989; Hornik, 1991) guarantees that a single hidden layer with a suitable non-linearity and sufficient width $d_{ff}$ can approximate any continuous function on a compact domain arbitrarily well. Placing such an MLP after every attention sublayer therefore ensures each block can, in principle, compute an arbitrary non-linear transformation of the (already contextualized) token vector, rather than being limited to affine recombination. Together these facts justify the architectural necessity of the FFN: it is not an arbitrary add-on but the component responsible for the network's non-linear representational capacity, applied at the finest possible granularity — one token at a time — so that its cost and expressive power scale with $n$ rather than $n^2$.

Take a toy setting with $d_{model} = 4$, $d_{ff} = 8$, and a 3-token sequence output from the attention + Add&Norm step, $Z \in \mathbb{R}^{3\times 4}$:

Suppose $W_1$ and $b_1$ (for illustration, using a simplified $4\times 8$ weight and a bias of $-0.1$ on every hidden unit) produce the pre-activation matrix $H_{pre} = ZW_1 + b_1$. We do not need the exact numeric entries of $W_1$ to see the mechanism — only that each row of $Z$ is multiplied against the same $W_1$, independently:

The key observation, verified by the shapes above: computing $y_2$ never reads $z_1$ or $z_3$. If we permuted the rows of $Z$ (swapped token 1 and token 3), the FFN output rows would simply permute the same way — $y_1$ and $y_3$ would swap, with no numerical change to their values. This is the formal signature of "position-wise": the FFN commutes with row permutations of its input, i.e. it is a permutation-equivariant, in fact a completely position-decoupled, operation. Contrast this with what would happen if we permuted the rows of $Z$ before the attention sublayer instead: the attention weights $\alpha_{ij}$ depend on pairwise dot products across positions, so permuting rows there changes which tokens attend to which, and the output values genuinely differ in a way that depends on the new arrangement of context, not merely a relabeling. Numerically, using ReLU as $\sigma$, note that token 2's pre-activation had a strongly negative first entry ($-1.4$), which gets zeroed — that hidden unit contributes nothing to $y_2$, while it may have contributed heavily to $y_1$ (whose corresponding entry was $0.9$, i.e. active). This is precisely the non-linear, content-dependent gating that a linear attention-only stack cannot replicate: different tokens activate different subsets of the $d_{ff}=8$ hidden units, letting the network compute qualitatively different functions per token content, not just different weighted averages of the same value vectors.

The FFN sublayer's mathematical ancestry runs through decades of neural network research on multilayer perceptrons and universal approximation, before being crystallized as a specific architectural component in the Transformer.

Every major Transformer-based system — Multi-Head Attention-based encoders like BERT, decoder-only language models like the GPT family, encoder-decoder systems like T5, and Vision Transformers applied to images — includes a position-wise FFN in each block, making it one of the most reused computational primitives in modern deep learning. Because the FFN dominates parameter count, it is also the primary target of parameter-efficiency and scaling research: Mixture-of-Experts (MoE) architectures such as Switch Transformer and Mixtral replace the single dense FFN with a bank of expert FFNs and a routing mechanism, letting models scale total parameters far beyond what a dense FFN could afford while keeping the compute per token roughly fixed. Techniques like LoRA and adapter modules frequently target the FFN's weight matrices for efficient fine-tuning, since they hold the bulk of the model's task-specific capacity. Understanding the FFN is also a prerequisite for understanding why the Transformer Encoder Stack scales the way it does in both parameters and FLOPs, and it directly informs engineering decisions in production inference systems, where FFN matrix multiplies are the dominant target for kernel fusion, quantization, and hardware-specific optimization (e.g. batching FFN computation across all tokens and all sequences in a batch into one large matrix multiply for maximal GPU utilization).

Ask the tutor: How does replacing the dense FFN with a Mixture-of-Experts layer change the complexity and training dynamics of a Transformer block?