Transformer Architecture

1. Overview

The Transformer (Vaswani et al., 2017) replaced recurrence with pure attention, enabling unprecedented parallelisation and long-range dependency modelling. It is the foundation of modern NLP, vision, and multimodal systems.

2. Scaled Dot-Product Attention

2.1 Definition

Given queries , keys , and values :

where is applied row-wise.

2.2 Why Scaling by ?

For random vectors , the dot product has:

Without scaling, the variance grows with dimension, pushing softmax into regions with extremely small gradients. Dividing by normalises:

2.3 Complexity Analysis

Time complexity: for the matrix multiplication, plus for softmax.

Space complexity: for the attention matrix .

For self-attention (), this is .

2.4 Attention as Kernel Regression

The attention output can be viewed as kernel regression with a data-dependent kernel:

This is a non-parametric estimate of the value function, where the kernel is the softmax kernel.

3. Multi-Head Attention

3.1 Formulation

Multi-head attention projects queries, keys, and values into different subspaces and computes attention independently:

where

with projection matrices:

Typically .

3.2 Parameter Count

For :

3.3 Information-Theoretic Perspective

Multi-head attention allows the model to attend to different types of information simultaneously:

• Head diversity: Different heads learn different attention patterns (syntactic, semantic, positional)
• Rank decomposition: Each head operates in a -dimensional subspace
• Ensemble effect: Multiple heads provide redundancy and robustness

The attention patterns learned by different heads often specialise:

4. Positional Encoding

4.1 Sinusoidal Positional Encoding

The original Transformer uses fixed sinusoidal encodings:

Key property: The encoding of position can be expressed as a linear function of the encoding at position :

where is a rotation matrix.

4.2 Relative Positional Encoding

Shaw et al. (2018) introduce relative position biases:

where is a learnable relative position embedding.

4.3 RoPE (Rotary Position Embedding)

RoPE (Su et al., 2021) encodes position by rotating the query and key vectors:

where .

Key advantage: The attention score depends only on relative position:

5. Layer Normalisation

5.1 LayerNorm vs BatchNorm

LayerNorm:

where and .

BatchNorm:

where are computed over the batch dimension.

Key difference: LayerNorm is normalised over the feature dimension, BatchNorm over the batch dimension.

5.2 Pre-LN vs Post-LN

Post-LN (original):

Pre-LN (modern):

Pre-LN provides better gradient flow and enables training without warmup.

6. Feed-Forward Network

6.1 Standard FFN

where and , typically .

6.2 SwiGLU (PaLM, LLaMA)

where .

This introduces gating and improves performance, requiring adjustment of to maintain parameter count.

6.3 Complexity

For , this is , which dominates the transformer FLOPs.

7. Full Transformer Architecture

7.1 Encoder

Each encoder layer consists of:

1. Multi-head self-attention
2. Add & Norm (residual connection + LayerNorm)
3. Feed-forward network
4. Add & Norm

7.2 Decoder

Each decoder layer adds:

1. Masked multi-head self-attention (prevents attending to future tokens)
2. Multi-head cross-attention (attends to encoder output)

7.3 Parameter Count

For a transformer with layers, , and heads:

For :

8. Code Examples

