Chapter 8.3 - Second Linear Layer

Overview

After the GELU function, our data is all puffed up like a balloon because the first linear layer made it really wide. Now, we need to let the air out and squeeze it back down to its normal size so it can move on to the next step!

---

🎯 Why we do it

Rationale

The AI temporarily expanded the data so it could find hidden patterns (using GELU). But the rest of the model expects the data to be a very specific size. The second linear layer acts like a funnel, compressing the expanded data back into its original shape.

🛠️ How we do it

Methodology

It's just another matrix multiplication! If our puffed-up data has a size of 400, and our normal size is 100, we multiply it by a matrix that transforms it from 400 back to 100.

import torch.nn as nn
import torch

# Puffed up data (size 400)
puffed_data = torch.randn(1, 400)

# The Funnel (Second Linear Layer)
# Compresses 400 back to 100
funnel = nn.Linear(400, 100)

normal_data = funnel(puffed_data)
print("Back to normal size:", normal_data.shape)

---

💻 Code Implementation

Here is the exact PyTorch implementation for the concepts discussed above:

import torch
import torch.nn as nn

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)