🎓 A1Learner
🔍
0 of 71 lessons completed

PyTorch The Encoder Layer

An encoder layer is one transformer block that refines every vector in a sequence. It runs attention and a feed-forward network, and it adds each part's input back before a norm rescales the result.

The parts inside one layer

One layer holds attention, two linear layers, two norms and three dropouts. The activation is a plain function, so it never appears in this list.

import torch
from torch import nn
torch.manual_seed(0)
layer = nn.TransformerEncoderLayer(d_model=16, nhead=4, dim_feedforward=32, dropout=0.0, batch_first=True)
print([name for name, _ in layer.named_children()])

The feed-forward sizes

The first linear layer expands each vector from 16 numbers to 32. The second contracts it back to 16.

print("First feed-forward weight:", tuple(layer.linear1.weight.shape))
print("Second feed-forward weight:", tuple(layer.linear2.weight.shape))

Attention on its own

The attention part returns its output and the weights. The weights hold one row of scores for each position pair.

layer.eval()
x = torch.randn(2, 4, 16)
attn_out, attn_weights = layer.self_attn(x, x, x)
print("Attention output:", tuple(attn_out.shape))
print("Attention weights:", tuple(attn_weights.shape))

Example

import torch
from torch import nn

torch.manual_seed(0)

layer = nn.TransformerEncoderLayer(
    d_model=16, nhead=4, dim_feedforward=32,
    dropout=0.0, batch_first=True)
layer.eval()

x = torch.randn(2, 4, 16)
print("Input:", tuple(x.shape))

attn_out, attn_weights = layer.self_attn(x, x, x)
print("Attention output:", tuple(attn_out.shape))
print("Attention weights:", tuple(attn_weights.shape))

after_attn = layer.norm1(x + layer.dropout1(attn_out))
print("After residual and norm:", tuple(after_attn.shape))

ff = layer.linear2(layer.dropout(layer.activation(layer.linear1(after_attn))))
print("Feed-forward output:", tuple(ff.shape))

manual = layer.norm2(after_attn + layer.dropout2(ff))
print("Manual output:", tuple(manual.shape))

full = layer(x)
print("Full layer output:", tuple(full.shape))
print("Manual matches full layer:", torch.allclose(manual, full))

Key points

  • Self-attention: attention where queries, keys and values all come from the same sequence
  • Feed-forward network: two linear layers with an activation between them, applied to each position alone
  • Residual: a part's input, added back to that part's output
  • Use layer.self_attn(x, x, x) to run the layer's own attention by hand
  • An encoder layer is attention plus a feed-forward network, each wrapped in a residual and a norm

In this example, one encoder layer is built with width 16 and dropout set to 0. The layer is put in eval mode, so the manual stages and the full layer run the same maths. A batch of two sequences with four positions enters as x. The attention part runs first and returns an output and the weights. The first norm rescales x plus the attention output. The feed-forward network then expands each vector to 32 numbers and contracts it back to 16. The second norm rescales the sum of that output and its input. The manual result and the full layer result match, which the last line checks.

Note

The default order is post-norm, where each norm comes after its residual add

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does a residual connection add to a part's output?