{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# DINOv2 Tutorial — Notebook 6: Attention Layer I/O Visualization\n", "\n", "**Series:** DINOv2 Deep Dive | **Prev:** `05_vit_architecture.ipynb` | **Next:** `07_intermediate_layer_visualization.ipynb`\n", "\n", "---\n", "\n", "## What You Will Learn\n", "- Multi-head self-attention (MHSA) mechanics in detail\n", "- Q, K, V projections: shapes and visualizations\n", "- Attention weight matrices: what each head attends to\n", "- How to extract attention maps from every layer\n", "- Visualizing attention from CLS token to image patches\n", "- Rollout attention: aggregating attention across layers\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib.patches as mpatches\n", "from PIL import Image, ImageDraw\n", "from transformers import AutoImageProcessor, AutoModel\n", "\n", "MODEL_PATH = \"/mnt/mydata/huggingface/dinov2-base\"\n", "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "\n", "processor = AutoImageProcessor.from_pretrained(MODEL_PATH)\n", "model = AutoModel.from_pretrained(MODEL_PATH).to(DEVICE).eval()\n", "cfg = model.config\n", "\n", "# Create test image\n", "def make_test_image(size=224):\n", " img = Image.new('RGB', (size, size), 'white')\n", " draw = ImageDraw.Draw(img)\n", " h = size // 2\n", " draw.rectangle([0, 0, h, h], fill=(220, 50, 50 ))\n", " draw.rectangle([h, 0, size, h], fill=(50, 100, 220))\n", " draw.rectangle([0, h, h, size], fill=(50, 180, 50 ))\n", " draw.rectangle([h, h, size, size], fill=(220, 200, 50 ))\n", " r = size // 14\n", " for cx, cy in [(h//2,h//2),(h+h//2,h//2),(h//2,h+h//2),(h+h//2,h+h//2)]:\n", " draw.ellipse([cx-r,cy-r,cx+r,cy+r], fill='white')\n", " return img\n", "\n", "test_image = make_test_image()\n", "inputs = processor(images=test_image, return_tensors='pt')\n", "inputs = {k: v.to(DEVICE) for k, v in inputs.items()}\n", "\n", "# Forward pass with attention outputs\n", "with torch.no_grad():\n", " outputs = model(**inputs, output_attentions=True, output_hidden_states=True)\n", "\n", "print(f\"Number of attention layers: {len(outputs.attentions)}\")\n", "print(f\"Attention shape per layer: {outputs.attentions[0].shape}\")\n", "print(f\" = [batch, heads, seq_len, seq_len]\")\n", "print(f\" = [1, {cfg.num_attention_heads}, {(cfg.image_size//cfg.patch_size)**2+1}, {(cfg.image_size//cfg.patch_size)**2+1}]\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6.1 Multi-Head Self-Attention Mechanics\n", "\n", "For input $X \\in \\mathbb{R}^{N \\times d}$ (N=257 tokens, d=768):\n", "\n", "$$Q = XW_Q, \\quad K = XW_K, \\quad V = XW_V$$\n", "\n", "where $W_Q, W_K, W_V \\in \\mathbb{R}^{d \\times d}$.\n", "\n", "Split into $h=12$ heads, each with $d_h = d/h = 64$ dimensions:\n", "\n", "$$\\text{Attention}(Q_i, K_i, V_i) = \\text{softmax}\\left(\\frac{Q_i K_i^T}{\\sqrt{d_h}}\\right) V_i$$\n", "\n", "Concatenate all heads and project:\n", "\n", "$$\\text{MHSA}(X) = \\text{Concat}(\\text{head}_1, \\ldots, \\text{head}_h) W_O$$\n", "\n", "**Shapes:**\n", "- Q, K, V: `[B, 12, 257, 64]`\n", "- Attention weights: `[B, 12, 257, 257]`\n", "- Output: `[B, 257, 768]`" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Extract Q, K, V from attention layer ──────────────────────────────────────\n", "# Hook into the attention module to capture Q, K, V\n", "qkv_store = {}\n", "\n", "def qkv_hook(module, input, output):\n", " \"\"\"Hook to capture Q, K, V matrices.\"\"\"\n", " x = input[0] # [B, seq_len, hidden_dim]\n", " B, N, D = x.shape\n", " H = cfg.num_attention_heads\n", " head_dim = D // H\n", " \n", " # Compute Q, K, V\n", " qkv = module.query, module.key, module.value\n", " Q = module.query(x).reshape(B, N, H, head_dim).permute(0, 2, 1, 3)\n", " K = module.key(x).reshape(B, N, H, head_dim).permute(0, 2, 1, 3)\n", " V = module.value(x).reshape(B, N, H, head_dim).permute(0, 2, 1, 3)\n", " \n", " qkv_store['Q'] = Q.detach().cpu()\n", " qkv_store['K'] = K.detach().cpu()\n", " qkv_store['V'] = V.detach().cpu()\n", "\n", "# Hook into last layer attention\n", "last_attn = model.encoder.layer[-1].attention.attention\n", "hook = last_attn.register_forward_hook(qkv_hook)\n", "\n", "with torch.no_grad():\n", " _ = model(**inputs)\n", "\n", "hook.remove()\n", "\n", "Q = qkv_store['Q'] # [1, 12, 257, 64]\n", "K = qkv_store['K']\n", "V = qkv_store['V']\n", "\n", "print(\"Q, K, V shapes:\")\n", "print(f\" Q: {Q.shape} [batch, heads, seq_len, head_dim]\")\n", "print(f\" K: {K.shape}\")\n", "print(f\" V: {V.shape}\")\n", "print()\n", "print(f\"Q stats: mean={Q.mean():.4f}, std={Q.std():.4f}\")\n", "print(f\"K stats: mean={K.mean():.4f}, std={K.std():.4f}\")\n", "print(f\"V stats: mean={V.mean():.4f}, std={V.std():.4f}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Visualize Q, K, V matrices ────────────────────────────────────────────────\n", "n_grid = (cfg.image_size // cfg.patch_size) # 16\n", "seq_len = n_grid**2 + 1 # 257\n", "\n", "fig, axes = plt.subplots(3, 4, figsize=(18, 12))\n", "\n", "for row, (mat, name, cmap) in enumerate([\n", " (Q[0], 'Q (Query)', 'RdBu_r'),\n", " (K[0], 'K (Key)', 'PiYG'),\n", " (V[0], 'V (Value)', 'viridis'),\n", "]):\n", " # Show first 4 heads\n", " for col in range(4):\n", " head_mat = mat[col].numpy() # [257, 64]\n", " im = axes[row, col].imshow(head_mat, cmap=cmap, aspect='auto')\n", " axes[row, col].set_title(f'{name} — Head {col+1}\\n[{seq_len}, 64]',\n", " fontweight='bold', fontsize=9)\n", " axes[row, col].set_xlabel('Head dim (64)')\n", " if col == 0:\n", " axes[row, col].set_ylabel('Token index (257)')\n", " plt.colorbar(im, ax=axes[row, col], fraction=0.046)\n", "\n", "plt.suptitle('Q, K, V Matrices — Last Transformer Layer (first 4 of 12 heads)',\n", " fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Compute and visualize attention weights manually ──────────────────────────\n", "head_dim = cfg.hidden_size // cfg.num_attention_heads # 64\n", "scale = head_dim ** -0.5\n", "\n", "# Compute attention for head 0\n", "q0 = Q[0, 0] # [257, 64]\n", "k0 = K[0, 0] # [257, 64]\n", "v0 = V[0, 0] # [257, 64]\n", "\n", "# Scaled dot-product attention\n", "attn_scores = (q0 @ k0.T) * scale # [257, 257]\n", "attn_weights = torch.softmax(attn_scores, dim=-1) # [257, 257]\n", "\n", "print(f\"Attention scores shape: {attn_scores.shape}\")\n", "print(f\"Attention weights shape: {attn_weights.shape}\")\n", "print(f\"Attention weights sum per row: {attn_weights.sum(dim=-1)[:5]} (should be ≈1.0)\")\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(18, 6))\n", "\n", "# Raw scores\n", "im0 = axes[0].imshow(attn_scores.numpy(), cmap='RdBu_r', aspect='auto')\n", "axes[0].set_title('Attention Scores (before softmax)\\nQ·Kᵀ / √d_h', fontweight='bold')\n", "axes[0].set_xlabel('Key position'); axes[0].set_ylabel('Query position')\n", "plt.colorbar(im0, ax=axes[0], fraction=0.046)\n", "\n", "# Attention weights\n", "im1 = axes[1].imshow(attn_weights.numpy(), cmap='hot', aspect='auto')\n", "axes[1].set_title('Attention Weights (after softmax)\\nsoftmax(Q·Kᵀ / √d_h)', fontweight='bold')\n", "axes[1].set_xlabel('Key position'); axes[1].set_ylabel('Query position')\n", "plt.colorbar(im1, ax=axes[1], fraction=0.046)\n", "\n", "# CLS token attention (row 0)\n", "cls_attn = attn_weights[0, 1:].numpy() # [256] — CLS attending to patches\n", "cls_attn_grid = cls_attn.reshape(n_grid, n_grid)\n", "im2 = axes[2].imshow(cls_attn_grid, cmap='hot', interpolation='bilinear')\n", "axes[2].set_title('CLS Token Attention Map\\n(row 0 of attention matrix, reshaped to 16×16)',\n", " fontweight='bold')\n", "axes[2].axis('off')\n", "plt.colorbar(im2, ax=axes[2], fraction=0.046)\n", "\n", "plt.suptitle('Attention Computation — Head 0, Last Layer', fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6.2 All 12 Attention Heads — Last Layer" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Visualize all 12 heads of the last layer ──────────────────────────────────\n", "last_layer_attn = outputs.attentions[-1][0].cpu().numpy() # [12, 257, 257]\n", "\n", "fig, axes = plt.subplots(3, 4, figsize=(18, 14))\n", "\n", "for head_idx in range(12):\n", " ax = axes[head_idx // 4, head_idx % 4]\n", " \n", " # CLS token attention to patches\n", " cls_attn = last_layer_attn[head_idx, 0, 1:] # [256]\n", " cls_attn_grid = cls_attn.reshape(n_grid, n_grid)\n", " \n", " im = ax.imshow(cls_attn_grid, cmap='hot', interpolation='bilinear')\n", " ax.set_title(f'Head {head_idx+1}\\nentropy={(-cls_attn*np.log(cls_attn+1e-8)).sum():.2f}',\n", " fontweight='bold', fontsize=9)\n", " ax.axis('off')\n", " plt.colorbar(im, ax=ax, fraction=0.046)\n", "\n", "plt.suptitle('All 12 Attention Heads — Last Layer\\n'\n", " 'CLS Token Attention to Image Patches (16×16 grid)',\n", " fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "# Statistics\n", "print(\"Head statistics (CLS→patch attention):\")\n", "print(f\"{'Head':>6} {'Max':>8} {'Mean':>8} {'Entropy':>10} {'Focus':>10}\")\n", "print(\"-\" * 50)\n", "for h in range(12):\n", " cls_attn = last_layer_attn[h, 0, 1:]\n", " entropy = -(cls_attn * np.log(cls_attn + 1e-8)).sum()\n", " focus = cls_attn.max() / cls_attn.mean() # how focused vs uniform\n", " print(f\"{h+1:>6} {cls_attn.max():>8.4f} {cls_attn.mean():>8.4f} {entropy:>10.3f} {focus:>10.2f}x\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6.3 Attention Maps Across All Layers" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── CLS attention across all 12 layers ────────────────────────────────────────\n", "fig, axes = plt.subplots(3, 4, figsize=(18, 14))\n", "\n", "for layer_idx in range(12):\n", " ax = axes[layer_idx // 4, layer_idx % 4]\n", " \n", " layer_attn = outputs.attentions[layer_idx][0].cpu().numpy() # [12, 257, 257]\n", " \n", " # Average over all heads\n", " mean_cls_attn = layer_attn[:, 0, 1:].mean(axis=0) # [256]\n", " mean_cls_grid = mean_cls_attn.reshape(n_grid, n_grid)\n", " \n", " im = ax.imshow(mean_cls_grid, cmap='hot', interpolation='bilinear')\n", " ax.set_title(f'Layer {layer_idx+1}\\n(mean over 12 heads)',\n", " fontweight='bold', fontsize=9)\n", " ax.axis('off')\n", " plt.colorbar(im, ax=ax, fraction=0.046)\n", "\n", "plt.suptitle('CLS Token Attention Maps — All 12 Layers\\n'\n", " '(averaged over all 12 heads per layer)',\n", " fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6.4 Attention Rollout\n", "\n", "**Attention rollout** (Abnar & Zuidema, 2020) aggregates attention across all layers to get a global view of which input patches the CLS token ultimately attends to.\n", "\n", "$$A_{\\text{rollout}} = A_1 \\cdot A_2 \\cdots A_L$$\n", "\n", "where each $A_l$ is the attention matrix of layer $l$, with residual connections accounted for:\n", "\n", "$$\\hat{A}_l = 0.5 \\cdot A_l + 0.5 \\cdot I$$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Attention rollout ─────────────────────────────────────────────────────────\n", "def attention_rollout(attentions, discard_ratio=0.9):\n", " \"\"\"\n", " Compute attention rollout across all layers.\n", " attentions: list of [B, H, N, N] tensors\n", " \"\"\"\n", " result = torch.eye(attentions[0].shape[-1]) # identity\n", " \n", " for attn in attentions:\n", " # Average over heads\n", " attn_avg = attn[0].mean(dim=0) # [N, N]\n", " \n", " # Add residual connection\n", " attn_avg = attn_avg + torch.eye(attn_avg.shape[0])\n", " attn_avg = attn_avg / attn_avg.sum(dim=-1, keepdim=True)\n", " \n", " # Discard low attention values\n", " flat = attn_avg.flatten()\n", " threshold = flat.kthvalue(int(flat.shape[0] * discard_ratio)).values\n", " attn_avg[attn_avg < threshold] = 0\n", " attn_avg = attn_avg / (attn_avg.sum(dim=-1, keepdim=True) + 1e-8)\n", " \n", " result = torch.matmul(attn_avg, result)\n", " \n", " return result\n", "\n", "# Compute rollout\n", "attentions_cpu = [a.cpu() for a in outputs.attentions]\n", "rollout = attention_rollout(attentions_cpu, discard_ratio=0.9)\n", "\n", "# CLS token rollout to patches\n", "cls_rollout = rollout[0, 1:].numpy() # [256]\n", "cls_rollout_grid = cls_rollout.reshape(n_grid, n_grid)\n", "\n", "# Normalize\n", "cls_rollout_grid = (cls_rollout_grid - cls_rollout_grid.min()) / \\\n", " (cls_rollout_grid.max() - cls_rollout_grid.min() + 1e-8)\n", "\n", "fig, axes = plt.subplots(1, 4, figsize=(18, 5))\n", "\n", "# Original image\n", "axes[0].imshow(test_image)\n", "axes[0].set_title('Original Image', fontweight='bold')\n", "axes[0].axis('off')\n", "\n", "# Last layer attention (mean heads)\n", "last_attn_mean = outputs.attentions[-1][0].cpu().mean(dim=0)[0, 1:].numpy().reshape(n_grid, n_grid)\n", "last_attn_mean = (last_attn_mean - last_attn_mean.min()) / (last_attn_mean.max() - last_attn_mean.min() + 1e-8)\n", "axes[1].imshow(last_attn_mean, cmap='hot', interpolation='bilinear')\n", "axes[1].set_title('Last Layer Attention\\n(mean over heads)', fontweight='bold')\n", "axes[1].axis('off')\n", "\n", "# Rollout\n", "axes[2].imshow(cls_rollout_grid, cmap='hot', interpolation='bilinear')\n", "axes[2].set_title('Attention Rollout\\n(aggregated across all layers)', fontweight='bold')\n", "axes[2].axis('off')\n", "\n", "# Overlay on image\n", "import matplotlib.cm as cm\n", "img_np = np.array(test_image.resize((n_grid*14, n_grid*14))) / 255.0\n", "rollout_resized = np.array(Image.fromarray((cls_rollout_grid * 255).astype(np.uint8)).resize(\n", " (n_grid*14, n_grid*14), Image.BILINEAR)) / 255.0\n", "heatmap = cm.hot(rollout_resized)[:, :, :3]\n", "overlay = 0.5 * img_np + 0.5 * heatmap\n", "axes[3].imshow(overlay.clip(0, 1))\n", "axes[3].set_title('Rollout Overlay\\non Original Image', fontweight='bold')\n", "axes[3].axis('off')\n", "\n", "plt.suptitle('Attention Rollout: Aggregating Attention Across All 12 Layers',\n", " fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6.5 Patch-to-Patch Attention\n", "\n", "Beyond CLS→patch attention, we can also visualize **patch→patch** attention to understand how patches communicate with each other." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Patch-to-patch attention ───────────────────────────────────────────────────\n", "last_attn = outputs.attentions[-1][0].cpu().numpy() # [12, 257, 257]\n", "\n", "# Pick a few query patches and show what they attend to\n", "query_patches = [\n", " (n_grid//4, n_grid//4), # top-left quadrant\n", " (n_grid//4, 3*n_grid//4), # top-right quadrant\n", " (3*n_grid//4, n_grid//4), # bottom-left quadrant\n", " (n_grid//2, n_grid//2), # center\n", "]\n", "\n", "fig, axes = plt.subplots(2, 4, figsize=(18, 9))\n", "\n", "for col, (qr, qc) in enumerate(query_patches):\n", " query_idx = qr * n_grid + qc + 1 # +1 for CLS token offset\n", " \n", " # Mean over heads\n", " patch_attn = last_attn[:, query_idx, 1:].mean(axis=0) # [256]\n", " patch_attn_grid = patch_attn.reshape(n_grid, n_grid)\n", " \n", " # Show original image with query patch marked\n", " axes[0, col].imshow(test_image)\n", " rect = mpatches.Rectangle(\n", " (qc * cfg.patch_size, qr * cfg.patch_size),\n", " cfg.patch_size, cfg.patch_size,\n", " linewidth=3, edgecolor='red', facecolor='red', alpha=0.5\n", " )\n", " axes[0, col].add_patch(rect)\n", " axes[0, col].set_title(f'Query patch [{qr},{qc}]\\n(red = query)', fontweight='bold', fontsize=9)\n", " axes[0, col].axis('off')\n", " \n", " # Attention map\n", " im = axes[1, col].imshow(patch_attn_grid, cmap='hot', interpolation='bilinear')\n", " axes[1, col].set_title(f'Attention from patch [{qr},{qc}]\\n(mean over 12 heads)',\n", " fontweight='bold', fontsize=9)\n", " axes[1, col].axis('off')\n", " plt.colorbar(im, ax=axes[1, col], fraction=0.046)\n", "\n", "plt.suptitle('Patch-to-Patch Attention — Last Layer\\n'\n", " 'What does each patch attend to?',\n", " fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6.6 Full Attention Matrix Visualization" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ── Full 257×257 attention matrix ─────────────────────────────────────────────\n", "fig, axes = plt.subplots(2, 3, figsize=(18, 12))\n", "\n", "# Show full attention matrix for 3 heads\n", "for col, head_idx in enumerate([0, 5, 11]):\n", " attn_full = last_attn[head_idx] # [257, 257]\n", " \n", " im = axes[0, col].imshow(attn_full, cmap='hot', aspect='auto')\n", " axes[0, col].set_title(f'Full Attention Matrix — Head {head_idx+1}\\n[257×257]',\n", " fontweight='bold')\n", " axes[0, col].set_xlabel('Key position (0=CLS, 1-256=patches)')\n", " axes[0, col].set_ylabel('Query position')\n", " # Mark CLS token\n", " axes[0, col].axhline(0.5, color='cyan', linewidth=1, alpha=0.7)\n", " axes[0, col].axvline(0.5, color='cyan', linewidth=1, alpha=0.7)\n", " plt.colorbar(im, ax=axes[0, col], fraction=0.046)\n", "\n", "# Attention entropy per head (measure of focus)\n", "entropies = []\n", "for h in range(12):\n", " attn_h = last_attn[h] # [257, 257]\n", " entropy = -(attn_h * np.log(attn_h + 1e-8)).sum(axis=-1) # [257]\n", " entropies.append(entropy)\n", "entropies = np.array(entropies) # [12, 257]\n", "\n", "# CLS token entropy per head\n", "cls_entropies = entropies[:, 0] # [12]\n", "axes[1, 0].bar(range(12), cls_entropies, color=plt.cm.tab20(np.linspace(0, 1, 12)))\n", "axes[1, 0].set_title('CLS Token Attention Entropy per Head\\n(higher = more spread out attention)',\n", " fontweight='bold')\n", "axes[1, 0].set_xlabel('Head index'); axes[1, 0].set_ylabel('Entropy')\n", "axes[1, 0].set_xticks(range(12)); axes[1, 0].set_xticklabels([f'H{i+1}' for i in range(12)])\n", "axes[1, 0].grid(True, alpha=0.3)\n", "\n", "# Mean entropy per token position\n", "mean_entropy_per_pos = entropies.mean(axis=0) # [257]\n", "axes[1, 1].plot(mean_entropy_per_pos, 'b-', linewidth=1.5)\n", "axes[1, 1].axvline(0, color='red', linestyle='--', label='CLS token')\n", "axes[1, 1].set_title('Mean Attention Entropy per Token Position\\n(averaged over 12 heads)',\n", " fontweight='bold')\n", "axes[1, 1].set_xlabel('Token position (0=CLS, 1-256=patches)')\n", "axes[1, 1].set_ylabel('Entropy')\n", "axes[1, 1].legend(); axes[1, 1].grid(True, alpha=0.3)\n", "\n", "# Attention to CLS token (how much do patches attend to CLS?)\n", "attn_to_cls = last_attn[:, 1:, 0].mean(axis=0) # [256] — mean over heads\n", "attn_to_cls_grid = attn_to_cls.reshape(n_grid, n_grid)\n", "im3 = axes[1, 2].imshow(attn_to_cls_grid, cmap='hot', interpolation='bilinear')\n", "axes[1, 2].set_title('Patch → CLS Attention\\n(how much do patches attend to CLS?)',\n", " fontweight='bold')\n", "axes[1, 2].axis('off')\n", "plt.colorbar(im3, ax=axes[1, 2], fraction=0.046)\n", "\n", "plt.suptitle('Full Attention Matrix Analysis — Last Layer', fontsize=14, fontweight='bold')\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6.7 Summary\n", "\n", "| Component | Shape | Description |\n", "|-----------|-------|-------------|\n", "| Input X | [B, 257, 768] | Token sequence |\n", "| Q, K, V | [B, 12, 257, 64] | Per-head projections |\n", "| Attn scores | [B, 12, 257, 257] | Q·Kᵀ / √64 |\n", "| Attn weights | [B, 12, 257, 257] | softmax(scores) |\n", "| Attn output | [B, 12, 257, 64] | weights · V |\n", "| MHSA output | [B, 257, 768] | concat + project |\n", "\n", "**Key observations:**\n", "- Different heads specialize in different spatial patterns\n", "- Early layers: local attention; later layers: global attention\n", "- CLS token attends broadly to semantically important patches\n", "- Attention rollout aggregates information flow across all layers\n", "\n", "---\n", "**Next:** [`07_intermediate_layer_visualization.ipynb`](07_intermediate_layer_visualization.ipynb)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 5 }