Inside a Transformer Block

The Transformer Block: A Modular Architecture

A transformer block is a reusable unit stacked many times to build models like BERT, GPT, and modern language models. Unlike RNNs, which process sequences step-by-step, a transformer processes an entire sequence in parallel. Each block contains four main components:

1. Self-attention: allows each position to attend to all other positions, computing a weighted combination of values.
2. Multi-head attention: runs multiple self-attention heads in parallel, each focusing on different representation subspaces.
3. Residual connections and layer norm: Add & Norm layers that stabilize deep networks.
4. Position-wise feed-forward: two fully connected layers (with a hidden layer in between) applied independently to each position.

The design is symmetric and modular: stack N identical blocks, and you have a deep transformer.

Scaled Dot-Product Attention

Self-attention computes a weighted sum of values (V) based on the similarity between queries (Q) and keys (K). The scaled dot-product attention formula is:

Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V

Breaking down each component:

1. Similarity computation: Q · K^T

• Q (query): shape (seq_len, d_k) — what each position is "looking for"
• K (key): shape (seq_len, d_k) — what each position "offers"
• Q · K^T: (seq_len, seq_len) matrix where element [i, j] = how much position i matches position j

For example, if Q[0] = [1, 0] and K[2] = [1, 0], their dot product is 1 (high similarity).

2. Scaling by √d_k:

scores = Q · K^T / √d_k

Why divide by √d_k? Without scaling, dot products can become very large (if d_k = 512 and all values are ~1, the sum is ~512). Large scores cause softmax to produce very sharp distributions (nearly one-hot), which limits gradient flow during backpropagation. Scaling by √d_k keeps scores in a reasonable range, typically [-1, 1] assuming normalized inputs.

3. Softmax for normalization:

attention_weights = softmax(scores)

This converts scores to a probability distribution: each row sums to 1, and all values are in [0, 1].

4. Weighted value aggregation:

output = attention_weights · V

Each position's output is a weighted combination of all values, where weights come from attention_weights. Position i's output is:

output[i] = Σ_j attention_weights[i, j] · V[j]

