Building a FeedForward Module

The FeedForward module maintains the shape of the input tensor, so the output has the same dimensions as the input.

Programming

Predicting the Shape Before Running the Program

This lesson takes a practical route. Instead of starting with a definition of a feed forward block, you will follow a short program that creates a FeedForward module, sends it a batch of token embeddings, and prints the shape of what comes back. The program is the anchor for everything that follows, so keep its sequence in mind as the trace unfolds.

The central fact of this practice is that the FeedForward module maintains the shape of the tensor it receives. Once the input has passed through the module, the output has exactly the same dimensions as the input. Before you look at any output, commit to an answer. The input is a batch of two samples, each containing three tokens, with 768 values per token. Will passing it through the module leave that shape alone, or will the module reshape it into something new?

What do you think happens?

After the batch of shape (2, 3, 768) passes through the FeedForward module, what shape does the output have?

The module reshapes the tensor into a new shape
The output keeps the same shape, (2, 3, 768)

Reveal answer

Answer:

The output keeps the same shape, (2, 3, 768).

The FeedForward module maintains the shape of the input tensor, so the three dimensions of the output match the three dimensions of the input: 2 samples, 3 tokens each, and 768 values per token.

Reading a Batched Token Tensor

The input in this practice is created with torch.rand(2, 3, 768). This tensor has three dimensions, and each one answers a different question. The first dimension gives the number of samples in the batch, which is 2. The second gives the number of tokens in each sample, which is 3. The last gives the number of values used to represent each token, which is 768. In this practice, 768 is the embedding size the module was configured for.

Which axis of the tensor is the batch, which is the token position, and which is the embedding, for a tensor of shape (2, 3, 768)?

Reading the dimensions from left to right, as batch, tokens, and embedding, gives you a reliable way to describe any tensor in this lesson. Keep that mental map ready, because the output is judged against it. To locate a single value, you name one index on each axis: which sample, which token, and which position within that token's embedding.

Key point

A shape of (2, 3, 768) means 2 samples, 3 tokens per sample, and 768 values per token. Read the numbers left to right as batch, tokens, and embedding.

Tracing the Program from Input to Printed Output

1
A tensor is created with torch.rand(2, 3, 768).
2
The FeedForward module is built from the GPT_CONFIG_124M configuration, which sets the token embedding size to 768.
3
The tensor is passed to the module through a call, and the module returns a new tensor.
4
The returned tensor is stored under the name out.
5
Printing out.shape reports the dimensions of that returned tensor.

python

import torch

batch = torch.rand(2, 3, 768)
ffn = FeedForward(GPT_CONFIG_124M)
out = ffn(batch)
print(out.shape)

import torch

batch = torch.rand(2, 3, 768)
ffn = FeedForward(GPT_CONFIG_124M)
out = ffn(batch)
print(out.shape)

Notes on the code

batch holds 2 samples, 3 tokens each, 768 values per token
the module is built from the configuration, and its embedding size is 768
the call sends the batch through the module and returns a new tensor
out.shape reports the dimensions of the returned tensor

Output

torch.Size([2, 3, 768])

Watch which numbers stay fixed and which, if any, change between the line that builds the input and the line that prints the result. In this trace, the batch size of 2, the token count of 3, and the embedding size of 768 all remain the same. Nothing in the printed result differs from the input dimensions.

Shape Preservation Through the Module

The module sits between two tensors, and the test is whether their dimensions agree. The input enters with shape (2, 3, 768), and the output leaves with the same shape. Viewed as a check rather than a transformation, the diagram below shows that the dimensions match, which is the whole point of the printed comparison.

How do the batch, token, and embedding dimensions compare between the tensor entering the module and the tensor leaving it?

Key point

Because the FeedForward module maintains the shape of the tensor it receives, the output has the same batch, token, and embedding dimensions as the input.

Using Shape Matching as a Check

Because the output has the same dimensions as the input, you can compare the two directly. If the shapes agree, the code that follows the module can keep expecting the same layout of batch, tokens, and embedding. You do not need to change that downstream code to accommodate a new shape. A direct comparison of shapes is a quick way to confirm this.

python

print(out.shape == batch.shape)

print(out.shape == batch.shape)

Output (expected)

True

Notes on the code

compares the dimensions of the output with the dimensions of the input

Mistakes Beginners Make with Output Shapes

Common mistake
Assuming the module changes the number of tokens or the embedding size
What it looks like
Reading out.shape as (2, 3, 768) and expecting a different token count after the call
Why it fails
The FeedForward module maintains the shape of the input tensor, so all three dimensions stay the same.
Do this instead
Compare the printed shape with the input shape dimension by dimension, reading batch, tokens, and embedding in order.
Common mistake
Treating a matching shape as proof that the values are unchanged
What it looks like
Concluding that the output equals the input because out.shape == batch.shape prints True
Why it fails
Matching shapes confirm structure only and say nothing about the values inside the tensor.
Do this instead
Use the shape check to confirm layout, and do not read it as a check on the numbers.
Common mistake
Mixing up the axes of the tensor
What it looks like
Reading the 768 as the number of tokens instead of the number of values per token
Why it fails
The dimensions answer different questions: 2 is the number of samples, 3 is the number of tokens, and 768 is the embedding size.
Do this instead
Always read the shape left to right as batch, tokens, and embedding.

Predicting the Shape for a Different Batch

Worked example

Predicting the output shape for a new input

An input is created with torch.rand(4, 5, 768) and passed through a FeedForward module. What shape should out have?

1
Read the input dimensions
The first dimension is 4, so there are 4 samples. The second is 5, so there are 5 tokens per sample. The last is 768, the embedding size.
2
Apply the shape-preservation rule
The FeedForward module maintains the shape of the input tensor, so the output keeps all three dimensions unchanged.
3
State the predicted shape
The predicted shape of out is (4, 5, 768).

Result

out.shape should print torch.Size([4, 5, 768]), with the batch size, token count, and embedding size all matching the input.

Practice: Predicting and Checking a Shape

EASY

An input is created with torch.rand(1, 7, 768) and passed through a FeedForward module built from GPT_CONFIG_124M. Before running anything, predict the printed shape of out. Then state how many samples the batch holds, how many tokens each sample has, and what the embedding size is. Finally, write the one-line comparison that would confirm the output can be checked against the input.

Hints

Read the first dimension of the input as the number of samples.
The module keeps all three dimensions, so the printed shape should match the input shape.
A comparison of out.shape with the input's shape confirms the dimensions agree.

Your attempt (not saved)

Key Takeaways

A FeedForward module maintains the shape of the tensor it receives, so the output has the same dimensions as the input.
A batched token tensor of shape (2, 3, 768) reads as 2 samples, 3 tokens each, and 768 values per token, from left to right.
You can predict the printed output shape from the input shape by keeping every dimension unchanged.
Because shapes match, the output can be compared directly with the input and the downstream code does not need to change.
Matching shapes confirm structure only; they do not establish that the values inside the tensor are the same.

More concepts starting with B