Transformers & Self-Attention Mechanisms

1. Scaled Dot-Product Attention Mechanics

Given an input sequence mapped to Query (Q), Key (K), and Value (V) matrices of dimension d_k, the attention weights represent the similarity of each token to every other token:

2. Accelerated Attention: F.scaled_dot_product_attention

Standard naive self-attention materializes the full $N \times N$ attention matrix in GPU High-Bandwidth Memory (HBM), consuming $O(N^2)$ memory and creating an I/O memory bottleneck. PyTorch 2.0+ integrates FlashAttention-2 directly intotorch.nn.functional.scaled_dot_product_attention (SDPA):

import torch
import torch.nn.functional as F

# Batch=4, Heads=8, Sequence Length=1024, Head Dimension=64
B, H, S, D = 4, 8, 1024, 64
q = torch.randn(B, H, S, D, device="cuda", dtype=torch.float16)
k = torch.randn(B, H, S, D, device="cuda", dtype=torch.float16)
v = torch.randn(B, H, S, D, device="cuda", dtype=torch.float16)

# Hardware-accelerated FlashAttention-2 with fused CUDA kernels
# is_causal=True automatically constructs and applies the autoregressive triangular mask!
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
    output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

print(f"SDPA Output Shape: {output.shape}")  # torch.Size([4, 8, 1024, 64])

3. Building a Multi-Head Attention Layer from Scratch

Here is a pure PyTorch implementation of Multi-Head Attention showing projection splits, head permutations, and causal masking:

import torch
import torch.nn as nn
import math

class CausalSelfAttention(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int):
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        # Key, Query, Value projected in a single efficient matrix
        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, S, C = x.shape  # Batch, Sequence Length, Channels (Embedding Dim)

        # 1. Project Q, K, V in one GEMM: [B, S, 3 * C]
        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)

        # 2. Reshape into [B, Num_Heads, S, Head_Dim]
        q = q.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)

        # 3. Fast Causal Scaled Dot-Product Attention
        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

        # 4. Concatenate heads back: [B, S, C]
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, S, C)
        
        # 5. Output projection
        return self.out_proj(attn_out)

4. Assembling a Modern GPT Transformer Block

Modern architectures (LLaMA, GPT-4, Mistral) employ Pre-LayerNorm (normalizing inputs before attention and MLP layers) for vastly superior training stability compared to original Post-LN Transformers:

class TransformerBlock(nn.Module):
    """
    Standard Pre-LN Autoregressive Transformer Decoder Block.
    """
    def __init__(self, embed_dim: int, num_heads: int, mlp_ratio: int = 4):
        super().__init__()
        self.ln1 = nn.LayerNorm(embed_dim)
        self.attn = CausalSelfAttention(embed_dim, num_heads)
        
        self.ln2 = nn.LayerNorm(embed_dim)
        # Position-wise Feedforward Network with GELU activation
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, mlp_ratio * embed_dim),
            nn.GELU(),
            nn.Linear(mlp_ratio * embed_dim, embed_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-LN Residual Attention Stream
        x = x + self.attn(self.ln1(x))
        # Pre-LN Residual MLP Stream
        x = x + self.mlp(self.ln2(x))
        return x

# Instantiate & Verify Block
block = TransformerBlock(embed_dim=256, num_heads=4)
sample_tokens = torch.randn(2, 64, 256)  # Batch=2, Seq=64, Dim=256
block_output = block(sample_tokens)
print(f"Transformer Block Output Shape: {block_output.shape}")  # [2, 64, 256]