Transformer Architecture, Self-Attention & Decoder Language Models

1. Scaled Dot-Product Attention Mechanics

The Transformer architecture (Vaswani et al. 2017) discards sequential recurrence entirely, computing global contextual dependencies across all sequence tokens simultaneously via Self-Attention.

Given an input sequence representation $\mathbf{X} \in \mathbb{R}^{T \times d_{\text{model}}}$, linear projection matrices $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V \in \mathbb{R}^{d_{\text{model}} \times d_k}$ generate Queries ($\mathbf{Q}$), Keys ($\mathbf{K}$), and Values ($\mathbf{V}$):

The Scaled Dot-Product Attention computes similarity coefficients across all token pairs:

where $\mathbf{M} \in \mathbb{R}^{T \times T}$ is an optional causal mask where $M_{i,j} = -\infty$ for $j > i$.

Why Scale by $\frac{1}{\sqrt{d_k}}$?

Assume components of $\mathbf{q}_i, \mathbf{k}_j \in \mathbb{R}^{d_k}$ are independent random variables with zero mean and unit variance ($\mathbb{E}[q_m] = \mathbb{E}[k_m] = 0, \; \text{Var}(q_m) = \text{Var}(k_m) = 1$). The dot product $s = \mathbf{q}_i^T \mathbf{k}_j = \sum_{m=1}^{d_k} q_{i,m} k_{j,m}$ has variance:

For large embedding dimensions ($d_k \ge 64$), unscaled dot products exhibit extremely high variance ($\sigma = \sqrt{d_k}$), pushing arguments of the Softmax function into saturation regions where gradient magnitudes vanish ($\frac{\partial \text{softmax}(z_i)}{\partial z_j} \approx 0$). Scaling by $\frac{1}{\sqrt{d_k}}$ restores unit variance ($\text{Var}(\frac{s}{\sqrt{d_k}}) = 1$).

2. Multi-Head Attention & Layer Normalization Topology

2.1 Multi-Head Attention (MHA)

To capture heterogeneous relationship subspaces simultaneously, MHA projects queries, keys, and values into $H$ orthogonal attention heads:

     Multi-Head Attention & Pre-LN Transformer Block
     
                x_{l+1}
                  ▲
                  │
                 (+) Skip Connection
                  ▲
                  │
               [MLP / FFN]
                  ▲
                  │
              [LayerNorm]
                  ▲
                  │
                 (+) Skip Connection
                  ▲
                  │
        [Multi-Head Attention]
                  ▲
                  │
              [LayerNorm]
                  ▲
                  │
                 x_l

3. Positional Encodings: Sinusoidal vs. Rotary (RoPE)

Since attention operates as a permutation-equivariant set operator, explicit temporal order must be injected into token representations.

3.1 Absolute Sinusoidal Encoding

For token index $t$ and embedding dimension $i$:

3.2 Rotary Position Embedding (RoPE - Su et al. 2021)

RoPE encodes relative token distance directly into query-key inner products by rotating complex 2D embedding pairs by angle $m \theta_i$:

4. FlashAttention: IO-Aware Hardware Tiling

Standard attention computes and materializes the $N \times N$ attention matrix $\mathbf{S} = \mathbf{Q}\mathbf{K}^T$ in GPU High-Bandwidth Memory (HBM), incurring $O(N^2)$ memory reads/writes.

FlashAttention (Dao et al. 2022) computes exact softmax attention using SRAM tiling and online softmax rescaling, eliminating $N \times N$ HBM instantiation and reducing memory complexity from $O(N^2)$ to $O(N)$.

5. PyTorch Implementation: Canonical Pre-LN Causal Transformer Decoder Block

Below is a complete PyTorch implementation of a Canonical Causal Transformer Decoder Block (Pre-LN architecture) standard for modern Large Language Models (LLMs).

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

class CausalSelfAttention(nn.Module):
    """Multi-Head Causal Self-Attention with Pre-LN formulation."""
    def __init__(self, d_model: int = 256, n_heads: int = 8, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        
        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)
        
        # Reshape to (Batch, Heads, SeqLen, HeadDim)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        
        # Scaled dot-product attention with lower-triangular causal mask
        scale = 1.0 / math.sqrt(self.head_dim)
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale
        
        causal_mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
        attn_scores = attn_scores.masked_fill(~causal_mask, float('-inf'))
        
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        context = torch.matmul(attn_weights, v)  # (Batch, Heads, SeqLen, HeadDim)
        context = context.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(context)

class TransformerDecoderBlock(nn.Module):
    """Canonical Pre-LayerNorm Transformer Decoder Block."""
    def __init__(self, d_model: int = 256, n_heads: int = 8, mlp_ratio: int = 4, dropout: float = 0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        
        mlp_hidden = d_model * mlp_ratio
        self.mlp = nn.Sequential(
            nn.Linear(d_model, mlp_hidden),
            nn.GELU(),
            nn.Linear(mlp_hidden, d_model),
            nn.Dropout(dropout)
        )
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

# Verification: Forward pass on sequence of token embeddings
if __name__ == "__main__":
    sample_embeddings = torch.randn(8, 64, 256)  # Batch=8, Sequence=64, d_model=256
    block = TransformerDecoderBlock(d_model=256, n_heads=8)
    
    output_representation = block(sample_embeddings)
    print("Input Sequence Embedding Shape:", sample_embeddings.shape)
    print("Output Transformer Block Representation Shape:", output_representation.shape)