AI Foundations

Transformer Architecture (Deep Dive)

By Arpit Tripathi, Founder

A component-level walkthrough of how a Transformer block is assembled: positional encoding, the position-wise feed-forward sublayer, residual connections, layer normalization, how multi-head attention slots into the block, and how these stack into encoder and decoder layers.

Inside the Transformer block

This page assumes the high-level picture of a Transformer and focuses on how a single block actually computes, component by component. For the plain-language overview of what a Transformer is, see the Transformer page; here the goal is the mechanics. Introduced by Vaswani et al. in the 2017 paper "Attention Is All You Need" (arXiv:1706.03762), the architecture replaces recurrence and convolution with attention as the only mechanism for moving information between positions in a sequence.

Each block is built from two kinds of sub-layer: an attention sub-layer that mixes information across positions, and a position-wise feed-forward network that transforms each position independently. Both sub-layers are wrapped in a residual connection and followed (in the original design) by layer normalization. The original model stacks N = 6 identical blocks in the encoder and 6 in the decoder, with a model width of d_model = 512.

The defining property of this design is that the attention sub-layer computes its output for all positions at once through matrix multiplications, rather than walking the sequence one step at a time. That single choice is what lets a Transformer be trained efficiently on modern accelerators and is the reason it displaced recurrent networks for large-scale language modeling.

  • A block has two sub-layers: a self-attention sub-layer and a position-wise feed-forward network.
  • Both sub-layers use a residual connection plus layer normalization.
  • The original encoder and decoder each stack N = 6 identical blocks at width d_model = 512.
  • Attention mixes information across positions; the feed-forward network transforms each position on its own.

Scaled dot-product attention and the 1/sqrt(d_k) scaling

Attention compares a set of queries against a set of keys to decide how much each value should contribute to the output. Queries (Q), keys (K), and values (V) are matrices whose rows are vectors, one per position. The score between a query and a key is their dot product; these scores are normalized with a softmax and used as weights over the value vectors. Keys and queries have dimension d_k.

The crucial detail is the 1/sqrt(d_k) factor applied before the softmax. The paper's own justification is direct: for large values of d_k, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients, so the dot products are scaled by 1/sqrt(d_k) to counteract this. Intuitively, a dot product of d_k independent unit-variance terms has variance that grows with d_k, so its standard deviation grows like sqrt(d_k); dividing by sqrt(d_k) keeps the score distribution at a stable scale regardless of head width, which keeps softmax gradients usable during training.

In the decoder's self-attention, a causal mask sets the scores for future positions to negative infinity before the softmax, so a position can only attend to itself and earlier positions. This masking is what makes autoregressive generation possible without leaking information from the future.

Attention(Q, K, V) = softmax( (Q K^T) / sqrt(d_k) ) V
Scaled dot-product attention. Q, K, V are query, key, and value matrices; d_k is the key/query dimension. The softmax runs over each row, producing weights that sum to one over the value positions.
  • Attention weights are softmax(QK^T/sqrt(d_k)) applied to the value matrix V.
  • Without scaling, large d_k makes dot products large and saturates the softmax, killing gradients.
  • The variance of a dot product of d_k terms grows with d_k, so dividing by sqrt(d_k) stabilizes the scale.
  • A causal mask (-infinity on future scores) turns self-attention into autoregressive decoding.

Multi-head attention

Rather than running one attention function over the full d_model-dimensional vectors, the Transformer projects the queries, keys, and values h times with separate learned matrices, runs attention in parallel on each lower-dimensional projection, then concatenates the results and projects them once more. Each of these parallel attention functions is a head. In the base model h = 8, and because d_k = d_v = d_model / h = 64, the total compute is similar to a single full-width attention.

Splitting attention into heads lets the model attend to information from different representation subspaces at different positions simultaneously. One head can track syntactic agreement, another can follow coreference, another can attend to nearby tokens; a single averaged attention would blur these patterns together. The separate output projection W^O then recombines the heads into one vector per position.

Multi-head attention appears in three roles in the original architecture: encoder self-attention (every position attends to every position of the same input), masked decoder self-attention (each output position attends only to earlier output positions), and encoder-decoder cross-attention (decoder queries attend over the encoder's output keys and values). The mechanism is identical in all three; only the source of Q, K, and V and the masking differ.

MultiHead(Q,K,V) = Concat(head_1, ..., head_h) W^O, head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)
Multi-head attention. Each head projects Q, K, V with its own learned matrices, attends, and the concatenated heads are mixed by the output projection W^O.
python
import torch
import torch.nn.functional as F

