Add source: tmt/experiments/02_mesh_only.ipynb

tmt/experiments/02_mesh_only.ipynb

@@ -0,0 +1,52 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["# Experiment 02 — Ablation: Mesh Attention Only\n", "TMT with only Innovation 1 (mesh attention) to isolate its contribution."]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from tmt.model.config import TMTConfig\n",
+    "from tmt.model.model import TMTModel\n",
+    "from tmt.training.trainer import TMTTrainer, TrainConfig\n",
+    "from tmt.data.dataset import load_text_dataset\n",
+    "\n",
+    "# Disable innovations 2 and 3: set exit_threshold high, decay_rate=0\n",
+    "cfg = TMTConfig(\n",
+    "    vocab_size=50258, d_model=512, n_heads=8, n_layers=6,\n",
+    "    graph_k=8, decay_rate=0.0, exit_threshold=1.1,  # never exit early\n",
+    "    dual_stream=False, memory_anchors=0,\n",
+    "    ffn_stream_dim=256,\n",
+    ")\n",
+    "print(cfg)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loaders = load_text_dataset('wikitext-2', seq_len=256, batch_size=16)\n",
+    "train_cfg = TrainConfig(total_steps=2000, warmup_steps=200, use_wandb=False)\n",
+    "\n",
+    "trainer = TMTTrainer(cfg, train_cfg, loaders['train'], loaders.get('validation'))\n",
+    "trainer.train()\n",
+    "mesh_only_ppl = trainer.evaluate()\n",
+    "print(f'Mesh-only perplexity: {mesh_only_ppl:.2f}')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
+  "language_info": {"name": "python", "version": "3.10.0"}
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}