Project: Building a Mini Transformer Block

Build the mental model

A transformer block is really two sub-layers glued together with the same pattern twice: sub-layer, then add the block's input back to the output, then normalize. Self-attention (via nn.MultiheadAttention) lets every position in the sequence look at every other position and pull in relevant context, producing a context-mixed representation. A small feedforward network then processes each position independently, adding representational capacity attention alone doesn't have. Both sub-layers are wrapped as x = LayerNorm(x + sublayer(x)) rather than just x = sublayer(x). The naive alternative — stacking sub-layers without the residual connection — forces gradients to flow only through each sublayer's transformation during backpropagation, and in a deep stack of many blocks that gradient signal shrinks or explodes layer by layer. The addition in x + sublayer(x) gives gradients a direct, unimpeded path straight back to earlier layers regardless of how deep the stack is, which is why transformers with dozens of stacked blocks are trainable at all. LayerNorm then keeps the added values in a stable numeric range so training doesn't drift. This attention/add-norm/feedforward/add-norm shape, repeated N times, is the entire transformer encoder.

Connect it to a real scenario

This exact block is the building unit behind a genuinely useful Tutorial Platform feature: a learned lesson-recommendation model that treats a learner's recent lesson history as a sequence and uses self-attention to weigh which past lessons are most relevant to predicting what they should study next — more powerful than a fixed rule like 'next lesson in the same chapter.' Stack a handful of these blocks, average or pool the final sequence output, and feed it into a small classifier head over available lessons, and you have the architectural core of a real personalized recommendation engine, not just a toy shape-check exercise.

• A transformer block always follows the same shape: sub-layer, then add the input back, then normalize — twice per block.
• Residual connections give gradients a direct path backward, which is what makes deep stacks of blocks trainable at all.
• A transformer block preserves input shape exactly — it transforms representations without changing sequence length or embedding dimension.

Try the working example

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim):
        super().__init__()
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(embed_dim)
        self.ff = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embed_dim),
        )
        self.norm2 = nn.LayerNorm(embed_dim)

    def forward(self, x):
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        ff_out = self.ff(x)
        x = self.norm2(x + ff_out)
        return x

batch_size, seq_len, embed_dim = 4, 10, 32
x = torch.randn(batch_size, seq_len, embed_dim)

block = TransformerBlock(embed_dim=embed_dim, num_heads=4, ff_dim=64)
output = block(x)

print("Input shape:", x.shape)
print("Output shape:", output.shape)

Input shape: torch.Size([4, 10, 32])
Output shape: torch.Size([4, 10, 32]) — the output shape exactly matches the input shape, confirming that a transformer block leaves sequence length and embedding dimension unchanged.

PyTorch Docs — Language Modeling with nn.Transformer — Deep Learning