FeedForward is the position-wise nonlinear subcomponent of a transformer block. It applies a linear expansion, GELU activation, and linear contraction independently to each token representation, preserving the input shape while enabling richer intermediate representations. section 4.3 section 4.5
The module’s data flow is:
- Receive token embeddings with shape batch × tokens × embedding size. section 4.3
- Expand each embedding from 768 to 3,072 values with a linear layer. section 4.3
- Apply GELU, then contract the representation back to 768 dimensions. section 4.3
- Return the same batch, token, and embedding dimensions for residual transformer-block processing. section 4.5
Shape-preserving execution
A FeedForward instance configured for the 124-million-parameter GPT model accepts a random tensor shaped [2, 3, 768] and produces torch.Size([2, 3, 768]). Internally, every token vector expands to 3,072 values before returning to 768. section 4.3
Listing 4.4 A feed forward neural network module
class FeedForward(nn.Module):
def __init__(self, cfg):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
GELU(),
nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)
def forward(self, x):
return self.layers(x)
Use the implementation listing to inspect the PyTorch Sequential module and its two Linear layers surrounding GELU. section 4.3
