🎓 A1Learner
🔍
0 of 53 lessons completed

PyTorch The Feed-Forward Part

The feed-forward network is a small two-layer network that processes one token vector at a time. Its first layer widens the vector, and its second layer narrows the vector back.

Widen, then narrow

The first Linear layer widens each vector from 128 numbers to 512 numbers. The second Linear layer narrows the vector back to 128 numbers.

torch.manual_seed(0)
ff = FeedForward()
v = torch.randn(2, n_embd)
print("Input shape:", v.shape)
print("Hidden shape:", ff.net[0](v).shape)
print("Output shape:", ff(v).shape)

ReLU clears the negatives

ReLU is the activation that turns negative numbers into zero and keeps the rest unchanged. The hidden values pass through ReLU between the two Linear layers.

hidden = ff.net[0](v)
print("Negatives before ReLU:", (hidden < 0).sum().item())
print("Negatives after ReLU:", (ff.net[1](hidden) < 0).sum().item())

One row alone gives the same row

Two copies of one vector enter as two rows in a single batch. The two output rows match exactly.

pair = torch.cat([v[0:1], v[0:1]], dim=0)
out_pair = ff(pair)
print("Two copies of one row, outputs equal:", torch.equal(out_pair[0], out_pair[1]))

Example

import torch
from torch import nn
from torch.nn import functional as F

batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4

class FeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.ReLU(),
            nn.Linear(4 * n_embd, n_embd),
        )

    def forward(self, x):
        return self.net(x)

torch.manual_seed(0)
ff = FeedForward()
token = torch.randn(1, n_embd)
hidden = ff.net[0](token)
after_relu = ff.net[1](hidden)
back = ff.net[2](after_relu)

print("Token shape:", token.shape)
print("After widening shape:", hidden.shape)
print("After narrowing shape:", back.shape)
print("Negative numbers before ReLU:", (hidden < 0).sum().item())
print("Negative numbers after ReLU:", (after_relu < 0).sum().item())

pair = torch.cat([token, token], dim=0)
out = ff(pair)
print("Same vector at two positions, outputs equal:", torch.equal(out[0], out[1]))
print("Largest difference between the two rows:", (out[0] - out[1]).abs().max().item())

Key points

  • Position independence: the same input vector at two positions gives the same output row
  • The hidden width of 512 is four times n_embd, set by the 4 * n_embd in the class
  • The feed-forward network holds 131,712 parameters at the course width
  • Attention mixes positions before this part runs, and this part never mixes positions itself

In this example, one token vector of width 128 passes through the FeedForward class. The widening layer stretches the vector to 512 numbers. The count of negative numbers falls from 237 to zero at the ReLU step. The narrowing layer returns the vector to width 128. The same vector then enters twice as two rows of one batch. Both output rows match exactly, with a largest difference of 0.0. No information crosses between rows anywhere in the network.

Note

The widening factor of four is a design choice in the class, not a rule of the operation

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What shape does FeedForward return for an input of shape (2, 128)?