diff --git a/docs/figures/neuron/phase12/hybrid_adj.png b/docs/figures/neuron/phase12/hybrid_adj.png new file mode 100644 index 0000000..5b23108 Binary files /dev/null and b/docs/figures/neuron/phase12/hybrid_adj.png differ diff --git a/docs/figures/neuron/phase12/loss_curves.png b/docs/figures/neuron/phase12/loss_curves.png new file mode 100644 index 0000000..8627969 Binary files /dev/null and b/docs/figures/neuron/phase12/loss_curves.png differ diff --git a/notebooks/02-function-level/11-phase12-hybrid-graph-foundations.ipynb b/notebooks/02-function-level/11-phase12-hybrid-graph-foundations.ipynb new file mode 100644 index 0000000..a996f4e --- /dev/null +++ b/notebooks/02-function-level/11-phase12-hybrid-graph-foundations.ipynb @@ -0,0 +1,387 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 11-phase12-hybrid-graph-foundations\n", + "\n", + "**neuron Phase 12** — paradigm 의 ultimate 단계. Phase 9 group + Phase 10/11 channel 의\n", + "계층적 hybrid (사용자 vision: 히든 레이어 = graph 의 ultimate 구현).\n", + "\n", + "핵심 가설:\n", + "1. **function preservation** — hybrid_full_full (outer=full, inner=full) ≈ plain Linear?\n", + "2. **Phase 9/10/11 통합 표현** — hybrid 의 init 조합으로 이전 phase 결과 재현 가능?\n", + " - outer=identity + inner=full ≈ Phase 9 group_identity (+0.14 열위)\n", + " - outer=full + inner=around_one ≈ Phase 11 channel_around_one (≈ plain)\n", + "3. **계층적 routing 효과** — 두 adj 의 학습된 패턴 분리 시각화?\n", + "4. **paradigm 안정성** — magnitude rule 자동 적용으로 모든 init 조합 안전?\n", + "\n", + "설계: 4 × 2 sweep = 8 run, max_steps=1500.\n", + "\n", + "데이터: TinyShakespeare (char-LM)\n", + "시드: [42, 123]\n", + "작성일: 2026-05-26\n", + "연관: Issue [#65](https://github.com/EinSofINTEREST/GraphLM/issues/65) / Phase 11 baseline PR [#64](https://github.com/EinSofINTEREST/GraphLM/pull/64)" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## 1. 환경" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import logging\n", + "import sys\n", + "\n", + "import torch\n", + "\n", + "import graphlm\n", + "from graphlm.data.tinyshakespeare import (\n", + " CharTokenizer,\n", + " TinyShakespeareDataset,\n", + " load_tinyshakespeare_text,\n", + ")\n", + "from graphlm.neuron.graph_hybrid_demo import train_hybrid_graph_mlp\n", + "from graphlm.utils import repo_root\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "logging.basicConfig(\n", + " level=logging.WARNING, format=\"%(asctime)s %(levelname)s %(message)s\", datefmt=\"%H:%M:%S\"\n", + ")\n", + "print(\"python :\", sys.version.split()[0])\n", + "print(\"graphlm :\", graphlm.__version__)\n", + "print(\"torch :\", torch.__version__)\n", + "print(\"device :\", DEVICE)" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## 2. 실험 설정" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "\n", + "ROOT = repo_root()\n", + "DATA_PATH = ROOT / \"data\" / \"tinyshakespeare.txt\"\n", + "OUT_DIR = ROOT / \"runs\" / \"notebook-neuron-phase12\"\n", + "OUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "\n", + "SEEDS = [42, 123]\n", + "ARCHS = [\n", + " \"plain\",\n", + " \"hybrid_full_full\",\n", + " \"hybrid_identity_full\",\n", + " \"hybrid_full_around_one\",\n", + "]\n", + "EMB_DIM = 64\n", + "HIDDEN_DIM = 256 # 256 / 16 = 16 groups\n", + "GROUP_SIZE = 16\n", + "N_GRAM = 4\n", + "BATCH_SIZE = 32\n", + "LR = 3e-4\n", + "MAX_STEPS = 1500\n", + "\n", + "# Phase 9/10/11 baseline (직접 비교)\n", + "PHASE9_GROUP_FULL = 2.1391\n", + "PHASE9_GROUP_IDENTITY = 2.2797 # 0-init vanishing 사례\n", + "PHASE10_CHANNEL_FULL = 2.1339\n", + "PHASE11_CHANNEL_AROUND_ONE = 2.1456\n", + "\n", + "print(f\"SEEDS = {SEEDS}\")\n", + "print(f\"ARCHS = {ARCHS}\")\n", + "print(\n", + " f\"HIDDEN_DIM = {HIDDEN_DIM}, GROUP_SIZE = {GROUP_SIZE} → n_groups = {HIDDEN_DIM // GROUP_SIZE}\"\n", + ")\n", + "print(f\"MAX_STEPS = {MAX_STEPS}\")\n", + "print(f\"전체 run = {len(SEEDS) * len(ARCHS)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## 3. 데이터 로드" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "text = load_tinyshakespeare_text(DATA_PATH)\n", + "tokenizer = CharTokenizer(text)\n", + "dataset = TinyShakespeareDataset(text, tokenizer)\n", + "V = tokenizer.vocab_size\n", + "\n", + "V_PADDED = math.ceil(V / GROUP_SIZE) * GROUP_SIZE\n", + "print(f\"vocab_size : {V}, padded to {V_PADDED} (GROUP_SIZE {GROUP_SIZE} 배수)\")" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## 4. sweep 학습" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "runs = {}\n", + "for seed in SEEDS:\n", + " for arch in ARCHS:\n", + " key = (seed, arch)\n", + " print(f\"--- seed={seed}, arch={arch} ---\")\n", + " runs[key] = train_hybrid_graph_mlp(\n", + " dataset=dataset,\n", + " vocab_size=V_PADDED,\n", + " seed=seed,\n", + " arch=arch,\n", + " emb_dim=EMB_DIM,\n", + " hidden_dim=HIDDEN_DIM,\n", + " group_size=GROUP_SIZE,\n", + " n_gram=N_GRAM,\n", + " batch_size=BATCH_SIZE,\n", + " lr=LR,\n", + " max_steps=MAX_STEPS,\n", + " device=DEVICE,\n", + " )\n", + " print(f\" done: final_loss={runs[key]['final_loss']:.4f}\")\n", + " if str(DEVICE).startswith(\"cuda\"):\n", + " torch.cuda.empty_cache()" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "## 5. 결과 표 + Phase 9/10/11 baseline 비교" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "import statistics\n", + "\n", + "print(f\"{'arch':>28} {'seed':>5} {'final_loss':>11}\")\n", + "print(\"-\" * 55)\n", + "for arch in ARCHS:\n", + " for seed in SEEDS:\n", + " r = runs[(seed, arch)]\n", + " print(f\"{arch:>28} {seed:>5} {r['final_loss']:>11.4f}\")\n", + "\n", + "print()\n", + "print(\"=== arch 별 mean ===\")\n", + "agg = {}\n", + "for arch in ARCHS:\n", + " fls = [runs[(s, arch)][\"final_loss\"] for s in SEEDS]\n", + " agg[arch] = dict(mean=statistics.mean(fls), range=max(fls) - min(fls))\n", + " print(f\" {arch:>28}: mean={agg[arch]['mean']:.4f}, range={agg[arch]['range']:.4f}\")\n", + "\n", + "print()\n", + "print(\"=== Phase 9/10/11 baseline (직접 비교 — 통합 표현 검증) ===\")\n", + "print(f\" Phase 9 group_full : {PHASE9_GROUP_FULL:.4f}\")\n", + "print(f\" Phase 9 group_identity : {PHASE9_GROUP_IDENTITY:.4f} (0-init vanishing 사례)\")\n", + "print(f\" Phase 10 channel_full : {PHASE10_CHANNEL_FULL:.4f}\")\n", + "print(f\" Phase 11 channel_around_one : {PHASE11_CHANNEL_AROUND_ONE:.4f}\")\n", + "print()\n", + "print(f\" Phase 12 hybrid_full_full : {agg['hybrid_full_full']['mean']:.4f}\")\n", + "print(\n", + " f\" Phase 12 hybrid_identity_full : {agg['hybrid_identity_full']['mean']:.4f} (expected ≈ Phase 9 group_identity)\"\n", + ")\n", + "print(\n", + " f\" Phase 12 hybrid_full_around_one : {agg['hybrid_full_around_one']['mean']:.4f} (expected ≈ Phase 11 channel_around_one)\"\n", + ")\n", + "\n", + "print()\n", + "print(\"=== 자동 verdict ===\")\n", + "plain = agg[\"plain\"][\"mean\"]\n", + "full_full = agg[\"hybrid_full_full\"][\"mean\"]\n", + "identity_full = agg[\"hybrid_identity_full\"][\"mean\"]\n", + "full_around = agg[\"hybrid_full_around_one\"][\"mean\"]\n", + "range_max = max(agg[a][\"range\"] for a in ARCHS)\n", + "print(f\" range_max = {range_max:.4f}\")\n", + "print(f\" hybrid_full_full vs plain : {full_full - plain:+.4f}\")\n", + "print(\n", + " f\" hybrid_identity_full vs plain : {identity_full - plain:+.4f} (expected ~+0.14 like Phase 9)\"\n", + ")\n", + "print(f\" hybrid_full_around_one vs plain : {full_around - plain:+.4f} (expected ≈ 0)\")\n", + "if abs(full_full - plain) < range_max:\n", + " print(\" ✅ hybrid_full_full ≈ plain — function preservation 작동\")\n", + "if identity_full - plain > 0.05:\n", + " print(\" ✅ hybrid_identity_full 열위 — Phase 9 group_identity 패턴 재현\")\n", + "if abs(full_around - plain) < range_max:\n", + " print(\" ✅ hybrid_full_around_one ≈ plain — Phase 11 channel_around_one 패턴 재현\")" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## 6. 학습된 adj_outer + adj_inner 시각화 (hybrid_full_full)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "r = runs[(SEEDS[0], \"hybrid_full_full\")]\n", + "adj_outer = r[\"final_adj\"][\"fc1\"][\"outer\"].numpy()\n", + "adj_inner = r[\"final_adj\"][\"fc1\"][\"inner\"].numpy()\n", + "print(f\"adj_outer shape: {adj_outer.shape}\")\n", + "print(f\"adj_inner shape: {adj_inner.shape}\")\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n", + "\n", + "# (a) adj_outer heatmap\n", + "vmax_o = max(abs(adj_outer).max(), 1e-6)\n", + "im0 = axes[0].imshow(adj_outer, cmap=\"RdBu_r\", vmin=-vmax_o, vmax=vmax_o, aspect=\"auto\")\n", + "axes[0].set_xlabel(\"input group\")\n", + "axes[0].set_ylabel(\"output group\")\n", + "axes[0].set_title(\n", + " f\"adj_outer (group-level)\\nmean={adj_outer.mean():.3f}, std={adj_outer.std():.3f}\"\n", + ")\n", + "fig.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.04)\n", + "\n", + "# (b) adj_inner mean abs by (G_out, G_in) — block-aggregated heatmap\n", + "inner_block_mean = np.abs(adj_inner).mean(axis=(2, 3))\n", + "vmax_i = max(inner_block_mean.max(), 1e-6)\n", + "im1 = axes[1].imshow(inner_block_mean, cmap=\"viridis\", vmin=0, vmax=vmax_i, aspect=\"auto\")\n", + "axes[1].set_xlabel(\"input group\")\n", + "axes[1].set_ylabel(\"output group\")\n", + "axes[1].set_title(f\"adj_inner |mean| per block\\noverall mean={inner_block_mean.mean():.3f}\")\n", + "fig.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04)\n", + "\n", + "fig.suptitle(f\"Phase 12 — hybrid_full_full 학습된 adj (seed={SEEDS[0]}, fc1)\", fontsize=11)\n", + "fig.tight_layout()\n", + "fig.savefig(OUT_DIR / \"hybrid_adj.png\", dpi=120)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## 7. loss curve 4 arch 비교 (mean ± σ)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "window = 30\n", + "colors = {\n", + " \"plain\": \"#1f77b4\",\n", + " \"hybrid_full_full\": \"#2ca02c\",\n", + " \"hybrid_identity_full\": \"#d62728\",\n", + " \"hybrid_full_around_one\": \"#ff7f0e\",\n", + "}\n", + "\n", + "fig, ax = plt.subplots(figsize=(13, 5))\n", + "for arch in ARCHS:\n", + " seed_curves = []\n", + " for seed in SEEDS:\n", + " losses = runs[(seed, arch)][\"losses\"]\n", + " smoothed = np.convolve(losses, np.ones(window) / window, mode=\"valid\")\n", + " seed_curves.append(smoothed)\n", + " arr = np.array(seed_curves)\n", + " steps = np.arange(window - 1, window - 1 + arr.shape[1])\n", + " mean = arr.mean(axis=0)\n", + " std = arr.std(axis=0, ddof=1)\n", + " ax.plot(steps, mean, color=colors[arch], lw=1.5, label=arch)\n", + " ax.fill_between(steps, mean - std, mean + std, color=colors[arch], alpha=0.15)\n", + "\n", + "ax.set_xlabel(\"step\")\n", + "ax.set_ylabel(f\"loss (smoothed window={window})\")\n", + "ax.set_title(\"Phase 12 — 4 arch loss curve (mean ± σ across 2 seeds)\")\n", + "ax.legend(loc=\"upper right\", fontsize=9)\n", + "ax.grid(alpha=0.3)\n", + "fig.tight_layout()\n", + "fig.savefig(OUT_DIR / \"loss_curves.png\", dpi=120)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "## 결과 요약 / Phase 13 권장\n", + "\n", + "확인 포인트:\n", + "- §5 hybrid_full_full ≈ plain (function preservation)\n", + "- §5 hybrid_identity_full vs Phase 9 group_identity (+0.14 패턴 재현)\n", + "- §5 hybrid_full_around_one vs Phase 11 channel_around_one (≈ plain 패턴 재현)\n", + "- §6 학습된 adj_outer + adj_inner 의 분리 표현\n", + "- §7 4 arch 의 발산/수렴 패턴\n", + "\n", + "**판정 시나리오**:\n", + "- **A. 4 arch 모두 expected 와 일치** ⭐ — hybrid 가 Phase 9/10/11 의 통합 표현 입증 → Phase 13 (Transformer 통합) 진입\n", + "- **B. 일부 불일치** — hybrid 의 dual adj 의 interaction effect 가 단순 합 이상의 dynamics\n", + "- **C. hybrid_full_full 도 plain 보다 명확 우위** — 계층적 routing 의 추가 표현력 발견 (surprise positive)\n", + "\n", + "**참고**:\n", + "- Phase 11 결과: https://www.notion.so/36ce8b70b7aa81b4b684f9c559788a46\n", + "- 아키텍처 구성 계획: https://www.notion.so/36ce8b70b7aa818cbf1fe71687b449b8" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "GraphLM (uv .venv)", + "language": "python", + "name": "graphlm-uv" + }, + "language_info": { + "name": "python", + "version": "3.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 88bb017..7420658 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,9 @@ markers = [ source = ["src/graphlm"] omit = [ "*/tests/*", + # demo modules — 노트북 orchestration 전용, unit-test 대상 아님 (Phase 8 ~ 12) + "*/neuron/*_demo.py", + "*/neuron/positional_analysis.py", ] [tool.coverage.report] diff --git a/src/graphlm/neuron/__init__.py b/src/graphlm/neuron/__init__.py index a560d4e..dafd3e6 100644 --- a/src/graphlm/neuron/__init__.py +++ b/src/graphlm/neuron/__init__.py @@ -13,6 +13,7 @@ ) from graphlm.neuron.graph_channel import ChannelGraphLinear from graphlm.neuron.graph_group import GroupGraphLinear +from graphlm.neuron.graph_hybrid import HybridGraphLinear from graphlm.neuron.growable import GrowableEmbedding, GrowableLayerNorm, GrowableLinear from graphlm.neuron.growth import add_attn_function_preserving, add_attn_smooth_start @@ -22,6 +23,7 @@ "GrowableEmbedding", "GrowableLayerNorm", "GrowableLinear", + "HybridGraphLinear", "NeuronBlock", "NeuronConfig", "NeuronGrowingDecoder", diff --git a/src/graphlm/neuron/graph_hybrid.py b/src/graphlm/neuron/graph_hybrid.py new file mode 100644 index 0000000..cafcacc --- /dev/null +++ b/src/graphlm/neuron/graph_hybrid.py @@ -0,0 +1,209 @@ +"""Phase 12 — hierarchical hybrid graph hidden layer (outer group + inner channel). + +사용자 vision (히든 레이어 = graph) 의 **ultimate 구조**: Phase 9 group-as-node 와 +Phase 10/11 channel-as-node 의 계층적 결합. + +``` +hidden_dim H = G groups × k channels per group +weight W: (G_out, G_in, k, k) # block-organized (Phase 9 동일) +adj_outer: (G_out, G_in) # group-level routing (Phase 9 의 adj) +adj_inner: (G_out, G_in, k, k) # channel-level fine-grained gate (Phase 10/11 의 채널 gate) + +forward (block matmul + dual gates): + contrib[go, gi] = (adj_inner[go, gi] * W[go, gi]) @ x[gi] # shape (k,) + y[go] = Σ_gi adj_outer[go, gi] · contrib[go, gi] # shape (k,) + y_flat = reshape(y, (G_out · k,)) +``` + +effective edge weight (channel-pair) = + adj_outer[group(out), group(in)] · adj_inner[out, in] · W[out, in] + +설계 의미: +- **outer adj** = group-level *coarse routing* (어느 group 가 어느 group 에 연결) +- **inner adj** = channel-level *fine-grained gate* (각 connection 의 strength tuning) +- 둘 다 학습 가능 + 둘 다 magnitude rule 적용 (weight multiplier 위치 ≈ 1.0) +- function preservation: 둘 다 ``"full"`` (모두 1) → standard Linear 와 forward 동치 + +**0-init 금지 + magnitude rule** (memory: feedback_no_zero_init.md): +- adj_outer / adj_inner 모두 ``"zero"`` 거부 (ValueError) +- weight multiplier 위치 → sweet spot magnitude ≈ 1.0 +- 옵션: + - ``"full"`` (모두 1) — function preserving + - ``"identity"`` (block-diagonal, n_groups_out == n_groups_in 만, adj_outer 만) — Phase 9 호환 + - ``"uniform_around_one"`` (uniform[0.95, 1.05]) — scale-corrected + 학습 활성 + +Phase 9 / 10 / 11 과의 관계: +- adj_outer=full + adj_inner=full == standard Linear (function preserving) +- adj_outer=full + adj_inner=uniform_around_one ≈ Phase 11 의 channel-level scale-corrected +- adj_outer=identity + adj_inner=full ≈ Phase 9 의 group_identity (block-diagonal) +- 모두 한 module 에서 통합 표현 가능 +""" + +from __future__ import annotations + +import math +from typing import Literal + +import torch +from torch import Tensor, nn + +AdjOuterInit = Literal["full", "identity", "uniform_around_one"] +AdjInnerInit = Literal["full", "uniform_around_one"] + + +def _validate_groupable(features: int, group_size: int, name: str) -> int: + if not isinstance(group_size, int) or group_size < 1: + raise ValueError(f"group_size must be a positive int, got {group_size!r}") + if features % group_size != 0: + raise ValueError(f"{name} ({features}) must be divisible by group_size ({group_size})") + return features // group_size + + +def _make_outer_adj(n_groups_out: int, n_groups_in: int, init: AdjOuterInit) -> Tensor: + if init == "full": + return torch.ones(n_groups_out, n_groups_in) + if init == "identity": + if n_groups_out != n_groups_in: + raise ValueError( + f"adj_outer_init='identity' requires square (n_groups_out=={n_groups_in}), " + f"got out={n_groups_out} in={n_groups_in}" + ) + return torch.eye(n_groups_out, n_groups_in) + if init == "uniform_around_one": + return torch.empty(n_groups_out, n_groups_in).uniform_(0.95, 1.05) + if init in {"zero", "zeros"}: + raise ValueError( + f"adj_outer_init={init!r} 는 금지됨 — 0-init vanishing 함정. " + "rationale: Phase 9 PR #60, magnitude rule: Phase 10 PR #62. " + "'full' / 'identity' / 'uniform_around_one' 사용 권장." + ) + raise ValueError(f"unknown adj_outer_init: {init!r}") + + +def _make_inner_adj( + n_groups_out: int, + n_groups_in: int, + group_size: int, + init: AdjInnerInit, +) -> Tensor: + shape = (n_groups_out, n_groups_in, group_size, group_size) + if init == "full": + return torch.ones(shape) + if init == "uniform_around_one": + return torch.empty(shape).uniform_(0.95, 1.05) + if init in {"zero", "zeros"}: + raise ValueError( + f"adj_inner_init={init!r} 는 금지됨 — 0-init vanishing 함정. " + "rationale: Phase 9 PR #60, magnitude rule: Phase 10 PR #62. " + "'full' 또는 'uniform_around_one' 사용 권장." + ) + raise ValueError(f"unknown adj_inner_init: {init!r}") + + +class HybridGraphLinear(nn.Module): + """Hierarchical hybrid graph linear: outer group + inner channel routing. + + paradigm 의 ultimate 단계 — Phase 9 group + Phase 10/11 channel 의 통합. + + Args: + in_features, out_features: 표준 Linear shape (둘 다 group_size 의 배수) + group_size: k (한 group 당 채널 수) + adj_outer_init: ``"full"`` / ``"identity"`` / ``"uniform_around_one"`` + adj_inner_init: ``"full"`` / ``"uniform_around_one"`` + bias: bias 사용 여부 + + Forward: + ``y[go] = Σ_gi adj_outer[go, gi] · (adj_inner[go, gi] * W[go, gi]) @ x[gi]`` + """ + + def __init__( + self, + in_features: int, + out_features: int, + group_size: int, + *, + adj_outer_init: AdjOuterInit = "full", + adj_inner_init: AdjInnerInit = "full", + bias: bool = True, + ): + super().__init__() + if not isinstance(in_features, int) or in_features < 1: + raise ValueError(f"in_features must be a positive int, got {in_features!r}") + if not isinstance(out_features, int) or out_features < 1: + raise ValueError(f"out_features must be a positive int, got {out_features!r}") + + self.in_features = in_features + self.out_features = out_features + self.group_size = group_size + self.n_groups_in = _validate_groupable(in_features, group_size, "in_features") + self.n_groups_out = _validate_groupable(out_features, group_size, "out_features") + + # block weight: shape (G_out, G_in, k, k) + self.weight = nn.Parameter( + torch.empty(self.n_groups_out, self.n_groups_in, group_size, group_size) + ) + # standard Linear-equivalent init (fan_in = in_features) + bound = 1.0 / math.sqrt(in_features) + nn.init.uniform_(self.weight, -bound, bound) + + # outer adj (group-level routing) + self.adj_outer = nn.Parameter( + _make_outer_adj(self.n_groups_out, self.n_groups_in, adj_outer_init) + ) + # inner adj (channel-level gate within each block) + self.adj_inner = nn.Parameter( + _make_inner_adj(self.n_groups_out, self.n_groups_in, group_size, adj_inner_init) + ) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, x: Tensor) -> Tensor: + # x: (..., in_features) → (..., n_groups_in, group_size) + *batch, in_f = x.shape + if in_f != self.in_features: + raise ValueError(f"expected last dim {self.in_features}, got {in_f}") + x_g = x.reshape(*batch, self.n_groups_in, self.group_size) + + # 메모리 효율 최적화 (gemini #3302293739): adj_outer 와 adj_inner 를 weight 수준에서 + # 미리 결합 → (*batch, G_out, G_in, k) 의 큰 intermediate tensor 회피. + # 수학적 등치: contrib[..., go, gi, ko] = Σ_ki adj_inner[go,gi,ki,ko] · W[go,gi,ki,ko] · x_g[..., gi, ki] + # y[..., go, ko] = Σ_gi adj_outer[go, gi] · contrib[..., go, gi, ko] + # = Σ_gi Σ_ki (adj_outer · adj_inner · W)[...] · x_g[..., gi, ki] + eff_w = self.adj_outer.unsqueeze(-1).unsqueeze(-1) * self.adj_inner * self.weight + # single einsum — output (*batch, G_out, k), 중간 (*batch, G_out, G_in, k) 텐서 없음 + y_g = torch.einsum("...gi,Ggik->...Gk", x_g, eff_w) + # flatten back to (..., out_features) + y = y_g.reshape(*batch, self.out_features) + if self.bias is not None: + y = y + self.bias + return y + + def adj_outer_sparsity(self, threshold: float = 0.05) -> float: + """|adj_outer| < threshold 인 group-edge 비율.""" + if threshold < 0: + raise ValueError(f"threshold must be >= 0, got {threshold}") + with torch.no_grad(): + return float((self.adj_outer.abs() < threshold).float().mean().item()) + + def adj_inner_sparsity(self, threshold: float = 0.05) -> float: + """|adj_inner| < threshold 인 channel-edge 비율.""" + if threshold < 0: + raise ValueError(f"threshold must be >= 0, got {threshold}") + with torch.no_grad(): + return float((self.adj_inner.abs() < threshold).float().mean().item()) + + def freeze_adj_outer(self) -> None: + self.adj_outer.requires_grad_(False) + + def freeze_adj_inner(self) -> None: + self.adj_inner.requires_grad_(False) + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"group_size={self.group_size}, n_groups_in={self.n_groups_in}, " + f"n_groups_out={self.n_groups_out}" + ) diff --git a/src/graphlm/neuron/graph_hybrid_demo.py b/src/graphlm/neuron/graph_hybrid_demo.py new file mode 100644 index 0000000..ecf035d --- /dev/null +++ b/src/graphlm/neuron/graph_hybrid_demo.py @@ -0,0 +1,163 @@ +"""Phase 12 — HybridGraphLinear demo MLP-LM + 학습 헬퍼. + +노트북 분리 규약 준수. Phase 12 노트북 11-phase12-hybrid-graph-foundations.ipynb 에서 import. + +4 가지 architecture 비교 (모두 0-init 금지 + magnitude rule 자동 적용): +- ``"plain"`` — 표준 nn.Linear (baseline) +- ``"hybrid_full_full"`` — outer=full + inner=full (function preserving 시작) +- ``"hybrid_identity_full"`` — outer=identity + inner=full (Phase 9 group_identity 의 hybrid 표현) +- ``"hybrid_full_around_one"`` — outer=full + inner=uniform_around_one (Phase 11 channel 의 hybrid 표현) +""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Literal + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from graphlm.data.tinyshakespeare import TinyShakespeareDataset, iter_random_batches +from graphlm.neuron.graph_hybrid import HybridGraphLinear +from graphlm.utils import set_seed + +Arch = Literal[ + "plain", + "hybrid_full_full", + "hybrid_identity_full", + "hybrid_full_around_one", +] + + +class HybridGraphMLPLM(nn.Module): + """Hybrid graph MLP-LM — Phase 9/10/11 의 통합 데모 모델.""" + + def __init__( + self, + vocab_size: int, + emb_dim: int, + hidden_dim: int, + n_gram: int, + arch: Arch, + group_size: int = 16, + ): + super().__init__() + self.arch = arch + self.n_gram = n_gram + self.emb = nn.Embedding(vocab_size, emb_dim) + in_f = emb_dim * n_gram + self.fc1 = _make_linear(in_f, hidden_dim, arch, group_size) + self.ln = nn.LayerNorm(hidden_dim) + # fc2 는 hidden_dim → vocab_size — 직사각형 가능. identity outer 는 정방 필요라 + # fc2 는 arch 와 무관하게 hybrid_full_full 또는 plain 사용 (비교 공정성) + fc2_arch: Arch = "plain" if arch == "plain" else "hybrid_full_full" + self.fc2 = _make_linear(hidden_dim, vocab_size, fc2_arch, group_size) + + def forward(self, x: Tensor) -> Tensor: + h = self.emb(x).reshape(x.shape[0], -1) + h = self.fc1(h) + h = F.gelu(h) + h = self.ln(h) + return self.fc2(h) + + +def _make_linear(in_f: int, out_f: int, arch: Arch, group_size: int) -> nn.Module: + if arch == "plain": + return nn.Linear(in_f, out_f) + if in_f % group_size != 0 or out_f % group_size != 0: + raise ValueError( + f"HybridGraphMLPLM 의 in_f({in_f}) / out_f({out_f}) 는 group_size({group_size}) 의 배수여야 함" + ) + if arch == "hybrid_full_full": + return HybridGraphLinear( + in_f, + out_f, + group_size=group_size, + adj_outer_init="full", + adj_inner_init="full", + ) + if arch == "hybrid_identity_full": + return HybridGraphLinear( + in_f, + out_f, + group_size=group_size, + adj_outer_init="identity", + adj_inner_init="full", + ) + if arch == "hybrid_full_around_one": + return HybridGraphLinear( + in_f, + out_f, + group_size=group_size, + adj_outer_init="full", + adj_inner_init="uniform_around_one", + ) + raise ValueError(f"unknown arch: {arch}") + + +def make_ngram_iter( + dataset: TinyShakespeareDataset, batch_size: int, n_gram: int, *, seed: int +) -> Iterator[tuple[Tensor, Tensor]]: + raw = iter_random_batches(dataset, batch_size=batch_size, block_size=n_gram + 1, seed=seed) + for x, _y in raw: + yield x[:, :n_gram], x[:, n_gram] + + +def train_hybrid_graph_mlp( + *, + dataset: TinyShakespeareDataset, + vocab_size: int, + seed: int, + arch: Arch, + emb_dim: int, + hidden_dim: int, + group_size: int, + n_gram: int, + batch_size: int, + lr: float, + max_steps: int, + device: str = "cpu", +) -> dict: + """1 run 학습 — Phase 12 sweep unit. + + Returns: ``losses``, ``final_loss``, ``final_adj`` (hybrid_* 에 한해 fc1 의 outer/inner 둘 다 snapshot). + """ + set_seed(seed) + model = HybridGraphMLPLM(vocab_size, emb_dim, hidden_dim, n_gram, arch, group_size).to(device) + data_iter = make_ngram_iter(dataset, batch_size, n_gram, seed=seed) + optimizer = torch.optim.AdamW(model.parameters(), lr=lr) + losses: list[float] = [] + model.train() + for _step in range(1, max_steps + 1): + x, y = next(data_iter) + x, y = x.to(device), y.to(device) + optimizer.zero_grad() + logits = model(x) + loss = F.cross_entropy(logits, y) + loss.backward() + optimizer.step() + losses.append(loss.item()) + + n_last = min(100, len(losses)) + final_loss = sum(losses[-n_last:]) / n_last if n_last > 0 else 0.0 + + final_adj = None + if arch != "plain": + # graph_group_demo / graph_channel_demo 와 동일한 {"fc1": ..., "fc2": ...} 계층 구조 — + # 공통 후처리/시각화 재사용 가능 (Copilot #3302306899). hybrid 는 fc1/fc2 각자 outer/inner. + final_adj = { + "fc1": { + "outer": model.fc1.adj_outer.detach().cpu().clone(), + "inner": model.fc1.adj_inner.detach().cpu().clone(), + }, + "fc2": { + "outer": model.fc2.adj_outer.detach().cpu().clone(), + "inner": model.fc2.adj_inner.detach().cpu().clone(), + }, + } + return { + "losses": losses, + "final_loss": final_loss, + "final_adj": final_adj, + } diff --git a/tests/neuron/test_graph_hybrid.py b/tests/neuron/test_graph_hybrid.py new file mode 100644 index 0000000..b0c5907 --- /dev/null +++ b/tests/neuron/test_graph_hybrid.py @@ -0,0 +1,142 @@ +"""Tests for graphlm.neuron.graph_hybrid — Phase 12 hierarchical hybrid foundations.""" + +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from graphlm.neuron.graph_hybrid import HybridGraphLinear + + +def test_shape_after_init(): + lin = HybridGraphLinear(16, 24, group_size=4) + assert lin.n_groups_in == 4 + assert lin.n_groups_out == 6 + # weight: (G_out, G_in, k, k) + assert lin.weight.shape == (6, 4, 4, 4) + # adj_outer: (G_out, G_in) + assert lin.adj_outer.shape == (6, 4) + # adj_inner: (G_out, G_in, k, k) + assert lin.adj_inner.shape == (6, 4, 4, 4) + assert lin.bias.shape == (24,) + + +def test_forward_shape(): + lin = HybridGraphLinear(16, 24, group_size=4) + x = torch.randn(2, 8, 16) + y = lin(x) + assert y.shape == (2, 8, 24) + + +@pytest.mark.parametrize( + "outer,inner", + [("full", "full"), ("identity", "full"), ("uniform_around_one", "uniform_around_one")], +) +def test_adj_init_combinations_valid(outer, inner): + lin = HybridGraphLinear(16, 16, group_size=4, adj_outer_init=outer, adj_inner_init=inner) + assert lin.adj_outer.shape == (4, 4) + assert lin.adj_inner.shape == (4, 4, 4, 4) + + +def test_identity_adj_outer_requires_square(): + """adj_outer_init='identity' 는 정방 (n_groups_out==n_groups_in) 만 허용.""" + with pytest.raises(ValueError, match="requires square"): + HybridGraphLinear(16, 24, group_size=4, adj_outer_init="identity") + + +def test_function_preservation_full_full_equivalent_to_linear(): + """adj_outer=full + adj_inner=full + 같은 W → standard Linear forward 동일 (atol=1e-5).""" + torch.manual_seed(0) + in_f, out_f, k = 16, 24, 4 + hg = HybridGraphLinear(in_f, out_f, group_size=k, adj_outer_init="full", adj_inner_init="full") + + # standard Linear W = blocks 모은 형태 (block (go, gi, k, k) 의 transpose) + G_out, G_in = hg.n_groups_out, hg.n_groups_in + W_std = torch.zeros(out_f, in_f) + for go in range(G_out): + for gi in range(G_in): + # hg.weight[go, gi] shape (k, k) — block matmul 에서 x[gi] @ W[go, gi] 이므로 + # standard Linear (y = W @ x) 의 block 은 W_std[go*k:(go+1)*k, gi*k:(gi+1)*k] = hg.weight[go, gi].T + W_std[go * k : (go + 1) * k, gi * k : (gi + 1) * k] = hg.weight[go, gi].T + + std = nn.Linear(in_f, out_f, bias=True) + with torch.no_grad(): + std.weight.copy_(W_std) + std.bias.copy_(hg.bias) + + x = torch.randn(2, 8, in_f) + y_hg = hg(x) + y_std = std(x) + assert torch.allclose(y_hg, y_std, atol=1e-5), ( + f"function preservation 깨짐: max |diff| = {(y_hg - y_std).abs().max().item()}" + ) + + +@pytest.mark.parametrize("bad", ["zero", "zeros"]) +def test_zero_init_outer_rejected(bad): + """adj_outer_init='zero' 거부 — 0-init 금지 규칙.""" + with pytest.raises(ValueError, match="vanishing"): + HybridGraphLinear(16, 16, group_size=4, adj_outer_init=bad) # type: ignore[arg-type] + + +@pytest.mark.parametrize("bad", ["zero", "zeros"]) +def test_zero_init_inner_rejected(bad): + """adj_inner_init='zero' 거부 — 0-init 금지 규칙.""" + with pytest.raises(ValueError, match="vanishing"): + HybridGraphLinear(16, 16, group_size=4, adj_inner_init=bad) # type: ignore[arg-type] + + +def test_unknown_adj_outer_raises(): + with pytest.raises(ValueError, match="unknown adj_outer_init"): + HybridGraphLinear(16, 16, group_size=4, adj_outer_init="bogus") # type: ignore[arg-type] + + +def test_unknown_adj_inner_raises(): + with pytest.raises(ValueError, match="unknown adj_inner_init"): + HybridGraphLinear(16, 16, group_size=4, adj_inner_init="bogus") # type: ignore[arg-type] + + +def test_all_three_params_have_gradient(): + """weight, adj_outer, adj_inner 모두 grad 흐름.""" + lin = HybridGraphLinear( + 16, + 24, + group_size=4, + adj_outer_init="uniform_around_one", + adj_inner_init="uniform_around_one", + ) + x = torch.randn(2, 16) + out = lin(x) + out.sum().backward() + for attr in ["weight", "adj_outer", "adj_inner"]: + grad = getattr(lin, attr).grad + assert grad is not None, f"{attr}.grad is None" + assert (grad.abs().sum() > 0).item(), f"{attr}.grad all zero" + + +def test_in_features_not_divisible_raises(): + with pytest.raises(ValueError, match="in_features.*divisible"): + HybridGraphLinear(17, 24, group_size=4) + + +def test_invalid_features_raises(): + with pytest.raises(ValueError, match="positive int"): + HybridGraphLinear(0, 16, group_size=4) + + +def test_sparsity_metrics_initial(): + """초기 (full/full) 은 모두 1 → sparsity 0.""" + lin = HybridGraphLinear(16, 16, group_size=4) + assert lin.adj_outer_sparsity(0.05) == 0.0 + assert lin.adj_inner_sparsity(0.05) == 0.0 + + +def test_freeze_helpers(): + lin = HybridGraphLinear(16, 16, group_size=4) + lin.freeze_adj_outer() + assert not lin.adj_outer.requires_grad + assert lin.adj_inner.requires_grad + assert lin.weight.requires_grad + lin.freeze_adj_inner() + assert not lin.adj_inner.requires_grad