Feed-Forward Networks in Transformers: Architecture

Self-attention lets tokens gather contextual information from across the sequence. But attention alone is limited: it computes weighted averages of value vectors, a fundamentally linear operation. After softmax normalization, each output position is a convex combination of value vectors. No matter how many attention heads you stack, no matter how carefully you tune the query and key projections, the mapping from input to output remains linear in the values. To learn complex functions of language, transformers need something more.

The feed-forward network (FFN) provides exactly this missing ingredient. Sitting in every transformer block alongside the self-attention sublayer, it applies a nonlinear transformation to each token's representation independently. Think of it as the transformer's "thinking" layer: after attention has routed information between positions and assembled contextually-aware representations, the FFN processes each of those representations through a small neural network that can learn curved, high-dimensional decision boundaries. This is where the actual computational work of language understanding happens.

Every transformer block contains two main components: a self-attention sublayer and a feed-forward sublayer. While attention handles inter-token communication, the FFN handles per-token computation. The division of labor is elegant and deliberate. Attention routes information: it reads from the full sequence and assembles each position's context. The FFN transforms information: it takes each position's assembled context and applies a learned nonlinear function to produce a richer representation. Together, they enable transformers to learn the rich, hierarchical, compositional representations that power modern language AI.

Understanding the FFN matters beyond academic interest. In most transformer architectures, the FFN contains significantly more parameters than the attention mechanism, often two-thirds or more of each layer's total parameters. It also dominates computational cost for short sequences. This combination, large parameter count and heavy compute, makes FFN design and optimization one of the most important engineering problems in modern LLM deployment. Techniques like mixture of experts, activation sparsity, and quantization all target the FFN specifically.

This chapter examines the FFN in detail. You'll learn its two-layer architecture and understand why the hidden dimension is expanded. You'll see how position independence enables massive parallelization and how the FFN acts as an associative memory that stores factual knowledge in its weights. You'll work through a numerical example step-by-step, trace through a complete implementation, and calculate the substantial parameter count that makes FFNs the largest component of most transformer models. By the end, you'll understand what the FFN computes, why it's designed the way it is, and what would be lost without it.

To understand why transformers need the feed-forward network, consider what attention alone provides. Self-attention computes weighted averages of value vectors, where the weights come from query-key similarities. This is powerful for gathering contextual information: a token can attend to any position in the sequence and assemble a representation that reflects its surrounding context. But no matter how sophisticated the attention patterns, the operation remains a convex combination of value vectors. Each output is literally a weighted sum of inputs. This is a linear function.

Linear functions have severe limitations for language understanding. They can only rotate and scale the input space, then translate it. They cannot learn the curved decision boundaries that distinguish "bank" (financial institution) from "bank" (riverbank). They cannot capture the complex feature interactions that tell you whether "not bad" is negative or positive in context. They cannot represent the conditional relationships that make "the cat sat on the mat" grammatically fine but "the cat sat on the cat" slightly odd. For a model to approximate arbitrary functions of language, it needs the ability to carve up representation space in nonlinear ways.

The key insight is that any sufficiently wide network with a nonlinear activation function can approximate arbitrary continuous functions. This is the universal approximation theorem, a foundational result in neural network theory. The FFN uses this result by introducing a nonlinear activation between two linear layers. The first layer projects the input into a higher-dimensional space. The activation function applies a nonlinearity. The second layer projects back. The result is a function that can approximate far more complex relationships than any linear operation could.

There is a critical design choice embedded in how the FFN is applied: instead of mixing information across positions like attention does, the FFN processes each position independently. If your sequence has 100 tokens, the FFN applies the same transformation to each of those 100 representations separately, using identical weights for all positions. This is called a position-wise or position-independent transformation, and it fundamentally shapes the FFN's role in the transformer.

This division of labor is elegant. Attention handles inter-position communication: it routes information between tokens, allowing the model to build representations that depend on context. The FFN handles intra-position computation: it transforms each position's representation using a learned nonlinear function. By separating these concerns, the transformer achieves both contextual awareness (from attention) and expressive power (from the FFN).

Why share the same network across all positions? Two reasons justify this choice. First, language exhibits translation invariance in a deep sense: the grammatical patterns that help understand "the cat" are equally useful whether those words appear at the beginning, middle, or end of a sentence. A transformation that converts a noun phrase's representation into something more useful should work wherever that noun phrase appears. Second, parameter sharing dramatically reduces model size. Instead of learning separate networks for each of the potentially thousands of positions in a sequence, we learn one network that generalizes across all positions. This is a major efficiency win.

