{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# The Functional Triad: Queries, Keys, and Values\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",
        "In traditional recurrent neural networks, a single hidden state must simultaneously encode token identity, store historical context, and act as the comparator against neighboring positions. Compressing these duties into an unspecialized representation creates an informational compromise.\n",
        "\n",
        "To resolve this limitation, the Transformer decomposes each token's representation into three distinct linear projections:\n",
        "\n",
        "- **Query ($Q$):** Specifies what information the current position requires.\n",
        "- **Key ($K$):** Advertises what information the position contains for indexing.\n",
        "- **Value ($V$):** Carries the substantive features to be blended into the output.\n",
        "\n",
        "Crucially, self-attention **never compares values directly**. It evaluates alignment between queries and keys to calculate normalized attention weights, and then applies those weights over values:\n",
        "\n",
        "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V$$\n",
        "\n",
        "This laboratory implements the complete functional triad in pure PyTorch and NumPy, demonstrating:\n",
        "\n",
        "1. **Linear Projections:** $x_i W_Q, x_i W_K, x_i W_V$.\n",
        "2. **The Variance Scaling Factor:** Why dividing by $\\sqrt{d_k}$ prevents softmax saturation.\n",
        "3. **Pairwise Compatibility Matrix:** $S = QK^T / \\sqrt{d_k}$.\n",
        "4. **Row-Stochastic Weights:** $A = \\text{softmax}(S, \\text{dim}=-1)$.\n",
        "5. **Value Aggregation:** $Y = AV$.\n",
        "6. **Multi-Head Tensor Extension:** Parallel execution across batch $B$ and heads $h$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import math\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "# Set seeds for reproducible tensor values\n",
        "torch.manual_seed(42)\n",
        "np.random.seed(42)\n",
        "\n",
        "print(f\"PyTorch version: {torch.__version__}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Linear Projections: Creating Queries, Keys, and Values\n",
        "\n",
        "Let an input sequence contain $T$ tokens, where each token has embedding dimension $C$ ($d_{\\text{model}}$ in Vaswani et al., 2017):\n",
        "\n",
        "$$X \\in \\mathbb{R}^{T \\times C}$$\n",
        "\n",
        "To assign functional roles, the input is projected through three learned weight matrices:\n",
        "\n",
        "$$W_Q \\in \\mathbb{R}^{C \\times d_k}, \\qquad W_K \\in \\mathbb{R}^{C \\times d_k}, \\qquad W_V \\in \\mathbb{R}^{C \\times d_v}$$\n",
        "\n",
        "yielding queries $Q \\in \\mathbb{R}^{T \\times d_k}$, keys $K \\in \\mathbb{R}^{T \\times d_k}$, and values $V \\in \\mathbb{R}^{T \\times d_v}$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "T = 4        # Sequence length (4 tokens)\n",
        "C = 12       # Input embedding dimension (d_model)\n",
        "d_k = 4      # Projection dimension for Queries and Keys\n",
        "d_v = 4      # Projection dimension for Values\n",
        "\n",
        "# Simulated input embeddings X in R^(T x C)\n",
        "X = torch.randn(T, C)\n",
        "\n",
        "# Linear projection layers without bias (matching standard formulation)\n",
        "W_q = nn.Linear(C, d_k, bias=False)\n",
        "W_k = nn.Linear(C, d_k, bias=False)\n",
        "W_v = nn.Linear(C, d_v, bias=False)\n",
        "\n",
        "# Compute Queries, Keys, and Values\n",
        "Q = W_q(X)  # (T, d_k)\n",
        "K = W_k(X)  # (T, d_k)\n",
        "V = W_v(X)  # (T, d_v)\n",
        "\n",
        "print(f\"Input X shape:   {X.shape}\")\n",
        "print(f\"Queries Q shape: {Q.shape}  [Searching space in R^{d_k}]\")\n",
        "print(f\"Keys K shape:    {K.shape}  [Indexing space in R^{d_k}]\")\n",
        "print(f\"Values V shape:  {V.shape}  [Payload space in R^{d_v}]\")\n",
        "\n",
        "assert Q.shape == (T, d_k)\n",
        "assert K.shape == (T, d_k)\n",
        "assert V.shape == (T, d_v)\n",
        "print(\"[OK] Linear projections verified.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 2. Why Scale by $\\sqrt{d_k}$? The Variance Proof\n",
        "\n",
        "Vaswani et al. (2017) note that if components of $q$ and $k$ are independent random variables with mean 0 and variance 1, their dot product:\n",
        "\n",
        "$$q^T k = \\sum_{m=1}^{d_k} q_m k_m$$\n",
        "\n",
        "has mean $\\mathbb{E}[q^T k] = 0$ and variance:\n",
        "\n",
        "$$\\text{Var}(q^T k) = \\sum_{m=1}^{d_k} \\text{Var}(q_m k_m) = d_k$$\n",
        "\n",
        "Consequently, the standard deviation is $\\sigma = \\sqrt{d_k}$. For large $d_k$ (e.g. $d_k = 64$ or $512$), dot products grow large in magnitude, pushing the softmax function into saturated regions where gradients vanish:\n",
        "\n",
        "$$\\frac{\\partial \\alpha_i}{\\partial s_j} = \\alpha_i (\\delta_{ij} - \\alpha_j) \\approx 0$$\n",
        "\n",
        "Scaling by $\\frac{1}{\\sqrt{d_k}}$ restores the variance back to 1."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Empirical validation of variance scaling\n",
        "dim_sizes = [16, 64, 256, 512]\n",
        "num_samples = 10000\n",
        "\n",
        "print(\"Empirical Variance of Dot Products:\")\n",
        "print(f\"{'Dimension d_k':<15} | {'Theoretical Var':<17} | {'Unscaled Empirical Var':<24} | {'Scaled (x 1/sqrt(d_k)) Var'}\")\n",
        "print(\"-\" * 85)\n",
        "\n",
        "for d in dim_sizes:\n",
        "    q_sample = torch.randn(num_samples, d)\n",
        "    k_sample = torch.randn(num_samples, d)\n",
        "    \n",
        "    # Unscaled dot product: sum across dimension d\n",
        "    dot_unscaled = (q_sample * k_sample).sum(dim=-1)\n",
        "    # Scaled dot product\n",
        "    dot_scaled = dot_unscaled / math.sqrt(d)\n",
        "    \n",
        "    print(f\"{d:<15} | {d:<17} | {dot_unscaled.var().item():<24.2f} | {dot_scaled.var().item():.2f}\")\n",
        "\n",
        "# Visual demonstration of Softmax Saturation\n",
        "d_test = 64\n",
        "q_ex = torch.randn(1, d_test)\n",
        "k_ex = torch.randn(10, d_test)\n",
        "\n",
        "unscaled_scores = (q_ex @ k_ex.T).squeeze(0)\n",
        "scaled_scores   = unscaled_scores / math.sqrt(d_test)\n",
        "\n",
        "softmax_unscaled = torch.softmax(unscaled_scores, dim=-1)\n",
        "softmax_scaled   = torch.softmax(scaled_scores, dim=-1)\n",
        "\n",
        "plt.figure(figsize=(8, 3.8), dpi=120)\n",
        "plt.plot(softmax_unscaled.numpy(), 'o--', color='#1c1812', lw=1.8, label=f'Unscaled (Peak = {softmax_unscaled.max():.2f}, Saturated)')\n",
        "plt.plot(softmax_scaled.numpy(),   's-',  color='#c9b994', lw=2.0, label=f'Scaled by $1/\\\\sqrt{{{d_test}}}$ (Smooth, Gradient Friendly)')\n",
        "plt.xticks(range(10), [f'Key {i}' for i in range(10)], fontsize=9)\n",
        "plt.ylabel('Attention Weight $\\\\alpha_j$', fontsize=10)\n",
        "plt.title(f'Softmax Saturation Effect ($d_k = {d_test}$)', fontsize=12, pad=10)\n",
        "plt.legend(frameon=True, facecolor='#faf1de', edgecolor='#c9b994', fontsize=9)\n",
        "plt.grid(True, linestyle=':', alpha=0.4)\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 3. Compatibility Scores: $S = \\frac{QK^T}{\\sqrt{d_k}}$\n",
        "\n",
        "For all $T$ positions simultaneously, inner products between every query $q_i$ and every key $k_j$ are computed in a single matrix multiplication:\n",
        "\n",
        "$$S = \\frac{Q K^T}{\\sqrt{d_k}}, \\qquad S \\in \\mathbb{R}^{T \\times T}$$\n",
        "\n",
        "Each element $S_{ij}$ measures how well query $i$ aligns with key $j$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Compute scaled dot-product compatibility scores\n",
        "S = (Q @ K.T) / math.sqrt(d_k)\n",
        "\n",
        "print(f\"Compatibility Scores Matrix S shape: {S.shape}  [T x T]\")\n",
        "print(\"Sample Score Matrix S:\")\n",
        "print(np.round(S.detach().numpy(), 3))\n",
        "\n",
        "assert S.shape == (T, T)\n",
        "print(\"[OK] Compatibility score matrix computed successfully.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Row-Wise Softmax Normalization: $A = \\text{softmax}(S, \\text{dim}=-1)$\n",
        "\n",
        "To convert compatibility scores into convex weights, the softmax is applied along the sequence dimension (rows of $S$):\n",
        "\n",
        "$$A_{ij} = \\frac{\\exp(S_{ij})}{\\sum_{m=1}^T \\exp(S_{im})}, \\qquad A \\in \\mathbb{R}^{T \\times T}$$\n",
        "\n",
        "Every row of $A$ forms a probability distribution over the sequence: $\\sum_{j=1}^T A_{ij} = 1$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Softmax over dim=-1 (columns) normalizes each row to 1.0\n",
        "A = torch.softmax(S, dim=-1)\n",
        "\n",
        "row_sums = A.sum(dim=-1)\n",
        "print(f\"Attention weights matrix A shape: {A.shape}\")\n",
        "print(f\"Row sums of A: {row_sums.tolist()}\")\n",
        "\n",
        "assert torch.allclose(row_sums, torch.ones(T)), \"Every row must sum to 1.0!\"\n",
        "\n",
        "# Heatmap visualization\n",
        "plt.figure(figsize=(5.5, 4.5), dpi=120)\n",
        "cax = plt.imshow(A.detach().numpy(), cmap='YlOrBr', vmin=0, vmax=1.0)\n",
        "plt.colorbar(cax, label='Attention Weight $A_{ij}$')\n",
        "\n",
        "labels = [f'Pos {i}' for i in range(T)]\n",
        "plt.xticks(range(T), labels, fontsize=9)\n",
        "plt.yticks(range(T), labels, fontsize=9)\n",
        "plt.xlabel('Keys (Source Position $j$)', fontsize=10)\n",
        "plt.ylabel('Queries (Target Position $i$)', fontsize=10)\n",
        "plt.title('Self-Attention Weights Matrix $A$', fontsize=11, pad=10)\n",
        "\n",
        "for i in range(T):\n",
        "    for j in range(T):\n",
        "        val = A[i, j].item()\n",
        "        text_color = 'white' if val > 0.4 else 'black'\n",
        "        plt.text(j, i, f\"{val:.2f}\", ha='center', va='center', color=text_color, fontsize=9)\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 5. Value Mixture: $Y = A @ V$\n",
        "\n",
        "Finally, the attention weights are multiplied by the value matrix $V$:\n",
        "\n",
        "$$Y = A V, \\qquad Y \\in \\mathbb{R}^{T \\times d_v}$$\n",
        "\n",
        "Every output vector $y_i$ is a convex combination of the values, weighted by row $i$ of $A$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Layer output Y = A @ V\n",
        "Y = A @ V\n",
        "\n",
        "print(f\"Value matrix V shape:   {V.shape}  [T x d_v]\")\n",
        "print(f\"Output matrix Y shape: {Y.shape}  [T x d_v]\")\n",
        "\n",
        "assert Y.shape == (T, d_v), \"Output shape must preserve (T, d_v)!\"\n",
        "\n",
        "# Verify that row 0 of Y equals manual sum over V\n",
        "y0_manual = sum(A[0, j] * V[j] for j in range(T))\n",
        "assert torch.allclose(Y[0], y0_manual), \"Output must match explicit convex combination!\"\n",
        "print(\"[OK] Value aggregation verified.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 6. Multi-Head Batched Execution\n",
        "\n",
        "In multi-head attention (Vaswani et al., 2017), the computation is vectorized across $B$ batch sequences and $h$ heads simultaneously:\n",
        "\n",
        "$$Q, K \\in \\mathbb{R}^{B \\times h \\times T \\times d_k}, \\qquad V \\in \\mathbb{R}^{B \\times h \\times T \\times d_v}$$\n",
        "\n",
        "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left( \\frac{Q K^T}{\\sqrt{d_k}} \\right) V \\in \\mathbb{R}^{B \\times h \\times T \\times d_v}$$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class ScaledDotProductAttention(nn.Module):\n",
        "    def __init__(self, scale=True):\n",
        "        super().__init__()\n",
        "        self.scale = scale\n",
        "\n",
        "    def forward(self, q, k, v, mask=None):\n",
        "        # q, k: (B, h, T, d_k)\n",
        "        # v:    (B, h, T, d_v)\n",
        "        d_k = q.size(-1)\n",
        "        \n",
        "        # Pairwise scores: (B, h, T, d_k) @ (B, h, d_k, T) -> (B, h, T, T)\n",
        "        scores = torch.matmul(q, k.transpose(-2, -1))\n",
        "        if self.scale:\n",
        "            scores = scores / math.sqrt(d_k)\n",
        "            \n",
        "        if mask is not None:\n",
        "            scores = scores.masked_fill(mask == 0, -1e9)\n",
        "            \n",
        "        # Row-wise softmax along last dimension\n",
        "        attn_weights = torch.softmax(scores, dim=-1)\n",
        "        \n",
        "        # Aggregate values: (B, h, T, T) @ (B, h, T, d_v) -> (B, h, T, d_v)\n",
        "        output = torch.matmul(attn_weights, v)\n",
        "        return output, attn_weights\n",
        "\n",
        "# Test module with multi-head batch tensors\n",
        "B, h, T, d_k, d_v = 2, 8, 16, 64, 64\n",
        "\n",
        "q_mha = torch.randn(B, h, T, d_k)\n",
        "k_mha = torch.randn(B, h, T, d_k)\n",
        "v_mha = torch.randn(B, h, T, d_v)\n",
        "\n",
        "attention_layer = ScaledDotProductAttention(scale=True)\n",
        "out_mha, weights_mha = attention_layer(q_mha, k_mha, v_mha)\n",
        "\n",
        "print(f\"Batched Q shape:       {q_mha.shape}\")\n",
        "print(f\"Batched K shape:       {k_mha.shape}\")\n",
        "print(f\"Batched V shape:       {v_mha.shape}\")\n",
        "print(f\"Attention Weights A:   {weights_mha.shape}  [(B, h, T, T)]\")\n",
        "print(f\"Attention Output Y:    {out_mha.shape}  [(B, h, T, d_v)]\")\n",
        "\n",
        "assert weights_mha.shape == (B, h, T, T)\n",
        "assert out_mha.shape == (B, h, T, d_v)\n",
        "assert torch.allclose(weights_mha.sum(dim=-1), torch.ones(B, h, T))\n",
        "print(\"[OK] Multi-head scaled dot-product attention module verified successfully.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 7. Dimensional and Conceptual Audit Summary\n",
        "\n",
        "| Object | Mathematical Domain | Coordinate Shape | Operational Purpose |\n",
        "| :--- | :--- | :--- | :--- |\n",
        "| **Input ($X$)** | Sequence embeddings | $(B, T, C)$ | Continuous token vectors entering the layer. |\n",
        "| **Queries ($Q$)** | Inquiry space | $(B, h, T, d_k)$ | $X W_Q$; specifies what information each position requires. |\n",
        "| **Keys ($K$)** | Indexing space | $(B, h, T, d_k)$ | $X W_K$; advertises what features each position contains. |\n",
        "| **Values ($V$)** | Payload space | $(B, h, T, d_v)$ | $X W_V$; the vectors to be blended into the output. |\n",
        "| **Scores ($S$)** | Compatibility | $(B, h, T, T)$ | $QK^T / \\sqrt{d_k}$; alignment normalized for unit variance. |\n",
        "| **Weights ($A$)** | Row-stochastic matrix | $(B, h, T, T)$ | $\\text{softmax}(S)$; convex mixing coefficients summing to 1. |\n",
        "| **Output ($Y$)** | Contextualized vectors | $(B, h, T, d_v)$ | $AV$; convex combinations bounded inside $\\text{Conv}(v_j)$."
      ]
    }
  ],
  "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.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}