Ground-Up VLA: Feed-Forward Sublayer

From GPU basics to robot foundation models

Modern deep learning, in a simplified manner, can be represented by three types of components (see the architecture post for the structure):

• Input

• Transformers

• Output

So far, we’ve primarily focused on understanding the transformer. A single transformer block (at least, in its simplest form) can be broken down into the “self-attention sublayer” and “feed-forward sublayer”. So far, we’ve covered the self-attention sublayer, first with the post on self-attention itself, and then with one on LayerNorm and residual addition. That leaves the feed-forward sublayer, which we’ll cover in this post. Luckily, the feed-forward sublayer re-uses LayerNorm and residual addition, so we’ve already covered part of what we need to know, and we can focus on the feed-forward network.

To recap, the self-attention sublayer (at least, the variant we are focusing on) looks like the following:

1. Input: P

2. L₁ = LayerNorm(P)

3. A₁ = SelfAttention(L₁)

4. Output: R₁ = ResidualAdd(P, A₁) = P + A₁

The self-attention sublayer is typically followed by a feed-forward sublayer, which looks like:

1. Input: R₁

2. L₂ = LayerNorm(R₁)

3. F₁ = FeedForward(L₂)

4. Output: R₂ = ResidualAdd(R₁, F₁) = R₁ + F₁

This post will focus on two key questions:

• Why do we need a feed-forward sublayer?

• How does the FeedForward(…) function in the above sequence work?

Motivation

We talked about how self-attention was important for a model to understand how to interpret each token of an input sequence (for example, how to interpret the word “tab” in “my brain has too many tabs open”) to develop a better, holistic understanding of the sequence as a whole. So, if self-attention already achieves this, then why would we put in this separate thing called a feed-forward sublayer?

Suppose you’re a researcher or a product designer refining your ideas. First, you go out, talk to people, and collect information from many sources. Then you sit alone and think through what you’ve learned, updating your views. In a transformer block, the self-attention sublayer is this ‘extroverted’ phase, where information is exchanged and mixed across tokens. The feed-forward sublayer is the ‘introverted’ phase, where each token, on its own, considers how to update its thinking about itself, given both what it originally was and what it just heard. In practice, this self-reflection is extremely useful in producing a practical output that would be considered high quality.

Mechanism

Now that we have an intuitive idea of why a feed-forward sublayer is useful, let’s consider what actually happens inside of one. The sublayer has three parts: (1) the LayerNorm, (2) the actual Feed-Forward Network (often called FF or FFN), and (3) a residual addition operation. The LayerNorm and residual addition pieces are present for the same reasons as they are present in the self-attention sublayer, and since we covered this topic in the previous post, we won’t cover them in detail again. We’ll keep the focus on what goes on in the FFN and mention the other pieces in passing.

At a high level, the FeedForward(…) part of the feed-forward sublayer will first convert each input vector from size d_model x 1 to size d_ff x 1, where d_ff > d_model (later in the post, we’ll explain why). Then, a nonlinear activation function will operate, element-wise, on the newly expanded vector, and a final operation will convert the vector back down to having dimensions d_model x 1 before it is passed to the residual addition piece of the feed-forward sublayer. In slightly more mathematical terms:

1. Input: x

2. y = LayerNorm(x)

3. Layer of neurons that performs:

   1. Linear projection: h = W₁y + b₁

   2. Nonlinear activation: h’ = ϕ(h) — e.g., ϕ could be ReLU or GELU, and it is applied element-wise.

4. Layer of neurons that performs:

   1. Linear projection: z = W₂h’ + b₂

   2. Linear identity activation: z’ = z.

5. Output: z’ + x

More succinctly, Steps 3-4 can be written as FFN(y) = W₂ ϕ(W₁ y + b₁) + b₂.

The first linear projection

Let’s now consider just the first linear projection: W₁ y + b₁. Typically, this will take an individual input token and project it into higher-dimensional space. That’s not very intuitive, so let’s break it down.

Suppose we have an input token (i.e., word in a sentence) that is represented as a vector of size d_model x 1. In the self-attention sublayer, the token is always represented as a vector of size d_model x 1, even as we mix in inputs from other nearby tokens.