Think of the position-wise FFN as a shared "dictionary of transformations." After attention has looked up what each position means in context, the FFN looks up how to transform that meaning into a richer representation, using the same dictionary regardless of where the position appears in the sequence.

With this motivation in place, let's examine what the FFN computes. The architecture is surprisingly simple: two linear transformations with a nonlinear activation function sandwiched between them. But this simplicity is deceptive. The expand-transform-contract structure enables remarkably expressive transformations, and the specific choice of expansion factor shapes the model's capacity in ways that affect both learning and efficiency.

Consider a single token's representation: a vector $x$ with $d_{\text{model}}$ dimensions. This vector encodes everything the model currently knows about that token in its context, having passed through however many transformer layers preceded this one. The FFN transforms this vector in three stages:

1. Expand: Project $x$ into a higher-dimensional space
2. Transform: Apply nonlinearity to enable learning of curved boundaries
3. Contract: Project back to the original dimension

The complete formula captures all three stages in one expression. Given an input vector $x \in \mathbb{R}^{d_{\text{model}}}$, the FFN computes:

where:

• $x \in \mathbb{R}^{d_{\text{model}}}$: the input vector for a single position (the token's representation from the previous layer)
• $W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{ff}}$: the first projection matrix, which expands the dimension from $d_{\text{model}}$ to $d_{ff}$
• $b_1 \in \mathbb{R}^{d_{ff}}$: the first bias vector, added after the first linear transformation
• $\sigma$: a nonlinear activation function (e.g., ReLU, GELU) applied element-wise to the hidden representation
• $W_2 \in \mathbb{R}^{d_{ff} \times d_{\text{model}}}$: the second projection matrix, which contracts the dimension back from $d_{ff}$ to $d_{\text{model}}$
• $b_2 \in \mathbb{R}^{d_{\text{model}}}$: the second bias vector, added to produce the final output
• $d_{ff}$: the hidden (or intermediate) dimension, typically set to $4 \times d_{\text{model}}$

To understand this formula, let's read it from the inside out, following the order of operations.

Stage 1: The Expansion ($xW_1 + b_1$)

The input vector $x$ is multiplied by a weight matrix $W_1$ and shifted by a bias vector $b_1$. This is a standard linear transformation, but with a twist: the output dimension is larger than the input dimension. If $x$ has 512 dimensions, the result might have 2048 dimensions. We're projecting into a higher-dimensional space where the data is easier to manipulate. In this expanded space, the model has more "room" to organize features before applying nonlinearity.

Why does projecting to higher dimensions help? Think of it this way: in the original $d_{\text{model}}$-dimensional space, the token's representation may conflate many different features in a compressed form. By expanding to $d_{ff}$ dimensions, the first layer can disentangle these features, spreading them across a larger space where they are more separable. The nonlinearity can then act on these individual features independently, creating the curved boundaries that linear operations cannot.

Stage 2: The Nonlinearity ($\sigma(\cdot)$)

