{ "cells": [ { "cell_type": "markdown", "id": "e840227d", "metadata": {}, "source": [ "# Attention as a Weighted Mixture\n", "### The Geometry of Softmax, Convex Combinations, and Value Aggregation\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 *Attention Is All You Need* (Vaswani et al., 2017), the authors define attention as mapping a query and a set of key-value pairs to an output vector, computed as a **weighted sum of the values**:\n", "\n", "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V$$\n", "\n", "While implementations often focus heavily on query-key interactions, the operation that executes attention is the linear aggregation of values. Dot products and softmax normalization exist exclusively to calculate scalar coefficients. Once obtained, those coefficients govern how existing representations blend into updated representations.\n", "\n", "This notebook provides a self-contained, educational implementation in pure PyTorch and NumPy that analyzes attention as a linear aggregation operator over vectors, examining:\n", "\n", "1. **Value Vectors in $\\mathbb{R}^{d_v}$:** The substrate being combined.\n", "2. **Convex Combinations & The Convex Hull:** Non-negativity, unit-sum constraints, and geometric containment.\n", "3. **From Logits to Weights via Softmax:** Normalization, position competition, and temperature scaling.\n", "4. **The Minimal Matrix Formulation:** Row-stochastic attention matrices $A \\in \\mathbb{R}^{T \\times T}$ and $Y = AV \\in \\mathbb{R}^{T \\times d_v}$.\n", "5. **Multi-Head Batch Geometry:** The tensor mapping $(B, h, T, T) \\times (B, h, T, d_v) \\longrightarrow (B, h, T, d_v)$." ] }, { "cell_type": "code", "execution_count": null, "id": "6aa235a1", "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 random seed for reproducible tensor operations and figures\n", "torch.manual_seed(42)\n", "np.random.seed(42)\n", "\n", "print(f\"PyTorch version: {torch.__version__}\")" ] }, { "cell_type": "markdown", "id": "c9959673", "metadata": {}, "source": [ "## 1. Value Vectors in $\\mathbb{R}^{d_v}$ and Convex Combinations\n", "\n", "Consider a sequence of $T$ token positions. Each position holds a value vector $v_j \\in \\mathbb{R}^{d_v}$. In the base Transformer architecture of Vaswani et al. (2017), the value dimension is set to $d_v = d_{\\text{model}} / h = 512 / 8 = 64$.\n", "\n", "Given scalar weights $\\alpha_1, \\ldots, \\alpha_T$, attention combines these values into a single output vector $y$:\n", "\n", "$$y = \\sum_{j=1}^T \\alpha_j v_j$$\n", "\n", "For $y$ to act as a normalized mixture, the weights must satisfy:\n", "\n", "$$\\alpha_j \\ge 0, \\qquad \\sum_{j=1}^T \\alpha_j = 1$$\n", "\n", "Any linear combination satisfying these two properties is a **convex combination**." ] }, { "cell_type": "code", "execution_count": null, "id": "6d7bab17", "metadata": {}, "outputs": [], "source": [ "# Define T=4 sequence positions and value dimension dv=3 for inspection\n", "T = 4\n", "d_v = 3\n", "\n", "# Matrix of value vectors V in R^(T x dv)\n", "V = torch.tensor([\n", " [1.0, 0.0, 0.0], # Position 0\n", " [0.0, 2.0, 0.0], # Position 1\n", " [0.0, 0.0, 3.0], # Position 2\n", " [1.0, 1.0, 1.0], # Position 3\n", "], dtype=torch.float32)\n", "\n", "# Valid convex weights (non-negative and sum to 1)\n", "alpha = torch.tensor([0.40, 0.10, 0.20, 0.30])\n", "\n", "assert (alpha >= 0).all(), \"Weights must be non-negative!\"\n", "assert torch.isclose(alpha.sum(), torch.tensor(1.0)), \"Weights must sum to 1.0!\"\n", "\n", "# 1. Explicit loop summation: y = sum(alpha_j * v_j)\n", "y_sum = torch.zeros(d_v)\n", "for j in range(T):\n", " y_sum += alpha[j] * V[j]\n", "\n", "# 2. Vector-matrix product: y = alpha @ V [(1, T) @ (T, dv) -> (1, dv)]\n", "y_matrix = alpha.unsqueeze(0) @ V\n", "\n", "print(f\"Value matrix V shape: {V.shape}\")\n", "print(f\"Weights alpha: {alpha.tolist()} (Sum = {alpha.sum().item():.2f})\")\n", "print(f\"Output mixture vector: {y_matrix.squeeze(0).tolist()}\")\n", "\n", "assert torch.allclose(y_sum, y_matrix.squeeze(0)), \"Loop sum and matrix product must be mathematically identical!\"\n", "print(\"[OK] Loop sum and matrix multiplication verified.\")" ] }, { "cell_type": "markdown", "id": "e88da9f6", "metadata": {}, "source": [ "## 2. Geometric Confinement: The Convex Hull in $\\mathbb{R}^2$\n", "\n", "Because the weights $\\alpha_j$ are non-negative and sum to one, the output vector $y$ is mathematically confined within the **convex hull** spanned by the values:\n", "\n", "$$y \\in \\text{Conv}(v_1, \\ldots, v_T)$$\n", "\n", "This geometric property ensures that attention outputs remain anchored within the spatial boundaries of the sequence representations. To observe this directly, let us project value vectors onto $\\mathbb{R}^2$ and sample $1,000$ random attention mixtures using the Dirichlet distribution." ] }, { "cell_type": "code", "execution_count": null, "id": "deed0176", "metadata": {}, "outputs": [], "source": [ "# 4 value vectors in 2D space\n", "V_2d = torch.tensor([\n", " [1.0, 1.0], # Position 0\n", " [1.5, 4.0], # Position 1\n", " [5.0, 4.5], # Position 2\n", " [4.0, 0.5] # Position 3\n", "], dtype=torch.float32)\n", "\n", "# Sample 1,000 random convex weight vectors using a symmetric Dirichlet distribution\n", "num_samples = 1000\n", "alpha_samples = torch.distributions.Dirichlet(torch.ones(4)).sample((num_samples,))\n", "mixtures_2d = alpha_samples @ V_2d # (1000, 4) @ (4, 2) -> (1000, 2)\n", "\n", "plt.figure(figsize=(7, 5), dpi=120)\n", "\n", "# Convex hull envelope polygon\n", "hull_indices = [0, 1, 2, 3, 0]\n", "plt.fill(V_2d[hull_indices, 0].numpy(), V_2d[hull_indices, 1].numpy(), \n", " color='#f0e4c8', alpha=0.55, label='Convex Hull $\\\\text{Conv}(v_1, \\\\ldots, v_T)$')\n", "plt.plot(V_2d[hull_indices, 0].numpy(), V_2d[hull_indices, 1].numpy(), \n", " '--', color='#5b5242', lw=1.5, label='Convex Boundary')\n", "\n", "# 1,000 random attention mixtures\n", "plt.scatter(mixtures_2d[:, 0].numpy(), mixtures_2d[:, 1].numpy(), \n", " color='#c9b994', alpha=0.35, s=12, label='1,000 Attention Mixtures')\n", "\n", "# Original value vectors\n", "plt.scatter(V_2d[:, 0].numpy(), V_2d[:, 1].numpy(), \n", " color='#1c1812', s=85, zorder=5, label='Original Values $v_j$')\n", "\n", "for j in range(4):\n", " plt.annotate(f\" $v_{j}$ ({V_2d[j,0]:.1f}, {V_2d[j,1]:.1f})\", \n", " (V_2d[j, 0].item() + 0.1, V_2d[j, 1].item()), \n", " fontsize=11, fontweight='bold', color='#1c1812')\n", "\n", "plt.title('Geometric Confinement: Attention Outputs Lie Inside the Convex Hull', fontsize=12, pad=12)\n", "plt.xlabel('Dimension 1 ($d_1$)', fontsize=10)\n", "plt.ylabel('Dimension 2 ($d_2$)', fontsize=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", "id": "1712f96b", "metadata": {}, "source": [ "## 3. From Logits to Weights: Softmax and Temperature Scaling\n", "\n", "Scoring mechanisms produce unbounded real-valued logits $s_j \\in \\mathbb{R}$. Softmax converts these logits into normalized probabilities:\n", "\n", "$$\\alpha_j = \\frac{\\exp(s_j / \\tau)}{\\sum_{k=1}^T \\exp(s_k / \\tau)}$$\n", "\n", "Here, $\\tau > 0$ controls the entropy of the competition:\n", "- **Low temperature ($\\tau \\to 0^+$):** The distribution sharpens toward an argmax one-hot vector (hard selection / pointer).\n", "- **Standard temperature ($\\tau = 1.0$):** Standard softmax competition.\n", "- **High temperature ($\\tau \\to \\infty$):** The distribution flattens toward uniform averaging $\\alpha_j \\approx 1/T$ (mean pooling)." ] }, { "cell_type": "code", "execution_count": null, "id": "f72d318f", "metadata": {}, "outputs": [], "source": [ "# Raw logits for 4 token positions\n", "logits = torch.tensor([1.2, 4.5, 2.0, 0.8], dtype=torch.float32)\n", "\n", "def custom_softmax(s, temperature=1.0):\n", " # Subtract max for numerical stability (prevents floating point overflow in exp)\n", " scaled = s / temperature\n", " stabilized = scaled - torch.max(scaled)\n", " exps = torch.exp(stabilized)\n", " return exps / torch.sum(exps)\n", "\n", "# Evaluate across three temperatures\n", "alpha_cold = custom_softmax(logits, temperature=0.25) # Sharp / Peaky\n", "alpha_std = custom_softmax(logits, temperature=1.00) # Standard\n", "alpha_hot = custom_softmax(logits, temperature=5.00) # Smooth / Uniform\n", "\n", "positions = [f\"Pos {j}\" for j in range(len(logits))]\n", "x_bar = np.arange(len(positions))\n", "width = 0.25\n", "\n", "plt.figure(figsize=(7, 4), dpi=120)\n", "plt.bar(x_bar - width, alpha_cold.numpy(), width, label='Low Temp ($\\tau=0.25$, Peaky)', color='#1c1812')\n", "plt.bar(x_bar, alpha_std.numpy(), width, label='Standard Temp ($\\tau=1.00$)', color='#5b5242')\n", "plt.bar(x_bar + width, alpha_hot.numpy(), width, label='High Temp ($\\tau=5.00$, Smooth)', color='#c9b994')\n", "\n", "plt.xticks(x_bar, positions, fontsize=10)\n", "plt.ylabel('Attention Weight $\\alpha_j$', fontsize=10)\n", "plt.title('Softmax Competition Under Temperature Scaling', fontsize=12, pad=12)\n", "plt.legend(frameon=True, facecolor='#faf1de', edgecolor='#c9b994', fontsize=9)\n", "plt.grid(True, axis='y', linestyle=':', alpha=0.4)\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "assert torch.allclose(alpha_std, torch.softmax(logits, dim=-1)), \"Custom softmax matches PyTorch built-in!\"\n", "print(\"[OK] Numerically stable custom softmax validated.\")" ] }, { "cell_type": "markdown", "id": "1e5efb4d", "metadata": {}, "source": [ "## 4. Minimal Matrix Formulation: $Y = AV$\n", "\n", "In a full Transformer sequence, every position $i$ acts as a query, generating its own set of attention weights over all positions $j$.\n", "\n", "Stacking these weight vectors forms the **attention matrix**:\n", "\n", "$$A \\in \\mathbb{R}^{T \\times T}, \\qquad A_{ij} \\ge 0, \\quad \\sum_{j=1}^T A_{ij} = 1$$\n", "\n", "In linear algebra, a matrix with non-negative real entries where each row sums to one is called a **row-stochastic matrix**.\n", "\n", "The entire sequence transformation is computed as:\n", "\n", "$$Y = AV, \\qquad Y \\in \\mathbb{R}^{T \\times d_v}$$\n", "\n", "Notice that $Y$ preserves sequence length $T$ and feature dimension $d_v$ exactly." ] }, { "cell_type": "code", "execution_count": null, "id": "6e441c99", "metadata": {}, "outputs": [], "source": [ "T = 5\n", "d_v = 4\n", "\n", "# Value matrix V in R^(T x dv)\n", "V = torch.randn(T, d_v)\n", "\n", "# Logit matrix S in R^(T x T)\n", "S = torch.randn(T, T)\n", "\n", "# Softmax applied strictly over the last dimension (dim=-1) so each row sums to 1\n", "A = torch.softmax(S, dim=-1)\n", "\n", "# Row-stochastic verification\n", "row_sums = A.sum(dim=-1)\n", "print(f\"Attention Matrix A shape: {A.shape}\")\n", "print(f\"Row sums of A: {row_sums.tolist()}\")\n", "assert torch.allclose(row_sums, torch.ones(T)), \"Every row of A must sum to 1.0!\"\n", "\n", "# Layer output Y = A @ V\n", "Y = A @ V\n", "print(f\"Output Matrix Y shape: {Y.shape} [Preserves (T, dv)]\")\n", "assert Y.shape == (T, d_v), \"Output Y must match shape (T, dv)!\"" ] }, { "cell_type": "markdown", "id": "33bd11ce", "metadata": {}, "source": [ "## 5. Visualizing the Attention Heatmap\n", "\n", "Let us inspect the row-stochastic matrix $A$. Each row $i$ reveals how target position $i$ distributes its convex weights across source positions $j$." ] }, { "cell_type": "code", "execution_count": null, "id": "45096117", "metadata": {}, "outputs": [], "source": [ "plt.figure(figsize=(6, 5), dpi=120)\n", "cax = plt.imshow(A.numpy(), cmap='YlOrBr', vmin=0, vmax=1.0)\n", "plt.colorbar(cax, label='Attention Weight $A_{ij}$')\n", "\n", "labels = [f\"Token {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 / Values (Source Position $j$)', fontsize=10)\n", "plt.ylabel('Queries (Target Position $i$)', fontsize=10)\n", "plt.title('Row-Stochastic Attention Matrix $A \\\\in \\\\mathbb{R}^{T \\\\times T}$', fontsize=12, pad=12)\n", "\n", "# Annotate values in each cell\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", "id": "001703af", "metadata": {}, "source": [ "## 6. Multi-Head Batch Tensor Formulation\n", "\n", "In multi-head attention (Vaswani et al., 2017), this computation executes in parallel across batch sequences $B$ and heads $h$:\n", "\n", "$$A \\in \\mathbb{R}^{B \\times h \\times T \\times T}, \\qquad V \\in \\mathbb{R}^{B \\times h \\times T \\times d_v}$$\n", "\n", "The matrix multiplication operator `@` automatically broadcasts over the batch ($B$) and head ($h$) dimensions, computing $B \\times h$ simultaneous matrix multiplications in parallel:\n", "\n", "$$Y = AV \\in \\mathbb{R}^{B \\times h \\times T \\times d_v}$$" ] }, { "cell_type": "code", "execution_count": null, "id": "74d91a85", "metadata": {}, "outputs": [], "source": [ "B = 2 # Batch size\n", "h = 3 # Number of attention heads\n", "T = 5 # Sequence length\n", "d_v = 4 # Dimension per head\n", "\n", "# Value tensor V in shape (B, h, T, dv)\n", "V_mha = torch.randn(B, h, T, d_v)\n", "\n", "# Unnormalized logits matrix S in shape (B, h, T, T)\n", "S_mha = torch.randn(B, h, T, T)\n", "A_mha = torch.softmax(S_mha, dim=-1)\n", "\n", "# Batched matrix multiplication\n", "Y_mha = A_mha @ V_mha\n", "\n", "print(f\"Batched Attention Matrix A shape: {A_mha.shape}\")\n", "print(f\"Batched Value Matrix V shape: {V_mha.shape}\")\n", "print(f\"Batched Output Matrix Y shape: {Y_mha.shape}\")\n", "\n", "# Verify row sums across all batches and heads\n", "row_sums_mha = A_mha.sum(dim=-1)\n", "assert torch.allclose(row_sums_mha, torch.ones(B, h, T)), \"All attention rows across all batches/heads must sum to 1.0!\"\n", "assert Y_mha.shape == (B, h, T, d_v), \"Output tensor shape must match (B, h, T, dv)!\"\n", "print(\"[OK] Multi-head batched tensor assertions verified successfully.\")" ] }, { "cell_type": "markdown", "id": "c4fe40af", "metadata": {}, "source": [ "## 7. Dimensional and Conceptual Audit\n", "\n", "| Object | Mathematical Domain | Tensor Shape | Function in Mechanism |\n", "| :--- | :--- | :--- | :--- |\n", "| **Values ($V$)** | Feature representations | $(B, h, T, d_v)$ | The vectors being blended; lives in $\\mathbb{R}^{d_v}$. |\n", "| **Logits ($S$)** | Unbounded compatibility | $(B, h, T, T)$ | Raw real scores from query-key inner products ($-\\infty, \\infty$). |\n", "| **Weights ($A$)** | Row-stochastic matrix | $(B, h, T, T)$ | Non-negative coefficients summing to 1 per row. |\n", "| **Output ($Y$)** | Contextualized vectors | $(B, h, T, d_v)$ | Convex combinations bounded inside $\\text{Conv}(v_1, \\ldots, v_T)$. |\n", "\n", "---\n", "\n", "### Core Pedagogical Takeaways\n", "\n", "1. **Attention is an aggregation operator:** Before queries and keys compute compatibility, attention is fundamentally a convex mixture of vectors.\n", "2. **Convex combinations preserve bounds:** Because weights $\\alpha_j \\ge 0$ and $\\sum \\alpha_j = 1$, outputs remain geometrically anchored inside the convex hull spanned by the values.\n", "3. **Attention preserves representation space:** The output $Y$ retains the exact feature dimension $d_v$ of $V$. Attention updates semantic information without altering spatial coordinates.\n", "4. **Attention is not a full causal explanation:** High attention weights indicate numerical reliance within an isolated layer; they do not replace circuit analysis of multi-layer representations." ] } ], "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 }