Feed-Forward Networks in Transformers

A TransformerBlock class builds its attention module, feed-forward network, two LayerNorm modules, and a dropout module from a configuration dictionary, but the forward method alone defines execution order.

Programming

Why the Order of Operations Matters

A common assumption among learners is that a transformer block is simply attention followed by a feed-forward network, with normalization placed wherever it fits. Under that assumption, the order of the operations looks like a minor detail. The code for a transformer block shows that the order matters a great deal. The same two components, attention and feed-forward, can be wrapped by normalization in two different ways, and the implementation commits to one of them.

This article corrects that belief by following the block through four stages: how it is constructed from a configuration dictionary, how the forward pass executes step by step, what the shortcut connections contribute, and how Pre-LayerNorm differs from Post-LayerNorm. Along the way you will learn a practical method for locating the step where a block's output first goes wrong.

Predicting the First Operation

Before reading further, commit to an answer. Think about a transformer block that receives an input tensor x and must produce an output. Which operation touches x first?

What do you think happens?

In a transformer block that receives input x, which operation touches x first?

The attention mechanism
The feed-forward network
A LayerNorm module (norm1)
The dropout module

Reveal answer

Answer:

A LayerNorm module (norm1)

The input is first saved as a shortcut, and then normalized by norm1 before it reaches attention. If you guessed attention, you are in good company, because attention is the most prominent component. The source code, however, places normalization at the front of the attention half.

Keep this ordering in mind. The rest of the article depends on it, because the forward method is what determines which operation runs when.

What the Constructor Builds

The TransformerBlock class is a PyTorch module that inherits from nn.Module. Its constructor does not hard-code sizes. Instead, it reads everything it needs from a configuration dictionary called cfg, which is passed in once when the block is created. The source gives the GPT_CONFIG_124M dictionary as an example of such a configuration.

Which modules does the constructor build from the config dictionary, and how is that separate from the order in which forward calls them?

Inside the constructor, the block creates five components. The attention module, stored as self.att, is a MultiHeadAttention built from the embedding dimension, the context length, the number of heads, the dropout rate, and a flag for query, key, and value bias. The feed-forward module, stored as self.ff, is a FeedForward network built directly from cfg. Two LayerNorm modules, norm1 and norm2, are each sized to the embedding dimension. Finally, a dropout module called drop_shortcut is created from the dropout rate.

Following the Forward Pass

The forward method is where the block does its work, and it is where most debugging questions are answered. The block handles two halves in sequence, and each half follows the same pattern. For the attention half, the block first saves the input as a shortcut, then normalizes x, runs attention, applies dropout, and adds the saved input back. For the feed-forward half, the block saves the current value as a new shortcut, normalizes it, runs the feed-forward network, applies dropout, and adds the shortcut back.

python

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        # self.att, self.ff, norm1, norm2, drop_shortcut
        # are created here, sized from cfg
        ...

    def forward(self, x):
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut

        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut
        return x

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        # self.att, self.ff, norm1, norm2, drop_shortcut
        # are created here, sized from cfg
        ...

    def forward(self, x):
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut

        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut
        return x

Notes on the code

The constructor sketch is elided; only the forward order is the point here.
Line order in forward defines execution order, not the order of construction.

In what exact order does data flow through the block's forward method, and which steps feed into which?

Read the diagram as two passes over the same data. The first pass takes the block input and produces an attention contribution, which is added back to that input. The second pass takes the resulting value and produces a feed-forward contribution, which is again added back to what it started from. The output of the block is therefore the result of the second addition, not the raw output of either component.

What a Shortcut Adds

A shortcut connection is a rule that adds the input of a block to the block's output. The source states that this feature helps gradients flow through the network during training and improves the learning of deep models. In the code, a shortcut is nothing more than a saved reference to a tensor, followed by a plus operation that combines it with the processed value.

How does the input bypass a sublayer and get added back to its output, so that the original signal stays available alongside the processed one?

Pre-LayerNorm Versus Post-LayerNorm

Definition

Pre-LayerNorm applies layer normalization before each of the two components. Post-LayerNorm applies it after them. These are different designs, not two names for the same arrangement.