def multi_head_attention(x, w_q, w_k, w_v, w_o, num_heads, mask=None):
    # x: (batch, seq, d_model). w_*: (d_model, d_model).
    B, T, d_model = x.shape
    d_k = d_model // num_heads
    # Project, then split d_model into (num_heads, d_k).
    q = (x @ w_q).view(B, T, num_heads, d_k).transpose(1, 2)
    k = (x @ w_k).view(B, T, num_heads, d_k).transpose(1, 2)
    v = (x @ w_v).view(B, T, num_heads, d_k).transpose(1, 2)
    # Scaled dot-product scores: (B, heads, T, T).
    scores = (q @ k.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    out = weights @ v                       # (B, heads, T, d_k)
    out = out.transpose(1, 2).reshape(B, T, d_model)
    return out @ w_o                        # final output projection
Minimal multi-head attention forward pass in PyTorch, showing the per-head split, the 1/sqrt(d_k) scaling, optional masking, concatenation via reshape, and the output projection.
  • h parallel heads each have their own Q, K, V projection matrices.
  • Head outputs are concatenated and passed through a single output projection W^O.
  • Base config: h = 8 heads with d_k = d_v = d_model/h = 64, so total cost matches one full-width head.
  • Used as encoder self-attention, masked decoder self-attention, and encoder-decoder cross-attention.

Positional encoding: sinusoidal, learned, and rotary

Attention is permutation invariant: it computes the same set of weighted sums no matter how the input positions are ordered, because it treats the input as a set of vectors with no notion of sequence order. A model that needs to know that "dog bites man" differs from "man bites dog" must therefore be given position information explicitly. The Transformer injects it by adding a positional encoding vector to each token embedding before the first block.

The original encoding is a fixed sinusoid. Each dimension of the encoding is a sine or cosine whose wavelength forms a geometric progression, so dimension i oscillates at frequency 1/10000^(2i/d_model). Even dimensions use sine and odd dimensions use cosine. A useful property is that the encoding for a position offset by a fixed amount is a linear function of the encoding at the original position, which gives the model an easy way to attend by relative offset, and the fixed form extrapolates to sequence lengths not seen in training.

Two widely used alternatives have since become common. Learned positional embeddings replace the fixed sinusoid with a trainable vector per position, which performs comparably but does not extrapolate beyond the trained length. Rotary position embedding (RoPE), from the RoFormer paper (arXiv:2104.09864), instead rotates the query and key vectors by an angle proportional to position, which encodes absolute position while making the attention score depend only on the relative distance between two tokens. RoPE is the dominant choice in many current large language models.

PE(pos, 2i) = sin( pos / 10000^(2i/d_model) ) PE(pos, 2i+1) = cos( pos / 10000^(2i/d_model) )
Original sinusoidal positional encoding. pos is the token position and i indexes the embedding dimension. Even dimensions get sine, odd dimensions get cosine, with wavelengths in geometric progression.
python
import torch

def sinusoidal_positional_encoding(seq_len, d_model):
    pos = torch.arange(seq_len).unsqueeze(1)                  # (seq_len, 1)
    i = torch.arange(0, d_model, 2)                           # even indices
    div = 10000.0 ** (i / d_model)                            # 10000^(2i/d_model)
    pe = torch.zeros(seq_len, d_model)
    pe[:, 0::2] = torch.sin(pos / div)                        # even dims -> sin
    pe[:, 1::2] = torch.cos(pos / div)                        # odd dims  -> cos
    return pe                                                 # add to embeddings
Constructing the original sinusoidal positional encoding in PyTorch. The result is added to the token embeddings before the first Transformer block.
  • Self-attention is permutation invariant, so position must be supplied separately.
  • The original encoding adds fixed sinusoids of geometrically spaced frequencies to the embeddings.
  • Learned positional embeddings train one vector per slot but do not extrapolate past the trained length.
  • Rotary position embedding (RoPE) rotates queries and keys so attention depends on relative distance.

Feed-forward network, residuals, and layer normalization

After the attention sub-layer, each block applies a position-wise feed-forward network: the same two-layer multilayer perceptron is run on every position independently. It projects up to a wider inner dimension, applies a nonlinearity, then projects back down. In the base model the width is d_model = 512 and the inner layer is d_ff = 2048, four times wider. This sub-layer adds per-position nonlinear capacity that attention alone, being a weighted average of values, does not provide.

Each sub-layer sits inside a residual connection, so its input is added back to its output. Residual connections give gradients a short path through the deep stack and let each sub-layer learn a refinement rather than a full replacement of its input. In the original paper the output of each sub-layer is LayerNorm(x + Sublayer(x)), a placement called post-LN, where normalization is applied after the residual is added.

Later work moved the normalization inside the residual branch, computing x + Sublayer(LayerNorm(x)), known as pre-LN. The paper "On Layer Normalization in the Transformer Architecture" (arXiv:2002.04745) showed that with pre-LN the gradients are well-behaved at initialization, which lets deep Transformers train stably without the learning-rate warmup that post-LN models need. Most large models trained today use the pre-LN arrangement for this reason.

FFN(x) = max(0, x W_1 + b_1) W_2 + b_2
Position-wise feed-forward network from the original paper: a linear up-projection, a ReLU nonlinearity, then a linear down-projection, applied to each position separately.
  • The feed-forward network is a two-layer MLP applied identically and independently at each position.
  • Base dimensions: d_model = 512 with a wider inner layer d_ff = 2048.
  • Residual connections add each sub-layer's input back to its output for stable gradient flow.
  • Post-LN normalizes after the residual; pre-LN normalizes before the sub-layer and trains more stably.

Encoder, decoder, and why it parallelizes

An encoder block contains self-attention followed by a feed-forward network. A decoder block adds a third sub-layer: masked self-attention over the generated output, then encoder-decoder cross-attention where the decoder's queries attend over the encoder's keys and values, then the feed-forward network. Three configurations are common. Encoder-only models (such as BERT) use bidirectional self-attention for understanding tasks like classification. Decoder-only models (the GPT family) use masked self-attention for autoregressive generation. Encoder-decoder models keep both stacks for sequence-to-sequence tasks like translation.

The reason the architecture scales is parallelism. A recurrent network computes the hidden state at position t from the state at position t-1, so a length-n sequence requires n sequential steps that cannot be overlapped. Self-attention instead computes the output for every position as a few large matrix multiplications, which means a constant number of sequential operations regardless of sequence length. On hardware built for dense linear algebra, this turns a length-n dependency chain into one parallel pass during training.

The cost of that parallelism is that attention compares every position with every other position, so compute and memory grow quadratically with sequence length n. This O(n^2) cost is the central pressure behind later work on efficient attention, key-value caching during generation, and long-context techniques, but it is the price for removing the sequential bottleneck that limited recurrent models.

  • Encoder block: self-attention then feed-forward; decoder block adds masked self-attention and cross-attention.
  • Encoder-only (BERT), decoder-only (GPT), and encoder-decoder are the three common configurations.
  • Self-attention runs a constant number of sequential operations per layer, unlike an RNN's n-step chain.
  • The tradeoff is O(n^2) compute and memory in sequence length, driving efficient-attention research.

Key takeaways

  • A Transformer block stacks an attention sub-layer and a position-wise feed-forward network, each wrapped in a residual connection and layer normalization.
  • Scaled dot-product attention divides scores by sqrt(d_k) so large key dimensions do not saturate the softmax and crush its gradients.
  • Multi-head attention runs h smaller attention functions in parallel, each in its own subspace, then concatenates and projects them so the model attends to several patterns at once.
  • Because attention is permutation invariant, position is supplied explicitly through sinusoidal, learned, or rotary (RoPE) positional encodings.
  • The original design used post-LN, LayerNorm(x + Sublayer(x)); pre-LN moves normalization inside the residual branch and trains more stably without warmup.
  • Self-attention processes a whole sequence in parallel, which is why Transformers replaced RNNs, at the cost of O(n^2) compute in sequence length.

Frequently asked questions

For large key dimension d_k, dot products between queries and keys grow large in magnitude and push the softmax into regions with extremely small gradients. Dividing the scores by sqrt(d_k) keeps them at a stable scale so training stays well-conditioned. The choice matches the fact that the standard deviation of such a dot product grows like sqrt(d_k).
Splitting attention into h heads lets the model attend to information from different representation subspaces at different positions at the same time. One head can track grammatical agreement while another follows long-range references, patterns a single averaged head would blur together. The heads are concatenated and mixed by an output projection so the block produces one combined vector per position.
Self-attention treats its input as an unordered set, so it computes the same result no matter how the tokens are permuted. Without added position information the model could not tell "dog bites man" from "man bites dog". Positional encodings, whether fixed sinusoids, learned vectors, or rotary embeddings, inject the order the attention mechanism otherwise ignores.
Post-LN, the original arrangement, applies layer normalization after the residual addition as LayerNorm(x + Sublayer(x)). Pre-LN instead normalizes the input before each sub-layer, computing x + Sublayer(LayerNorm(x)). Pre-LN gives well-behaved gradients at initialization and lets deep models train without learning-rate warmup, which is why most large models use it.
Attention mixes information across positions by taking weighted averages of value vectors, but it applies no per-position nonlinear transform on its own. The feed-forward network runs the same two-layer MLP independently at every position, projecting up to a wider inner dimension (2048 in the base model) and back down. It supplies the per-position nonlinear capacity that attention lacks.
An RNN computes each hidden state from the previous one, forcing n sequential steps for a length-n sequence that cannot be parallelized. Self-attention computes all positions together as a few matrix multiplications, a constant number of sequential operations regardless of length. On accelerators built for dense linear algebra this trains far faster, at the cost of compute that grows quadratically with sequence length.

Put the idea into practice

MemX is an AI memory agent built on these ideas: store anything, skip the folders, and find it again by asking in plain English.

Try MemX Free