In contrast, the first linear projection of a FF sublayer will make an input token of size d_model x 1 be represented as an input token of size (4d_model) x 1. It’s like going from describing something with just its visible traits like color and shape (a snowy owl is white-feathered and looks a bit round), to then also including some representation of what it sounds like (owls “hoot”). They represent the same underlying concept, but the more-featured variant provides richer detail. Similarly, letting the FF sublayer convert an input token from dimensions of d_model x 1 to dimensions of d_ff x 1, where d_ff is often around 4d_model, allows the model to develop a more fine-grained and nuanced understanding of the original input.

Let’s revisit the expression: W₁ y + b_1. Based on what we’ve just discussed, we know that y is of size d_model x 1. The output of this projection should be of size d_ff x 1, so that means that W₁ should be of size d_ff x d_model for the matrix multiplication to work. Similarly, b₁ should be of size d_ff x 1.

At the end of the day, however, we don’t really want to have such a high-dimensional representation going through all parts of the network, so we’ll bring everything back to vectors of size d_model x 1, which is exactly what the second linear projection will achieve. Note, however, that the second linear projection is not an “undoing” of the dimension-expansion done in this first linear projection, because of the nonlinear activation function between these two pieces.

Why don’t we lose information? Well, in some way, we can. But, continuing our owl example, you can encode the understanding from the more feature-rich space into the less feature-rich space. For instance, the sound of an owl hooting can be encoded with the information “owls hoot”. It’s not quite the same, but it does capture the essence. And that’s a reasonable tradeoff, because it allows the pipeline to develop a deeper understanding of an input without causing dimensional explosion as we stack transformer layers atop one another, each with their own feed-forward layer.

The post-projection nonlinear activation

The nonlinear activation function that is paired with the first linear projection in a feed-forward sublayer is essential to developing that deeper understanding of what a particular token means. And the fact that it’s nonlinear is essential, otherwise the first and second linear layers would together constitute a single linear operation (because W₁ and W₂ could be multiplied together to produce a single weights matrix). This is the same underlying concept that we discussed when we first talked about what a neuron is.

These nonlinear activation functions let the pipeline model complex relationships in the newly feature-expanded representation of each incoming token, and the essence of what is produced as the output of the nonlinear activation gets preserved even when the second linear projection layer brings a token’s representation from d_ff x 1 back down to d_model x 1.

The second linear projection

At this point, we’ve covered the intuition of the second linear projection, but we’ll recap it here. The output after data passes through the nonlinear activation function discussed above is h’ = ϕ(W₁ y + b₁). The second linear projection will let us produce output vectors of the same size (d_model x 1) as the original input vectors when it performs the operation z = W₂h’ + b₂.

The vector h’ is going to be of size d_ff x 1, which means that W₂ will have to be of size d_model x d_ff for the output to be the same size as the original input vector. And this will also mean that b₂ must be of size d_model x 1.

The activation function for this second projection is just the identity. I.e., we directly pass on the output of the second linear projection to the residual add portion of the feed-forward sublayer before exiting the transformer block. This is because the key thing for this sublayer to do is to develop a deeper understanding of the relationship between features of individual token vectors. That was achieved by the first linear projection and its corresponding nonlinear activation function.

The second linear projection distills the essence of that understanding and passes that along directly, without trying to add its own “layer of understanding” via a nonlinear activation function.

Recap

After the second linear projection operation completes, we have a transformed representation of the original input token. This new representation has performed a kind of “self-reflection” operation by first expanding the input token’s representation to a larger feature space, thinking about how the features relate to one another, and distilling the results of that thought process back down into a more compact form that can be passed along.

A basic transformer block consists of two parts: the self-attention sublayer and the feed-forward sublayer. At this point, we’ve discussed both sublayers in detail, and so have a conceptual understanding of what a transformer block is used for in deep learning networks! Deep learning networks can often have dozens of transformer blocks stacked on top of one another, and are particularly good for modeling really difficult-to-model problems that are highly nonlinear, because they don’t require a human to sit down and think about what the exact mathematical formulae for super complex processes must be.

Next, we’ll consider the problem of how to go from raw data like “my brain has too many tabs open” to vector representations of each word/token in the sequence, such that a neural network can understand it.