Where does the LayerNorm sit relative to the attention and feed-forward sublayers in each variant, and which one does this block use?

The before-and-after view isolates the single difference that matters for each half. The same operations appear in both arrangements, but their order is reversed. The block described here is the Pre-LayerNorm form, because norm1 runs before the attention module and norm2 runs before the feed-forward module. The source notes that older architectures, including the original transformer model, applied layer normalization after the self-attention and feed-forward networks, which is the Post-LayerNorm arrangement.

Worked Example: Tracing a Block Input

Worked example

Tracing which operation handles each value

Given a block input x, name the operation that produces each intermediate value in the forward method, in order.

1
Save the shortcut
shortcut = x. The original input is kept as a reference for the first addition.
2
Normalize
x = norm1(x). This is the first operation applied to the input, before any attention.
3
Attend and apply dropout
x = att(x), then x = drop_shortcut(x). The attention contribution is produced and passed through dropout.
4
First addition
x = x + shortcut. The attention contribution is added back to the original input.
5
Second half
shortcut = x, then x = norm2(x), x = ff(x), x = drop_shortcut(x), and finally x = x + shortcut. The output is the result of this second addition.

Result

The block output is the feed-forward contribution added to the value that was present after the first addition, not the raw output of attention or feed-forward alone.

Common Mistakes When Reading a Transformer Block

Common mistake
Assuming attention runs first because it is the most prominent component
What it looks like
Reasoning that the input meets attention before anything else.
Why it fails
The forward method saves the input as a shortcut and passes it through norm1 before attention.
Do this instead
Read the forward method line by line and treat its order as the definition of execution.
Common mistake
Treating the constructor as the place where execution order is set
Why it fails
The constructor only creates self.att, self.ff, norm1, norm2, and drop_shortcut. It does not say which runs first.
Do this instead
Look to forward for the sequence of calls.
Common mistake
Treating the shortcut addition as optional decoration
Why it fails
The source says shortcut connections help gradients flow and improve learning in deep models, and the block output depends on the second addition.
Do this instead
Confirm each add line is present and that it adds the correct saved value.
Common mistake
Confusing Pre-LayerNorm with Post-LayerNorm
What it looks like
Describing the block as normalizing after attention and feed-forward.
Why it fails
This block is Pre-LayerNorm, so norm1 runs before attention and norm2 runs before feed-forward.
Do this instead
Check whether each LayerNorm call comes before the sublayer call it feeds.
Common mistake
Assuming the two dropout applications are different modules, when the block uses one drop_shortcut module for both halves

Finding Where Output Diverges

Suppose a block produces an output that does not match what a colleague expects. The fastest way to locate the cause is to ask a question about each step in the forward pass: is the value at this point what I expect it to be? Because the block has a fixed sequence of steps, the first step where the expected and actual values disagree marks the place where the divergence starts.

If the output looks wrong, at which step in the forward sequence should I inspect intermediate values to find the mismatch?

MEDIUM

A block's output is wrong. The value after the first addition looks correct, but the value after the second addition does not. Which half of the forward pass should you investigate first, and which intermediate values would you compare?

Hints

Count the addition steps in the forward order and find which one is failing.
If the value after the first addition is correct, the error lies in the steps that follow it.
The second half runs norm2, the feed-forward network, dropout, and then the second addition.

Your attempt (not saved)

Key Takeaways

A TransformerBlock class builds its attention module, feed-forward network, two LayerNorm modules, and a dropout module from a configuration dictionary, but the forward method alone defines execution order.
In the forward pass, the input is saved as a shortcut and normalized by norm1 before attention, dropout, and the first addition; the same pattern then runs with norm2 and the feed-forward network, ending with the second addition.
Shortcut connections add the block input to the output, which the source says helps gradients flow and improves learning in deep models.
This block uses Pre-LayerNorm, where normalization comes before each sublayer, and this is distinct from Post-LayerNorm, where normalization comes after.
To locate an output mismatch, check intermediate values in forward order and investigate at the first step where the expected and actual values disagree.

Related concepts

More concepts starting with F