8.1 Multi-Head Attention Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    """Multi-head attention with optional causal masking."""
    
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        self.dropout = nn.Dropout(dropout)
        self.scale = math.sqrt(self.d_k)
    
    def forward(self, query, key, value, mask=None):
        """
        Forward pass.
        
        Parameters
        ----------
        query : Tensor (batch, seq_len, d_model)
        key : Tensor (batch, seq_len, d_model)
        value : Tensor (batch, seq_len, d_model)
        mask : Tensor (batch, 1, 1, seq_len) or (batch, 1, seq_len, seq_len)
        
        Returns
        -------
        output : Tensor (batch, seq_len, d_model)
        attn_weights : Tensor (batch, num_heads, seq_len, seq_len)
        """
        batch_size = query.size(0)
        
        # Linear projections and reshape
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # Scaled dot-product attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # Apply attention to values
        context = torch.matmul(attn_weights, V)
        
        # Reshape and project
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.W_o(context)
        
        return output, attn_weights
    
    def compute_attention_patterns(self, query, key, value, mask=None):
        """Analyse attention patterns across heads."""
        _, attn_weights = self.forward(query, key, value, mask)
        
        patterns = {}
        for h in range(self.num_heads):
            head_attn = attn_weights[0, h]  # first batch element
            
            # Compute metrics
            patterns[f'head_{h}'] = {
                'entropy': -(head_attn * (head_attn + 1e-10).log()).sum(dim=-1).mean().item(),
                'max_attention': head_attn.max(dim=-1).values.mean().item(),
                'avg_distance': self._compute_avg_distance(head_attn),
            }
        
        return patterns
    
    def _compute_avg_distance(self, attn_matrix):
        """Compute average attention distance."""
        seq_len = attn_matrix.size(0)
        positions = torch.arange(seq_len, device=attn_matrix.device).float()
        
        # Weighted average of distances
        distances = torch.abs(positions.unsqueeze(0) - positions.unsqueeze(1))
        avg_dist = (attn_matrix * distances).sum(dim=-1).mean()
        
        return avg_dist.item()


# Example: Create and test multi-head attention
mha = MultiHeadAttention(d_model=512, num_heads=8)
x = torch.randn(2, 10, 512)  # batch=2, seq_len=10, d_model=512

output, attn_weights = mha(x, x, x)
print(f"Output shape: {output.shape}")
print(f"Attention shape: {attn_weights.shape}")

# Analyse patterns
patterns = mha.compute_attention_patterns(x, x, x)
for head, metrics in patterns.items():
    print(f"{head}: entropy={metrics['entropy']:.3f}, "
          f"max_attn={metrics['max_attention']:.3f}, "
          f"avg_dist={metrics['avg_distance']:.3f}")

8.2 Positional Encoding

class SinusoidalPositionalEncoding(nn.Module):
    """Fixed sinusoidal positional encoding."""
    
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """Add positional encoding to input."""
        return x + self.pe[:, :x.size(1)]


class RotaryPositionalEncoding(nn.Module):
    """Rotary Position Embedding (RoPE)."""
    
    def __init__(self, d_model, max_len=5000, base=10000):
        super().__init__()
        self.d_model = d_model
        
        # Compute rotation frequencies
        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
        self.register_buffer('inv_freq', inv_freq)
        
        # Precompute rotary embeddings
        self._build_cache(max_len)
    
    def _build_cache(self, max_len):
        """Precompute cos and sin caches."""
        t = torch.arange(max_len, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        
        self.register_buffer('cos_cached', emb.cos())
        self.register_buffer('sin_cached', emb.sin())
    
    def _rotate_half(self, x):
        """Rotate half the hidden dims."""
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)
    
    def forward(self, q, k, position_ids=None):
        """
        Apply rotary embedding to queries and keys.
        
        Parameters
        ----------
        q, k : Tensor (batch, num_heads, seq_len, d_k)
        position_ids : Tensor (batch, seq_len), optional
        
        Returns
        -------
        q_rot, k_rot : Tensor (batch, num_heads, seq_len, d_k)
        """
        seq_len = q.size(2)
        
        if position_ids is None:
            cos = self.cos_cached[:seq_len].unsqueeze(0).unsqueeze(0)
            sin = self.sin_cached[:seq_len].unsqueeze(0).unsqueeze(0)
        else:
            cos = self.cos_cached[position_ids].unsqueeze(1)
            sin = self.sin_cached[position_ids].unsqueeze(1)
        
        q_rot = q * cos + self._rotate_half(q) * sin
        k_rot = k * cos + self._rotate_half(k) * sin
        
        return q_rot, k_rot


# Example: Compare positional encodings
d_model = 64
seq_len = 100

