§3.4Feed-Forward, Layer Normalization & Positional Embeddings

Attention is only half an encoder layer. This section adds the other half — the feed-forward network — plus the two tricks that let us stack layers deep (layer normalization + skip connections) and the fix for attention’s order-blindness (positional embeddings).

The feed-forward layer

The second sublayer is a plain two-layer network — but a position-wise one: it processes each token embedding independently, the same MLP applied at every position. (Some call it a 1-D convolution with kernel size 1.)

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.linear_1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.linear_2 = nn.Linear(config.intermediate_size, config.hidden_size)
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, x):
        x = self.linear_1(x)
        x = self.gelu(x)
        x = self.linear_2(x)
        x = self.dropout(x)
        return x

A rule of thumb is to make the hidden size of the first layer four times the embedding size, with a GELU gelu Gaussian Error Linear Unit, the smooth activation function used in the transformer's feed-forward sublayer. defined in ch. 3 — open in glossary activation in between:

Layer normalization and skip connections

To train a deep stack we need two standard tricks. Layer normalization layer normalization A normalization that rescales each input to zero mean and unit variance across its features, stabilizing deep-network training. defined in ch. 3 — open in glossary rescales each input to zero mean and unit variance. A skip connection skip connection A residual connection that adds a sublayer's input to its output, letting gradients flow and easing training of deep stacks. defined in ch. 3 — open in glossary (residual) passes a tensor to the next layer unchanged and adds it to the processed tensor, giving gradients a clean path backward.

The only real choice is where to put the norm. Step through the two options:

• Post-LN post-layer normalization Placing layer normalization between the skip connections (the original Transformer arrangement); powerful but tricky to train, often needing warm-up. defined in ch. 3 — open in glossary (original paper) puts the norm between the skip connections; gradients can diverge, so it usually needs learning-rate warm-up learning rate warm-up Gradually increasing the learning rate from a small value to its maximum early in training, needed to stabilize post-LN transformers. defined in ch. 3 — open in glossary .
• Pre-LN pre-layer normalization Placing layer normalization inside the span of the skip connections; the most common arrangement, much more stable to train. defined in ch. 3 — open in glossary puts the norm inside the skip span; it’s much more stable and is what we’ll use.

With that decided, a full encoder layer is just these blocks wired together:

class TransformerEncoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.layer_norm_1 = nn.LayerNorm(config.hidden_size)
        self.layer_norm_2 = nn.LayerNorm(config.hidden_size)
        self.attention = MultiHeadAttention(config)
        self.feed_forward = FeedForward(config)

    def forward(self, x):
        hidden_state = self.layer_norm_1(x)
        x = x + self.attention(hidden_state)                 # skip connection
        x = x + self.feed_forward(self.layer_norm_2(x))      # skip connection
        return x

Click through the assembled layer — note that its output shape equals its input shape, which is exactly what lets us stack N of them:

encoder_layer = TransformerEncoderLayer(config)
inputs_embeds.shape, encoder_layer(inputs_embeds).size()
# (torch.Size([1, 5, 768]), torch.Size([1, 5, 768]))

Positional embeddings

There’s a catch: as built, the encoder layer is totally invariant to token order — attention and the feed-forward layer are permutation equivariant permutation equivariant A property of attention and feed-forward layers whereby permuting the input permutes the output identically — the reason positional information must be added. defined in ch. 3 — open in glossary , so “dog bites man” and “man bites dog” would come out identical. We fix this by adding a positional embedding positional embedding A position-dependent vector added to token embeddings so the otherwise order-blind attention mechanism can model token order. defined in ch. 3 — open in glossary — a position-dependent vector — to each token embedding.

The simplest, most common approach is a learnable table, indexed by position instead of token id, summed with the token embeddings:

class Embeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings,
                                                config.hidden_size)
        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout()

    def forward(self, input_ids):
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0)
        token_embeddings = self.token_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        embeddings = token_embeddings + position_embeddings   # just add them
        embeddings = self.layer_norm(embeddings)
        return self.dropout(embeddings)

An alternative uses fixed sinusoidal patterns — each dimension is a sine or cosine of a different frequency, so every position gets a unique, smooth fingerprint (and no parameters to learn):

The three main families:

Putting the embeddings and the encoder layers together gives the full encoder:

class TransformerEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embeddings = Embeddings(config)
        self.layers = nn.ModuleList([TransformerEncoderLayer(config)
                                     for _ in range(config.num_hidden_layers)])

    def forward(self, x):
        x = self.embeddings(x)
        for layer in self.layers:
            x = layer(x)
        return x
# encoder(inputs.input_ids).size()  ->  torch.Size([1, 5, 768])

The encoder emits one hidden state per token — a flexible output we can adapt for masked-language-modeling, question answering (Ch. 7), or, next, sequence classification. Let’s attach a head and meet the decoder.