This means position i attends to all positions, but allocates most weight to similar positions (high attention_weights[i, j] if position j has a key matching position i's query).

At each position i, attention_weights[i, :] is a probability distribution over all positions j. High attention weights mean position i "looks at" position j more strongly. This is learned: Q, K, V come from learned linear projections of the input.

Implementing Scaled Dot-Product Attention

Here is a numpy implementation to see the mechanics:

import numpy as np

def scaled_dot_product_attention(query, key, value, mask=None):
    """
    Compute scaled dot-product attention.

    Args:
        query: (seq_len, d_k)
        key: (seq_len, d_k)
        value: (seq_len, d_k)
        mask: optional (seq_len, seq_len) boolean mask (True = positions to ignore)

    Returns:
        output: (seq_len, d_k) weighted value sum
        attention_weights: (seq_len, seq_len) softmax weights
    """
    d_k = query.shape[-1]

    # Compute similarity scores
    scores = np.matmul(query, key.T) / np.sqrt(d_k)
    # scores: (seq_len, seq_len)

    # Apply mask if provided (useful for causal/autoregressive masks)
    if mask is not None:
        scores = np.where(mask, -np.inf, scores)

    # Softmax to get attention weights
    attention_weights = np.exp(scores - scores.max(axis=1, keepdims=True))
    # Subtract max for numerical stability (softmax is invariant to adding constants)
    attention_weights = attention_weights / attention_weights.sum(axis=1, keepdims=True)
    # attention_weights: (seq_len, seq_len)

    # Weighted combination of values
    output = np.matmul(attention_weights, value)
    # output: (seq_len, d_k)

    return output, attention_weights

# Example: 4-token sequence, d_k=2
seq_len, d_k = 4, 2

# Toy Q, K, V (in practice, these are learned projections of input)
query = np.array([
    [1.0, 0.0],
    [0.5, 0.5],
    [0.0, 1.0],
    [0.5, 1.0]
])

key = np.array([
    [1.0, 0.0],
    [0.5, 0.5],
    [0.0, 1.0],
    [0.5, 1.0]
])

value = np.array([
    [2.0, 0.0],
    [1.5, 1.5],
    [0.0, 2.0],
    [1.5, 2.0]
])

output, weights = scaled_dot_product_attention(query, key, value)
print("Attention weights shape:", weights.shape)  # (4, 4)
print("Attention weights (row 0, where does position 0 attend?):")
print(weights[0, :])
# Should be softmax of similarities [1.0, 0.7, 0.0, 1.4] / sqrt(2)
# Position 0 will attend more strongly to positions with similar queries

print("\nOutput shape:", output.shape)  # (4, 2)
print("Output (position 0):", output[0, :])
# Weighted combination of all values, with weights from attention_weights[0, :]

The key insight: each position independently computes attention weights over all positions, then computes a context vector (weighted average of values). Positions with similar queries and keys get high attention.

Multi-Head Attention

Instead of computing attention once, compute it h times in parallel on different subspaces:

def multi_head_attention(query, key, value, num_heads, d_model, mask=None):
    """
    Multi-head attention.

    Args:
        query, key, value: (seq_len, d_model)
        num_heads: number of attention heads
        d_model: total dimension (must be divisible by num_heads)
        mask: optional mask

    Returns:
        output: (seq_len, d_model) concatenated head outputs
        weights: list of (seq_len, seq_len) attention weights for each head
    """
    d_k = d_model // num_heads
    batch_heads = []
    weights_list = []

    for head in range(num_heads):
        # Project to subspace for this head (simplified: just slice)
        q_head = query[:, head*d_k:(head+1)*d_k]
        k_head = key[:, head*d_k:(head+1)*d_k]
        v_head = value[:, head*d_k:(head+1)*d_k]

        # Apply attention to this subspace
        head_output, head_weights = scaled_dot_product_attention(
            q_head, k_head, v_head, mask
        )
        batch_heads.append(head_output)
        weights_list.append(head_weights)

    # Concatenate all head outputs
    output = np.concatenate(batch_heads, axis=1)  # (seq_len, d_model)
    return output, weights_list

# Example: 4 tokens, d_model=4, 2 heads (each head operates on d_k=2)
mha_output, mha_weights = multi_head_attention(
    query, key, value, num_heads=2, d_model=4
)
print("Multi-head output shape:", mha_output.shape)  # (4, 4)
print(f"Number of attention weight matrices: {len(mha_weights)}")  # 2

Each head learns to attend to different aspects. One head might focus on syntax, another on semantics. Concatenating them gives a richer representation than a single attention head.

Residual Connections and Layer Normalization

A transformer block computes:

x_attn = MultiHeadAttention(x)
x_after_attn = x + x_attn                           # Residual add
x_after_norm1 = LayerNorm(x_after_attn)              # Normalize

x_ffn = FeedForward(x_after_norm1)
x_after_ffn = x_after_norm1 + x_ffn                 # Another residual add
x_output = LayerNorm(x_after_ffn)                    # Final normalize

(Modern transformers often use "pre-norm" instead, normalizing before the sublayer: x_attn = MultiHeadAttention(LayerNorm(x)), which is slightly different but achieves similar effect.)

Residual connections (skip connections)

The + operation is a residual connection or skip connection. Why is this critical?

Without residuals, gradients backpropagate through many matrix multiplications:

∂L/∂x^(0) = (∂L/∂x^(n)) · (∂x^(n)/∂x^(n-1)) · ... · (∂x^(1)/∂x^(0))

This is a product of n Jacobian matrices. If each Jacobian has singular values < 1, the product shrinks exponentially — vanishing gradient. This is why RNNs struggle with long sequences, and why pre-LSTM networks struggled with deep layers.

Residual connections introduce an additive path:

∂L/∂x^(0) includes: ... + (∂L/∂x^(n)) · (identity matrix) = (∂L/∂x^(n))

The gradient for the residual term is the identity, so gradients can flow directly through. This allows training of very deep networks (BERT has 12+ layers, GPT-3 has 96 layers).

Layer normalization

Layer normalization (LayerNorm) normalizes each feature independently across the sequence:

x_normalized[i, j] = (x[i, j] - mean_j(x[:, j])) / √(var_j(x[:, j]) + ε)

where the mean and variance are computed across the sequence dimension for each feature j.

Why normalize?

• Gradient stability: Normalizes activations to zero mean and unit variance, preventing internal covariate shift (activation distributions drifting during training).
• Training speed: Allows use of higher learning rates.
• No batch dependence: Unlike batch norm, layer norm doesn't depend on batch statistics, making it suitable for variable-length sequences and inference.

Transformers typically normalize after adding residuals, ensuring stable gradients through deep stacks.

Comparison: Residual connections vs. other approaches

| Component | Benefit | Cost | Use Case | |---|---|---|---| | Residual connections | Direct gradient flow; enables deep networks | Slight computational overhead (addition) | Essential for > 6 layers | | Layer normalization | Gradient stability; train faster; batch-independent | Slightly slower than batch norm; learnable affine parameters | All modern transformers | | Batch normalization | Gradient stability; can reduce learning rate | Depends on batch statistics; issues with small batches | CNNs, some RNNs (less in transformers) | | Dropout | Regularization; prevents overfitting | Slower training; need proper tuning (p=0.1-0.5) | Large models; overfitting-prone tasks |

# PyTorch layer norm
ln = np.nn.LayerNorm(normalized_shape=4)
# Normalizes the last dimension (d_model = 4)

x = np.random.randn(4, 4)  # (seq_len, d_model)
x_norm = ln(x)
# Each of the 4 positions is normalized independently

A transformer block's forward pass is:

def transformer_block_forward(x, attn_layer, ffn_layer, norm1, norm2):
    # Self-attention with residual and norm
    attn_out = attn_layer(x)
    x = norm1(x + attn_out)  # "Add & Norm"

    # Feed-forward with residual and norm
    ffn_out = ffn_layer(x)
    x = norm2(x + ffn_out)

    return x

The Position-Wise Feed-Forward Network

After attention, each position independently passes through the same two-layer feed-forward network:

FFN(x) = max(0, x * W_1 + b_1) * W_2 + b_2

• W_1: (d_model, d_ff), e.g., (768, 3072)
• W_2: (d_ff, d_model), e.g., (3072, 768)
• ReLU (or GELU) activations introduce nonlinearity

This is applied independently at each position, so it is very parallel.

def position_wise_ffn(x, d_model, d_ff):
    """
    Args:
        x: (seq_len, d_model)

    Returns:
        (seq_len, d_model)
    """
    # Expand to d_ff (e.g., 4*d_model)
    x_expanded = np.dot(x, np.random.randn(d_model, d_ff))
    x_expanded = np.maximum(0, x_expanded)  # ReLU

    # Project back to d_model
    output = np.dot(x_expanded, np.random.randn(d_ff, d_model))

    return output

# Example
x_ffn = position_wise_ffn(mha_output, d_model=4, d_ff=16)
print("FFN output shape:", x_ffn.shape)  # (4, 4)

PyTorch Transformer

PyTorch provides efficient implementations:

import torch
import torch.nn as nn

# A single transformer block
transformer_block = nn.TransformerEncoderLayer(
    d_model=256,
    nhead=8,  # 8 attention heads
    dim_feedforward=1024,  # Feed-forward hidden size
    dropout=0.1,
    batch_first=True  # (batch, seq, features)
)

# Input: (batch, seq_len, d_model)
x = torch.randn(2, 10, 256)

# Forward pass through one block
output = transformer_block(x)
print(f"Input shape: {x.shape}")   # (2, 10, 256)
print(f"Output shape: {output.shape}")  # (2, 10, 256)

# Stack multiple blocks
transformer_encoder = nn.TransformerEncoder(
    encoder_layer=transformer_block,
    num_layers=6  # 6 blocks stacked
)

deep_output = transformer_encoder(x)
print(f"Deep output shape: {deep_output.shape}")  # (2, 10, 256)

A 6-layer transformer encoder (like the encoder half of the original "Attention Is All You Need" model) is built by stacking 6 identical blocks.

Why Transformers Parallelize Better Than RNNs

RNNs process sequences step-by-step: h_t depends on h_, so computation is inherently sequential. Training an RNN on a 1000-token sequence requires 1000 sequential forward passes, then 1000 backward passes. GPUs cannot parallelize this well because they must wait for each previous hidden state.

Transformers process all positions simultaneously: each position computes attention to all other positions in parallel. A transformer processes a 1000-token sequence in one parallel forward pass. This is much faster on modern hardware (GPUs, TPUs) and enables training on much larger datasets.

Complexity comparison: RNN vs. Transformer

| Metric | RNN (vanilla) | LSTM/GRU | Transformer | |---|---|---|---| | Sequential depth | O(seq_len) | O(seq_len) | O(1) (fully parallel) | | Time complexity per step | O(hidden² + hidden×input) | O(4×hidden² + hidden×input) (4 gates) | O(seq_len² × d_model) | | Total training time for seq_len=T | O(T × hidden²) | O(T × hidden²) | O(seq_len² × d_model) = O(T² × d_model) | | Memory (forward pass) | O(seq_len × hidden) | O(seq_len × hidden) | O(seq_len² × d_model) |

For short sequences (T < 100):

• RNN: ~1000× faster per step; total time is reasonable
• Transformer: slower per step but fully parallel; total time is similar or faster due to parallelization

For long sequences (T = 1000):

• RNN: sequential; trains slowly, requires many iterations; still uses O(T × hidden) memory
• Transformer: O(T²) memory and compute become prohibitive; but data parallelism across TPUs/GPUs is efficient

Real-world impact on training time

Training a model with 12 transformer blocks on a 512-token sequence with batch size 32:

• RNN (optimistic estimate): 512 sequential steps × 12 layers × complex computations = ~100,000 operations per token sequentially. With GPU parallelization across batch, roughly 3-5s per batch.
• Transformer: All 512 tokens processed in parallel; (512² × 768 × 12) / (GPU FLOPs) ≈ 1-2s per batch (higher FLOPs but greater parallelism).

Transformers become 2-3× faster for long sequences due to parallelization, despite higher theoretical complexity.

The trade-off: transformers use O(seq_len²) memory and compute for self-attention (the (seq_len, seq_len) similarity matrix). For very long sequences (>10k tokens), this becomes problematic, consuming 100GB+ GPU memory. Solutions include:

• Linear attention: approximate attention with lower complexity (e.g., Performer, BigBird)
• Sparse attention: attend only to nearby positions (e.g., local attention in Longformer)
• Hierarchical methods: summarize earlier tokens then attend (e.g., Reformer)

A Minimal Transformer Classifier

Here is a tiny transformer-based text classifier:

class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers, output_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=d_model*4,
                dropout=0.1,
                batch_first=True
            ),
            num_layers=num_layers
        )
        self.fc = nn.Linear(d_model, output_dim)

    def forward(self, x):
        """
        Args:
            x: (batch, seq_len) of token indices

        Returns:
            logits: (batch, output_dim)
        """
        # Embed tokens
        embedded = self.embedding(x)  # (batch, seq_len, d_model)

        # Pass through transformer
        transformed = self.transformer(embedded)  # (batch, seq_len, d_model)

        # Mean pooling over sequence
        pooled = transformed.mean(dim=1)  # (batch, d_model)

        # Classify
        logits = self.fc(pooled)  # (batch, output_dim)

        return logits