The activation function $\sigma$ is applied element-wise to the expanded representation. Common choices include ReLU (which zeros out negative values) and GELU (a smoother alternative that we'll examine in detail in the next chapter). This is where the magic happens: the nonlinearity allows the network to learn curved decision boundaries that would be impossible with linear transformations alone.

The element-wise nature of the activation matters. Because it operates on each hidden dimension independently, it creates a kind of gating: some hidden dimensions activate strongly and contribute to the output, while others are suppressed to near zero. Different input patterns activate different subsets of hidden dimensions, which is exactly how the network implements different transformations for different types of inputs.

Stage 3: The Contraction ($(\cdot)W_2 + b_2$)

Finally, the transformed representation is projected back to the original dimension via weight matrix $W_2$ and bias $b_2$. The output has the same dimensionality as the input, which is essential for the residual connection that adds the FFN output back to its input. The second linear layer combines the activated hidden dimensions into a compressed $d_{\text{model}}$-dimensional output, essentially summarizing the computation done in the expanded space.

Why does this formula make sense? Notice that without the nonlinearity $\sigma$, the composition of two linear maps is itself a linear map: $x W_1 W_2 + \text{const}$ is just another linear function of $x$. The nonlinearity is the only thing that prevents the two layers from collapsing into a single linear layer. The biases also allow the network to shift the activation function's operating range, making it easier to learn the right threshold for each hidden dimension.

The expansion ratio $d_{ff} / d_{\text{model}}$ is a key hyperparameter. The original transformer used a 4x expansion: for $d_{\text{model}} = 512$, the hidden dimension was $d_{ff} = 2048$. This ratio has become a de facto standard, though modern architectures sometimes use different values, especially when combined with gated variants, as we'll explore later.

With the formula understood, translating it to code reveals its simplicity. The entire FFN is just two matrix multiplications with a nonlinearity in between:

The FFN preserves the dimensionality of its input: a 512-dimensional vector goes in, and a 512-dimensional vector comes out. This is essential for the residual connection that adds the FFN output back to its input. The residual connection (which adds the original input $x$ to the FFN output) requires both tensors to have the same shape, so the FFN's dimension-preserving property is architecturally necessary rather than merely convenient.

Notice that we initialize the weights using Xavier/Glorot initialization, scaling by $\sqrt{2 / (d_{in} + d_{out})}$. This initialization keeps the variance of activations stable across layers, preventing the exploding or vanishing gradients that plagued deep networks before careful initialization became standard practice. The biases are initialized to zero, which is standard; the weight matrices carry all the initial structure.

The most striking aspect of the FFN architecture is the dimension expansion. The first linear layer projects from $d_{\text{model}}$ to $d_{ff}$, typically with $d_{ff} = 4 \times d_{\text{model}}$. For a model with $d_{\text{model}} = 512$, this means expanding to 2048 dimensions before projecting back down. For GPT-3 with $d_{\text{model}} = 12288$, the hidden dimension reaches 49152, nearly fifty thousand neurons per FFN layer.

This expansion might seem wasteful. Why create a 2048-dimensional intermediate representation just to immediately compress it back to 512 dimensions? The answer lies in the expressiveness of the network, and it connects to some deep results in neural network theory.

A fundamental result in approximation theory shows that wider hidden layers can approximate more complex functions. Consider what happens without expansion: if $d_{ff} = d_{\text{model}}$, the hidden layer has the same dimensionality as the input. While the formula still applies, the network has limited capacity to decompose and recombine features. Each hidden neuron is a linear combination of the input features, gated by nonlinearity. With only $d_{\text{model}}$ neurons, the network can detect $d_{\text{model}}$ different patterns. With $4 \times d_{\text{model}}$ neurons, it can detect four times as many patterns and combine them in far richer ways.

Think of the expansion as temporarily working in a higher-dimensional space where the data is easier to manipulate. In the expanded 2048-dimensional space, the activation function can create very specific, highly discriminating patterns. One hidden neuron might activate only when the input represents a verb in past tense. Another might activate for proper nouns that appear after articles. These fine-grained detectors would be impossible to implement in 512 dimensions because there simply isn't enough room to disentangle all the features. The expansion provides the working space to learn them, and the contraction summarizes the results back into the original dimensionality for the residual connection.

The 4x expansion factor was established in the original "Attention Is All You Need" paper and has become a standard choice. It represents a pragmatic balance between expressiveness and efficiency: large enough to provide substantial nonlinear capacity, small enough to remain computationally tractable. Modern models like LLaMA use slightly different ratios (around 2.7x) when using gated linear units (GLUs), which effectively increase the hidden dimension through gating, but the underlying logic of expanding before contracting remains the same.

A important property of the FFN is that it processes each position independently. Unlike attention, where every position can influence every other position through the query-key-value mechanism, the FFN applies an identical transformation to each position in isolation. This has practical consequences for both computation and interpretation, and it reflects a deep design decision about the division of responsibilities in the transformer architecture.

When we say the FFN is position-independent, we mean that the output at position $i$ depends only on the input at position $i$. The FFN does not know or care about adjacent tokens. It does not look at the whole sequence. It simply applies the same transformation $\text{FFN}(x) = \sigma(xW_1 + b_1)W_2 + b_2$ to whatever vector it receives, and that vector could have come from any position in the sequence.

This might seem like a limitation. After all, language is deeply contextual. The meaning of "bank" depends on whether "river" or "money" appeared nearby. But here is the key insight: by the time the FFN sees the representation at position $i$, that representation has already been processed by the self-attention sublayer in the same block. Attention has already incorporated contextual information. The FFN is operating on a context-enriched representation, not on a raw token embedding. Its job is to transform that already-contextualized representation further, not to gather more context.

Think of the transformer block as a two-phase processor: attention gathers context (inter-position), and the FFN processes the gathered context (intra-position). The FFN can focus entirely on the "what to do with this information" question because attention has already handled the "what context is relevant" question.

Let's verify this independence empirically:

One of the most illuminating ways to understand the FFN is to see it as a neural network layer and as an associative memory: a structure that stores key-value pairs in its weights and retrieves the appropriate values when presented with matching inputs. This interpretation, developed by researchers at Tel Aviv University and other institutions, helps explain how transformers store and retrieve factual knowledge, and why simply scaling up the FFN increases a model's capacity to "know" things.

The idea is elegant. Consider the first layer's weight matrix $W_1$. Each column of $W_1$ can be thought of as a "key" that matches certain input patterns. When we compute $xW_1 + b_1$, we are essentially computing the dot product between the input $x$ and each of the $d_{ff}$ keys. Each element of the resulting vector measures how well the input matches the corresponding key.

To make this precise: the pre-activation value for hidden dimension $i$ measures the alignment between the current input and the $i$-th key pattern. Given an input $x$ and the $i$-th column of $W_1$ (call it $\mathbf{k}_i$), this alignment is:

where:

• $h_i$: the pre-activation value for hidden dimension $i$, measuring how well the input matches key $i$
• $x$: the input vector ($d_{\text{model}}$-dimensional), representing the current token's contextual state
• $\mathbf{k}_i$: the $i$-th column of $W_1$, encoding the pattern that key $i$ is tuned to detect
• $b_{1,i}$: the $i$-th element of the bias vector $b_1$, which acts as a threshold, biasing the key toward activation or suppression

When the input $x$ aligns well with key $\mathbf{k}_i$ (high dot product), the corresponding hidden dimension activates strongly. After ReLU, only positive activations survive, so keys that match the current input "fire" while keys that don't match are silenced. This is exactly how an associative memory works: present a query, retrieve only the memories that match.

The second layer's weight matrix $W_2$ contains the "values" associated with each key. Each row of $W_2$ (call it $\mathbf{v}_i$ for the $i$-th row) is the value vector that gets added to the output when hidden dimension $i$ is active. The final FFN output is a weighted sum of these values, where the weights are the hidden activations:

where:

• $h_i^+ = \text{ReLU}(h_i) = \max(0, h_i)$: the post-activation strength of the $i$-th key match
• $\mathbf{v}_i$: the $i$-th row of $W_2$, encoding the value retrieved when key $i$ fires
• The sum runs over all $d_{ff}$ key-value pairs, but in practice only the active (nonzero) $h_i^+$ values contribute

This memory interpretation explains several important empirical observations about transformers. Researchers have found that specific neurons in FFN layers activate for particular concepts: individual neurons that fire consistently for "the Eiffel Tower," for "programming languages," or for "past tense verbs." These neurons are not hand-designed. They emerge from training, as the model discovers that certain hidden dimensions reliably encode certain patterns. The FFN learns to store factual associations in its weights, and the attention mechanism retrieves them by routing relevant input patterns through the FFN.

This interpretation also explains a peculiar phenomenon called "knowledge editing" in language models. When researchers want to update a model's stored facts (for instance, changing "The president of France is X" to a newer value), they can sometimes accomplish this by modifying specific rows of FFN weight matrices. The change propagates naturally because those rows encode the associated value that gets retrieved when the "president of France" key fires.

The memory interpretation also clarifies why larger FFN dimensions store more knowledge. More hidden dimensions means more key-value pairs, which means more distinct patterns the model can recognize and more distinct responses it can retrieve. This is why scaling the FFN dimension is one of the primary ways to increase a model's factual capacity.

Let's visualize this interpretation with a concrete example:

The heatmap reveals the key-value memory character directly. Each column represents one hidden dimension (one key-value pair). The white cells are zero, meaning that key did not fire for that input. The blue cells indicate firing keys, with darker blue indicating stronger activation. Notice that no two input samples have exactly the same activation pattern: different inputs retrieve different combinations of stored associations. This input-dependent sparse access is what makes the FFN such a powerful memory structure.

This sparsity also has practical implications for efficiency. In a model with large $d_{ff}$, only a fraction of hidden dimensions are nonzero for any given input. Methods that exploit this sparsity can skip computation for inactive dimensions, reducing the compute needed well below the theoretical maximum. Recent work on ReLU-based sparse FFNs has shown that sparsity rates of 90% or higher can be achieved with minimal quality loss, meaning only 10% of hidden dimensions are needed for most tokens.

Feed-forward networks are the largest component of transformer models by parameter count. In a standard transformer, the FFN contains significantly more parameters than the attention mechanism. This matters for model scaling and efficiency optimization, and it helps explain why so many recent research directions (mixture of experts, weight sharing, factorization) target the FFN specifically.

For a single FFN layer, we count the total number of learnable parameters by summing the sizes of all weight matrices and bias vectors. The calculation is straightforward. Matrix $W_1$ has shape $d_{\text{model}} \times d_{ff}$, contributing $d_{\text{model}} \times d_{ff}$ parameters. Matrix $W_2$ has shape $d_{ff} \times d_{\text{model}}$, contributing an equal $d_{ff} \times d_{\text{model}}$ parameters. The two bias vectors contribute $d_{ff}$ and $d_{\text{model}}$ parameters respectively. Adding these together:

where:

• $d_{\text{model}} \times d_{ff}$: the number of elements in weight matrix $W_1$ (rows $\times$ columns)
• $d_{ff} \times d_{\text{model}}$: the number of elements in weight matrix $W_2$ (rows $\times$ columns, equal to the count in $W_1$)
• $d_{ff}$: the number of elements in bias vector $b_1$
• $d_{\text{model}}$: the number of elements in bias vector $b_2$
• The factor of 2 in front of $d_{\text{model}} \times d_{ff}$ accounts for both weight matrices $W_1$ and $W_2$ having the same total element count

With the standard expansion factor $d_{ff} = 4 \times d_{\text{model}}$, we can simplify this expression. Substituting $d_{ff} = 4 d_{\text{model}}$:

For large models where $d_{\text{model}}$ is in the hundreds or thousands, the quadratic term dominates and the linear bias terms become negligible, giving approximately $8 \times d_{\text{model}}^2$ parameters per FFN layer.

Why does this formula make sense? Notice that the parameter count grows quadratically with $d_{\text{model}}$: if you double the model dimension, you quadruple the FFN parameter count. This is the same quadratic scaling that makes attention expensive with respect to sequence length, but here the quadratic factor is in the model dimension rather than sequence length. Doubling $d_{\text{model}}$ from 512 to 1024 increases FFN parameters from roughly 2 million to 8 million per layer, a 4x increase for a 2x dimension increase.

Beyond parameter count, we need to consider computational cost, measured in floating-point operations (FLOPs). Understanding FFN compute requirements helps explain why these layers dominate inference time for short sequences, and why long-context models face a different bottleneck.

The FFN's computational cost per token follows directly from its structure. For each token, the FFN performs two matrix-vector multiplications. Each multiply-add operation consists of one multiplication and one addition, counting as 2 floating-point operations. Working through the arithmetic:

The first layer computes $xW_1$ where $x$ is a $d_{\text{model}}$-dimensional vector and $W_1$ is $d_{\text{model}} \times d_{ff}$. This requires $d_{\text{model}} \times d_{ff}$ multiply-adds, so $2 \times d_{\text{model}} \times d_{ff}$ FLOPs. The second layer computes $hW_2$ where $h$ is $d_{ff}$-dimensional and $W_2$ is $d_{ff} \times d_{\text{model}}$, requiring another $2 \times d_{ff} \times d_{\text{model}}$ FLOPs. The total FLOPs per token for the FFN is therefore:

where:

• $d_{\text{model}}$: the input and output dimension of the FFN
• $d_{ff}$: the hidden dimension (typically $4 \times d_{\text{model}}$)
• The factor of 4 comes from: 2 layers $\times$ 2 FLOPs per multiply-add operation

For a sequence of $n$ tokens, each token is processed independently, so the total FFN cost scales linearly with sequence length. Because the FFN has no interaction between positions, the computation at each position is identical and independent:

where $n$ is the number of tokens in the sequence. This is $O(n)$ scaling: double the sequence length, double the FLOPs.

This linear scaling contrasts sharply with attention, which has quadratic complexity $O(n^2 \cdot d_k)$ due to the $n \times n$ attention matrix computation. For every pair of positions, attention must compute a query-key dot product. With $n$ positions, there are $n^2$ such pairs. As sequences grow longer, attention's quadratic term eventually dominates.

The crossover point depends on model dimensions. For a model with $d_{\text{model}} = 768$ and $d_{ff} = 3072$, let's compute when attention FLOPs equal FFN FLOPs. FFN FLOPs scale as $4n \times 768 \times 3072 \approx 9.4 \times 10^6 \times n$. Attention's quadratic term (for the attention matrix alone) scales as $4n^2 \times 768 \approx 3 \times 10^3 \times n^2$. Setting these equal: $9.4 \times 10^6 \times n = 3 \times 10^3 \times n^2$, giving $n \approx 3000$ tokens. At sequence lengths below 3000, FFN dominates. Above 3000, attention takes over.

This crossover point has important practical implications. Short-context applications like question answering and classification, as well as code completion, typically operate with sequences well below the crossover point, meaning the FFN is the primary compute bottleneck. Research into FFN efficiency (quantization, pruning, low-rank approximation) offers the greatest gains here. Long-context applications like document summarization, long-form generation, and retrieval-augmented generation operate above the crossover, making attention efficiency (sparse attention, linear attention, sliding window attention) the priority.

The formula $\text{FFN}(x) = \sigma(xW_1 + b_1)W_2 + b_2$ is compact, but its compactness can obscure what happens at each step. To understand the FFN, let's trace through a concrete computation with actual numbers. We'll use intentionally small dimensions ($d_{\text{model}} = 3$ and $d_{ff} = 4$) so you can follow every multiplication and addition by hand if you like, verifying each step against the formulas.

The goal of this example is threefold: to see exactly how the input vector is transformed at each stage, to observe how ReLU zeros out negative activations to introduce nonlinearity, and to verify that the output dimension matches the input dimension, ready for the residual connection. We'll use simple, hand-readable weight values to keep the arithmetic transparent.

We have chosen an input with mixed positive and negative values ($x = [1.0, -0.5, 0.8]$), which is representative of what you'd see in practice after layer normalization. The weight matrices use a mix of positive and negative values. This ensures some hidden dimensions will have positive pre-activations (and survive ReLU) while others will be zeroed out.

The first operation computes $xW_1 + b_1$. Our 3-dimensional input vector gets multiplied by a $3 \times 4$ weight matrix, producing a 4-dimensional hidden representation. Each element of this output is a weighted sum of the input elements plus a bias term. Specifically, the $j$-th element of the pre-activation is $x[0] \cdot W_1[0,j] + x[1] \cdot W_1[1,j] + x[2] \cdot W_1[2,j] + b_1[j]$: a dot product between the input and the $j$-th column of $W_1$.

Notice the dimension change: a 3-dimensional vector goes in, and a 4-dimensional vector comes out. This expansion is the "higher-dimensional working space" we discussed earlier, where the network has more room to manipulate the representation before applying nonlinearity. The result is called the pre-activation because we haven't applied the nonlinearity yet. Some of these pre-activation values are positive and some are negative, which sets up the selective filtering that ReLU will perform in the next step.

Here is where the FFN gains its expressive power. The ReLU activation function applies a simple rule: keep positive values unchanged, but set negative values to zero. Mathematically, $\text{ReLU}(z) = \max(0, z)$. This is one of the simplest possible nonlinear functions, yet it is sufficient to make the FFN a universal approximator when the hidden dimension is large enough.

Finally, we project back to the original dimension. The 4-dimensional hidden vector is multiplied by a $4 \times 3$ weight matrix, and a 3-dimensional bias is added. This contraction reduces the dimension and combines the contributions of all active hidden dimensions into an output vector. Each output dimension receives a weighted sum of the active hidden dimensions, where the weights are the corresponding elements of $W_2$.

Think of the second linear layer as the "summarization" step. The first layer asked "which features of this input are relevant?" The nonlinearity filtered out irrelevant dimensions. The second layer asks "given these relevant features, what transformation should we apply to the representation?"

The output has the same dimension as the input (3 elements), which is essential. In the full transformer, this output will be added to the original input via a residual connection, and both must have the same shape. The FFN has transformed the representation while preserving its dimensionality, enabling the residual addition that stabilizes training in deep networks.

Let's verify that our step-by-step calculation matches the complete FFN function applied directly:

The step-by-step and function-based computations match exactly. This worked example demonstrates the complete journey: a 3-dimensional input expands to 4 dimensions, passes through nonlinearity (with some dimensions zeroed out), and contracts back to 3 dimensions. The transformation is nonlinear (different inputs will activate different subsets of hidden dimensions, producing qualitatively different transformations), position-independent (the same computation would apply to any token's representation), and dimension-preserving (the output shape matches the input shape for the residual connection).

Having traced through the mathematics by hand, we can now build a reusable FFN module that encapsulates everything we've learned. This implementation follows patterns used in production transformer libraries: it initializes weights using proper scaling (Xavier/Glorot initialization), supports optional biases, and handles both single vectors and batched sequences without any special casing.

naturally the class structure maps onto the mathematical formula. The __init__ method creates the four learnable parameter tensors ($W_1$, $b_1$, $W_2$, $b_2$). The __call__ method implements the formula $\sigma(xW_1 + b_1)W_2 + b_2$ in three readable lines. The num_parameters method computes the formula we derived earlier. The entire forward pass is just three matrix multiplications (including the bias additions). This reflects the fundamental simplicity of the FFN despite its expressive power.

The module correctly handles both single vectors and batched sequences. When a 2D input with shape $(n, d_{\text{model}})$ is provided (a batch of $n$ token representations), the matrix multiplication automatically broadcasts across all $n$ positions simultaneously. This is the source of the FFN's computational efficiency: instead of calling the function once per token in a loop, the entire sequence is processed in a single pair of batch matrix multiplications.

The parameter count matches our theoretical formula: $2 \times 512 \times 2048 + 2048 + 512 = 2{,}099{,}200 + 2{,}560 = 2{,}101{,}760$, close to 2.1 million parameters just for this single FFN layer. A 12-layer GPT-2 Small model has 12 such layers, plus attention layers, adding up to the 117 million total parameters commonly cited.

The feed-forward network is conceptually simple: two linear layers with a nonlinearity between them. Yet this simplicity masks significant computational cost, and understanding those costs is essential for anyone working with large language models in practice.

The FFN's most significant limitation is its memory footprint. Because FFN weights are static after training, they must reside in memory at inference time. For a model like LLaMA-70B with a hidden dimension of 8192 and an FFN dimension of 28672 (approximately 3.5x), each FFN layer holds about 2 $\times$ 8192 $\times$ 28672 $\approx$ 470 million parameters. Across 80 layers, the FFN alone stores roughly 37 billion parameters, each requiring at least 2 bytes in FP16 format, meaning over 74 GB of memory just for FFN weights. This is why running large models requires high-memory GPUs or specialized hardware.

Weight quantization is the most common approach to this memory problem. By representing weights in lower precision (INT8, INT4, or even INT2), the memory footprint shrinks proportionally. A model quantized to INT4 requires one-quarter the memory of its FP16 counterpart, enabling deployment on consumer hardware. The challenge is maintaining model quality: aggressive quantization can degrade generation quality, particularly for reasoning tasks that require precise numerical computations in the FFN's hidden layers. Research into quantization-aware training and calibration has made significant progress, but the tradeoff between compression and quality remains an active area of work.

Sparsity offers another path forward. The ReLU activation naturally creates sparse hidden representations, as negative pre-activations become zero. Researchers have exploited this by identifying which hidden dimensions will be active for a given input and computing only those, skipping computation for dimensions that would be zeroed anyway. This requires predicting active neurons before the full computation, which introduces overhead, but the net savings can be substantial for highly sparse models. The GELU and SiLU activations that have replaced ReLU in many modern models are smoother alternatives that do not produce exact zeros, trading the sparsity benefit for better gradient flow during training.

The Mixture of Experts (MoE) architecture takes the sparsity idea further. Instead of a single FFN per layer, an MoE layer has multiple "expert" FFNs and a router that directs each token to only a small subset (typically 2 of 8, or 2 of 64). The model scales parameters without proportionally scaling compute, because only a fraction of experts process each token. Models like Mixtral-8x7B and GPT-4 (reportedly) use MoE to achieve parameter counts that would otherwise be computationally prohibitive. The key challenge is load balancing: ensuring that all experts are used approximately equally across the training data, preventing some experts from specializing in everything while others specialize in nothing.

Despite its computational weight, the FFN's role is needed and irreplaceable. It provides the nonlinearity that maps attention's linear weighted averages into expressive function approximation. It is the model's "memory," storing factual associations in its weights that attention retrieves based on context. Its position-wise nature enables massive parallelization that makes transformer training tractable at the scale of billions of parameters and trillions of training tokens.

The interplay between attention and FFN defines transformer expressiveness. Attention routes information between positions, creating context-dependent representations. The FFN transforms those representations position-by-position, adding computational depth and nonlinear capacity. Together, they enable the hierarchical, compositional language understanding that powers modern NLP. Neither component alone would be sufficient: attention without the FFN is a purely linear contextual aggregation; the FFN without attention is a powerful token-level classifier with no awareness of surrounding context.

The feed-forward network is the workhorse of transformer computation, applying identical transformations to each position independently. This chapter covered its architecture and efficiency, then examined its interpretation within the broader transformer block.

Key takeaways:

• Two-layer architecture: The FFN formula $\text{FFN}(x) = \sigma(xW_1 + b_1)W_2 + b_2$ consists of an expansion layer ($d_{\text{model}} \to d_{ff}$), a nonlinear activation $\sigma$ (ReLU, GELU, etc.), and a contraction layer ($d_{ff} \to d_{\text{model}}$). The standard expansion factor is $d_{ff} = 4 \times d_{\text{model}}$, establishing a working space four times larger than the model dimension where nonlinear transformations are easier to learn.

• Position independence: Each position is processed separately with shared weights. This enables parallel computation and cleanly separates the FFN's role (per-position nonlinear transformation) from attention's role (inter-position information routing). The position independence is not a limitation but a design choice: by the time the FFN processes each token, attention has already incorporated contextual information into that token's representation.

• Key-value memory interpretation: The FFN can be viewed as an associative memory where $W_1$ columns are keys, hidden activations are match scores, and $W_2$ rows are values. Different inputs activate different subsets of hidden dimensions, creating sparse, input-dependent retrievals. This explains how transformers store factual knowledge and why larger FFN dimensions increase a model's capacity to "know" things.

• Dominant parameter count: The FFN contains approximately two-thirds of a transformer block's parameters. With 4x expansion, the total is approximately $8 \times d_{\text{model}}^2$ parameters per layer (from the formula $2 \times d_{\text{model}} \times d_{ff} + d_{ff} + d_{\text{model}}$). This quadratic scaling with $d_{\text{model}}$ makes FFN optimization critical for model efficiency.

• Linear computational scaling: FFN compute cost is $4 \times n \times d_{\text{model}} \times d_{ff}$ FLOPs, scaling linearly with sequence length $n$. This contrasts with attention's quadratic $O(n^2)$ scaling. For short sequences, FFN dominates compute; for long sequences, attention becomes the bottleneck.

• Crossover point: Around 1000-2000 tokens, FFN and attention have comparable computational costs for typical model sizes. This crossover point shapes optimization strategy: short-context systems prioritize FFN compression, while long-context systems prioritize attention efficiency.

The next chapter examines activation functions used in FFNs in depth, comparing ReLU, GELU, SiLU/Swish, and their gated variants (GLU, SwiGLU), and exploring why modern models have moved beyond the original ReLU choice toward smoother, more expressive alternatives.

When implementing or configuring feed-forward networks in transformers, these parameters control capacity and efficiency:

• d_model (input/output dimension): The embedding dimension that the FFN preserves. Typical values range from 256 (small models) to 12288 (GPT-3 scale). This dimension must match the attention layer output and determines the FFN's interface with the rest of the transformer.

• d_ff (hidden dimension): The expanded dimension of the intermediate representation. The standard choice is $d_{ff} = 4 \times d_{\text{model}}$, though modern architectures like LLaMA use ratios around 2.7x when combined with gated linear units. Larger values increase expressiveness but proportionally increase parameters and compute.

• activation: The nonlinear function applied element-wise after the first linear layer. ReLU was used in the original transformer, but GELU has become the standard for encoder models (BERT, RoBERTa) and SiLU/Swish for decoder models (LLaMA, GPT-NeoX). The choice affects gradient flow, sparsity patterns, and training stability.

• use_bias: Whether to include bias terms $b_1$ and $b_2$. Some modern architectures (LLaMA, PaLM) omit biases entirely to reduce parameters and simplify quantization. The impact on model quality is typically minimal.

• dropout (not shown in our implementation): Dropout rate applied to the hidden representation after activation. Values of 0.1-0.2 are common during training to prevent overfitting. Set to 0 during inference.

• initialization: Weight initialization scale affects training stability. Xavier/Glorot initialization (scaling by $\sqrt{2/(d_{in} + d_{out})}$) is standard. Some architectures use scaled initialization for residual paths to maintain signal magnitude through deep networks.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about feed-forward networks in transformers.