{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# DINOv2 Tutorial — Notebook 10: Attention Head Specialization\n",
                "\n",
                "**Series:** DINOv2 Deep Dive | **Prev:** `09_cls_token_analysis.ipynb` | **Next:** `11_image_similarity_retrieval.ipynb`\n",
                "\n",
                "---\n",
                "\n",
                "## What You Will Learn\n",
                "- How different attention heads specialize in different spatial patterns\n",
                "- Visualizing all 12 heads across all 12 layers\n",
                "- Head entropy analysis: focused vs diffuse attention\n",
                "- Identifying heads that attend to edges, textures, objects, backgrounds\n",
                "- Head redundancy and diversity analysis\n",
                "- The famous DINOv2 attention maps that reveal semantic segmentation\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.cm as cm\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",
                "n_grid   = cfg.image_size // cfg.patch_size\n",
                "n_heads  = cfg.num_attention_heads\n",
                "n_layers = cfg.num_hidden_layers\n",
                "\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",
                "with torch.no_grad():\n",
                "    outputs = model(**inputs, output_attentions=True)\n",
                "\n",
                "attentions = [a[0].cpu().numpy() for a in outputs.attentions]\n",
                "print(f\"Attention tensors: {len(attentions)} layers x {attentions[0].shape}\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 10.1 All Heads x All Layers: Complete Attention Map Grid"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, axes = plt.subplots(n_layers, n_heads, figsize=(24, 24))\n",
                "\n",
                "for layer_idx in range(n_layers):\n",
                "    for head_idx in range(n_heads):\n",
                "        ax = axes[layer_idx, head_idx]\n",
                "        cls_attn = attentions[layer_idx][head_idx, 0, 1:]\n",
                "        cls_attn_grid = cls_attn.reshape(n_grid, n_grid)\n",
                "        ax.imshow(cls_attn_grid, cmap='hot', interpolation='bilinear')\n",
                "        ax.axis('off')\n",
                "        if layer_idx == 0:\n",
                "            ax.set_title(f'H{head_idx+1}', fontsize=8, fontweight='bold')\n",
                "        if head_idx == 0:\n",
                "            ax.set_ylabel(f'L{layer_idx+1}', fontsize=8, fontweight='bold',\n",
                "                          rotation=0, labelpad=20)\n",
                "\n",
                "plt.suptitle('All Attention Heads x All Layers\\n'\n",
                "             'CLS Token Attention to Image Patches\\n'\n",
                "             'Rows=Layers (1-12), Columns=Heads (1-12)',\n",
                "             fontsize=14, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 10.2 Head Entropy Analysis"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "entropy_matrix = np.zeros((n_layers, n_heads))\n",
                "focus_matrix   = np.zeros((n_layers, n_heads))\n",
                "\n",
                "for layer_idx in range(n_layers):\n",
                "    for head_idx in range(n_heads):\n",
                "        cls_attn = attentions[layer_idx][head_idx, 0, 1:]\n",
                "        entropy  = -(cls_attn * np.log(cls_attn + 1e-8)).sum()\n",
                "        focus    = cls_attn.max() / (cls_attn.mean() + 1e-8)\n",
                "        entropy_matrix[layer_idx, head_idx] = entropy\n",
                "        focus_matrix[layer_idx, head_idx]   = focus\n",
                "\n",
                "fig, axes = plt.subplots(1, 3, figsize=(18, 6))\n",
                "\n",
                "im0 = axes[0].imshow(entropy_matrix, cmap='RdYlGn_r', aspect='auto')\n",
                "axes[0].set_xticks(range(n_heads))\n",
                "axes[0].set_xticklabels([f'H{i+1}' for i in range(n_heads)])\n",
                "axes[0].set_yticks(range(n_layers))\n",
                "axes[0].set_yticklabels([f'L{i+1}' for i in range(n_layers)])\n",
                "for i in range(n_layers):\n",
                "    for j in range(n_heads):\n",
                "        axes[0].text(j, i, f'{entropy_matrix[i,j]:.1f}',\n",
                "                     ha='center', va='center', fontsize=7)\n",
                "plt.colorbar(im0, ax=axes[0])\n",
                "axes[0].set_title('Attention Entropy\\n(lower=focused, higher=diffuse)', fontweight='bold')\n",
                "axes[0].set_xlabel('Head'); axes[0].set_ylabel('Layer')\n",
                "\n",
                "im1 = axes[1].imshow(focus_matrix, cmap='hot', aspect='auto')\n",
                "axes[1].set_xticks(range(n_heads))\n",
                "axes[1].set_xticklabels([f'H{i+1}' for i in range(n_heads)])\n",
                "axes[1].set_yticks(range(n_layers))\n",
                "axes[1].set_yticklabels([f'L{i+1}' for i in range(n_layers)])\n",
                "plt.colorbar(im1, ax=axes[1])\n",
                "axes[1].set_title('Attention Focus (max/mean)\\n(higher=more focused)', fontweight='bold')\n",
                "axes[1].set_xlabel('Head'); axes[1].set_ylabel('Layer')\n",
                "\n",
                "mean_entropy_per_layer = entropy_matrix.mean(axis=1)\n",
                "axes[2].plot(range(1, n_layers+1), mean_entropy_per_layer, 'b-o', linewidth=2, markersize=8)\n",
                "axes[2].fill_between(range(1, n_layers+1),\n",
                "                     entropy_matrix.min(axis=1),\n",
                "                     entropy_matrix.max(axis=1),\n",
                "                     alpha=0.2, color='blue', label='Min-Max range')\n",
                "axes[2].set_title('Mean Attention Entropy per Layer\\n(with min-max range across heads)',\n",
                "                  fontweight='bold')\n",
                "axes[2].set_xlabel('Layer'); axes[2].set_ylabel('Entropy')\n",
                "axes[2].set_xticks(range(1, n_layers+1))\n",
                "axes[2].legend(); axes[2].grid(True, alpha=0.3)\n",
                "\n",
                "plt.suptitle('Attention Head Entropy Analysis', fontsize=14, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 10.3 Head Specialization: Spatial Statistics"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "last_layer_attn = attentions[-1]\n",
                "\n",
                "head_stats = []\n",
                "for h in range(n_heads):\n",
                "    cls_attn  = last_layer_attn[h, 0, 1:]\n",
                "    attn_grid = cls_attn.reshape(n_grid, n_grid)\n",
                "    rows = np.arange(n_grid)\n",
                "    cols = np.arange(n_grid)\n",
                "    row_w = attn_grid.sum(axis=1)\n",
                "    col_w = attn_grid.sum(axis=0)\n",
                "    com_row = (rows * row_w).sum() / (row_w.sum() + 1e-8)\n",
                "    com_col = (cols * col_w).sum() / (col_w.sum() + 1e-8)\n",
                "    h_mid = n_grid // 2\n",
                "    head_stats.append({\n",
                "        'entropy':   -(cls_attn * np.log(cls_attn + 1e-8)).sum(),\n",
                "        'max_attn':  cls_attn.max(),\n",
                "        'com_row':   com_row,\n",
                "        'com_col':   com_col,\n",
                "        'q_tl': attn_grid[:h_mid, :h_mid].sum(),\n",
                "        'q_tr': attn_grid[:h_mid, h_mid:].sum(),\n",
                "        'q_bl': attn_grid[h_mid:, :h_mid].sum(),\n",
                "        'q_br': attn_grid[h_mid:, h_mid:].sum(),\n",
                "        'attn_grid': attn_grid,\n",
                "    })\n",
                "\n",
                "fig, axes = plt.subplots(3, 4, figsize=(18, 14))\n",
                "\n",
                "for h, stats in enumerate(head_stats):\n",
                "    ax = axes[h // 4, h % 4]\n",
                "    im = ax.imshow(stats['attn_grid'], cmap='hot', interpolation='bilinear')\n",
                "    ax.plot(stats['com_col'], stats['com_row'], 'c+', markersize=12, markeredgewidth=2)\n",
                "    h_mid = n_grid // 2\n",
                "    for (tr, tc, key) in [\n",
                "        (h_mid//2,       h_mid//2,       'q_tl'),\n",
                "        (h_mid//2,       h_mid+h_mid//2, 'q_tr'),\n",
                "        (h_mid+h_mid//2, h_mid//2,       'q_bl'),\n",
                "        (h_mid+h_mid//2, h_mid+h_mid//2, 'q_br'),\n",
                "    ]:\n",
                "        ax.text(tc, tr, f'{stats[key]:.2f}', ha='center', va='center',\n",
                "                fontsize=8, color='cyan', fontweight='bold')\n",
                "    ax.set_title(f'Head {h+1}\\nH={stats[\"entropy\"]:.2f}, max={stats[\"max_attn\"]:.3f}',\n",
                "                 fontweight='bold', fontsize=9)\n",
                "    ax.axis('off')\n",
                "    plt.colorbar(im, ax=ax, fraction=0.046)\n",
                "\n",
                "plt.suptitle('Last Layer: All 12 Heads with Spatial Statistics\\n'\n",
                "             'Cyan + = center of mass | Numbers = quadrant attention sum',\n",
                "             fontsize=13, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 10.4 Head Redundancy and Diversity"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, axes = plt.subplots(3, 4, figsize=(18, 14))\n",
                "\n",
                "for layer_idx in range(n_layers):\n",
                "    ax = axes[layer_idx // 4, layer_idx % 4]\n",
                "    head_maps = np.array([attentions[layer_idx][h, 0, 1:] for h in range(n_heads)])\n",
                "    head_maps_norm = head_maps / (np.linalg.norm(head_maps, axis=1, keepdims=True) + 1e-8)\n",
                "    head_sim = head_maps_norm @ head_maps_norm.T\n",
                "    im = ax.imshow(head_sim, cmap='RdYlGn', vmin=0, vmax=1)\n",
                "    ax.set_xticks(range(n_heads))\n",
                "    ax.set_xticklabels([f'H{i+1}' for i in range(n_heads)], fontsize=6)\n",
                "    ax.set_yticks(range(n_heads))\n",
                "    ax.set_yticklabels([f'H{i+1}' for i in range(n_heads)], fontsize=6)\n",
                "    mask = ~np.eye(n_heads, dtype=bool)\n",
                "    mean_sim = head_sim[mask].mean()\n",
                "    ax.set_title(f'Layer {layer_idx+1}\\nmean sim={mean_sim:.3f}', fontweight='bold', fontsize=9)\n",
                "    plt.colorbar(im, ax=ax, fraction=0.046)\n",
                "\n",
                "plt.suptitle('Head Similarity per Layer\\n'\n",
                "             'How similar are different heads? (lower = more diverse)',\n",
                "             fontsize=14, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 10.5 The Famous DINOv2 Attention Maps\n",
                "\n",
                "One of DINOv2's most striking properties: the **last layer attention maps** from the CLS token reveal **semantic segmentation** without any supervision.\n",
                "\n",
                "Different heads attend to different semantic regions of the image."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "def make_complex_image(size=224):\n",
                "    img = Image.new('RGB', (size, size), (135, 206, 235))\n",
                "    draw = ImageDraw.Draw(img)\n",
                "    draw.rectangle([0, size*2//3, size, size], fill=(80, 160, 80))\n",
                "    draw.rectangle([40, size//3, size//2, size*2//3], fill=(180, 120, 80))\n",
                "    draw.polygon([(20, size//3), (size//2+20, size//3), (size//4+10, size//6)],\n",
                "                 fill=(160, 60, 60))\n",
                "    draw.rectangle([size//4-15, size//2, size//4+15, size*2//3], fill=(100, 60, 30))\n",
                "    draw.rectangle([55, size//3+20, 95, size//3+60], fill=(200, 230, 255))\n",
                "    draw.rectangle([size*3//4-10, size//2, size*3//4+10, size*2//3], fill=(100, 70, 40))\n",
                "    draw.ellipse([size*3//4-40, size//4, size*3//4+40, size//2+20], fill=(50, 140, 50))\n",
                "    draw.ellipse([size*3//4, 20, size*3//4+40, 60], fill=(255, 220, 50))\n",
                "    return img\n",
                "\n",
                "complex_image = make_complex_image()\n",
                "inp = processor(images=complex_image, return_tensors='pt')\n",
                "inp = {k: v.to(DEVICE) for k, v in inp.items()}\n",
                "with torch.no_grad():\n",
                "    out = model(**inp, output_attentions=True)\n",
                "\n",
                "last_attn_complex = out.attentions[-1][0].cpu().numpy()\n",
                "\n",
                "fig, axes = plt.subplots(3, 5, figsize=(22, 14))\n",
                "\n",
                "axes[0, 0].imshow(complex_image)\n",
                "axes[0, 0].set_title('Original Image', fontweight='bold')\n",
                "axes[0, 0].axis('off')\n",
                "\n",
                "mean_attn = last_attn_complex[:, 0, 1:].mean(axis=0).reshape(n_grid, n_grid)\n",
                "mean_attn_norm = (mean_attn - mean_attn.min()) / (mean_attn.max() - mean_attn.min() + 1e-8)\n",
                "mean_up = np.array(Image.fromarray((mean_attn_norm*255).astype(np.uint8)).resize(\n",
                "    (224, 224), Image.BILINEAR)) / 255.0\n",
                "heatmap = cm.hot(mean_up)[:, :, :3]\n",
                "img_np = np.array(complex_image) / 255.0\n",
                "axes[0, 1].imshow((0.5*img_np + 0.5*heatmap).clip(0, 1))\n",
                "axes[0, 1].set_title('Mean Attention\\n(all heads)', fontweight='bold')\n",
                "axes[0, 1].axis('off')\n",
                "\n",
                "head_positions = [(0,2),(0,3),(0,4),(1,0),(1,1),(1,2),(1,3),(1,4),(2,0),(2,1)]\n",
                "for h_idx, (row, col) in enumerate(head_positions):\n",
                "    if h_idx >= n_heads:\n",
                "        break\n",
                "    cls_attn  = last_attn_complex[h_idx, 0, 1:].reshape(n_grid, n_grid)\n",
                "    attn_norm = (cls_attn - cls_attn.min()) / (cls_attn.max() - cls_attn.min() + 1e-8)\n",
                "    attn_up   = np.array(Image.fromarray((attn_norm*255).astype(np.uint8)).resize(\n",
                "        (224, 224), Image.BILINEAR)) / 255.0\n",
                "    heatmap   = cm.hot(attn_up)[:, :, :3]\n",
                "    overlay   = (0.4*img_np + 0.6*heatmap).clip(0, 1)\n",
                "    entropy   = -(last_attn_complex[h_idx, 0, 1:] *\n",
                "                  np.log(last_attn_complex[h_idx, 0, 1:]+1e-8)).sum()\n",
                "    axes[row, col].imshow(overlay)\n",
                "    axes[row, col].set_title(f'Head {h_idx+1}\\nH={entropy:.2f}',\n",
                "                              fontweight='bold', fontsize=9)\n",
                "    axes[row, col].axis('off')\n",
                "\n",
                "for row, col in [(2,2),(2,3),(2,4)]:\n",
                "    axes[row, col].axis('off')\n",
                "\n",
                "plt.suptitle('DINOv2 Last Layer Attention Maps\\n'\n",
                "             'Different heads attend to different semantic regions',\n",
                "             fontsize=14, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 10.6 Head Importance via Attention Variance"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "head_variances = np.zeros((n_layers, n_heads))\n",
                "\n",
                "for layer_idx in range(n_layers):\n",
                "    for head_idx in range(n_heads):\n",
                "        cls_attn = attentions[layer_idx][head_idx, 0, 1:]\n",
                "        head_variances[layer_idx, head_idx] = cls_attn.var()\n",
                "\n",
                "fig, axes = plt.subplots(1, 3, figsize=(18, 6))\n",
                "\n",
                "im0 = axes[0].imshow(head_variances * 1000, cmap='hot', aspect='auto')\n",
                "axes[0].set_xticks(range(n_heads))\n",
                "axes[0].set_xticklabels([f'H{i+1}' for i in range(n_heads)])\n",
                "axes[0].set_yticks(range(n_layers))\n",
                "axes[0].set_yticklabels([f'L{i+1}' for i in range(n_layers)])\n",
                "plt.colorbar(im0, ax=axes[0])\n",
                "axes[0].set_title('Attention Map Variance x1000\\n(higher = more discriminative)', fontweight='bold')\n",
                "axes[0].set_xlabel('Head'); axes[0].set_ylabel('Layer')\n",
                "\n",
                "k = 3\n",
                "top_heads_per_layer = head_variances.argsort(axis=1)[:, -k:]\n",
                "top_head_grid = np.zeros((n_layers, n_heads))\n",
                "for layer_idx in range(n_layers):\n",
                "    for rank, head_idx in enumerate(top_heads_per_layer[layer_idx]):\n",
                "        top_head_grid[layer_idx, head_idx] = k - rank\n",
                "\n",
                "im1 = axes[1].imshow(top_head_grid, cmap='YlOrRd', aspect='auto', vmin=0, vmax=k)\n",
                "axes[1].set_xticks(range(n_heads))\n",
                "axes[1].set_xticklabels([f'H{i+1}' for i in range(n_heads)])\n",
                "axes[1].set_yticks(range(n_layers))\n",
                "axes[1].set_yticklabels([f'L{i+1}' for i in range(n_layers)])\n",
                "plt.colorbar(im1, ax=axes[1])\n",
                "axes[1].set_title(f'Top-{k} Most Discriminative Heads per Layer',\n",
                "                  fontweight='bold')\n",
                "axes[1].set_xlabel('Head'); axes[1].set_ylabel('Layer')\n",
                "\n",
                "head_freq = np.zeros(n_heads)\n",
                "for layer_idx in range(n_layers):\n",
                "    for head_idx in top_heads_per_layer[layer_idx]:\n",
                "        head_freq[head_idx] += 1\n",
                "\n",
                "axes[2].bar(range(n_heads), head_freq,\n",
                "            color=cm.hot(head_freq / (head_freq.max() + 1e-8)))\n",
                "axes[2].set_title(f'Head Frequency in Top-{k} Across All Layers\\n(higher = consistently important)',\n",
                "                  fontweight='bold')\n",
                "axes[2].set_xlabel('Head index'); axes[2].set_ylabel(f'Times in top-{k}')\n",
                "axes[2].set_xticks(range(n_heads))\n",
                "axes[2].set_xticklabels([f'H{i+1}' for i in range(n_heads)])\n",
                "axes[2].grid(True, alpha=0.3)\n",
                "\n",
                "plt.suptitle('Head Importance Analysis', fontsize=14, fontweight='bold')\n",
                "plt.tight_layout()\n",
                "plt.show()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## 10.7 Summary\n",
                "\n",
                "| Observation | Explanation |\n",
                "|-------------|-------------|\n",
                "| Early layers: focused attention | Local feature detection |\n",
                "| Late layers: diffuse attention | Global context aggregation |\n",
                "| Different heads: different regions | Specialization for different semantic parts |\n",
                "| High-entropy heads | Attend broadly (background, global context) |\n",
                "| Low-entropy heads | Attend to specific objects/regions |\n",
                "| Last layer reveals segmentation | Semantic features emerge from SSL training |\n",
                "\n",
                "**The key DINOv2 insight:** Self-supervised training with multi-crop and iBOT forces the model to develop **semantically meaningful attention patterns** — different heads learn to attend to different semantic regions without any segmentation supervision.\n",
                "\n",
                "---\n",
                "**Next:** [`11_image_similarity_retrieval.ipynb`](11_image_similarity_retrieval.ipynb)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "name": "python",
            "version": "3.10.0"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}