# Sinusoidal
sin_pe = SinusoidalPositionalEncoding(d_model)
x = torch.randn(1, seq_len, d_model)
x_sin = sin_pe(x)

# RoPE
rope = RotaryPositionalEncoding(d_model)
q = torch.randn(1, 8, seq_len, d_model // 8)
k = torch.randn(1, 8, seq_len, d_model // 8)
q_rot, k_rot = rope(q, k)

print(f"Sinusoidal PE - input: {x.shape}, output: {x_sin.shape}")
print(f"RoPE - q: {q.shape}, q_rot: {q_rot.shape}")

8.3 Full Transformer Encoder

class TransformerEncoderLayer(nn.Module):
    """Single transformer encoder layer with pre-LN."""
    
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )
        
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        """
        Forward pass with pre-LN.
        
        Parameters
        ----------
        x : Tensor (batch, seq_len, d_model)
        mask : Tensor (batch, 1, 1, seq_len)
        
        Returns
        -------
        output : Tensor (batch, seq_len, d_model)
        """
        # Pre-LN self-attention
        residual = x
        x_norm = self.norm1(x)
        attn_out, _ = self.self_attn(x_norm, x_norm, x_norm, mask)
        x = residual + self.dropout1(attn_out)
        
        # Pre-LN FFN
        residual = x
        x_norm = self.norm2(x)
        ffn_out = self.ffn(x_norm)
        x = residual + self.dropout2(ffn_out)  # Note: dropout already in ffn
        
        return x


class TransformerEncoder(nn.Module):
    """Full transformer encoder."""
    
    def __init__(self, vocab_size, d_model, num_heads, d_ff, 
                 num_layers, max_len=5000, dropout=0.1):
        super().__init__()
        
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = SinusoidalPositionalEncoding(d_model, max_len)
        self.dropout = nn.Dropout(dropout)
        
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
        self.norm = nn.LayerNorm(d_model)
    
    def forward(self, x, mask=None):
        """
        Forward pass.
        
        Parameters
        ----------
        x : Tensor (batch, seq_len) - token indices
        mask : Tensor (batch, 1, 1, seq_len) - padding mask
        
        Returns
        -------
        output : Tensor (batch, seq_len, d_model)
        """
        # Embedding and positional encoding
        x = self.embedding(x) * math.sqrt(self.d_model)
        x = self.pos_encoding(x)
        x = self.dropout(x)
        
        # Encoder layers
        for layer in self.layers:
            x = layer(x, mask)
        
        x = self.norm(x)
        
        return x


# Example: Create transformer encoder
encoder = TransformerEncoder(
    vocab_size=30000,
    d_model=512,
    num_heads=8,
    d_ff=2048,
    num_layers=6
)

# Forward pass
x = torch.randint(0, 30000, (2, 100))  # batch=2, seq_len=100
output = encoder(x)

print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")

# Count parameters
total_params = sum(p.numel() for p in encoder.parameters())
print(f"Total parameters: {total_params:,}")

9. Summary

1. Scaled dot-product attention computes weighted averages of values based on query-key similarity, with complexity.

2. Multi-head attention allows simultaneous attention to different representation subspaces, improving model capacity.

3. Positional encoding injects sequence order information; RoPE is now standard for its relative position awareness.

4. Layer normalisation stabilises training; pre-LN is preferred for deep transformers.

5. Feed-forward networks provide the bulk of parameters and non-linearity, with SwiGLU becoming the modern standard.

References

• Vaswani, A., et al. (2017). Attention is all you need. NeurIPS.
• Shaw, P., et al. (2018). Self-attention with relative position representations. NAACL.
• Su, J., et al. (2021). RoFormer: Enhanced transformer with rotary position embedding. arXiv.
• Layer normalisation: Ba, J., Kiros, J., & Hinton, G. (2016). Layer normalization. arXiv.
• SwiGLU: Shazeer, N. (2020). GLU variants improve transformer. arXiv.