[FEAT#77] neuron Phase 16b — Net2Net function-preserving FFN grow (cross-shape expansion)

notebooks/02-function-level/16-phase16b-net2net-grow.ipynb

@@ -0,0 +1,399 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# 16-phase16b-net2net-grow\n",
+    "\n",
+    "**neuron Phase 16b** — paradigm 의 dynamic phase 두 번째 갈래. Phase 16a 가 *within-shape DST* (sparsity 재할당, parameter 수 동일) 였다면, Phase 16b 는 **cross-shape expansion** — 학습 중 `ffn_dim` 을 확장하여 *parameter 수 자체 증가*. **function preservation** 보장 (Net2Net-style: 새 input weight = 0).\n",
+    "\n",
+    "핵심 가설:\n",
+    "1. **function preservation** — grow 직후 forward output 이 grow 전과 정확히 동일?\n",
+    "2. **grown ≤ large-baseline + threshold** — 작게 시작 + 중간에 grow 한 모델이 처음부터 큰 모델 근방까지 학습?\n",
+    "3. **grown < small-baseline** — grow 가 capacity 부족 모델보다 명확히 좋음?\n",
+    "4. **all-finite** — grow 후 학습 안정성?\n",
+    "\n",
+    "설계: 3 mode × 2 seed = 6 run.\n",
+    "- `small_baseline`: ffn=128 끝까지 (small capacity)\n",
+    "- `grown`: ffn=128 시작 → step 750 에서 256 으로 grow\n",
+    "- `large_baseline`: ffn=256 끝까지 (large capacity)\n",
+    "\n",
+    "arch 고정: `hybrid_around_one_around_one` + `use_full_graph=True` (Phase 14 최저 loss 구조)\n",
+    "데이터: TinyShakespeare (char-LM, block_size=64)\n",
+    "시드: [42, 123]\n",
+    "작성일: 2026-05-27\n",
+    "연관: Issue [#77](https://github.com/EinSofINTEREST/GraphLM/issues/77) (Phase 16 main [#75](https://github.com/EinSofINTEREST/GraphLM/issues/75)) / Phase 16a PR [#78](https://github.com/EinSofINTEREST/GraphLM/pull/78)"
+   ]
+  },
+  {
+   "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",
+    "HIDDEN_DIM = 128\n",
+    "N_HEADS = 4\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",
+    "GROW_AT_STEP = MAX_STEPS // 2  # 750\n",
+    "SEEDS = [42, 123]\n",
+    "ARCH = \"hybrid_around_one_around_one\"\n",
+    "\n",
+    "FFN_SMALL = 128\n",
+    "FFN_LARGE = 256\n",
+    "\n",
+    "# 3 mode 정의\n",
+    "MODES = [\n",
+    "    {\n",
+    "        \"name\": \"small_baseline\",\n",
+    "        \"ffn_dim\": FFN_SMALL,\n",
+    "        \"grow_at_step\": None,\n",
+    "        \"grow_ffn_target\": None,\n",
+    "    },\n",
+    "    {\n",
+    "        \"name\": \"grown\",\n",
+    "        \"ffn_dim\": FFN_SMALL,\n",
+    "        \"grow_at_step\": GROW_AT_STEP,\n",
+    "        \"grow_ffn_target\": FFN_LARGE,\n",
+    "    },\n",
+    "    {\n",
+    "        \"name\": \"large_baseline\",\n",
+    "        \"ffn_dim\": FFN_LARGE,\n",
+    "        \"grow_at_step\": None,\n",
+    "        \"grow_ffn_target\": None,\n",
+    "    },\n",
+    "]\n",
+    "print(f\"\\nGrow 시점: step {GROW_AT_STEP} / {MAX_STEPS}, ffn_dim {FFN_SMALL} → {FFN_LARGE}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "## 2. Sweep 실행 (3 mode × 2 seed = 6 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=mode[\"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",
+    "            grow_at_step=mode[\"grow_at_step\"],\n",
+    "            grow_ffn_target=mode[\"grow_ffn_target\"],\n",
+    "            seed=seed,\n",
+    "            device=device,\n",
+    "        )\n",
+    "        out = train_hybrid_transformer_lm(cfg)\n",
+    "        results[key] = out\n",
+    "        ev = out[\"grow_event\"]\n",
+    "        print(\n",
+    "            f\"  final_loss = {out['final_loss']:.4f}  (ppl = {safe_perplexity(out['final_loss']):.2f})\"\n",
+    "            f\"  params = {out['final_param_count']:,}\"\n",
+    "            f\"  grow_event = {ev}\"\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "## 3. 결과 표 + 자동 verdict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"{'mode':>16s} {'seed':>6s} {'final_loss':>12s} {'perplexity':>12s} {'params':>10s}\")\n",
+    "print(\"-\" * 65)\n",
+    "for (name, seed), out in results.items():\n",
+    "    fl = out[\"final_loss\"]\n",
+    "    print(\n",
+    "        f\"{name:>16s} {seed:>6d} {fl:>12.4f} {safe_perplexity(fl):>12.2f} \"\n",
+    "        f\"{out['final_param_count']:>10,}\"\n",
+    "    )\n",
+    "\n",
+    "# mode 별 평균\n",
+    "print(\"\\n== Mode summary (mean ± σ across seeds) ==\")\n",
+    "summary = {}\n",
+    "for mode in MODES:\n",
+    "    name = mode[\"name\"]\n",
+    "    vals = [results[(name, s)][\"final_loss\"] for s in SEEDS]\n",
+    "    params = [results[(name, s)][\"final_param_count\"] 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(params))\n",
+    "    print(\n",
+    "        f\"  {name:>16s} {m:.4f} ± {sd:.4f}  (ppl ≈ {safe_perplexity(m):.2f})  \"\n",
+    "        f\"params={statistics.mean(params):,.0f}\"\n",
+    "    )\n",
+    "\n",
+    "# 자동 verdict\n",
+    "print(\"\\n== Verdict ==\")\n",
+    "small_loss = summary[\"small_baseline\"][0]\n",
+    "grown_loss = summary[\"grown\"][0]\n",
+    "large_loss = summary[\"large_baseline\"][0]\n",
+    "\n",
+    "# 1. all-finite\n",
+    "all_finite = all(math.isfinite(out[\"final_loss\"]) for out in results.values())\n",
+    "verdict_1 = \"PASS\" if all_finite else \"FAIL\"\n",
+    "print(f\"1. all-finite (grow stability): {all_finite}  [{verdict_1}]\")\n",
+    "\n",
+    "# 2. grown < small_baseline — grow 가 capacity 부족 모델보다 명확히 좋음\n",
+    "diff_grown_small = grown_loss - small_loss\n",
+    "verdict_2 = \"PASS\" if diff_grown_small < 0 else \"FAIL\"\n",
+    "print(f\"2. grown < small_baseline: diff = {diff_grown_small:+.4f}  [{verdict_2}]\")\n",
+    "\n",
+    "# 3. grown ≤ large_baseline + 0.05 — grown 이 처음부터 큰 모델 근방까지 학습\n",
+    "diff_grown_large = grown_loss - large_loss\n",
+    "verdict_3 = \"PASS\" if diff_grown_large <= 0.05 else \"FAIL\"\n",
+    "print(f\"3. grown ≤ large_baseline + 0.05: diff = {diff_grown_large:+.4f}  [{verdict_3}]\")\n",
+    "\n",
+    "# 4. grow 직후 function preservation 검증 — 학습 curve 의 step 750 근방에 spike 없음\n",
+    "# (정량 검증은 별도 unit test 가 atol=1e-5 로 보장; 여기서는 학습 curve smoothness 만 추정)\n",
+    "for seed in SEEDS:\n",
+    "    losses = results[(\"grown\", seed)][\"losses\"]\n",
+    "    # step 750 직전 100step 평균 vs step 750 직후 1 step 비교 (drift 측정)\n",
+    "    before = sum(losses[max(0, GROW_AT_STEP - 100) : GROW_AT_STEP]) / 100\n",
+    "    after = losses[GROW_AT_STEP] if len(losses) > GROW_AT_STEP else losses[-1]\n",
+    "    spike = after - before\n",
+    "    print(\n",
+    "        f\"   seed={seed}: grow 직전 100 step 평균 = {before:.4f}, grow 직후 1 step = {after:.4f}, spike = {spike:+.4f}\"\n",
+    "    )\n",
+    "\n",
+    "# grow 이벤트 정보\n",
+    "print(\"\\n== Grow events ==\")\n",
+    "for (name, seed), out in results.items():\n",
+    "    if out[\"grow_event\"] is not None:\n",
+    "        ev = out[\"grow_event\"]\n",
+    "        print(\n",
+    "            f\"  {name} seed={seed}: step={ev['step']}, layers_grown={ev['n_layers_grown']}, \"\n",
+    "            f\"ffn {ev['old_ffn_dim']} → {ev['new_ffn_dim']}, n_new_groups={ev['n_new_groups']}\"\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9",
+   "metadata": {},
+   "source": [
+    "## 4. Loss curve 시각화 — grow step 표시"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(1, 1, figsize=(12, 6))\n",
+    "colors = {\n",
+    "    \"small_baseline\": \"tab:red\",\n",
+    "    \"grown\": \"tab:blue\",\n",
+    "    \"large_baseline\": \"tab:green\",\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.15)\n",
+    "\n",
+    "# grow step 수직선\n",
+    "ax.axvline(GROW_AT_STEP, color=\"black\", linestyle=\":\", alpha=0.5, label=f\"grow @ {GROW_AT_STEP}\")\n",
+    "\n",
+    "ax.set_xlabel(\"step\")\n",
+    "ax.set_ylabel(f\"loss (rolling mean w={window})\")\n",
+    "ax.set_title(\n",
+    "    f\"Phase 16b — Net2Net grow (ffn {FFN_SMALL} → {FFN_LARGE}): small / grown / large (mean ± σ over 2 seeds)\"\n",
+    ")\n",
+    "ax.legend(loc=\"upper right\")\n",
+    "ax.grid(alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "out_dir = Path(\"../../runs/notebook-neuron-phase16b\")\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. Parameter count 비교"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(1, 1, figsize=(8, 5))\n",
+    "names = [m[\"name\"] for m in MODES]\n",
+    "params = [summary[n][2] for n in names]\n",
+    "losses = [summary[n][0] for n in names]\n",
+    "stds = [summary[n][1] for n in names]\n",
+    "\n",
+    "color_list = [colors[n] for n in names]\n",
+    "ax.errorbar(params, losses, yerr=stds, marker=\"o\", capsize=4, linewidth=1.5, color=\"tab:gray\")\n",
+    "for p, lo, n in zip(params, losses, names, strict=True):\n",
+    "    ax.annotate(n, (p, lo), xytext=(8, -8), textcoords=\"offset points\", fontsize=10)\n",
+    "\n",
+    "ax.set_xlabel(\"final parameter count\")\n",
+    "ax.set_ylabel(\"final loss (last 100 mean ± σ)\")\n",
+    "ax.set_title(\"Phase 16b — params vs loss: grown vs static baselines\")\n",
+    "ax.grid(alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "fig.savefig(out_dir / \"params_vs_loss.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.show()\n",
+    "print(f\"saved: {out_dir / 'params_vs_loss.png'}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13",
+   "metadata": {},
+   "source": [
+    "## 6. 결론 / 다음 단계\n",
+    "\n",
+    "(셀 출력 보고 사용자가 채울 영역)\n",
+    "\n",
+    "- grow 직후 function preservation 확인 (spike 없음)?\n",
+    "- grown 의 final loss 가 large_baseline 근방까지 회복?\n",
+    "- 16a (DST) 와 16b (grow) 의 의미 비교 — 16b 의 capacity 증가가 더 효과적?\n",
+    "\n",
+    "**Phase 17 후보**:\n",
+    "- 16a + 16b 결합 — grow + shrink 동시 dynamic\n",
+    "- layer-wise 차등 (attention vs FFN 별 다른 grow / prune 정책)\n",
+    "- hidden_dim 까지 grow (downstream layer 모두 영향 — 더 invasive)"
+   ]
+  }
+ ],
+ "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
+}