[FEAT#76] neuron Phase 16a — RigL/SET edge regrow (constant sparsity DST)

notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb

@@ -0,0 +1,339 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# 15-phase16a-rigl-set-regrow\n",
+    "\n",
+    "**neuron Phase 16a** — paradigm 의 dynamic phase 두 번째 단계. Phase 15 의 *static prune* 을 확장하여 **iterative prune + regrow** 로 학습 중 topology 가 진화하는 DST (Dynamic Sparse Training) 구현.\n",
+    "\n",
+    "핵심 가설:\n",
+    "1. **RigL ≥ static prune?** — 같은 50% sparsity 에서 RigL DST 가 단순 prune 보다 낮은 loss?\n",
+    "2. **RigL ≥ SET?** — gradient-guided regrow 가 random regrow 보다 우수?\n",
+    "3. **constant sparsity 유지** — 각 DST cycle 후 sparsity 일정?\n",
+    "4. **all-finite** — DST cycle 도중 학습 안정성?\n",
+    "\n",
+    "설계: full graph block (`hybrid_around_one_around_one`, Phase 14 최저 loss 구조) 위에서 **4 mode × 2 seed = 8 run**.\n",
+    "- `dense`: prune 없음 (baseline)\n",
+    "- `static_50`: 50% prune at step 750, regrow 없음 (Phase 15)\n",
+    "- `dst_set_50`: 50% prune at step 750 + random regrow (SET) every 50 steps, swap 10%\n",
+    "- `dst_rigl_50`: 50% prune at step 750 + gradient-based regrow (RigL) every 50 steps, swap 10%\n",
+    "- dst_end_step = 1300 (마지막 200 step 은 stabilize)\n",
+    "\n",
+    "데이터: TinyShakespeare (char-LM, block_size=64)\n",
+    "시드: [42, 123]\n",
+    "작성일: 2026-05-27\n",
+    "연관: Issue [#76](https://github.com/EinSofINTEREST/GraphLM/issues/76) (Phase 16 main [#75](https://github.com/EinSofINTEREST/GraphLM/issues/75)) / Phase 15 PR [#72](https://github.com/EinSofINTEREST/GraphLM/pull/72)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "## 0. 환경 / 의존성"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations\n",
+    "\n",
+    "import math\n",
+    "import statistics\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import torch\n",
+    "\n",
+    "from graphlm.data.tinyshakespeare import (\n",
+    "    CharTokenizer,\n",
+    "    TinyShakespeareDataset,\n",
+    "    load_tinyshakespeare_text,\n",
+    ")\n",
+    "from graphlm.neuron.hybrid_transformer_demo import (\n",
+    "    HybridTransformerTrainConfig,\n",
+    "    train_hybrid_transformer_lm,\n",
+    ")\n",
+    "from graphlm.utils import safe_perplexity\n",
+    "\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "print(f\"device: {device}\")\n",
+    "print(f\"torch: {torch.__version__}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "## 1. Config + 데이터"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text = load_tinyshakespeare_text()\n",
+    "tokenizer = CharTokenizer(text)\n",
+    "dataset = TinyShakespeareDataset(text, tokenizer)\n",
+    "vocab_size = tokenizer.vocab_size\n",
+    "print(f\"vocab_size = {vocab_size}, dataset size = {len(dataset)}\")\n",
+    "\n",
+    "# Phase 15 와 동일 hyperparameter (공정 비교)\n",
+    "HIDDEN_DIM = 128\n",
+    "N_HEADS = 4\n",
+    "FFN_DIM = 256\n",
+    "N_LAYERS = 4\n",
+    "GROUP_SIZE = 16\n",
+    "BLOCK_SIZE = 64\n",
+    "BATCH_SIZE = 32\n",
+    "LR = 3e-4\n",
+    "MAX_STEPS = 1500\n",
+    "PRUNE_AT_STEP = MAX_STEPS // 2  # 750\n",
+    "DST_PERIOD = 50  # 매 50 step 마다 DST cycle\n",
+    "DST_SWAP_FRACTION = 0.1  # alive 의 10% swap\n",
+    "DST_END_STEP = 1300  # 마지막 200 step stabilize\n",
+    "SEEDS = [42, 123]\n",
+    "ARCH = \"hybrid_around_one_around_one\"\n",
+    "\n",
+    "# 4 mode 정의\n",
+    "MODES = [\n",
+    "    {\"name\": \"dense\", \"prune_fraction\": 0.0, \"prune_at_step\": None, \"regrow_method\": None},\n",
+    "    {\n",
+    "        \"name\": \"static_50\",\n",
+    "        \"prune_fraction\": 0.5,\n",
+    "        \"prune_at_step\": PRUNE_AT_STEP,\n",
+    "        \"regrow_method\": None,\n",
+    "    },\n",
+    "    {\n",
+    "        \"name\": \"dst_set_50\",\n",
+    "        \"prune_fraction\": 0.5,\n",
+    "        \"prune_at_step\": PRUNE_AT_STEP,\n",
+    "        \"regrow_method\": \"random\",\n",
+    "    },\n",
+    "    {\n",
+    "        \"name\": \"dst_rigl_50\",\n",
+    "        \"prune_fraction\": 0.5,\n",
+    "        \"prune_at_step\": PRUNE_AT_STEP,\n",
+    "        \"regrow_method\": \"rigl\",\n",
+    "    },\n",
+    "]\n",
+    "print(\n",
+    "    f\"\\nDST cycle: period={DST_PERIOD}, swap_fraction={DST_SWAP_FRACTION}, end_step={DST_END_STEP}\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "## 2. Sweep 실행 (4 mode × 2 seed = 8 run)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "results = {}\n",
+    "for mode in MODES:\n",
+    "    for seed in SEEDS:\n",
+    "        key = (mode[\"name\"], seed)\n",
+    "        print(f\"\\n== mode={mode['name']} seed={seed} ==\")\n",
+    "        cfg = HybridTransformerTrainConfig(\n",
+    "            dataset=dataset,\n",
+    "            vocab_size=vocab_size,\n",
+    "            hidden_dim=HIDDEN_DIM,\n",
+    "            n_heads=N_HEADS,\n",
+    "            ffn_dim=FFN_DIM,\n",
+    "            n_layers=N_LAYERS,\n",
+    "            group_size=GROUP_SIZE,\n",
+    "            arch=ARCH,\n",
+    "            use_full_graph=True,\n",
+    "            block_size=BLOCK_SIZE,\n",
+    "            batch_size=BATCH_SIZE,\n",
+    "            lr=LR,\n",
+    "            max_steps=MAX_STEPS,\n",
+    "            prune_at_step=mode[\"prune_at_step\"],\n",
+    "            prune_fraction=mode[\"prune_fraction\"],\n",
+    "            regrow_method=mode[\"regrow_method\"],\n",
+    "            dst_period=DST_PERIOD if mode[\"regrow_method\"] else None,\n",
+    "            dst_swap_fraction=DST_SWAP_FRACTION,\n",
+    "            dst_end_step=DST_END_STEP if mode[\"regrow_method\"] else None,\n",
+    "            seed=seed,\n",
+    "            device=device,\n",
+    "        )\n",
+    "        out = train_hybrid_transformer_lm(cfg)\n",
+    "        results[key] = out\n",
+    "        n_cycles = len(out[\"dst_cycles\"])\n",
+    "        print(\n",
+    "            f\"  final_loss = {out['final_loss']:.4f}  (ppl = {safe_perplexity(out['final_loss']):.2f})\"\n",
+    "            f\"  sparsity = {out['final_sparsity']:.3f}  dst_cycles = {n_cycles}\"\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "## 3. 결과 표 + 자동 verdict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": "print(\n    f\"{'mode':>12s} {'seed':>6s} {'final_loss':>12s} {'perplexity':>12s} \"\n    f\"{'sparsity':>10s} {'dst_cycles':>12s}\"\n)\nprint(\"-\" * 75)\nfor (name, seed), out in results.items():\n    fl = out[\"final_loss\"]\n    print(\n        f\"{name:>12s} {seed:>6d} {fl:>12.4f} {safe_perplexity(fl):>12.2f} \"\n        f\"{out['final_sparsity']:>10.3f} {len(out['dst_cycles']):>12d}\"\n    )\n\n# mode 별 평균\nprint(\"\\n== Mode summary (mean ± σ across seeds) ==\")\nsummary = {}\nfor mode in MODES:\n    name = mode[\"name\"]\n    vals = [results[(name, s)][\"final_loss\"] for s in SEEDS]\n    sparsities = [results[(name, s)][\"final_sparsity\"] for s in SEEDS]\n    m = statistics.mean(vals)\n    sd = statistics.stdev(vals) if len(vals) > 1 else 0.0\n    summary[name] = (m, sd, statistics.mean(sparsities))\n    print(\n        f\"  {name:>12s} {m:.4f} ± {sd:.4f}  (ppl ≈ {safe_perplexity(m):.2f})  \"\n        f\"sparsity={statistics.mean(sparsities):.3f}\"\n    )\n\n# 자동 verdict\nprint(\"\\n== Verdict ==\")\ndense_loss = summary[\"dense\"][0]\nstatic_loss = summary[\"static_50\"][0]\nset_loss = summary[\"dst_set_50\"][0]\nrigl_loss = summary[\"dst_rigl_50\"][0]\n\n# 1. constant sparsity — static 은 final, DST 는 모든 cycle 별 sparsity 검증\n# (CodeRabbit #3308022917 — final 만 보면 cycle 중간 drift 미감지)\nstatic_ok = abs(summary[\"static_50\"][2] - 0.5) < 0.02\ndst_cycles_ok = all(\n    abs(c[\"sparsity_after\"] - 0.5) < 0.02\n    for mode_name in (\"dst_set_50\", \"dst_rigl_50\")\n    for seed in SEEDS\n    for c in results[(mode_name, seed)][\"dst_cycles\"]\n)\nsp_ok = static_ok and dst_cycles_ok\nverdict_1 = \"PASS\" if sp_ok else \"FAIL\"\nprint(\n    f\"1. constant sparsity (static final + 모든 DST cycle, target 0.5 ±0.02): \"\n    f\"{sp_ok}  [{verdict_1}]\"\n)\n\n# 2. RigL ≥ static — DST RigL 이 static prune 보다 ≤ (낮거나 같음)\ndiff_rigl_static = rigl_loss - static_loss\nverdict_2 = \"PASS\" if diff_rigl_static <= 0.02 else \"FAIL\"\nprint(f\"2. RigL ≤ static + 0.02: diff = {diff_rigl_static:+.4f}  [{verdict_2}]\")\n\n# 3. RigL ≥ SET — gradient-guided 가 random 보다 ≤\ndiff_rigl_set = rigl_loss - set_loss\nverdict_3 = \"PASS\" if diff_rigl_set <= 0.02 else \"FAIL\"\nprint(f\"3. RigL ≤ SET + 0.02: diff = {diff_rigl_set:+.4f}  [{verdict_3}]\")\n\n# 4. all-finite\nall_finite = all(math.isfinite(out[\"final_loss\"]) for out in results.values())\nverdict_4 = \"PASS\" if all_finite else \"FAIL\"\nprint(f\"4. all-finite (DST stability): {all_finite}  [{verdict_4}]\")\n\n# DST cycle 정보 (RigL seed=42 기준 샘플)\nrigl_42_cycles = results[(\"dst_rigl_50\", 42)][\"dst_cycles\"]\nprint(f\"\\n== RigL seed=42 의 DST cycles ({len(rigl_42_cycles)}건) ==\")\nfor c in rigl_42_cycles[:3]:\n    print(f\"  step={c['step']} swap={c['total_swap']} sparsity_after={c['sparsity_after']:.3f}\")\nif len(rigl_42_cycles) > 6:\n    print(f\"  ... (중략 {len(rigl_42_cycles) - 6}건) ...\")\nfor c in rigl_42_cycles[-3:]:\n    print(f\"  step={c['step']} swap={c['total_swap']} sparsity_after={c['sparsity_after']:.3f}\")"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9",
+   "metadata": {},
+   "source": [
+    "## 4. Loss curve 시각화 — prune + DST cycle 표시"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(1, 1, figsize=(12, 6))\n",
+    "colors = {\n",
+    "    \"dense\": \"tab:gray\",\n",
+    "    \"static_50\": \"tab:red\",\n",
+    "    \"dst_set_50\": \"tab:orange\",\n",
+    "    \"dst_rigl_50\": \"tab:blue\",\n",
+    "}\n",
+    "window = 50\n",
+    "\n",
+    "for mode in MODES:\n",
+    "    name = mode[\"name\"]\n",
+    "    losses_per_seed = [results[(name, s)][\"losses\"] for s in SEEDS]\n",
+    "    smoothed = []\n",
+    "    for losses in losses_per_seed:\n",
+    "        smoothed.append(\n",
+    "            [\n",
+    "                sum(losses[max(0, i - window + 1) : i + 1]) / min(i + 1, window)\n",
+    "                for i in range(len(losses))\n",
+    "            ]\n",
+    "        )\n",
+    "    arr = torch.tensor(smoothed)\n",
+    "    mean = arr.mean(dim=0)\n",
+    "    std = arr.std(dim=0)\n",
+    "    steps = list(range(len(mean)))\n",
+    "    ax.plot(steps, mean, label=name, color=colors[name], linewidth=1.5)\n",
+    "    ax.fill_between(steps, mean - std, mean + std, color=colors[name], alpha=0.12)\n",
+    "\n",
+    "# prune step + DST end 수직선\n",
+    "ax.axvline(PRUNE_AT_STEP, color=\"black\", linestyle=\":\", alpha=0.5, label=f\"prune @ {PRUNE_AT_STEP}\")\n",
+    "ax.axvline(DST_END_STEP, color=\"gray\", linestyle=\":\", alpha=0.5, label=f\"DST end @ {DST_END_STEP}\")\n",
+    "\n",
+    "ax.set_xlabel(\"step\")\n",
+    "ax.set_ylabel(f\"loss (rolling mean w={window})\")\n",
+    "ax.set_title(\"Phase 16a — DST (RigL vs SET vs static): all 50% sparsity (mean ± σ over 2 seeds)\")\n",
+    "ax.legend(loc=\"upper right\")\n",
+    "ax.grid(alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "out_dir = Path(\"../../runs/notebook-neuron-phase16a\")\n",
+    "out_dir.mkdir(parents=True, exist_ok=True)\n",
+    "fig.savefig(out_dir / \"loss_curves.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "print(f\"saved: {out_dir / 'loss_curves.png'}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "11",
+   "metadata": {},
+   "source": [
+    "## 5. DST cycle sparsity 추적 (RigL)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
+    "for seed in SEEDS:\n",
+    "    cycles = results[(\"dst_rigl_50\", seed)][\"dst_cycles\"]\n",
+    "    steps = [c[\"step\"] for c in cycles]\n",
+    "    sparsities = [c[\"sparsity_after\"] for c in cycles]\n",
+    "    ax.plot(steps, sparsities, marker=\"o\", label=f\"RigL seed={seed}\", linewidth=1.5)\n",
+    "\n",
+    "ax.axhline(0.5, color=\"black\", linestyle=\"--\", alpha=0.5, label=\"target sparsity 0.5\")\n",
+    "ax.set_xlabel(\"step\")\n",
+    "ax.set_ylabel(\"effective sparsity (mask=0 비율)\")\n",
+    "ax.set_title(\"Phase 16a — DST cycle 후 sparsity (constant 유지 검증)\")\n",
+    "ax.legend()\n",
+    "ax.grid(alpha=0.3)\n",
+    "ax.set_ylim(0.45, 0.55)\n",
+    "plt.tight_layout()\n",
+    "fig.savefig(out_dir / \"sparsity_trace.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "print(f\"saved: {out_dir / 'sparsity_trace.png'}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13",
+   "metadata": {},
+   "source": [
+    "## 6. 결론 / 다음 단계\n",
+    "\n",
+    "(셀 출력 보고 사용자가 채울 영역)\n",
+    "\n",
+    "- RigL vs SET vs static 의 final loss ranking?\n",
+    "- DST 가 static prune 보다 의미 있는 개선?\n",
+    "- constant sparsity 정확히 유지되는지?\n",
+    "\n",
+    "**Phase 16b 진입 시 가설**:\n",
+    "- 16a 의 *constant sparsity DST* 가 paradigm 안에서 작동 입증 → 16b 의 *capacity expansion (Net2Net)* 도 비슷한 통합 메커니즘 (edge_mask) 위에서 작동 가능?\n",
+    "- 16a + 16b 결합 → *grow + shrink 동시 dynamic*"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "GraphLM (uv .venv)",
+   "language": "python",
+   "name": "graphlm-uv-venv"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}