Feed-Forward Networks › Computational Cost

Michael BrenndoerferUpdated June 10, 20255 min read
Language AI HandbookAll chapters
Part 15 of 67Transformer Blocks
Chapter 107Feed-Forward Networks › Computational Cost
Reading section 7 of 13 · Computational Cost · 5 minSection 7 of 13 · 5 min

Book contents

Section 7 of 13 · 5 min read

Feed-Forward Networks › Computational CostOne section · Sections in this chapter7 / 13 ↓
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

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 xW1xW_1 where xx is a dmodeld_{\text{model}}-dimensional vector and W1W_1 is dmodel×dffd_{\text{model}} \times d_{ff}. This requires dmodel×dffd_{\text{model}} \times d_{ff} multiply-adds, so 2×dmodel×dff2 \times d_{\text{model}} \times d_{ff} FLOPs. The second layer computes hW2hW_2 where hh is dffd_{ff}-dimensional and W2W_2 is dff×dmodeld_{ff} \times d_{\text{model}}, requiring another 2×dff×dmodel2 \times d_{ff} \times d_{\text{model}} FLOPs. The total FLOPs per token for the FFN is therefore:

FLOPsFFN per token=2×(dmodel×dff)+2×(dff×dmodel)=4×dmodel×dff\text{FLOPs}_{\text{FFN per token}} = 2 \times (d_{\text{model}} \times d_{ff}) + 2 \times (d_{ff} \times d_{\text{model}}) = 4 \times d_{\text{model}} \times d_{ff}

where:

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

For a sequence of nn 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:

Total FLOPsFFN=4×n×dmodel×dff\text{Total FLOPs}_{\text{FFN}} = 4 \times n \times d_{\text{model}} \times d_{ff}

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

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

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

In[19]:
Code
def ffn_flops(n, d_model, d_ff):
    """Calculate FFN FLOPs for sequence length n."""
    return 4 * n * d_model * d_ff


def attention_flops(n, d_model):
    """Calculate attention FLOPs (simplified)."""
    # Q, K, V projections: 3 * 2 * n * d_model^2
    # QK^T: 2 * n^2 * d_model
    # Attention @ V: 2 * n^2 * d_model
    # Output projection: 2 * n * d_model^2
    projection_flops = 4 * 2 * n * d_model * d_model
    attention_matrix_flops = 4 * n * n * d_model
    return projection_flops + attention_matrix_flops


# Compare across sequence lengths
sequence_lengths = [128, 512, 1024, 2048, 4096, 8192]
d_model_test = 768
d_ff_test = 3072
Out[20]:
Console
FLOPs comparison: FFN vs Attention (GPT-2 Small)

  Seq Length       FFN FLOPs      Attn FLOPs     FFN/Attn
------------------------------------------------------------
         128   1,207,959,552     654,311,424         1.85x
         512   4,831,838,208   3,221,225,472         1.50x
        1024   9,663,676,416   8,053,063,680         1.20x
        2048  19,327,352,832  22,548,578,304         0.86x
        4096  38,654,705,664  70,866,960,384         0.55x
        8192  77,309,411,328 244,813,135,872         0.32x

For short sequences (128-512 tokens), the FFN contributes over twice the FLOPs of attention. As sequences grow to 4096 or 8192 tokens, attention's quadratic cost catches up and eventually dominates. The crossover point around 1024-2048 tokens is where attention and FFN have comparable costs for this model configuration.

Out[21]:
Visualization
Line plot showing FFN FLOPs as a straight line and attention FLOPs as a curve that crosses over and grows faster at longer sequences.
Computational cost in FLOPs for FFN (red) and attention (blue) as sequence length increases, shown on a log-log scale. FFN cost grows linearly while attention cost grows quadratically. The vertical dotted line marks the crossover point where attention surpasses the FFN in compute. Below this crossover, optimizing the FFN yields greater efficiency gains; above it, attention efficiency becomes the priority.

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.