diff --git a/notebooks/06_VLAv3_Research_Engineer_Colab.ipynb b/notebooks/06_VLAv3_Research_Engineer_Colab.ipynb new file mode 100644 index 0000000..2f528b3 --- /dev/null +++ b/notebooks/06_VLAv3_Research_Engineer_Colab.ipynb @@ -0,0 +1,707 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# VLAv3 Research Notebook (Colab T4 Ready)\n", + "\n", + "This notebook implements and compares four Transformer variants:\n", + "1. **Softmax causal attention**\n", + "2. **Linear attention (ELU+1)**\n", + "3. **DeltaNet-style recurrent attention**\n", + "4. **VLAv3 recurrent attention**\n", + "\n", + "It also runs the three requested experiments with detailed logging and paper-grade plots:\n", + "- **Experiment 1 \u2014 Scaling behavior** (forward time + tokens/sec vs sequence length)\n", + "- **Experiment 2 \u2014 Stability behavior** (track $\\|S_t\\|$ and $\\|A_t\\|$ over time)\n", + "- **Experiment 3 \u2014 Behavior comparison** (associative retrieval accuracy vs length)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 0) Environment + reproducibility =====\n", + "import os, math, time, json, random\n", + "from dataclasses import dataclass, asdict\n", + "from pathlib import Path\n", + "from typing import Dict, List, Tuple, Optional\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "\n", + "SEED = 42\n", + "random.seed(SEED)\n", + "np.random.seed(SEED)\n", + "torch.manual_seed(SEED)\n", + "if torch.cuda.is_available():\n", + " torch.cuda.manual_seed_all(SEED)\n", + "\n", + "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n", + "DTYPE = torch.float32\n", + "\n", + "ROOT = Path('/content' if Path('/content').exists() else '.')\n", + "OUT = ROOT / 'vlav3_paper_artifacts'\n", + "OUT.mkdir(parents=True, exist_ok=True)\n", + "(OUT / 'plots').mkdir(exist_ok=True)\n", + "(OUT / 'logs').mkdir(exist_ok=True)\n", + "\n", + "print('Device:', DEVICE)\n", + "print('Torch :', torch.__version__)\n", + "print('Out :', OUT)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 1) Config =====\n", + "@dataclass\n", + "class ModelConfig:\n", + " d_model: int = 128\n", + " n_heads: int = 4\n", + " n_layers: int = 2\n", + " d_ff: int = 256\n", + " dropout: float = 0.0\n", + " vocab_size: int = 128\n", + " max_seq_len: int = 4096\n", + "\n", + "@dataclass\n", + "class TrainConfig:\n", + " batch_size: int = 32\n", + " steps: int = 400\n", + " lr: float = 3e-4\n", + " weight_decay: float = 1e-2\n", + " grad_clip: float = 1.0\n", + " log_every: int = 20" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 2) Attention implementations =====\n", + "\n", + "def _split_heads(x, n_heads):\n", + " B, T, D = x.shape\n", + " H = n_heads\n", + " Dh = D // H\n", + " return x.view(B, T, H, Dh).transpose(1, 2) # B,H,T,Dh\n", + "\n", + "def _merge_heads(x):\n", + " B, H, T, Dh = x.shape\n", + " return x.transpose(1, 2).contiguous().view(B, T, H * Dh)\n", + "\n", + "class SoftmaxCausalAttention(nn.Module):\n", + " def __init__(self, d_model, n_heads, dropout=0.0):\n", + " super().__init__()\n", + " assert d_model % n_heads == 0\n", + " self.d_model, self.n_heads = d_model, n_heads\n", + " self.qkv = nn.Linear(d_model, 3 * d_model)\n", + " self.o = nn.Linear(d_model, d_model)\n", + " self.drop = nn.Dropout(dropout)\n", + "\n", + " def forward(self, x, need_state=False):\n", + " B, T, D = x.shape\n", + " qkv = self.qkv(x)\n", + " q, k, v = qkv.chunk(3, dim=-1)\n", + " q, k, v = _split_heads(q, self.n_heads), _split_heads(k, self.n_heads), _split_heads(v, self.n_heads)\n", + " scale = 1.0 / math.sqrt(q.shape[-1])\n", + " scores = torch.matmul(q, k.transpose(-2, -1)) * scale\n", + " mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)\n", + " scores = scores.masked_fill(mask, float('-inf'))\n", + " p = F.softmax(scores, dim=-1)\n", + " p = self.drop(p)\n", + " y = torch.matmul(p, v)\n", + " y = _merge_heads(y)\n", + " y = self.o(y)\n", + " if need_state:\n", + " # Softmax has no recurrent S/A states\n", + " return y, {'S_norm': None, 'A_norm': None}\n", + " return y\n", + "\n", + "class LinearAttention(nn.Module):\n", + " \"\"\"Causal linear attention with ELU+1 features.\"\"\"\n", + " def __init__(self, d_model, n_heads, eps=1e-6):\n", + " super().__init__()\n", + " assert d_model % n_heads == 0\n", + " self.d_model, self.n_heads, self.eps = d_model, n_heads, eps\n", + " self.q = nn.Linear(d_model, d_model)\n", + " self.k = nn.Linear(d_model, d_model)\n", + " self.v = nn.Linear(d_model, d_model)\n", + " self.o = nn.Linear(d_model, d_model)\n", + "\n", + " def forward(self, x, need_state=False):\n", + " B, T, D = x.shape\n", + " H = self.n_heads\n", + " Dh = D // H\n", + "\n", + " q = F.elu(_split_heads(self.q(x), H)) + 1.0\n", + " k = F.elu(_split_heads(self.k(x), H)) + 1.0\n", + " v = _split_heads(self.v(x), H)\n", + "\n", + " S = torch.zeros(B, H, Dh, Dh, device=x.device, dtype=x.dtype)\n", + " z = torch.zeros(B, H, Dh, device=x.device, dtype=x.dtype)\n", + "\n", + " ys = []\n", + " S_norm = []\n", + " for t in range(T):\n", + " kt = k[:, :, t, :]\n", + " vt = v[:, :, t, :]\n", + " qt = q[:, :, t, :]\n", + " S = S + torch.einsum('bhd,bhe->bhde', vt, kt)\n", + " z = z + kt\n", + " yt = torch.einsum('bhde,bhe->bhd', S, qt)\n", + " den = torch.einsum('bhd,bhd->bh', z, qt).unsqueeze(-1).clamp(min=self.eps)\n", + " ys.append(yt / den)\n", + " if need_state:\n", + " S_norm.append(S.norm(dim=(-2,-1)).mean().item())\n", + "\n", + " y = torch.stack(ys, dim=2)\n", + " y = _merge_heads(y)\n", + " y = self.o(y)\n", + " if need_state:\n", + " return y, {'S_norm': S_norm, 'A_norm': None}\n", + " return y\n", + "\n", + "class DeltaNetAttention(nn.Module):\n", + " \"\"\"Simple DeltaNet-style recurrent fast-weight update with learned gate.\"\"\"\n", + " def __init__(self, d_model, n_heads, eps=1e-6):\n", + " super().__init__()\n", + " assert d_model % n_heads == 0\n", + " self.d_model, self.n_heads, self.eps = d_model, n_heads, eps\n", + " self.q = nn.Linear(d_model, d_model)\n", + " self.k = nn.Linear(d_model, d_model)\n", + " self.v = nn.Linear(d_model, d_model)\n", + " self.g = nn.Linear(d_model, d_model) # gate\n", + " self.o = nn.Linear(d_model, d_model)\n", + "\n", + " def forward(self, x, need_state=False):\n", + " B, T, D = x.shape\n", + " H = self.n_heads\n", + " Dh = D // H\n", + "\n", + " q = F.elu(_split_heads(self.q(x), H)) + 1.0\n", + " k = F.elu(_split_heads(self.k(x), H)) + 1.0\n", + " v = _split_heads(self.v(x), H)\n", + " g = torch.sigmoid(_split_heads(self.g(x), H))\n", + "\n", + " S = torch.zeros(B, H, Dh, Dh, device=x.device, dtype=x.dtype)\n", + " A = torch.zeros(B, H, Dh, Dh, device=x.device, dtype=x.dtype)\n", + "\n", + " ys, S_norm, A_norm = [], [], []\n", + " for t in range(T):\n", + " kt = k[:, :, t, :]\n", + " vt = v[:, :, t, :]\n", + " qt = q[:, :, t, :]\n", + " gt = g[:, :, t, :]\n", + "\n", + " delta = torch.einsum('bhd,bhe->bhde', vt, kt)\n", + " S = gt.unsqueeze(-1) * S + delta\n", + " A = 0.99 * A + torch.einsum('bhd,bhe->bhde', kt, kt)\n", + "\n", + " yt = torch.einsum('bhde,bhe->bhd', S, qt)\n", + " den = torch.einsum('bhde,bhd,bhe->bh', A + torch.eye(Dh, device=x.device).view(1,1,Dh,Dh)*1e-3, qt, qt).unsqueeze(-1)\n", + " ys.append(yt / den.clamp(min=self.eps))\n", + "\n", + " if need_state:\n", + " S_norm.append(S.norm(dim=(-2,-1)).mean().item())\n", + " A_norm.append(A.norm(dim=(-2,-1)).mean().item())\n", + "\n", + " y = torch.stack(ys, dim=2)\n", + " y = _merge_heads(y)\n", + " y = self.o(y)\n", + " if need_state:\n", + " return y, {'S_norm': S_norm, 'A_norm': A_norm}\n", + " return y\n", + "\n", + "class VLAv3Attention(nn.Module):\n", + " \"\"\"Stable VLAv3-style recurrent update with normalized alpha and key.\"\"\"\n", + " def __init__(self, d_model, n_heads, lambda_0=10.0, stab_eps=1e-5):\n", + " super().__init__()\n", + " assert d_model % n_heads == 0\n", + " self.d_model, self.n_heads = d_model, n_heads\n", + " self.lambda_0, self.stab_eps = lambda_0, stab_eps\n", + " self.q = nn.Linear(d_model, d_model)\n", + " self.k = nn.Linear(d_model, d_model)\n", + " self.v = nn.Linear(d_model, d_model)\n", + " self.u = nn.Linear(d_model, d_model, bias=False)\n", + " self.o = nn.Linear(d_model, d_model)\n", + "\n", + " def forward(self, x, need_state=False):\n", + " B, T, D = x.shape\n", + " H = self.n_heads\n", + " Dh = D // H\n", + "\n", + " q = F.elu(_split_heads(self.q(x), H)) + 1.0\n", + " k_raw = _split_heads(self.k(x), H)\n", + " k = F.elu(k_raw) + 1.0\n", + " v = _split_heads(self.v(x), H)\n", + " u = F.normalize(_split_heads(self.u(x), H), p=2, dim=-1)\n", + "\n", + " eye = torch.eye(Dh, device=x.device, dtype=x.dtype).view(1,1,Dh,Dh)\n", + " A = (1.0 / self.lambda_0) * eye.repeat(B, H, 1, 1)\n", + " S = torch.zeros(B, H, Dh, Dh, device=x.device, dtype=x.dtype)\n", + "\n", + " ys, S_norm, A_norm = [], [], []\n", + " for t in range(T):\n", + " ut = u[:, :, t, :] / math.sqrt(Dh)\n", + " z = torch.einsum('bhde,bhe->bhd', A, ut)\n", + " dot = torch.einsum('bhd,bhd->bh', ut, z)\n", + " delta = (1.0 + dot).clamp(min=self.stab_eps)\n", + " A = A - torch.einsum('bhd,bhe->bhde', z, z) / delta.unsqueeze(-1).unsqueeze(-1)\n", + "\n", + " kt = F.normalize(k[:, :, t, :], p=2, dim=-1)\n", + " alpha = torch.einsum('bhde,bhe->bhd', A, kt)\n", + " alpha = F.normalize(alpha, p=2, dim=-1)\n", + " vt = v[:, :, t, :]\n", + "\n", + " pred = torch.einsum('bhde,bhe->bhd', S, kt)\n", + " err = vt - pred\n", + " S = S + torch.einsum('bhd,bhe->bhde', err, alpha)\n", + "\n", + " qt = q[:, :, t, :]\n", + " num = torch.einsum('bhde,bhe->bhd', S, qt)\n", + " den = torch.einsum('bhde,bhd,bhe->bh', A, qt, qt).unsqueeze(-1).clamp(min=self.stab_eps)\n", + " ys.append(num / den)\n", + "\n", + " if need_state:\n", + " S_norm.append(S.norm(dim=(-2,-1)).mean().item())\n", + " A_norm.append(A.norm(dim=(-2,-1)).mean().item())\n", + "\n", + " y = torch.stack(ys, dim=2)\n", + " y = _merge_heads(y)\n", + " y = self.o(y)\n", + " if need_state:\n", + " return y, {'S_norm': S_norm, 'A_norm': A_norm}\n", + " return y" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 3) Transformer stack with pluggable attention =====\n", + "\n", + "ATTN_REGISTRY = {\n", + " 'softmax': SoftmaxCausalAttention,\n", + " 'linear': LinearAttention,\n", + " 'deltanet': DeltaNetAttention,\n", + " 'vla': VLAv3Attention,\n", + "}\n", + "\n", + "class TransformerBlock(nn.Module):\n", + " def __init__(self, cfg: ModelConfig, attn_type: str):\n", + " super().__init__()\n", + " self.ln1 = nn.LayerNorm(cfg.d_model)\n", + " self.ln2 = nn.LayerNorm(cfg.d_model)\n", + " self.attn = ATTN_REGISTRY[attn_type](cfg.d_model, cfg.n_heads, dropout=cfg.dropout) if attn_type=='softmax' else ATTN_REGISTRY[attn_type](cfg.d_model, cfg.n_heads)\n", + " self.ff = nn.Sequential(\n", + " nn.Linear(cfg.d_model, cfg.d_ff),\n", + " nn.GELU(),\n", + " nn.Linear(cfg.d_ff, cfg.d_model),\n", + " )\n", + "\n", + " def forward(self, x, need_state=False):\n", + " if need_state:\n", + " y, st = self.attn(self.ln1(x), need_state=True)\n", + " x = x + y\n", + " x = x + self.ff(self.ln2(x))\n", + " return x, st\n", + " x = x + self.attn(self.ln1(x))\n", + " x = x + self.ff(self.ln2(x))\n", + " return x\n", + "\n", + "class TinyTransformerLM(nn.Module):\n", + " def __init__(self, cfg: ModelConfig, attn_type: str):\n", + " super().__init__()\n", + " self.cfg = cfg\n", + " self.token = nn.Embedding(cfg.vocab_size, cfg.d_model)\n", + " self.pos = nn.Embedding(cfg.max_seq_len, cfg.d_model)\n", + " self.blocks = nn.ModuleList([TransformerBlock(cfg, attn_type) for _ in range(cfg.n_layers)])\n", + " self.ln_f = nn.LayerNorm(cfg.d_model)\n", + " self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)\n", + " self.head.weight = self.token.weight\n", + "\n", + " def forward(self, idx, need_state=False):\n", + " B, T = idx.shape\n", + " pos = torch.arange(T, device=idx.device).unsqueeze(0)\n", + " x = self.token(idx) + self.pos(pos)\n", + " state = None\n", + " for i, block in enumerate(self.blocks):\n", + " if need_state and i == len(self.blocks) - 1:\n", + " x, state = block(x, need_state=True)\n", + " else:\n", + " x = block(x)\n", + " x = self.ln_f(x)\n", + " logits = self.head(x)\n", + " return logits, state" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 4) Synthetic associative retrieval task =====\n", + "\n", + "def make_associative_batch(batch_size: int, seq_len: int, vocab_size: int, n_pairs: int, device: str):\n", + " \"\"\"\n", + " Format per sample (length = 2*n_pairs + 1 + n_queries):\n", + " k1 v1 k2 v2 ... q1 q2 ...\n", + " Targets only at query positions: corresponding values.\n", + " \"\"\"\n", + " sep = vocab_size - 1\n", + " n_queries = seq_len - (2*n_pairs + 1)\n", + " assert n_queries > 0\n", + "\n", + " x = torch.full((batch_size, seq_len), sep, dtype=torch.long, device=device)\n", + " y = torch.full((batch_size, seq_len), -100, dtype=torch.long, device=device)\n", + "\n", + " for b in range(batch_size):\n", + " keys = torch.randperm(vocab_size - 1, device=device)[:n_pairs]\n", + " vals = torch.randint(0, vocab_size - 1, (n_pairs,), device=device)\n", + " for i in range(n_pairs):\n", + " x[b, 2*i] = keys[i]\n", + " x[b, 2*i + 1] = vals[i]\n", + " x[b, 2*n_pairs] = sep\n", + "\n", + " q_idx = torch.randint(0, n_pairs, (n_queries,), device=device)\n", + " q_keys = keys[q_idx]\n", + " q_vals = vals[q_idx]\n", + " x[b, 2*n_pairs + 1:] = q_keys\n", + " y[b, 2*n_pairs + 1:] = q_vals\n", + " return x, y\n", + "\n", + "@torch.no_grad()\n", + "def eval_assoc_accuracy(model, seq_len, cfg, n_pairs=16, batches=20):\n", + " model.eval()\n", + " all_acc = []\n", + " for _ in range(batches):\n", + " x, y = make_associative_batch(cfg.batch_size, seq_len, model.cfg.vocab_size, n_pairs, DEVICE)\n", + " logits, _ = model(x)\n", + " pred = logits.argmax(-1)\n", + " mask = y != -100\n", + " acc = (pred[mask] == y[mask]).float().mean().item()\n", + " all_acc.append(acc)\n", + " return float(np.mean(all_acc))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 5) Training + logging helpers =====\n", + "\n", + "def train_assoc_model(attn_type: str, mcfg: ModelConfig, tcfg: TrainConfig, seq_len: int, n_pairs: int = 16):\n", + " model = TinyTransformerLM(mcfg, attn_type).to(DEVICE)\n", + " opt = torch.optim.AdamW(model.parameters(), lr=tcfg.lr, weight_decay=tcfg.weight_decay)\n", + "\n", + " logs = []\n", + " model.train()\n", + " t0 = time.time()\n", + " for step in range(1, tcfg.steps + 1):\n", + " x, y = make_associative_batch(tcfg.batch_size, seq_len, mcfg.vocab_size, n_pairs, DEVICE)\n", + " logits, _ = model(x)\n", + " loss = F.cross_entropy(logits.view(-1, mcfg.vocab_size), y.view(-1), ignore_index=-100)\n", + "\n", + " opt.zero_grad(set_to_none=True)\n", + " loss.backward()\n", + " torch.nn.utils.clip_grad_norm_(model.parameters(), tcfg.grad_clip)\n", + " opt.step()\n", + "\n", + " with torch.no_grad():\n", + " pred = logits.argmax(-1)\n", + " mask = y != -100\n", + " acc = (pred[mask] == y[mask]).float().mean().item()\n", + "\n", + " if step % tcfg.log_every == 0 or step == 1 or step == tcfg.steps:\n", + " row = {\n", + " 'attn_type': attn_type,\n", + " 'step': step,\n", + " 'loss': float(loss.item()),\n", + " 'train_acc': float(acc),\n", + " 'elapsed_sec': float(time.time() - t0),\n", + " }\n", + " logs.append(row)\n", + " print(f\"[{attn_type:8s}] step={step:4d} loss={row['loss']:.4f} acc={row['train_acc']:.4f} elapsed={row['elapsed_sec']:.1f}s\")\n", + "\n", + " return model, pd.DataFrame(logs)\n", + "\n", + "\n", + "def benchmark_forward(model, seq_len, batch_size=8, warmup=10, iters=30):\n", + " model.eval()\n", + " x = torch.randint(0, model.cfg.vocab_size, (batch_size, seq_len), device=DEVICE)\n", + "\n", + " for _ in range(warmup):\n", + " _ = model(x)\n", + " if DEVICE == 'cuda':\n", + " torch.cuda.synchronize()\n", + "\n", + " t0 = time.perf_counter()\n", + " for _ in range(iters):\n", + " _ = model(x)\n", + " if DEVICE == 'cuda':\n", + " torch.cuda.synchronize()\n", + " t1 = time.perf_counter()\n", + "\n", + " sec = (t1 - t0) / iters\n", + " toks_per_sec = (batch_size * seq_len) / max(sec, 1e-9)\n", + " return sec, toks_per_sec" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 6) Experiment 1 \u2014 scaling =====\n", + "\n", + "SEQ_LENS = [256, 512, 1024, 2048]\n", + "MODEL_TYPES = ['softmax', 'linear', 'deltanet', 'vla']\n", + "\n", + "mcfg = ModelConfig(d_model=128, n_heads=4, n_layers=2, d_ff=256, vocab_size=128, max_seq_len=max(SEQ_LENS))\n", + "\n", + "scaling_rows = []\n", + "for name in MODEL_TYPES:\n", + " model = TinyTransformerLM(mcfg, name).to(DEVICE)\n", + " for L in SEQ_LENS:\n", + " fwd_sec, tps = benchmark_forward(model, seq_len=L, batch_size=8, warmup=10, iters=40)\n", + " row = {\n", + " 'model': name,\n", + " 'seq_len': L,\n", + " 'forward_sec': fwd_sec,\n", + " 'tokens_per_sec': tps,\n", + " }\n", + " scaling_rows.append(row)\n", + " print(f\"[scaling] model={name:8s} L={L:4d} time={fwd_sec*1e3:8.3f} ms tok/s={tps:10.1f}\")\n", + "\n", + "scaling_df = pd.DataFrame(scaling_rows)\n", + "scaling_df.to_csv(OUT / 'logs' / 'exp1_scaling.csv', index=False)\n", + "scaling_df.head()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Plot Experiment 1\n", + "plt.figure(figsize=(12,4))\n", + "plt.subplot(1,2,1)\n", + "for m, g in scaling_df.groupby('model'):\n", + " plt.plot(g['seq_len'], g['forward_sec']*1e3, marker='o', label=m)\n", + "plt.xlabel('Sequence length')\n", + "plt.ylabel('Forward pass time (ms)')\n", + "plt.title('Experiment 1: Runtime vs Sequence length')\n", + "plt.grid(True, alpha=0.3)\n", + "plt.legend()\n", + "\n", + "plt.subplot(1,2,2)\n", + "for m, g in scaling_df.groupby('model'):\n", + " plt.plot(g['seq_len'], g['tokens_per_sec'], marker='o', label=m)\n", + "plt.xlabel('Sequence length')\n", + "plt.ylabel('Tokens / sec')\n", + "plt.title('Experiment 1: Throughput vs Sequence length')\n", + "plt.grid(True, alpha=0.3)\n", + "plt.legend()\n", + "\n", + "plt.tight_layout()\n", + "plt.savefig(OUT / 'plots' / 'exp1_scaling.png', dpi=220)\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 7) Experiment 2 \u2014 stability =====\n", + "\n", + "STAB_LEN = 2048\n", + "x = torch.randint(0, mcfg.vocab_size, (4, STAB_LEN), device=DEVICE)\n", + "\n", + "stab_rows = []\n", + "for name in MODEL_TYPES:\n", + " model = TinyTransformerLM(mcfg, name).to(DEVICE).eval()\n", + " with torch.no_grad():\n", + " _, st = model(x, need_state=True)\n", + " S_norm = st['S_norm'] if st is not None else None\n", + " A_norm = st['A_norm'] if st is not None else None\n", + "\n", + " if S_norm is None:\n", + " # preserve row format for paper tables\n", + " for t in range(STAB_LEN):\n", + " stab_rows.append({'model': name, 't': t+1, 'S_norm': np.nan, 'A_norm': np.nan})\n", + " else:\n", + " for t, s in enumerate(S_norm, start=1):\n", + " a = np.nan if A_norm is None else A_norm[t-1]\n", + " stab_rows.append({'model': name, 't': t, 'S_norm': s, 'A_norm': a})\n", + "\n", + "stability_df = pd.DataFrame(stab_rows)\n", + "stability_df.to_csv(OUT / 'logs' / 'exp2_stability.csv', index=False)\n", + "stability_df.head()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Plot Experiment 2\n", + "fig, axes = plt.subplots(1, 2, figsize=(12,4))\n", + "\n", + "for m, g in stability_df.groupby('model'):\n", + " if g['S_norm'].notna().any():\n", + " axes[0].plot(g['t'], g['S_norm'], label=m)\n", + "axes[0].set_title('Experiment 2: ||S_t|| over timestep')\n", + "axes[0].set_xlabel('t')\n", + "axes[0].set_ylabel('||S_t||')\n", + "axes[0].grid(True, alpha=0.3)\n", + "axes[0].legend()\n", + "\n", + "for m, g in stability_df.groupby('model'):\n", + " if g['A_norm'].notna().any():\n", + " axes[1].plot(g['t'], g['A_norm'], label=m)\n", + "axes[1].set_title('Experiment 2: ||A_t|| over timestep')\n", + "axes[1].set_xlabel('t')\n", + "axes[1].set_ylabel('||A_t||')\n", + "axes[1].grid(True, alpha=0.3)\n", + "axes[1].legend()\n", + "\n", + "plt.tight_layout()\n", + "plt.savefig(OUT / 'plots' / 'exp2_stability.png', dpi=220)\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 8) Experiment 3 \u2014 behavior comparison (associative retrieval) =====\n", + "\n", + "EVAL_LENS = [256, 512, 1024]\n", + "TRAIN_LEN = 256\n", + "\n", + "train_cfg = TrainConfig(batch_size=32, steps=300, lr=3e-4, log_every=25)\n", + "\n", + "behavior_rows = []\n", + "train_logs = {}\n", + "for name in MODEL_TYPES:\n", + " model, log_df = train_assoc_model(name, mcfg, train_cfg, seq_len=TRAIN_LEN, n_pairs=16)\n", + " train_logs[name] = log_df\n", + " log_df.to_csv(OUT / 'logs' / f'exp3_trainlog_{name}.csv', index=False)\n", + "\n", + " for L in EVAL_LENS:\n", + " acc = eval_assoc_accuracy(model, seq_len=L, cfg=train_cfg, n_pairs=16, batches=25)\n", + " behavior_rows.append({'model': name, 'seq_len': L, 'accuracy': acc})\n", + " print(f\"[behavior] model={name:8s} eval_len={L:4d} acc={acc:.4f}\")\n", + "\n", + "behavior_df = pd.DataFrame(behavior_rows)\n", + "behavior_df.to_csv(OUT / 'logs' / 'exp3_behavior.csv', index=False)\n", + "behavior_df.head()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Plot Experiment 3\n", + "plt.figure(figsize=(6.5,4.5))\n", + "for m, g in behavior_df.groupby('model'):\n", + " plt.plot(g['seq_len'], g['accuracy'], marker='o', label=m)\n", + "plt.xlabel('Sequence length')\n", + "plt.ylabel('Accuracy')\n", + "plt.title('Experiment 3: Associative retrieval accuracy vs length')\n", + "plt.grid(True, alpha=0.3)\n", + "plt.ylim(0.0, 1.0)\n", + "plt.legend()\n", + "plt.tight_layout()\n", + "plt.savefig(OUT / 'plots' / 'exp3_behavior.png', dpi=220)\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# ===== 9) Paper-ready summary tables + JSON export =====\n", + "\n", + "summary = {\n", + " 'device': DEVICE,\n", + " 'seed': SEED,\n", + " 'model_config': asdict(mcfg),\n", + " 'train_config': asdict(train_cfg),\n", + " 'experiment_1_scaling': scaling_df.to_dict(orient='records'),\n", + " 'experiment_2_stability_tail': (\n", + " stability_df.groupby('model')[['S_norm','A_norm']].tail(1).to_dict(orient='records')\n", + " ),\n", + " 'experiment_3_behavior': behavior_df.to_dict(orient='records'),\n", + "}\n", + "\n", + "with open(OUT / 'logs' / 'paper_summary.json', 'w') as f:\n", + " json.dump(summary, f, indent=2)\n", + "\n", + "print('Saved:')\n", + "print('-', OUT / 'plots' / 'exp1_scaling.png')\n", + "print('-', OUT / 'plots' / 'exp2_stability.png')\n", + "print('-', OUT / 'plots' / 'exp3_behavior.png')\n", + "print('-', OUT / 'logs' / 'paper_summary.json')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Notes for Colab-T4 runs\n", + "\n", + "- Keep `d_model=128`, `n_layers=2` for quick iteration.\n", + "- For publication-quality runs:\n", + " - Increase benchmark repetitions (`iters`) in Experiment 1.\n", + " - Increase training steps in Experiment 3 (e.g., 2k\u201310k).\n", + " - Run 3 seeds and report mean \u00b1 std.\n", + "- This notebook already exports CSV logs and PNG plots to `vlav3_paper_artifacts/`." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file