# Example: vocab=1000, d_model=128, 2 attention heads, 2 layers, binary classification
classifier = TransformerClassifier(
    vocab_size=1000,
    d_model=128,
    nhead=2,
    num_layers=2,
    output_dim=2
)

x = torch.randint(0, 1000, (8, 20))  # Batch of 8, sequence length 20
logits = classifier(x)
print(f"Logits shape: {logits.shape}")  # (8, 2)

# Check model size
total_params = sum(p.numel() for p in classifier.parameters())
print(f"Total parameters: {total_params}")

This is a tiny model (around 450k parameters) compared to BERT (110M+), but it demonstrates the full architecture: embeddings, transformer blocks, and a classification head.

Common mistake: not handling variable-length sequences

Transformers can handle variable-length sequences with padding and attention masks. If you have sequences of different lengths in a batch, pad them to the same length and pass a mask to the transformer telling it to ignore padding tokens:

# Example: batch with different lengths, padded to max_len=20
x_padded = torch.randint(0, 1000, (8, 20))
padding_mask = x_padded == 0  # Assuming 0 is the padding token

output = transformer_encoder(x_padded, src_key_padding_mask=padding_mask)
# The mask tells the transformer to ignore padding positions

Another mistake: forgetting that self-attention is all-to-all. If you want a causal (autoregressive) transformer where position i can only attend to positions <= i, you need to pass an attn_mask (a (seq_len, seq_len) boolean matrix):

causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
# Upper triangular matrix: True (ignore) for positions j > i

output = transformer_encoder(x, src_mask=~causal_mask)
# Note: True = ignore, so we use ~causal_mask

This is critical for language model pretraining (next token prediction), where the model must not see future tokens during training.