Transformer Architecture

Build the mental model

The transformer architecture, introduced in 'Attention Is All You Need', replaced recurrence entirely with self-attention — every position in a sequence attends to every other position in that same sequence — followed by a feedforward network, stacked into repeated blocks. Because there's no recurrence, there's no requirement to process token 1 before token 2 before token 3; every position can be computed simultaneously on a GPU, which is why transformers train dramatically faster than RNNs on the same hardware, RNNs being forced into sequential, step-by-step computation. But this parallelism has a cost: self-attention alone treats the input as an unordered set — swapping two tokens' positions wouldn't change the output — so the model has no inherent sense of word order. Positional encoding fixes this by injecting a pattern (often sinusoidal, varying by position and dimension) directly into the token embeddings before they enter the attention layers, giving the model position information to work with. Multi-head attention runs several attention operations in parallel, each with its own learned query/key/value projections, so different heads can specialize — one tracking syntactic relationships, another tracking topical similarity — and their outputs are concatenated.

Connect it to a real scenario

A search ranking feature on the Tutorial Platform needs to judge how relevant a lesson is to a user's query, and both the query text and the lesson text are sequences where word order changes meaning ('learn Rust before Go' vs 'learn Go before Rust'). A transformer encoder processes the whole query in parallel — fast enough to rank search results in real time — while positional encoding keeps 'before' and 'after' meaningful, and multiple attention heads can separately capture which words are the programming languages versus which word expresses the ordering relationship, combining both signals into a single relevance score.

• Transformers replace recurrence with stacked self-attention + feedforward blocks — all positions process in parallel, making training much faster than RNNs
• Self-attention alone has no sense of order, so positional encoding (a position pattern added to token embeddings) is required
• Multi-head attention runs several attention operations in parallel, letting each head specialize in a different kind of relationship (syntax, topic) simultaneously

Try the working example

import torch
import torch.nn as nn

torch.manual_seed(0)

embed_dim, num_heads, seq_len, batch = 16, 4, 5, 2
mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
x = torch.randn(batch, seq_len, embed_dim)
attn_out, attn_weights = mha(x, x, x)
print("multi-head attention output shape:", attn_out.shape)

def sinusoidal_positional_encoding(seq_len, dim):
    pos = torch.arange(seq_len).unsqueeze(1)
    i = torch.arange(dim).unsqueeze(0)
    angle_rates = 1 / torch.pow(10000, (2 * (i // 2)) / dim)
    angles = pos * angle_rates
    pe = torch.zeros(seq_len, dim)
    pe[:, 0::2] = torch.sin(angles[:, 0::2])
    pe[:, 1::2] = torch.cos(angles[:, 1::2])
    return pe

pe = sinusoidal_positional_encoding(seq_len, embed_dim)
token_embeddings = torch.randn(seq_len, embed_dim)
combined = token_embeddings + pe
print("positional encoding shape:", pe.shape)
print("token embedding + positional encoding shape:", combined.shape)

multi-head attention output shape: torch.Size([2, 5, 16])
positional encoding shape: torch.Size([5, 16])
token embedding + positional encoding shape: torch.Size([5, 16])

Since nn.MultiheadAttention is called as self-attention (Q=K=V=x), the output shape matches the input shape (batch, seq_len, embed_dim). The positional encoding is added element-wise to the token embedding, so the shape stays the same while position information is now baked in.

Wikipedia — Transformer (deep learning architecture) — Deep Learning