{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实验 02：从公式实现因果注意力\n",
    "\n",
    "**运行条件：** CPU，约 1 分钟，内存低于 1 GB。\n",
    "\n",
    "实验逐步构造 Q、K、V、缩放点积、因果遮罩和注意力输出，并与 PyTorch 标准实现核对。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "torch.manual_seed(7)\n",
    "batch, heads, seq_len, head_dim = 1, 2, 6, 8\n",
    "q = torch.randn(batch, heads, seq_len, head_dim)\n",
    "k = torch.randn(batch, heads, seq_len, head_dim)\n",
    "v = torch.randn(batch, heads, seq_len, head_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 缩放点积和因果遮罩\n",
    "\n",
    "缩放项 $\\sqrt{d_k}$ 控制点积方差；上三角遮罩阻止当前位置读取未来 token。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_scores = q @ k.transpose(-2, -1)\n",
    "scores = raw_scores / math.sqrt(head_dim)\n",
    "future = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool), diagonal=1)\n",
    "scores = scores.masked_fill(future, float('-inf'))\n",
    "weights = torch.softmax(scores, dim=-1)\n",
    "manual = weights @ v\n",
    "print('weights:', tuple(weights.shape))\n",
    "print('row sums:', weights.sum(dim=-1))\n",
    "assert torch.allclose(weights.sum(dim=-1), torch.ones_like(weights.sum(dim=-1)))\n",
    "assert torch.count_nonzero(weights.masked_select(future)) == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 和标准实现核对"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reference = F.scaled_dot_product_attention(q, k, v, is_causal=True)\n",
    "error = (manual - reference).abs().max().item()\n",
    "print('maximum absolute error:', error)\n",
    "assert error < 1e-5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 可视化注意力矩阵\n",
    "\n",
    "图中右上角必须保持为零，这是 decoder-only 模型自回归约束的直接表现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix = weights[0, 0].detach().numpy()\n",
    "fig, ax = plt.subplots(figsize=(5.6, 4.6))\n",
    "image = ax.imshow(matrix, vmin=0, vmax=matrix.max(), cmap='viridis')\n",
    "ax.set_xlabel('key position'); ax.set_ylabel('query position')\n",
    "ax.set_title('Causal attention weights — head 0')\n",
    "fig.colorbar(image, ax=ax); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 练习\n",
    "\n",
    "1. 删除缩放项，增大 `head_dim`，比较 softmax 熵。\n",
    "2. 把因果遮罩对角线改为 `diagonal=0`，解释为什么当前 token 也无法读取自己。\n",
    "3. 修改一个头的 Q，观察不同头如何形成不同注意力分布。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}