{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e5efb596",
   "metadata": {},
   "source": [
    "# Attention Is All You Need — Walkthrough\n",
    "\n",
    "**Paper:** [Attention Is All You Need](https://arxiv.org/abs/1706.03762) (Vaswani et al., 2017)\n",
    "\n",
    "This notebook walks through the Transformer architecture, connecting each paper section to code.\n",
    "All examples use **tiny dimensions** so everything runs instantly on CPU.\n",
    "The architecture is identical to the paper — only the numbers change.\n",
    "\n",
    "For each component:\n",
    "1. Paper excerpt (the relevant passage)\n",
    "2. Code (the implementation)\n",
    "3. Sanity check (verify shapes and behavior)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c8a4651",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from dataclasses import dataclass\n",
    "from typing import Optional\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# Toy dimensions — paper uses d_model=512, we use 64 for CPU walkthrough\n",
    "BATCH_SIZE = 2\n",
    "SRC_LEN = 10\n",
    "TGT_LEN = 8\n",
    "D_MODEL = 64      # Paper: 512\n",
    "N_HEADS = 4        # Paper: 8\n",
    "D_FF = 128         # Paper: 2048\n",
    "N_LAYERS = 2       # Paper: 6\n",
    "VOCAB_SIZE = 100   # Paper: 37000\n",
    "DROPOUT = 0.0      # Set to 0 for deterministic walkthrough\n",
    "\n",
    "print(f\"Using toy dimensions: d_model={D_MODEL}, n_heads={N_HEADS}, d_ff={D_FF}\")\n",
    "print(f\"d_k = d_model / n_heads = {D_MODEL // N_HEADS}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8346be9",
   "metadata": {},
   "source": [
    "## 1. Scaled Dot-Product Attention (§3.2.1)\n",
    "\n",
    "> \"An attention function can be described as mapping a query and a set of key-value\n",
    "> pairs to an output, where the query, keys, values, and output are all vectors.\n",
    "> The output is computed as a weighted sum of the values, where the weight assigned\n",
    "> to each value is computed by a compatibility function of the query with the\n",
    "> corresponding key.\"\n",
    ">\n",
    "> — §3.2, Attention Is All You Need\n",
    "\n",
    "**Equation 1:**\n",
    "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\frac{QK^T}{\\sqrt{d_k}}\\right)V$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a40f5933",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scaled_dot_product_attention(\n",
    "    query: torch.Tensor,\n",
    "    key: torch.Tensor,\n",
    "    value: torch.Tensor,\n",
    "    mask: Optional[torch.Tensor] = None,\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"§3.2.1, Eq. 1 — Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V\"\"\"\n",
    "    d_k = query.size(-1)\n",
    "    \n",
    "    # (batch, n_heads, seq_q, d_k) @ (batch, n_heads, d_k, seq_k) -> (batch, n_heads, seq_q, seq_k)\n",
    "    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)\n",
    "    \n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask == 0, float('-inf'))\n",
    "    \n",
    "    attn_weights = F.softmax(scores, dim=-1)  # (batch, n_heads, seq_q, seq_k)\n",
    "    return torch.matmul(attn_weights, value)   # (batch, n_heads, seq_q, d_v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee66cd76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sanity check: scaled dot-product attention\n",
    "d_k = D_MODEL // N_HEADS\n",
    "q = torch.randn(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)\n",
    "k = torch.randn(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)\n",
    "v = torch.randn(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)\n",
    "\n",
    "attn_out = scaled_dot_product_attention(q, k, v)\n",
    "assert attn_out.shape == (BATCH_SIZE, N_HEADS, SRC_LEN, d_k), \\\n",
    "    f\"Expected {(BATCH_SIZE, N_HEADS, SRC_LEN, d_k)}, got {attn_out.shape}\"\n",
    "print(f\"✓ Scaled dot-product attention: ({BATCH_SIZE}, {N_HEADS}, {SRC_LEN}, {d_k}) -> {attn_out.shape}\")\n",
    "\n",
    "# Verify attention weights sum to 1\n",
    "scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)\n",
    "weights = F.softmax(scores, dim=-1)\n",
    "weight_sums = weights.sum(dim=-1)\n",
    "assert torch.allclose(weight_sums, torch.ones_like(weight_sums), atol=1e-6), \\\n",
    "    \"Attention weights should sum to 1\"\n",
    "print(f\"✓ Attention weights sum to 1 (verified)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9585614f",
   "metadata": {},
   "source": [
    "## 2. Multi-Head Attention (§3.2.2)\n",
    "\n",
    "> \"Multi-head attention allows the model to jointly attend to information from\n",
    "> different representation subspaces at different positions. With a single attention\n",
    "> head, averaging inhibits this.\"\n",
    ">\n",
    "> \"MultiHead(Q, K, V) = Concat(head₁, ..., headₕ)W^O\n",
    ">  where headᵢ = Attention(QWᵢ^Q, KWᵢ^K, VWᵢ^V)\"\n",
    ">\n",
    "> \"We employ h = 8 parallel attention layers, or heads. For each of these we use\n",
    "> d_k = d_v = d_model/h = 64.\"\n",
    ">\n",
    "> — §3.2.2, Attention Is All You Need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab01cf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\"§3.2.2 — Multi-Head Attention.\"\"\"\n",
    "    \n",
    "    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):\n",
    "        super().__init__()\n",
    "        assert d_model % n_heads == 0\n",
    "        self.d_model = d_model\n",
    "        self.n_heads = n_heads\n",
    "        self.d_k = d_model // n_heads  # §3.2.2 — d_k = d_model/h\n",
    "        \n",
    "        # §3.2.2 — projection matrices W_Q, W_K, W_V, W_O\n",
    "        self.W_q = nn.Linear(d_model, d_model)\n",
    "        self.W_k = nn.Linear(d_model, d_model)\n",
    "        self.W_v = nn.Linear(d_model, d_model)\n",
    "        self.W_o = nn.Linear(d_model, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        batch_size = query.size(0)\n",
    "        \n",
    "        # Project and reshape: (batch, seq, d_model) -> (batch, n_heads, seq, d_k)\n",
    "        q = self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)\n",
    "        k = self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)\n",
    "        v = self.W_v(value).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)\n",
    "        \n",
    "        # Scaled dot-product attention\n",
    "        attn_out = scaled_dot_product_attention(q, k, v, mask=mask)\n",
    "        \n",
    "        # Concatenate and project: (batch, n_heads, seq, d_k) -> (batch, seq, d_model)\n",
    "        attn_out = attn_out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)\n",
    "        return self.W_o(attn_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61e16c15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sanity check: multi-head attention\n",
    "mha = MultiHeadAttention(D_MODEL, N_HEADS, dropout=0.0)\n",
    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
    "out = mha(x, x, x)  # self-attention\n",
    "assert out.shape == (BATCH_SIZE, SRC_LEN, D_MODEL), \\\n",
    "    f\"Expected {(BATCH_SIZE, SRC_LEN, D_MODEL)}, got {out.shape}\"\n",
    "print(f\"✓ MultiHeadAttention: (batch={BATCH_SIZE}, seq={SRC_LEN}, d_model={D_MODEL}) -> {out.shape}\")\n",
    "print(f\"  Parameters: {sum(p.numel() for p in mha.parameters()):,}\")\n",
    "print(f\"  Expected: 4 * d_model^2 + 4 * d_model (biases) = {4 * D_MODEL**2 + 4 * D_MODEL:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6191f70",
   "metadata": {},
   "source": [
    "## 3. Position-wise Feed-Forward Network (§3.3)\n",
    "\n",
    "> \"In addition to attention sub-layers, each of the layers in our encoder and decoder\n",
    "> contains a fully connected feed-forward network, which is applied to each position\n",
    "> separately and identically.\"\n",
    ">\n",
    "> \"FFN(x) = max(0, xW₁ + b₁)W₂ + b₂\" (Eq. 2)\n",
    ">\n",
    "> \"The dimensionality of input and output is d_model = 512, and the inner-layer has\n",
    "> dimensionality d_ff = 2048.\"\n",
    ">\n",
    "> — §3.3, Attention Is All You Need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a52d14af",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    \"\"\"§3.3, Eq. 2 — FFN(x) = max(0, xW_1 + b_1)W_2 + b_2\"\"\"\n",
    "    \n",
    "    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0):\n",
    "        super().__init__()\n",
    "        self.linear1 = nn.Linear(d_model, d_ff)\n",
    "        self.linear2 = nn.Linear(d_ff, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # §3.3, Eq. 2 — ReLU activation\n",
    "        x = self.linear1(x)     # (batch, seq, d_model) -> (batch, seq, d_ff)\n",
    "        x = F.relu(x)           # max(0, ...)\n",
    "        x = self.dropout(x)\n",
    "        x = self.linear2(x)     # (batch, seq, d_ff) -> (batch, seq, d_model)\n",
    "        return x\n",
    "\n",
    "# Sanity check\n",
    "ff = FeedForward(D_MODEL, D_FF)\n",
    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
    "out = ff(x)\n",
    "assert out.shape == (BATCH_SIZE, SRC_LEN, D_MODEL)\n",
    "print(f\"✓ FeedForward: {x.shape} -> {out.shape}\")\n",
    "print(f\"  Inner dimension: {D_FF} (expansion factor: {D_FF/D_MODEL}x)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92ca2960",
   "metadata": {},
   "source": [
    "## 4. Positional Encoding (§3.5)\n",
    "\n",
    "> \"Since our model contains no recurrence and no convolution, in order for the model\n",
    "> to make use of the order of the sequence, we must inject some information about the\n",
    "> relative or absolute position of the tokens in the sequence.\"\n",
    ">\n",
    "> \"We use sine and cosine functions of different frequencies:\n",
    ">   PE(pos, 2i) = sin(pos / 10000^(2i/d_model))\n",
    ">   PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))\"\n",
    ">\n",
    "> — §3.5, Attention Is All You Need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63a07ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    \"\"\"§3.5 — Sinusoidal positional encoding.\"\"\"\n",
    "    \n",
    "    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.0):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len).unsqueeze(1).float()\n",
    "        div_term = torch.exp(\n",
    "            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)\n",
    "        )\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)  # even: sin\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)  # odd: cos\n",
    "        pe = pe.unsqueeze(0)  # (1, max_len, d_model)\n",
    "        self.register_buffer('pe', pe)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:, :x.size(1)]\n",
    "        return self.dropout(x)\n",
    "\n",
    "# Sanity check\n",
    "pos_enc = PositionalEncoding(D_MODEL, max_len=100)\n",
    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
    "out = pos_enc(x)\n",
    "assert out.shape == x.shape\n",
    "print(f\"✓ PositionalEncoding: {x.shape} -> {out.shape} (added, not concatenated)\")\n",
    "\n",
    "# Verify PE values are bounded (sin/cos are in [-1, 1])\n",
    "pe_values = pos_enc.pe[0, :SRC_LEN, :]\n",
    "assert pe_values.min() >= -1.0 and pe_values.max() <= 1.0\n",
    "print(f\"✓ PE values bounded in [-1, 1]: min={pe_values.min():.3f}, max={pe_values.max():.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "067a45be",
   "metadata": {},
   "source": [
    "## 5. Encoder Layer (§3.1)\n",
    "\n",
    "> \"The encoder is composed of a stack of N = 6 identical layers. Each layer has two\n",
    "> sub-layers. The first is a multi-head self-attention mechanism, and the second is a\n",
    "> simple, position-wise fully connected feed-forward network.\"\n",
    ">\n",
    "> \"We employ a residual connection around each of the two sub-layers, followed by\n",
    "> layer normalization. That is, the output of each sub-layer is LayerNorm(x + Sublayer(x)).\"\n",
    ">\n",
    "> — §3.1, Attention Is All You Need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a7aa6d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    \"\"\"§3.1 — Single encoder layer: self-attention + feed-forward, with residual + norm.\"\"\"\n",
    "    \n",
    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.0, norm_eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
    "        self.feed_forward = FeedForward(d_model, d_ff, dropout)\n",
    "        # [UNSPECIFIED] LayerNorm epsilon not stated in paper — using 1e-6\n",
    "        self.norm1 = nn.LayerNorm(d_model, eps=norm_eps)\n",
    "        self.norm2 = nn.LayerNorm(d_model, eps=norm_eps)\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x, mask=None):\n",
    "        # §3.1 — Post-norm: LayerNorm(x + Sublayer(x))\n",
    "        attn_out = self.self_attn(x, x, x, mask)  # self-attention\n",
    "        x = self.norm1(x + self.dropout1(attn_out))\n",
    "        ff_out = self.feed_forward(x)\n",
    "        x = self.norm2(x + self.dropout2(ff_out))\n",
    "        return x\n",
    "\n",
    "# Sanity check\n",
    "enc_layer = EncoderLayer(D_MODEL, N_HEADS, D_FF)\n",
    "x = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
    "out = enc_layer(x)\n",
    "assert out.shape == (BATCH_SIZE, SRC_LEN, D_MODEL)\n",
    "print(f\"✓ EncoderLayer: {x.shape} -> {out.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b658381",
   "metadata": {},
   "source": [
    "## 6. Decoder Layer (§3.1)\n",
    "\n",
    "> \"In addition to the two sub-layers in each encoder layer, the decoder inserts a\n",
    "> third sub-layer, which performs multi-head attention over the output of the encoder\n",
    "> stack.\"\n",
    ">\n",
    "> \"We also modify the self-attention sub-layer in the decoder stack to prevent\n",
    "> positions from attending to subsequent positions. This masking, combined with the\n",
    "> fact that the output embeddings are offset by one position, ensures that the\n",
    "> predictions for position i can depend only on the known outputs at positions less\n",
    "> than i.\"\n",
    ">\n",
    "> — §3.1, §3.2.3, Attention Is All You Need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffdaf93e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    \"\"\"§3.1 — Single decoder layer: masked self-attn + cross-attn + FFN.\"\"\"\n",
    "    \n",
    "    def __init__(self, d_model, n_heads, d_ff, dropout=0.0, norm_eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
    "        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)\n",
    "        self.feed_forward = FeedForward(d_model, d_ff, dropout)\n",
    "        self.norm1 = nn.LayerNorm(d_model, eps=norm_eps)\n",
    "        self.norm2 = nn.LayerNorm(d_model, eps=norm_eps)\n",
    "        self.norm3 = nn.LayerNorm(d_model, eps=norm_eps)\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "        self.dropout3 = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x, memory, tgt_mask=None, memory_mask=None):\n",
    "        # Sub-layer 1: Masked self-attention\n",
    "        self_attn_out = self.self_attn(x, x, x, tgt_mask)\n",
    "        x = self.norm1(x + self.dropout1(self_attn_out))\n",
    "        # Sub-layer 2: Cross-attention over encoder output\n",
    "        cross_attn_out = self.cross_attn(x, memory, memory, memory_mask)\n",
    "        x = self.norm2(x + self.dropout2(cross_attn_out))\n",
    "        # Sub-layer 3: Feed-forward\n",
    "        ff_out = self.feed_forward(x)\n",
    "        x = self.norm3(x + self.dropout3(ff_out))\n",
    "        return x\n",
    "\n",
    "# Sanity check\n",
    "dec_layer = DecoderLayer(D_MODEL, N_HEADS, D_FF)\n",
    "x = torch.randn(BATCH_SIZE, TGT_LEN, D_MODEL)\n",
    "memory = torch.randn(BATCH_SIZE, SRC_LEN, D_MODEL)\n",
    "\n",
    "# Causal mask for decoder self-attention\n",
    "causal_mask = torch.tril(torch.ones(TGT_LEN, TGT_LEN, dtype=torch.bool)).unsqueeze(0).unsqueeze(0)\n",
    "out = dec_layer(x, memory, tgt_mask=causal_mask)\n",
    "assert out.shape == (BATCH_SIZE, TGT_LEN, D_MODEL)\n",
    "print(f\"✓ DecoderLayer: tgt {x.shape} + memory {memory.shape} -> {out.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9677217c",
   "metadata": {},
   "source": [
    "## 7. Full Transformer (§3)\n",
    "\n",
    "> \"Most competitive neural sequence transduction models have an encoder-decoder\n",
    "> structure. Here, the encoder maps an input sequence of symbol representations\n",
    "> (x₁, ..., xₙ) to a sequence of continuous representations z = (z₁, ..., zₙ).\n",
    "> Given z, the decoder then generates an output sequence (y₁, ..., yₘ) of symbols\n",
    "> one element at a time.\"\n",
    ">\n",
    "> — §2, Attention Is All You Need\n",
    "\n",
    "> \"In the embedding layers, we multiply those weights by √d_model.\"\n",
    "> \"In our model, we share the same weight matrix between the two embedding layers\n",
    "> and the pre-softmax linear transformation.\"\n",
    ">\n",
    "> — §3.4, Attention Is All You Need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb4adb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    \"\"\"§3 — Full Transformer encoder-decoder.\"\"\"\n",
    "    \n",
    "    def __init__(self, d_model, n_heads, d_ff, n_layers, vocab_size,\n",
    "                 max_len=5000, dropout=0.0, norm_eps=1e-6, tie_weights=True):\n",
    "        super().__init__()\n",
    "        \n",
    "        # §3.4 — Embeddings\n",
    "        self.src_embed = nn.Embedding(vocab_size, d_model)\n",
    "        self.tgt_embed = nn.Embedding(vocab_size, d_model)\n",
    "        if tie_weights:  # §3.4 — \"share the same weight matrix\"\n",
    "            self.tgt_embed.weight = self.src_embed.weight\n",
    "        \n",
    "        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)\n",
    "        self.embed_scale = math.sqrt(d_model)  # §3.4 — multiply by sqrt(d_model)\n",
    "        \n",
    "        # §3.1 — Encoder: N identical layers\n",
    "        self.encoder_layers = nn.ModuleList([\n",
    "            EncoderLayer(d_model, n_heads, d_ff, dropout, norm_eps)\n",
    "            for _ in range(n_layers)\n",
    "        ])\n",
    "        \n",
    "        # §3.1 — Decoder: N identical layers\n",
    "        self.decoder_layers = nn.ModuleList([\n",
    "            DecoderLayer(d_model, n_heads, d_ff, dropout, norm_eps)\n",
    "            for _ in range(n_layers)\n",
    "        ])\n",
    "        \n",
    "        # §3.4 — Output projection (tied with embeddings)\n",
    "        self.output_proj = nn.Linear(d_model, vocab_size, bias=False)\n",
    "        if tie_weights:\n",
    "            self.output_proj.weight = self.src_embed.weight\n",
    "    \n",
    "    def forward(self, src, tgt, src_mask=None, tgt_mask=None):\n",
    "        tgt_len = tgt.size(1)\n",
    "        \n",
    "        # Auto-generate causal mask if not provided\n",
    "        if tgt_mask is None:\n",
    "            tgt_mask = torch.tril(torch.ones(tgt_len, tgt_len, \n",
    "                         dtype=torch.bool, device=tgt.device)).unsqueeze(0).unsqueeze(0)\n",
    "        \n",
    "        # §3.4 — Embed and scale by sqrt(d_model)\n",
    "        src_emb = self.pos_encoding(self.src_embed(src) * self.embed_scale)\n",
    "        tgt_emb = self.pos_encoding(self.tgt_embed(tgt) * self.embed_scale)\n",
    "        \n",
    "        # Encode\n",
    "        memory = src_emb\n",
    "        for layer in self.encoder_layers:\n",
    "            memory = layer(memory, src_mask)\n",
    "        \n",
    "        # Decode\n",
    "        output = tgt_emb\n",
    "        for layer in self.decoder_layers:\n",
    "            output = layer(output, memory, tgt_mask, src_mask)\n",
    "        \n",
    "        # §3.4 — Project to vocabulary\n",
    "        logits = self.output_proj(output)  # (batch, tgt_len, vocab_size)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc02bf24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sanity check: full model forward pass\n",
    "model = Transformer(\n",
    "    d_model=D_MODEL, n_heads=N_HEADS, d_ff=D_FF,\n",
    "    n_layers=N_LAYERS, vocab_size=VOCAB_SIZE,\n",
    "    dropout=0.0,\n",
    ")\n",
    "\n",
    "src = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SRC_LEN))\n",
    "tgt = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, TGT_LEN))\n",
    "logits = model(src, tgt)\n",
    "\n",
    "assert logits.shape == (BATCH_SIZE, TGT_LEN, VOCAB_SIZE), \\\n",
    "    f\"Expected {(BATCH_SIZE, TGT_LEN, VOCAB_SIZE)}, got {logits.shape}\"\n",
    "print(f\"✓ Full Transformer: src {src.shape} + tgt {tgt.shape} -> logits {logits.shape}\")\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"  Parameters: {total_params:,}\")\n",
    "print(f\"  (Paper base model has ~65M params with d_model=512, vocab=37k)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e4799ba",
   "metadata": {},
   "source": [
    "## 8. Label-Smoothed Loss (§5.4)\n",
    "\n",
    "> \"During training, we employed label smoothing of value ε_ls = 0.1. This hurt\n",
    "> perplexity, as the model learns to be more unsure, but improved accuracy and\n",
    "> BLEU score.\"\n",
    ">\n",
    "> — §5.4, Attention Is All You Need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a59bd69",
   "metadata": {},
   "outputs": [],
   "source": [
    "def label_smoothed_cross_entropy(logits, targets, smoothing=0.1, pad_idx=0):\n",
    "    \"\"\"§5.4 — Label smoothing with ε_ls = 0.1.\"\"\"\n",
    "    vocab_size = logits.size(-1)\n",
    "    logits = logits.contiguous().view(-1, vocab_size)\n",
    "    targets = targets.contiguous().view(-1)\n",
    "    \n",
    "    log_probs = F.log_softmax(logits, dim=-1)\n",
    "    nll_loss = -log_probs.gather(dim=-1, index=targets.unsqueeze(-1)).squeeze(-1)\n",
    "    smooth_loss = -log_probs.mean(dim=-1)\n",
    "    \n",
    "    loss = (1.0 - smoothing) * nll_loss + smoothing * smooth_loss\n",
    "    non_pad = targets != pad_idx\n",
    "    return (loss * non_pad.float()).sum() / non_pad.float().sum().clamp(min=1.0)\n",
    "\n",
    "# Sanity check: loss computes and backprops\n",
    "targets = torch.randint(1, VOCAB_SIZE, (BATCH_SIZE, TGT_LEN))  # no padding (idx 0)\n",
    "loss = label_smoothed_cross_entropy(logits, targets, smoothing=0.1)\n",
    "loss.backward()\n",
    "print(f\"✓ Label-smoothed loss: {loss.item():.4f}\")\n",
    "print(f\"  (Compare: -log(1/V) = {-math.log(1/VOCAB_SIZE):.4f} for random predictions)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1174650f",
   "metadata": {},
   "source": [
    "## 9. Gradient Check\n",
    "\n",
    "Verify that gradients flow through all parameters of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da524fdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check gradients exist for all parameters\n",
    "no_grad_params = []\n",
    "for name, param in model.named_parameters():\n",
    "    if param.grad is None:\n",
    "        no_grad_params.append(name)\n",
    "\n",
    "if no_grad_params:\n",
    "    print(f\"⚠ Parameters without gradients: {no_grad_params}\")\n",
    "else:\n",
    "    print(f\"✓ All {sum(1 for _ in model.parameters())} parameter tensors have gradients\")\n",
    "    \n",
    "# Check gradient norms\n",
    "total_norm = 0.0\n",
    "for param in model.parameters():\n",
    "    if param.grad is not None:\n",
    "        total_norm += param.grad.data.norm(2).item() ** 2\n",
    "total_norm = total_norm ** 0.5\n",
    "print(f\"  Total gradient norm: {total_norm:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605c8821",
   "metadata": {},
   "source": [
    "## Common Pitfalls\n",
    "\n",
    "1. **Pre-norm vs post-norm**: The paper's Figure 1 shows post-norm (`LayerNorm(x + Sublayer(x))`), and §3.1 says \"followed by layer normalization\" which also implies post-norm. However, many modern reimplementations use pre-norm (`x + Sublayer(LayerNorm(x))`) because it's more stable for deep models. If you're having trouble training, try switching to pre-norm. [§3.1]\n",
    "\n",
    "2. **Attention scale factor**: Must divide by √d_k where d_k = d_model/n_heads, NOT √d_model. For the base model: √64 = 8, not √512 ≈ 22.6. This is a 3x difference that significantly affects attention sharpness. [§3.2.1, Eq. 1]\n",
    "\n",
    "3. **Positional encoding frequencies**: The div_term uses `exp(arange(0, d_model, 2) * (-log(10000) / d_model))`, stepping by 2. Using `arange(0, d_model)` (step 1) doubles the frequency count and gives wrong wavelengths. [§3.5]\n",
    "\n",
    "4. **Label smoothing**: ε_ls = 0.1 means the true class gets probability 1 - ε + ε/V ≈ 0.9 (not exactly 0.9). PyTorch's `CrossEntropyLoss(label_smoothing=0.1)` uses a slightly different formula — for large V this is negligible but worth knowing. [§5.4]\n",
    "\n",
    "5. **Learning rate schedule**: The Transformer schedule (Eq. 3) computes the LR from scratch — it's not a multiplier on a base LR. Set the optimizer LR to 1.0 and let the schedule compute the actual LR. Peaks at ~step 4000 for d_model=512. [§5.3, Eq. 3]\n",
    "\n",
    "6. **Embedding scaling**: The paper multiplies embedding weights by √d_model (§3.4). This is easy to forget and affects the relative magnitude of embeddings vs positional encodings.\n",
    "\n",
    "7. **Weight tying**: The paper shares weights between (a) source embedding, (b) target embedding, and (c) output projection. Missing this changes the parameter count and can affect quality. [§3.4]"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
