{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Anatomy of a Transformer's Tensor\n", "\n", "**Author:** Juan Simancas \n", "**Reference Paper:** *Attention Is All You Need* (Vaswani et al., 2017, [arXiv:1706.03762](https://arxiv.org/abs/1706.03762))\n", "\n", "---\n", "\n", "## Theoretical Motivation\n", "\n", "When studying the Transformer architecture, it is common to focus on the mathematical formulation of attention:\n", "\n", "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V$$\n", "\n", "However, implementing and understanding Transformers requires mastering **tensor geometry**: tracking the exact dimensional transformations that reshape continuous word representations into multi-head query, key, and value matrices.\n", "\n", "This notebook provides a minimal, self-contained educational implementation that walks step-by-step through every tensor transformation in PyTorch, highlighting:\n", "\n", "1. The initial coordinate system: $(B, T, C)$\n", "2. The fused linear projection: $C \\longrightarrow 3C$\n", "3. Separating roles into $Q$, $K$, and $V$\n", "4. The divisibility constraint: $C = H \\cdot D$\n", "5. Factoring channels into heads via `view`: $(B, T, C) \\longrightarrow (B, T, H, D)$\n", "6. Aligning heads before sequence positions via `transpose`: $(B, T, H, D) \\longrightarrow (B, H, T, D)$\n", "7. Computing pairwise attention score matrices: $(B, H, T, T)$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import math\n", "import torch\n", "import torch.nn as nn\n", "\n", "# Set random seed for reproducible tensor values\n", "torch.manual_seed(42)\n", "print(f\"PyTorch version: {torch.__version__}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. The Input Coordinate System: $(B, T, C)$\n", "\n", "After tokenization, vocabulary mapping, and the addition of positional embeddings, input text enters the Transformer as a 3D tensor:\n", "\n", "$$X \\in \\mathbb{R}^{B \\times T \\times C}$$\n", "\n", "- **$B$ (Batch Size):** Number of independent sequences processed simultaneously.\n", "- **$T$ (Sequence Length / Time):** Number of token positions in each sequence (denoted $n$ in Vaswani et al., 2017).\n", "- **$C$ (Channels / Embedding Dimension):** Feature vector size per token (denoted $d_{\\text{model}}$ in Vaswani et al., 2017)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Educational dimensions (kept small for clear inspection)\n", "B = 2 # 2 sentences in the batch\n", "T = 5 # 5 tokens per sentence\n", "C = 12 # 12 embedding channels per token (d_model)\n", "\n", "# Create synthetic input representation\n", "X = torch.randn(B, T, C)\n", "\n", "print(f\"Input tensor X shape: {X.shape} -> (Batch={B}, Time={T}, Channels={C})\")\n", "print(f\"Total numerical values: {X.numel()}\")\n", "print(f\"Vector for Sequence 0, Token 0 (first 4 channels): {X[0, 0, :4].round(decimals=3)}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. The Fused Linear Projection: $C \\longrightarrow 3C$\n", "\n", "To prepare for attention, each token must generate three representations:\n", "- **Query ($Q$):** What this token is seeking.\n", "- **Key ($K$):** What this token contains (used for matching).\n", "- **Value ($V$):** The contextual information carried by this token.\n", "\n", "Rather than executing three separate matrix multiplications ($XW^Q$, $XW^K$, $XW^V$), high-performance implementations fuse them into a single linear layer:\n", "\n", "$$W_{\\text{qkv}} \\in \\mathbb{R}^{C \\times 3C}$$\n", "\n", "$$QKV = X W_{\\text{qkv}} \\in \\mathbb{R}^{B \\times T \\times 3C}$$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Single fused projection mapping C -> 3C\n", "qkv_projection = nn.Linear(C, 3 * C, bias=False)\n", "\n", "# Project input tensor\n", "QKV = qkv_projection(X)\n", "\n", "print(f\"Projected QKV shape: {QKV.shape} -> (B={B}, T={T}, 3C={3 * C})\")\n", "print(\"Notice: Axes B and T remain untouched. The transformation occurs purely along the channel axis.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Separating Roles: $Q$, $K$, and $V$\n", "\n", "The resulting $3C$ channels represent $Q$, $K$, and $V$ stacked side by side. We partition this axis into three equal chunks of size $C$:\n", "\n", "$$Q, K, V \\in \\mathbb{R}^{B \\times T \\times C}$$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Split the last dimension (-1) into 3 equal chunks of size C\n", "Q, K, V = torch.chunk(QKV, chunks=3, dim=-1)\n", "\n", "print(f\"Q shape: {Q.shape} -> (B={B}, T={T}, C={C})\")\n", "print(f\"K shape: {K.shape} -> (B={B}, T={T}, C={C})\")\n", "print(f\"V shape: {V.shape} -> (B={B}, T={T}, C={C})\")\n", "\n", "# Verification: chunking is identical to contiguous slicing\n", "assert torch.equal(Q, QKV[:, :, :C])\n", "assert torch.equal(K, QKV[:, :, C:2*C])\n", "assert torch.equal(V, QKV[:, :, 2*C:])\n", "print(\"[OK] Slicing check passed: Q, K, and V correctly isolated.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. The Divisibility Constraint: $C = H \\cdot D$\n", "\n", "Multi-head attention divides the feature space into $H$ orthogonal subspaces (heads), each with dimension $D$ (denoted $d_k$ in Vaswani et al., 2017):\n", "\n", "$$C = H \\cdot D \\implies D = \\frac{C}{H}$$\n", "\n", "If $C$ is not divisible by $H$, uniform integer partitioning is mathematically impossible." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "H = 3 # Number of attention heads (h = 3)\n", "\n", "# Check divisibility\n", "assert C % H == 0, f\"Channels C={C} must be divisible by heads H={H}\"\n", "D = C // H # Dimension per head (d_k)\n", "\n", "print(f\"Divisibility condition satisfied: C ({C}) = H ({H}) * D ({D})\")\n", "\n", "print(\"\\n--- Failure Diagnostic Test ---\")\n", "invalid_C = 10\n", "invalid_H = 3\n", "test_tensor = torch.randn(B, T, invalid_C)\n", "\n", "try:\n", " # Attempting to reshape 10 channels into 3 integer heads\n", " test_tensor.view(B, T, invalid_H, -1)\n", "except RuntimeError as e:\n", " print(f\"Captured expected RuntimeError: {e}\")\n", " print(f\"Explanation: 10 numbers cannot be partitioned into 3 heads of equal integer size.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. Factoring Channels into Subspaces via `view`: $(B, T, H, D)$\n", "\n", "We now reshape the final axis from $C$ into $(H, D)$:\n", "\n", "$$(B, T, C) \\longrightarrow (B, T, H, D)$$\n", "\n", "This operation reinterprets the channel indices as belonging to $H$ separate heads of size $D$ without moving numbers in memory." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Reshape channel axis into (Heads, Head_Dim)\n", "Q_view = Q.view(B, T, H, D)\n", "K_view = K.view(B, T, H, D)\n", "V_view = V.view(B, T, H, D)\n", "\n", "print(f\"Q_view shape: {Q_view.shape} -> (B={B}, T={T}, H={H}, D={D})\")\n", "print(f\"Total values preserved: {Q.numel()} == {Q_view.numel()}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6. The Alignment Transposition: $(B, H, T, D)$\n", "\n", "The tensor is currently structured as $(B, T, H, D)$: for each token position, we have its $H$ heads.\n", "\n", "However, attention compares **all token positions across the sequence** within the **same head**. To enable batched matrix multiplication over all heads simultaneously, we swap axis 1 ($T$) and axis 2 ($H$):\n", "\n", "$$(B, T, H, D) \\xrightarrow{\\text{transpose}(1, 2)} (B, H, T, D)$$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Swap sequence length (axis 1) and heads (axis 2)\n", "Q_heads = Q_view.transpose(1, 2)\n", "K_heads = K_view.transpose(1, 2)\n", "V_heads = V_view.transpose(1, 2)\n", "\n", "print(f\"Q_heads shape: {Q_heads.shape} -> (B={B}, H={H}, T={T}, D={D})\")\n", "print(f\"Memory contiguous after transpose? {Q_heads.is_contiguous()}\")\n", "print(\"Note: Transposing changes logical strides. The accelerator now treats (B, H) as a batch of 2D matrices.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 7. Computing Pairwise Attention Scores: $(B, H, T, T)$\n", "\n", "With the tensor shaped as $(B, H, T, D)$, computing attention is a batched matrix multiplication between queries and transposed keys:\n", "\n", "$$\\text{Scores} = \\frac{Q K^T}{\\sqrt{D}}$$\n", "\n", "$$(B, H, T, D) \\times (B, H, D, T) \\longrightarrow (B, H, T, T)$$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Transpose the last two dimensions of K to prepare for dot products\n", "K_transposed = K_heads.transpose(-2, -1)\n", "print(f\"K_transposed shape: {K_transposed.shape} -> (B={B}, H={H}, D={D}, T={T})\")\n", "\n", "# Batched matrix multiplication scaled by sqrt(D)\n", "scores = torch.matmul(Q_heads, K_transposed) / math.sqrt(D)\n", "print(f\"Scores shape: {scores.shape} -> (Batch={B}, Heads={H}, Queries={T}, Keys={T})\")\n", "\n", "# Compute normalized attention distribution\n", "attention_weights = torch.softmax(scores, dim=-1)\n", "print(f\"Attention weights shape: {attention_weights.shape}\")\n", "\n", "print(\"\\n--- Attention Grid for Sequence 0, Head 0 (Rows sum to 1.0) ---\")\n", "print(attention_weights[0, 0].round(decimals=3))\n", "print(f\"Row sums check: {attention_weights[0, 0].sum(dim=-1)}\")" ] }, { "cell_type": "markdown", "id": "6ed7c338", "metadata": {}, "source": [ "## 8. Complete Shape Audit Summary\n", "\n", "| Step | Operation | Input Shape | Output Shape | What Changed | What Was Preserved |\n", "| :--- | :--- | :--- | :--- | :--- | :--- |\n", "| **1** | Token + Positional Embedding | Tokens | $(B, T, C)$ | Defined coordinate grid | $B$ examples, $T$ tokens, $C$ features |\n", "| **2** | Fused Linear Projection | $(B, T, C)$ | $(B, T, 3C)$ | $C \\longrightarrow 3C$ channels | $B$ and $T$ |\n", "| **3** | Split into Roles | $(B, T, 3C)$ | $3 \\times (B, T, C)$ | $3C \\longrightarrow C + C + C$ | $B$ and $T$ |\n", "| **4** | Factor Heads via `view` | $(B, T, C)$ | $(B, T, H, D)$ | $C \\longrightarrow H \\times D$ | $B$ and $T$ |\n", "| **5** | Reorder via `transpose(1, 2)` | $(B, T, H, D)$ | $(B, H, T, D)$ | Swapped order of $T$ and $H$ | Numerical values and $B, D$ |\n", "| **6** | Batched Dot-Product Attention | $(B, H, T, D) \\times (B, H, D, T)$ | $(B, H, T, T)$ | Scaled dot product over $D$ | $B, H$ batch sub-spaces |\n", "\n", "
\n", " \"Anatomy\n", "
\n", "\n", "*Every axis has a clear physical meaning. Tracking these dimensions transforms attention from an opaque algorithm into predictable matrix calculus.*" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.12" } }, "nbformat": 4, "nbformat_minor": 5 }