Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Sparse Models › Why Replace the Feed-Forward Network?

Michael BrenndoerferUpdated December 31, 20256 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 239Sparse Models › Why Replace the Feed-Forward Network?

Reading section 5 of 10 · Why Replace the Feed-Forward Network? · 6 minSection 5 of 10 · 6 min

Section 5 of 10 · 6 min read

Sparse Models › Why Replace the Feed-Forward Network?One section · Sections in this chapter5 / 10 ↓

Sections in this chapter10

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

You might wonder why sparse architectures typically replace the FFN rather than the attention mechanism. As we explored in Part XIII, attention's computational cost comes from the quadratic relationship between sequence length and compute. But there's another consideration: the role each component plays in the model. Understanding these distinct roles helps explain why the FFN is the natural candidate for sparsification.

Attention mechanisms handle information routing, determining which tokens should attend to which other tokens. This function is inherently about relationships within the input and benefits from seeing all tokens simultaneously. Making attention sparse (as in Longformer and BigBird from Part XVII) requires careful design to ensure important relationships aren't missed. The attention mechanism needs to maintain a global view of the sequence to effectively aggregate information across positions.

Feed-forward networks, by contrast, operate on each token position independently. They're often interpreted as key-value memories that store factual knowledge and learned transformations. This independence is important: since each token's FFN computation doesn't depend on other tokens in the sequence, we can freely route different tokens to different experts without disrupting information flow. This makes them natural candidates for specialization: different experts can store different types of knowledge without interfering with the information routing that attention gives.

The FFN also dominates the parameter count in typical transformers. In a standard architecture with $d_{ff} = 4 \times d_{model}$ , the FFN contains roughly two-thirds of each layer's parameters. This makes the FFN the natural target for sparsification. Converting it to a mixture of experts yields the largest efficiency gains. To see why the FFN dominates, let's calculate the parameter counts step by step, starting with the FFN:

\begin{aligned} \text{FFN parameters} &= 2 \times d_{\text{model}} \times d_{\text{ff}} && \text{(two weight matrices)} \\ &= 2 \times d_{\text{model}} \times (4 \times d_{\text{model}}) && \text{(substitute } d_{\text{ff}} = 4 \times d_{\text{model}}\text{)} \\ &= 8 \times d_{\text{model}}^2 && \text{(simplify)} \end{aligned}

For comparison, the attention mechanism contains four projection matrices (Q, K, V, and output projection), each with dimension $d_{\text{model}} \times d_{\text{model}}$ . These projections transform the input into queries, keys, and values for the attention computation, then project the attention output back to the model dimension. This gives:

\text{Attention parameters} = 4 \times d_{\text{model}} \times d_{\text{model}} = 4 \times d_{\text{model}}^2

where:

$4$ : the number of projection matrices in multi-head attention (query, key, value, and output projections)
$d_{\text{model}} \times d_{\text{model}}$ : the dimensions of each projection matrix, mapping from the model dimension to itself

To compute the FFN's share of total layer parameters, we divide FFN parameters by the sum of FFN and attention parameters. This ratio tells us what fraction of each layer's capacity resides in the feed-forward network:

\begin{aligned} \text{FFN fraction} &= \frac{\text{FFN parameters}}{\text{FFN parameters} + \text{Attention parameters}} && \text{(definition)} \\ &= \frac{8d_{\text{model}}^2}{8d_{\text{model}}^2 + 4d_{\text{model}}^2} && \text{(substitute parameter counts)} \\ &= \frac{8d_{\text{model}}^2}{12d_{\text{model}}^2} && \text{(combine denominators)} \\ &= \frac{8}{12} && (\text{cancel } d_{\text{model}}^2) \\ &= \frac{2}{3} && \text{(simplify fraction)} \end{aligned}

This means the FFN accounts for approximately 67% of each layer's parameters. This two-thirds dominance makes the FFN the natural target for sparsification: converting it to a mixture of experts allows us to scale the parameter count (by adding more experts) while keeping computational cost proportional to the number of active experts rather than total experts. Since the FFN already holds most of the parameters, this is where we get the biggest capacity gains from sparse computation. If we were to sparsify the attention mechanism instead, we would affect only one-third of the layer's parameters, yielding proportionally smaller efficiency gains.

In[12]:

Code

def analyze_layer_parameters(
    d_model: int = 4096, n_heads: int = 32, d_ff: int = 16384
):
    """Break down parameter distribution in a transformer layer."""

    d_head = d_model // n_heads

    # Self-attention parameters: Q, K, V projections and output projection
    attention_params = 4 * d_model * d_model  # 4 matrices of d_model × d_model

    # FFN parameters: two matrices d_model by d_ff and d_ff by d_model
    ffn_params = 2 * d_model * d_ff

    # Layer norms: small contribution, included for completeness
    norm_params = 4 * d_model  # 2 layer norms with scale and bias

    total = attention_params + ffn_params + norm_params

    return {
        "attention": attention_params,
        "ffn": ffn_params,
        "norms": norm_params,
        "total": total,
        "ffn_fraction": ffn_params / total,
    }

Out[13]:

Console

Attention: 67,108,864 params (33.3%)
FFN: 134,217,728 params (66.7%)
Total: 201,342,976 params

FFN accounts for 66.7% of layer parameters

Out[14]:

Visualization

Pie chart showing transformer layer parameter distribution with Feed-Forward Network in coral taking about two-thirds, Attention in blue taking about one-third, and Layer Norms as a negligible sliver, with total count annotated below. — Proportional breakdown of parameters in a standard transformer layer. The Feed-Forward Network (FFN) accounts for approximately two-thirds (66.7%) of total parameters, identifying it as the high-leverage target for sparsification strategies.

The FFN's dominance at approximately two-thirds of layer parameters makes it the natural target for sparsification. In typical transformer configurations, the FFN contains roughly twice as many parameters as the attention mechanism. This 2:1 ratio means that converting the FFN to a mixture of experts yields the largest efficiency gains. When we add more experts to the FFN, we scale the parameter count materially while keeping computational cost proportional to the number of active experts. Since the FFN already holds most parameters, this is where sparse computation gives the biggest capacity increase per unit of additional compute. Replacing the dense FFN with a sparse mixture of experts while keeping attention dense preserves the model's ability to route information effectively while sharply increasing its knowledge capacity.