diff --git a/notebooks/02-function-level/08-phase9-group-graph-foundations.ipynb b/notebooks/02-function-level/08-phase9-group-graph-foundations.ipynb new file mode 100644 index 0000000..90902a2 --- /dev/null +++ b/notebooks/02-function-level/08-phase9-group-graph-foundations.ipynb @@ -0,0 +1,333 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 08-phase9-group-graph-foundations\n", + "\n", + "**neuron Phase 9** — Hidden layer 자체를 graph 로 (group-as-node, 사용자 vision 의 architecture\n", + "구성 계획 B 진입).\n", + "\n", + "핵심 가설:\n", + "1. **function preservation** — GroupGraphLinear(adj=full) 가 standard Linear 와 forward 동치?\n", + "2. **adjacency 학습** — adj 파라미터에 gradient 가 흘러 routing 이 학습됨?\n", + "3. **identity init (block-diagonal)** 의 inductive bias — 시작부터 group-wise 독립 학습 vs full 보다 빠르거나 느림?\n", + "4. **adjacency 시각화** — 학습 후 어떤 group 들이 서로 강하게 연결되는가?\n", + "\n", + "설계: 3-way sweep × 2 seed = 6 run, max_steps=1500.\n", + "- arch ∈ {plain, group_full, group_identity}\n", + "- seed ∈ {42, 123}\n", + "\n", + "데이터: TinyShakespeare (char-LM)\n", + "시드: [42, 123]\n", + "작성일: 2026-05-26\n", + "연관: Issue [#59](https://github.com/EinSofINTEREST/GraphLM/issues/59) / Phase 8 baseline PR [#56](https://github.com/EinSofINTEREST/GraphLM/pull/56) / [아키텍처 구성 계획 (Notion)](https://www.notion.so/36ce8b70b7aa818cbf1fe71687b449b8)" + ] + }, + { + "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_group_demo import train_group_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": [ + "ROOT = repo_root()\n", + "DATA_PATH = ROOT / \"data\" / \"tinyshakespeare.txt\"\n", + "OUT_DIR = ROOT / \"runs\" / \"notebook-neuron-phase9\"\n", + "OUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "\n", + "SEEDS = [42, 123]\n", + "ARCHS = [\"plain\", \"group_full\", \"group_identity\"]\n", + "EMB_DIM = 64\n", + "HIDDEN_DIM = 256 # group_size 16 의 배수\n", + "GROUP_SIZE = 16 # → n_groups_in/out = 16 (256/16) — 시각화 적합\n", + "N_GRAM = 4\n", + "BATCH_SIZE = 32\n", + "LR = 3e-4\n", + "MAX_STEPS = 1500\n", + "\n", + "# vocab*N_GRAM*EMB_DIM = N_GRAM * EMB_DIM = 256 = HIDDEN_DIM = vocab_logits 입력\n", + "# vocab_size 는 65 라 GROUP_SIZE 의 배수가 아님 — group_*는 vocab 도 padding 또는 unequal split\n", + "# 단순화: vocab_size 도 GROUP_SIZE 배수가 되도록 pad — 노트북에서 wrapper 처리\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", + "# vocab_size 가 GROUP_SIZE 의 배수가 아니면 group_*는 fc2 의 out_features 가 안 맞음.\n", + "# 노트북 단순화: vocab_size 를 GROUP_SIZE 배수로 padding (예: 65 → 80) 한 가상 vocab 사용\n", + "import math\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 학습\n", + "\n", + "각 (seed, arch) 에 대해 1 run. plain 은 baseline, group_full 은 function-preserving 시작,\n", + "group_identity 는 block-diagonal sparse 시작." + ] + }, + { + "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_group_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. 결과 표 — arch × seed" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "import statistics\n", + "\n", + "print(f\"{'arch':>16} {'seed':>5} {'final_loss':>11}\")\n", + "print(\"-\" * 40)\n", + "for arch in ARCHS:\n", + " for seed in SEEDS:\n", + " r = runs[(seed, arch)]\n", + " print(f\"{arch:>16} {seed:>5} {r['final_loss']:>11.4f}\")\n", + "\n", + "print()\n", + "print(\"=== arch 별 mean ===\")\n", + "for arch in ARCHS:\n", + " fls = [runs[(s, arch)][\"final_loss\"] for s in SEEDS]\n", + " print(f\" {arch:>16}: mean={statistics.mean(fls):.4f}, range={max(fls) - min(fls):.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## 6. adjacency 학습 진화 — fc1 / fc2 의 학습된 adj heatmap" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "fig = plt.figure(figsize=(14, 9))\n", + "gs = fig.add_gridspec(2, 4, width_ratios=[1, 1, 1, 1])\n", + "\n", + "# row 0: group_full 의 fc1, fc2 (seed 42)\n", + "# row 1: group_identity 의 fc1, fc2 (seed 42)\n", + "for row_i, arch in enumerate([\"group_full\", \"group_identity\"]):\n", + " r = runs[(SEEDS[0], arch)]\n", + " if r[\"final_adj\"] is None:\n", + " continue\n", + " for col_i, layer_name in enumerate([\"fc1\", \"fc2\"]):\n", + " ax = fig.add_subplot(gs[row_i, col_i * 2 : col_i * 2 + 2])\n", + " adj = r[\"final_adj\"][layer_name].numpy()\n", + " vmax = max(abs(adj).max(), 1e-6)\n", + " im = ax.imshow(adj, cmap=\"RdBu_r\", vmin=-vmax, vmax=vmax, aspect=\"auto\")\n", + " ax.set_title(f\"{arch} — {layer_name} adj (seed={SEEDS[0]}) shape={adj.shape}\")\n", + " ax.set_xlabel(\"group_in\")\n", + " ax.set_ylabel(\"group_out\")\n", + " fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)\n", + "\n", + "fig.tight_layout()\n", + "fig.savefig(OUT_DIR / \"adjacency_heatmaps.png\", dpi=120)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## 7. loss curve 비교 (arch 별 mean ± σ across 2 seeds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "window = 30\n", + "colors = {\"plain\": \"#1f77b4\", \"group_full\": \"#2ca02c\", \"group_identity\": \"#ff7f0e\"}\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", + "ax.set_xlabel(\"step\")\n", + "ax.set_ylabel(f\"loss (smoothed window={window})\")\n", + "ax.set_title(f\"Phase 9 — plain Linear vs GroupGraphLinear (mean ± σ over {len(SEEDS)} 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 10 권장 방향\n", + "\n", + "확인 포인트:\n", + "- §5 final_loss — group_full 이 plain 과 비슷? (function preservation 가설 입증)\n", + "- §5 group_identity vs plain — block-diagonal sparse init 의 학습 성능 (inductive bias 효과)\n", + "- §6 adjacency heatmap — group_full 학습 후 어떤 group 간 연결이 강해졌나? group_identity 의 off-diagonal 학습 활성도?\n", + "- §7 loss curve — 세 arch 의 수렴 속도 비교\n", + "\n", + "**판정 시나리오**:\n", + "- **A. group_full ≈ plain** ⭐ — function preservation 입증. group routing 학습 가능성 확인\n", + "- **B. group_full < plain** — adjacency 학습이 plain Linear 보다 능력 강화 (drop-in 대체 후보)\n", + "- **C. group_identity ≈ group_full** — block-diagonal sparse init 도 충분 (Phase 10 의 sparsification 권장)\n", + "- **D. group_identity 명확 열위** — sparse 시작이 학습 능력 제한 — adjacency growth 메커니즘 필요\n", + "\n", + "**참고**:\n", + "- 아키텍처 구성 계획 (Notion): https://www.notion.so/36ce8b70b7aa818cbf1fe71687b449b8\n", + "- ML 용어집: https://www.notion.so/36ce8b70b7aa812298bbe1388e61b753" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "GraphLM (uv .venv)", + "language": "python", + "name": "graphlm-uv" + }, + "language_info": { + "name": "python", + "version": "3.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/graphlm/neuron/__init__.py b/src/graphlm/neuron/__init__.py index 9618d88..6a6dce8 100644 --- a/src/graphlm/neuron/__init__.py +++ b/src/graphlm/neuron/__init__.py @@ -11,10 +11,12 @@ NeuronGrowingDecoder, SinusoidalAlpha, ) +from graphlm.neuron.graph_group import GroupGraphLinear from graphlm.neuron.growable import GrowableEmbedding, GrowableLayerNorm, GrowableLinear from graphlm.neuron.growth import add_attn_function_preserving, add_attn_smooth_start __all__ = [ + "GroupGraphLinear", "GrowableEmbedding", "GrowableLayerNorm", "GrowableLinear", diff --git a/src/graphlm/neuron/graph_group.py b/src/graphlm/neuron/graph_group.py new file mode 100644 index 0000000..3ff4566 --- /dev/null +++ b/src/graphlm/neuron/graph_group.py @@ -0,0 +1,159 @@ +"""Phase 9 — group-as-node graph hidden layer foundations. + +채널을 ``group_size`` 단위로 묶어 graph node 로 다룬다. dense ``nn.Linear`` 대신 그룹 간 +adjacency 가 결정하는 block-sparse routing. + +``` +hidden_dim H 채널을 k 개씩 묶어 G = H/k 그룹 → G 개 node +within-group: dense (k × k) block weight +between-group: 학습된 adjacency adj[go, gi] 가 routing scalar +forward: y[go] = Σ_gi adj[go, gi] · (x[gi] @ W[go, gi]) +``` + +이 구조의 의미: +- node = k 채널 묶음 (graph 의 vertex) +- edge = group 간 block routing (k×k block weight × adj scalar) +- adjacency = 1 모두 + identity-ish weight init → standard Linear 와 forward 동치 (function preservation) +- adjacency 의 일부 entry 를 0 으로 sparsify → block-sparse routing emerges + +참고 ML 패턴: MoE (Switch Transformer) 의 expert routing, multi-head attention 의 head grouping, +DeepSeek-V3 의 block-sparse attention, grouped convolution. + +Phase 10+ 에서: +- A (channel-as-node) foundations + 계층적 hybrid (그룹 내 채널 graph nest) +- DARTS / L0 / Gumbel-softmax 로 adjacency sparsification 학습 +- Transformer 의 Q/K/V/O 통합 +""" + +from __future__ import annotations + +import math +from typing import Literal + +import torch +from torch import Tensor, nn + +AdjInit = Literal["full", "identity"] + + +def _validate_groupable(features: int, group_size: int, name: str) -> int: + # group_size 자체 검증 — 0 또는 음수 시 % 에서 ZeroDivisionError / 의미 없는 결과 회피 (Copilot #3301531298) + 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 + + +class GroupGraphLinear(nn.Module): + """Group-as-node graph linear layer. + + 표준 ``nn.Linear`` 와 입출력 shape 동일하나, 내부적으로 ``group_size`` 단위 block 으로 + 파라미터를 보관하고 그룹 간 routing 은 ``adj`` 가 결정. + + Args: + in_features: 입력 차원 — ``group_size`` 의 배수여야 함. + out_features: 출력 차원 — ``group_size`` 의 배수여야 함. + group_size: 한 그룹의 채널 수. 표준 head_dim (64). + adj_init: ``"full"`` 이면 adjacency 모두 1 (모든 그룹 routing 활성), + ``"identity"`` 이면 동일 그룹 index 끼리만 1, 나머지 0 (block-diagonal — pure + grouped operation, 가장 sparse 시작). + + Forward (입력 shape ``(..., in_features)`` → 출력 ``(..., out_features)``): + x 를 ``(..., n_groups_in, group_size)`` 로 reshape → + 각 출력 group ``go`` 에 대해 ``y[go] = Σ_gi adj[go, gi] · (x[gi] @ W[go, gi])`` → + ``(..., out_features)`` 로 reshape. + """ + + def __init__( + self, + in_features: int, + out_features: int, + group_size: int, + *, + adj_init: AdjInit = "full", + bias: bool = True, + ): + super().__init__() + 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 weights: shape (n_groups_out, n_groups_in, group_size, group_size) + self.weight = nn.Parameter( + torch.empty(self.n_groups_out, self.n_groups_in, group_size, group_size) + ) + # adjacency: shape (n_groups_out, n_groups_in) — continuous routing scalar + if adj_init == "full": + adj = torch.ones(self.n_groups_out, self.n_groups_in) + elif adj_init == "identity": + # identity 의 의미는 정방 (n_groups_out == n_groups_in) 일 때만 명확 — + # 직사각형이면 일부 입력/출력 group 이 isolated 됨 (Copilot #3301531337). 명시적 거부. + if self.n_groups_out != self.n_groups_in: + raise ValueError( + f"adj_init='identity' requires square (n_groups_out=={self.n_groups_in}), " + f"got out={self.n_groups_out} in={self.n_groups_in}" + ) + # torch.eye 는 (n, m) rectangular 직접 지원 — max+slice 보다 간결 (gemini #3301524134) + adj = torch.eye(self.n_groups_out, self.n_groups_in) + else: + raise ValueError(f"unknown adj_init: {adj_init}") + self.adj = nn.Parameter(adj) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_features)) + else: + self.register_parameter("bias", None) + + # standard nn.Linear 와 동일 스케일 init — fan_in = in_features 로 직접 계산. + # nn.init.kaiming_uniform_ 은 4D tensor 의 fan_in 을 잘못 계산 (n_groups_in * group_size²) — + # 실제 in_features (= n_groups_in × group_size) 와 다름 (gemini #3301524127 / Copilot #3301531321). + bound = 1.0 / math.sqrt(self.in_features) + nn.init.uniform_(self.weight, -bound, bound) + + 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}") + # reshape (= view 호환 + non-contiguous tensor 도 안전) — Copilot #3301531355 + x_g = x.reshape(*batch, self.n_groups_in, self.group_size) + + # block matmul: for each (go, gi) compute x[gi] @ W[go, gi] → (..., go, gi, group_size) + # einsum: ...gi,Goik (G=n_groups_out, g=n_groups_in 동일 index, i=in_chans, k=out_chans) + # 결과 (..., G, g, k) — go 별 gi-routed contributions + contrib = torch.einsum("...gi,Ggik->...Ggk", x_g, self.weight) + # adjacency-weighted sum over gi: y[go] = Σ_gi adj[go, gi] · contrib[go, gi] + y_g = torch.einsum("Gg,...Ggk->...Gk", self.adj, contrib) + # reshape (..., n_groups_out, group_size) → (..., out_features) + y = y_g.reshape(*batch, self.out_features) + if self.bias is not None: + y = y + self.bias + return y + + def freeze_adjacency(self) -> None: + """adjacency 를 학습 비대상으로 — Phase 10+ sparsification 학습 분리용.""" + self.adj.requires_grad_(False) + + def sparsify_adjacency(self, threshold: float) -> int: + """절대값이 ``threshold`` 미만인 adjacency entry 를 0 으로 강제 (in-place). + + Phase 10+ 의 hard sparsification 학습 결과 적용을 위한 stub. Returns 비활성화된 entry 수. + """ + if threshold < 0: + raise ValueError(f"threshold must be >= 0, got {threshold}") + with torch.no_grad(): + # masked_fill_ 가 mul_(mask.to(dtype)) 보다 깔끔 + 캐스팅 불필요 (gemini #3301524139) + zero_mask = self.adj.abs() < threshold + n_zeroed = int(zero_mask.sum().item()) + self.adj.masked_fill_(zero_mask, 0.0) + return n_zeroed + + 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_group_demo.py b/src/graphlm/neuron/graph_group_demo.py new file mode 100644 index 0000000..f687be1 --- /dev/null +++ b/src/graphlm/neuron/graph_group_demo.py @@ -0,0 +1,136 @@ +"""Phase 9 — GroupGraphLinear demo MLP-LM + 학습 헬퍼. + +노트북 분리 규약 준수 (.claude/rules/06-code-style.md). Phase 9 노트북 +08-phase9-group-graph-foundations.ipynb 는 여기서 import. + +세 가지 architecture 를 같은 학습 루프로 비교: +- ``"plain"`` — 표준 nn.Linear (baseline) +- ``"group_full"`` — GroupGraphLinear with adj_init="full" (function preserving 시작) +- ``"group_identity"`` — GroupGraphLinear with adj_init="identity" (block-diagonal, 가장 sparse) +""" + +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_group import GroupGraphLinear +from graphlm.utils import set_seed + +Arch = Literal["plain", "group_full", "group_identity"] + + +class GroupGraphMLPLM(nn.Module): + """Phase 8 의 GrowableMLPLM 의 graph 버전 — fc1 / fc2 가 architecture 별로 다름. + + Architecture: + emb (vocab × emb_dim) → flatten n_gram → fc1 → GELU → LayerNorm → fc2 → vocab logits + """ + + 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 (직사각형 — n_groups_out ≠ n_groups_in 가능). + # arch="group_identity" 의 의미는 정방에서만 정의되므로 fc2 는 항상 "full" 사용 — + # 비교 공정성 유지 (group_identity 비교의 차이는 fc1 에 한정) (Copilot #3301531369). + fc2_arch: Arch = "plain" if arch == "plain" else "group_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) + # group_*: in_f / out_f 가 group_size 의 배수가 아니면 명시적으로 ValueError — + # padding 은 caller (예: 노트북 V_PADDED) 책임이며 본 함수는 검증만 (Copilot #3301531359). + if in_f % group_size != 0 or out_f % group_size != 0: + raise ValueError( + f"GroupGraphMLPLM 의 in_f({in_f}) / out_f({out_f}) 는 group_size({group_size}) 의 배수여야 함" + ) + adj_init = "full" if arch == "group_full" else "identity" + return GroupGraphLinear(in_f, out_f, group_size=group_size, adj_init=adj_init) + + +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_group_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 9 sweep 의 단위. + + Returns dict with ``losses``, ``final_loss`` (last 100 avg), ``final_adj`` (학습된 adjacency + snapshots — plain 의 경우 None). + """ + set_seed(seed) + model = GroupGraphMLPLM(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 + + # adjacency snapshot (group_* 에만 있음) + final_adj = None + if arch != "plain": + final_adj = { + "fc1": model.fc1.adj.detach().cpu().clone(), + "fc2": model.fc2.adj.detach().cpu().clone(), + } + return { + "losses": losses, + "final_loss": final_loss, + "final_adj": final_adj, + } diff --git a/tests/neuron/test_graph_group.py b/tests/neuron/test_graph_group.py new file mode 100644 index 0000000..b107ac7 --- /dev/null +++ b/tests/neuron/test_graph_group.py @@ -0,0 +1,131 @@ +"""Tests for graphlm.neuron.graph_group — Phase 9 group-as-node foundations.""" + +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from graphlm.neuron.graph_group import GroupGraphLinear + + +def test_shape_after_init(): + lin = GroupGraphLinear(16, 24, group_size=4) + assert lin.n_groups_in == 4 + assert lin.n_groups_out == 6 + assert lin.weight.shape == (6, 4, 4, 4) + assert lin.adj.shape == (6, 4) + assert lin.bias.shape == (24,) + + +def test_forward_shape(): + lin = GroupGraphLinear(16, 24, group_size=4) + x = torch.randn(2, 8, 16) # (B, T, in_features) + y = lin(x) + assert y.shape == (2, 8, 24) + + +def test_in_features_not_divisible_raises(): + with pytest.raises(ValueError, match="in_features.*divisible"): + GroupGraphLinear(17, 24, group_size=4) + + +def test_out_features_not_divisible_raises(): + with pytest.raises(ValueError, match="out_features.*divisible"): + GroupGraphLinear(16, 25, group_size=4) + + +def test_full_adj_init_value(): + lin = GroupGraphLinear(16, 24, group_size=4, adj_init="full") + assert torch.allclose(lin.adj, torch.ones(6, 4)) + + +def test_identity_adj_init_value(): + lin = GroupGraphLinear(16, 16, group_size=4, adj_init="identity") + assert torch.allclose(lin.adj, torch.eye(4)) + + +def test_invalid_adj_init_raises(): + with pytest.raises(ValueError, match="unknown adj_init"): + GroupGraphLinear(16, 16, group_size=4, adj_init="bogus") # type: ignore[arg-type] + + +def test_function_preservation_equivalent_to_standard_linear(): + """adj=full + Linear 와 같은 W 로 init 하면 forward 결과 동일. + + 수학적으로: GroupGraphLinear (adj=1 모든 곳) = standard Linear (블록을 모두 모은 W). + """ + torch.manual_seed(0) + in_f, out_f, k = 16, 24, 4 + gg = GroupGraphLinear(in_f, out_f, group_size=k, adj_init="full") + # gg.weight: shape (G_out=6, G_in=4, k=4, k=4) — 각 block 이 standard W 의 block tile + # standard W 와 같은 effective W 를 만들려면: W_std[go*k:(go+1)*k, gi*k:(gi+1)*k] = gg.weight[go, gi].T + # (linear.weight 는 (out, in) shape, x @ W^T 형식) + G_out, G_in = gg.n_groups_out, gg.n_groups_in + W_std = torch.zeros(out_f, in_f) + for go in range(G_out): + for gi in range(G_in): + W_std[go * k : (go + 1) * k, gi * k : (gi + 1) * k] = gg.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_(gg.bias) + + x = torch.randn(2, 8, in_f) + y_gg = gg(x) + y_std = std(x) + assert torch.allclose(y_gg, y_std, atol=1e-5), ( + f"function preservation 깨짐: max |diff| = {(y_gg - y_std).abs().max().item()}" + ) + + +def test_identity_init_zero_off_diagonal_contribution(): + """adj=identity init 시 off-diagonal group 의 routing = 0 → block-diagonal forward.""" + torch.manual_seed(0) + in_f, out_f, k = 16, 16, 4 # n_groups_in = n_groups_out = 4 + lin = GroupGraphLinear(in_f, out_f, group_size=k, adj_init="identity") + x = torch.randn(2, in_f) + y = lin(x) + # 각 출력 그룹 [go*k:(go+1)*k] 는 *동일* 입력 그룹 [go*k:(go+1)*k] 만의 함수여야 함 + # → x 의 다른 group entry 변경이 영향 없음 확인 + x_perturbed = x.clone() + # group 1 (indices 4~7) 만 변경 + x_perturbed[:, 4:8] += 100 + y_perturbed = lin(x_perturbed) + # group 0 (indices 0~3) 출력은 안 변해야 + assert torch.allclose(y[:, 0:4], y_perturbed[:, 0:4], atol=1e-5) + # group 1 출력은 변해야 + assert not torch.allclose(y[:, 4:8], y_perturbed[:, 4:8], atol=1e-3) + + +def test_weight_and_adj_both_have_gradient(): + lin = GroupGraphLinear(16, 24, group_size=4) + x = torch.randn(2, 16) + out = lin(x) + out.sum().backward() + assert lin.weight.grad is not None + assert lin.adj.grad is not None + assert (lin.weight.grad.abs().sum() > 0).item() + assert (lin.adj.grad.abs().sum() > 0).item() + + +def test_freeze_adjacency(): + lin = GroupGraphLinear(16, 24, group_size=4) + lin.freeze_adjacency() + assert not lin.adj.requires_grad + # weight 는 여전히 학습 가능 + assert lin.weight.requires_grad + + +def test_sparsify_adjacency_zeros_below_threshold(): + lin = GroupGraphLinear(16, 24, group_size=4, adj_init="full") + # all-ones (=1.0) → threshold=1.5 면 모두 0 으로 강제 + n_zeroed = lin.sparsify_adjacency(threshold=1.5) + assert n_zeroed == 6 * 4 # 모든 entry + assert torch.allclose(lin.adj, torch.zeros_like(lin.adj)) + + +def test_sparsify_adjacency_negative_threshold_raises(): + lin = GroupGraphLinear(16, 24, group_size=4) + with pytest.raises(ValueError, match="threshold"): + lin.sparsify_adjacency(threshold=-0.1)