Feed-Forward Layer

The feed-forward layer helps an LLM to learn features of the text that it is trained on.

By feature, we mean a pattern or property of the text that the model finds useful for predicting what comes next.

Examples might include:

“This word is probably a person’s name.”

“The sentence is discussing something in the past.”

“This phrase expresses disagreement.”

“The current token is inside a quotation.”

“The subject is plural, so the following verb should probably be plural.”

“The context concerns France and its capital.”

At the beginning of training, the model’s internal transformations are mostly random. It repeatedly predicts the next token, measures its error, and adjusts its weights. Over many examples, these adjustments create internal patterns that respond to useful properties of text. That is what we mean when we say the model learns features.

There are many feed-forward layers in an LLM. Each one does the following:

Expand, then contract A feed-forward layer creates a wider hidden representation, then returns it to the original size.
Feed-forward layer expands and contracts vector dimensions Six input dimensions connect to twelve hidden dimensions, which connect to six output dimensions. The wider middle column represents expansion into a higher-dimensional space before contraction back to the original size.

This is a setup that is very common in neural networks which are aimed at learning complex features. Projecting into a higher-dimensional space gives the LLM more capacity to detect and transform complex features. Projecting back down then combines those features into a compact representation whose dimensions remain consistent throughout the LLM architecture.

In code this looks like the following:

import torch.nn as nn


class FeedForwardLayer(nn.Module):

    def __init__(self, model_dim, hidden_dim):
        super().__init__()
        self.expand = nn.Linear(model_dim, hidden_dim)
        self.activation = nn.GELU()
        self.project = nn.Linear(hidden_dim, model_dim)

    def forward(self, x):
        hidden = self.expand(x)
        activated = self.activation(hidden)
        return self.project(activated)

The linear layers here (nn.Linear) contain weights (random to start with) which determine how the input vectors are projected into the larger space, and how the results of that projection are projected back to the smaller space. These weights are then optimised during training.