diff --git a/docs/figures/neuron/phase10/adj_distribution.png b/docs/figures/neuron/phase10/adj_distribution.png new file mode 100644 index 0000000..68d61a5 Binary files /dev/null and b/docs/figures/neuron/phase10/adj_distribution.png differ diff --git a/docs/figures/neuron/phase10/adj_heatmaps.png b/docs/figures/neuron/phase10/adj_heatmaps.png new file mode 100644 index 0000000..791410b Binary files /dev/null and b/docs/figures/neuron/phase10/adj_heatmaps.png differ diff --git a/docs/figures/neuron/phase10/loss_curves.png b/docs/figures/neuron/phase10/loss_curves.png new file mode 100644 index 0000000..2001323 Binary files /dev/null and b/docs/figures/neuron/phase10/loss_curves.png differ diff --git a/notebooks/02-function-level/09-phase10-channel-graph-foundations.ipynb b/notebooks/02-function-level/09-phase10-channel-graph-foundations.ipynb new file mode 100644 index 0000000..5f8f55c --- /dev/null +++ b/notebooks/02-function-level/09-phase10-channel-graph-foundations.ipynb @@ -0,0 +1,386 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 09-phase10-channel-graph-foundations\n", + "\n", + "**neuron Phase 10** — channel-as-node graph hidden layer foundations (사용자 vision 의 본질 axis).\n", + "\n", + "핵심 가설:\n", + "1. **function preservation** — channel_full (adj=full) 가 plain Linear 와 forward 동치?\n", + "2. **adjacency 학습** — adj 파라미터에 gradient 흐름 + 학습된 per-edge importance?\n", + "3. **0-init 금지 + sweet spot 적용** — channel_uniform_small (adj ∈ [0.05, 0.15]) 시작이 plain 과 비교?\n", + "4. **post-training adj 분포** — Phase 5 의 implicit pruning 패턴이 edge-level 에서 재현?\n", + "5. **Phase 9 group 과의 비교** — channel-level granularity 가 group-level (Phase 9: 2.1378~2.1391) 보다 우위?\n", + "\n", + "설계: 3-way sweep × 2 seed = 6 run, max_steps=1500.\n", + "- arch ∈ {plain, channel_full, channel_uniform_small}\n", + "- seed ∈ {42, 123}\n", + "\n", + "데이터: TinyShakespeare (char-LM)\n", + "시드: [42, 123]\n", + "작성일: 2026-05-26\n", + "연관: Issue [#61](https://github.com/EinSofINTEREST/GraphLM/issues/61) / Phase 9 baseline PR [#60](https://github.com/EinSofINTEREST/GraphLM/pull/60) / [Phase 9 결과 PR #60](https://github.com/EinSofINTEREST/GraphLM/pull/60) — 0-init 금지 규칙" + ] + }, + { + "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_channel_demo import train_channel_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-phase10\"\n", + "OUT_DIR.mkdir(parents=True, exist_ok=True)\n", + "\n", + "SEEDS = [42, 123]\n", + "ARCHS = [\"plain\", \"channel_full\", \"channel_uniform_small\"]\n", + "EMB_DIM = 64\n", + "HIDDEN_DIM = 256 # channel granularity — group 무관, 어떤 dim 도 가능\n", + "N_GRAM = 4\n", + "BATCH_SIZE = 32\n", + "LR = 3e-4\n", + "MAX_STEPS = 1500\n", + "\n", + "# Phase 9 baseline (PR #60 결과)\n", + "PHASE9_PLAIN_MEAN = 2.1378\n", + "PHASE9_GROUP_FULL_MEAN = 2.1391\n", + "PHASE9_GROUP_IDENTITY_MEAN = 2.2797 # 0-init vanishing 사례\n", + "\n", + "print(f\"SEEDS = {SEEDS}\")\n", + "print(f\"ARCHS = {ARCHS}\")\n", + "print(f\"HIDDEN_DIM = {HIDDEN_DIM} (channel granularity)\")\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", + "print(f\"vocab_size : {V}\") # padding 불필요 — channel granularity 는 모든 dim 지원" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## 4. sweep 학습\n", + "\n", + "각 (seed, arch) 에 대해 1 run.\n", + "- plain: standard nn.Linear baseline\n", + "- channel_full: adj=full (function preservation, 'free' graph)\n", + "- channel_uniform_small: adj ∈ [0.05, 0.15] (Phase 2 sweet spot 패턴 적용, 0-init 회피)" + ] + }, + { + "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_channel_graph_mlp(\n", + " dataset=dataset,\n", + " vocab_size=V,\n", + " seed=seed,\n", + " arch=arch,\n", + " emb_dim=EMB_DIM,\n", + " hidden_dim=HIDDEN_DIM,\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 + Phase 9 baseline 비교" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "import statistics\n", + "\n", + "print(f\"{'arch':>26} {'seed':>5} {'final_loss':>11}\")\n", + "print(\"-\" * 50)\n", + "for arch in ARCHS:\n", + " for seed in SEEDS:\n", + " r = runs[(seed, arch)]\n", + " print(f\"{arch:>26} {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:>26}: mean={agg[arch]['mean']:.4f}, range={agg[arch]['range']:.4f}\")\n", + "\n", + "print()\n", + "print(\"=== Phase 9 (group-level) baseline 비교 ===\")\n", + "print(f\" Phase 9 plain : {PHASE9_PLAIN_MEAN:.4f}\")\n", + "print(f\" Phase 9 group_full : {PHASE9_GROUP_FULL_MEAN:.4f}\")\n", + "print(f\" Phase 9 group_identity : {PHASE9_GROUP_IDENTITY_MEAN:.4f} (0-init vanishing 사례)\")\n", + "print()\n", + "print(f\" Phase 10 plain (재현) : {agg['plain']['mean']:.4f}\")\n", + "print(f\" Phase 10 channel_full : {agg['channel_full']['mean']:.4f}\")\n", + "print(f\" Phase 10 channel_uniform : {agg['channel_uniform_small']['mean']:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## 6. 학습된 adj 분포 분석 — implicit pruning at edge-level?" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "# channel_* 의 fc1 adj 분포 vs init\n", + "fig, axes = plt.subplots(2, 2, figsize=(13, 8))\n", + "\n", + "for col_i, arch in enumerate([\"channel_full\", \"channel_uniform_small\"]):\n", + " r = runs[(SEEDS[0], arch)]\n", + " if r[\"final_adj\"] is None:\n", + " continue\n", + " for row_i, layer in enumerate([\"fc1\", \"fc2\"]):\n", + " ax = axes[row_i, col_i]\n", + " adj = r[\"final_adj\"][layer].numpy().flatten()\n", + " ax.hist(adj, bins=80, alpha=0.7, color=\"#1f77b4\")\n", + " ax.set_xlabel(\"adj value\")\n", + " ax.set_ylabel(\"count\")\n", + " ax.set_title(f\"{arch} — {layer} adj distribution (n={len(adj)})\")\n", + " ax.axvline(0, color=\"red\", linestyle=\"--\", lw=0.8, alpha=0.5)\n", + " ax.axvline(adj.mean(), color=\"green\", linestyle=\":\", lw=1, label=f\"mean={adj.mean():.3f}\")\n", + " ax.legend(fontsize=8)\n", + " ax.grid(alpha=0.3)\n", + "\n", + "fig.suptitle(f\"Phase 10 — channel-level adj 분포 (seed={SEEDS[0]})\", fontsize=11)\n", + "fig.tight_layout()\n", + "fig.savefig(OUT_DIR / \"adj_distribution.png\", dpi=120)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## 7. adj heatmap (fc1, channel × channel)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n", + "\n", + "for i, arch in enumerate([\"channel_full\", \"channel_uniform_small\"]):\n", + " r = runs[(SEEDS[0], arch)]\n", + " if r[\"final_adj\"] is None:\n", + " continue\n", + " adj = r[\"final_adj\"][\"fc1\"].numpy()\n", + " vmax = max(abs(adj).max(), 1e-6)\n", + " im = axes[i].imshow(adj, cmap=\"RdBu_r\", vmin=-vmax, vmax=vmax, aspect=\"auto\")\n", + " axes[i].set_xlabel(\"input channel\")\n", + " axes[i].set_ylabel(\"output channel\")\n", + " axes[i].set_title(f\"{arch} — fc1 adj (shape={adj.shape})\")\n", + " fig.colorbar(im, ax=axes[i], fraction=0.046, pad=0.04)\n", + "\n", + "fig.suptitle(f\"Phase 10 — channel adj heatmap (seed={SEEDS[0]}, fc1)\", fontsize=11)\n", + "fig.tight_layout()\n", + "fig.savefig(OUT_DIR / \"adj_heatmaps.png\", dpi=120)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "15", + "metadata": {}, + "source": [ + "## 8. loss curve 비교 (arch × mean ± σ across 2 seeds)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "window = 30\n", + "colors = {\n", + " \"plain\": \"#1f77b4\",\n", + " \"channel_full\": \"#2ca02c\",\n", + " \"channel_uniform_small\": \"#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", + "ax.axhline(\n", + " PHASE9_GROUP_FULL_MEAN,\n", + " color=\"gray\",\n", + " linestyle=\":\",\n", + " lw=1,\n", + " alpha=0.7,\n", + " label=f\"Phase 9 group_full ({PHASE9_GROUP_FULL_MEAN})\",\n", + ")\n", + "ax.set_xlabel(\"step\")\n", + "ax.set_ylabel(f\"loss (smoothed window={window})\")\n", + "ax.set_title(f\"Phase 10 — plain Linear vs ChannelGraphLinear (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": "17", + "metadata": {}, + "source": [ + "## 결과 요약 / Phase 11 권장 방향\n", + "\n", + "확인 포인트:\n", + "- §5 channel_full vs plain — function preservation 입증 (~0)?\n", + "- §5 channel_uniform_small vs plain — 0-init 회피한 sweet spot 패턴이 잘 작동? (열위 ≤ 0.05 면 OK)\n", + "- §6 adj 분포 — 학습 후 spread? 일부 edge 가 자연 약화 (implicit pruning)?\n", + "- §7 heatmap — sparse/structured 패턴 emerge?\n", + "- §8 Phase 9 group_full (2.1391) 와의 비교 — channel granularity 가 group 보다 우위/동등?\n", + "\n", + "**판정 시나리오**:\n", + "- **A. channel_full ≈ plain + uniform_small 도 비슷** ⭐ — function preservation + 0-init 회피 둘 다 입증, Phase 11 (hybrid) 진입\n", + "- **B. channel_full ≈ plain, channel_uniform_small 열위** — sweet spot 의 small init 이 fine-grained edge 에는 부족, 추가 sweep 필요\n", + "- **C. channel_uniform_small 우위** — implicit pruning at edge-level 이 실제 효과 입증 (sparsification 학습 motivation 강화)\n", + "\n", + "**참고**:\n", + "- 아키텍처 구성 계획 (Notion): https://www.notion.so/36ce8b70b7aa818cbf1fe71687b449b8\n", + "- Phase 9 결과: https://www.notion.so/36ce8b70b7aa8100b0acf756686d2e9f" + ] + } + ], + "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/src/graphlm/neuron/__init__.py b/src/graphlm/neuron/__init__.py index 6a6dce8..a560d4e 100644 --- a/src/graphlm/neuron/__init__.py +++ b/src/graphlm/neuron/__init__.py @@ -11,11 +11,13 @@ NeuronGrowingDecoder, SinusoidalAlpha, ) +from graphlm.neuron.graph_channel import ChannelGraphLinear 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__ = [ + "ChannelGraphLinear", "GroupGraphLinear", "GrowableEmbedding", "GrowableLayerNorm", diff --git a/src/graphlm/neuron/graph_channel.py b/src/graphlm/neuron/graph_channel.py new file mode 100644 index 0000000..2153017 --- /dev/null +++ b/src/graphlm/neuron/graph_channel.py @@ -0,0 +1,138 @@ +"""Phase 10 — channel-as-node graph hidden layer foundations. + +Hidden_dim H 채널을 graph 의 node 로 다룬다 (paradigm 의 finest unit). standard ``nn.Linear`` +의 weight matrix 의 각 entry 를 *edge* 로 재해석하고, 학습 가능 ``adj`` (per-edge gate) 를 +도입. + +``` +in_features H_in 채널 = 입력 nodes +out_features H_out 채널 = 출력 nodes +edge[out, in]: 학습 가능 scalar adj[out, in] +effective weight: W_eff = adj ⊙ W (elementwise product) +forward: y = W_eff @ x = (adj * W) @ x +``` + +설계 의미: +- node = 1 채널 (graph 의 vertex) +- edge = (in, out) 쌍의 connection — adj 가 routing strength, W 가 transformation +- adj=full (모두 1) → effective = W → standard Linear forward 동치 (function preserving) +- adj=uniform_small → Phase 2 sweet spot 패턴 channel-level edge 에 적용 (0-init 회피) + +**0-init 금지 규칙 적용** (rationale: Phase 9 결과 PR #60 + +[feedback_no_zero_init.md](https://www.notion.so/36ce8b70b7aa8100b0acf756686d2e9f) Notion 정리): +- Phase 1 dead block / Phase 7 amplitude vanishing / Phase 9 block-diagonal 의 3차 반복 발견 +- 본 모듈은 ``adj_init="zero"`` 옵션 명시적 거부 (ValueError) +- default 권장 = ``"full"`` (function preserving) 또는 ``"uniform_small"`` (sweet spot 패턴) + +Phase 9 (GroupGraphLinear) 와의 차이: +- Phase 9: group_size 단위 block 으로 묶어 (n_groups_out, n_groups_in) adjacency +- Phase 10: 채널 1개 단위 — (out, in) 전체 adjacency, paradigm 의 finest unit +""" + +from __future__ import annotations + +import math +from typing import Literal + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +AdjInit = Literal["full", "uniform_small"] + + +class ChannelGraphLinear(nn.Module): + """Channel-as-node graph linear layer (Phase 10 paradigm 의 finest unit foundation). + + Args: + in_features: 입력 채널 수 (graph 의 in-nodes). + out_features: 출력 채널 수 (graph 의 out-nodes). + adj_init: ``"full"`` (모두 1, function preserving) 또는 ``"uniform_small"`` + (uniform[0.05, 0.15], Phase 2 sweet spot 패턴). ``"zero"`` 등 0-init 옵션은 + 거부 (rationale: Phase 9 결과 PR #60 — block-diagonal sparse 시작이 +0.14 loss + 열위, vanishing gradient 함정). + bias: bias 사용 여부 (standard Linear 동일). + + Forward: + ``y = (adj * weight) @ x + bias`` + """ + + def __init__( + self, + in_features: int, + out_features: int, + *, + adj_init: AdjInit = "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 + + # standard Linear init for W (fan_in = in_features) + self.weight = nn.Parameter(torch.empty(out_features, in_features)) + bound = 1.0 / math.sqrt(in_features) + nn.init.uniform_(self.weight, -bound, bound) + + # adj init — 0-init 거부 (feedback_no_zero_init.md) + if adj_init == "full": + adj = torch.ones(out_features, in_features) + elif adj_init == "uniform_small": + # Phase 2 sweet spot (0.10) ± δ — gradient flow 확보 + edge importance 학습 시작 + adj = torch.empty(out_features, in_features).uniform_(0.05, 0.15) + elif adj_init in {"zero", "zeros"}: + raise ValueError( + f"adj_init={adj_init!r} 는 금지됨 — 0-init 은 vanishing gradient 함정 " + "(Phase 1 dead block / Phase 7 amplitude vanishing / Phase 9 block-diagonal " + "에서 3차 재현). rationale: Phase 9 PR #60. " + "'full' (function preserving) 또는 'uniform_small' (sweet spot 패턴) 사용 권장." + ) + else: + raise ValueError(f"unknown adj_init: {adj_init!r}") + self.adj = nn.Parameter(adj) + + if bias: + self.bias = nn.Parameter(torch.zeros(out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, x: Tensor) -> Tensor: + effective_w = self.adj * self.weight + return F.linear(x, effective_w, self.bias) + + def adj_sparsity(self, threshold: float = 0.05) -> float: + """현재 |adj| < threshold 인 edge 의 비율 (0~1).""" + if threshold < 0: + raise ValueError(f"threshold must be >= 0, got {threshold}") + with torch.no_grad(): + below = (self.adj.abs() < threshold).float().mean().item() + return float(below) + + def sparsify_adj(self, threshold: float) -> int: + """|adj| < threshold 인 entry 를 0 으로 강제 (in-place). returns 비활성화된 entry 수. + + ``self.adj.data.masked_fill_`` 사용 — leaf+requires_grad tensor 에서 in-place op 의 + autograd 안전성 명시적 보장 (gemini #3301728271 방어적 패턴). + """ + if threshold < 0: + raise ValueError(f"threshold must be >= 0, got {threshold}") + with torch.no_grad(): + zero_mask = self.adj.abs() < threshold + n_zeroed = int(zero_mask.sum().item()) + self.adj.data.masked_fill_(zero_mask, 0.0) + return n_zeroed + + def freeze_adjacency(self) -> None: + """adj 를 학습 비대상으로.""" + self.adj.requires_grad_(False) + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"adj_shape={tuple(self.adj.shape)}" + ) diff --git a/src/graphlm/neuron/graph_channel_demo.py b/src/graphlm/neuron/graph_channel_demo.py new file mode 100644 index 0000000..e2d2452 --- /dev/null +++ b/src/graphlm/neuron/graph_channel_demo.py @@ -0,0 +1,122 @@ +"""Phase 10 — ChannelGraphLinear demo MLP-LM + 학습 헬퍼. + +노트북 분리 규약 준수. Phase 10 노트북 09-phase10-channel-graph-foundations.ipynb 는 여기서 +import. Phase 9 의 ``graph_group_demo`` 의 channel-level 대응. + +3 가지 architecture 를 같은 학습 루프로 비교: +- ``"plain"`` — 표준 nn.Linear (baseline) +- ``"channel_full"`` — ChannelGraphLinear with adj_init="full" (function preserving 시작) +- ``"channel_uniform_small"`` — ChannelGraphLinear with adj_init="uniform_small" + (Phase 2 sweet spot 패턴 channel-level edge 에 적용) +""" + +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_channel import ChannelGraphLinear +from graphlm.utils import set_seed + +Arch = Literal["plain", "channel_full", "channel_uniform_small"] + + +class ChannelGraphMLPLM(nn.Module): + """Channel-as-node graph MLP-LM — Phase 9 의 GroupGraphMLPLM 의 channel-level 버전.""" + + def __init__( + self, + vocab_size: int, + emb_dim: int, + hidden_dim: int, + n_gram: int, + arch: Arch, + ): + 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) + self.ln = nn.LayerNorm(hidden_dim) + self.fc2 = _make_linear(hidden_dim, vocab_size, arch) + + 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) -> nn.Module: + if arch == "plain": + return nn.Linear(in_f, out_f) + if arch == "channel_full": + return ChannelGraphLinear(in_f, out_f, adj_init="full") + if arch == "channel_uniform_small": + return ChannelGraphLinear(in_f, out_f, adj_init="uniform_small") + 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_channel_graph_mlp( + *, + dataset: TinyShakespeareDataset, + vocab_size: int, + seed: int, + arch: Arch, + emb_dim: int, + hidden_dim: int, + n_gram: int, + batch_size: int, + lr: float, + max_steps: int, + device: str = "cpu", +) -> dict: + """1 run 학습 — Phase 10 sweep 의 단위. + + Returns dict: ``losses``, ``final_loss``, ``final_adj`` (channel_* 의 경우 fc1/fc2 의 adj snapshot). + """ + set_seed(seed) + model = ChannelGraphMLPLM(vocab_size, emb_dim, hidden_dim, n_gram, arch).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": + 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/src/graphlm/neuron/graph_group.py b/src/graphlm/neuron/graph_group.py index 3ff4566..05d635f 100644 --- a/src/graphlm/neuron/graph_group.py +++ b/src/graphlm/neuron/graph_group.py @@ -146,9 +146,10 @@ def sparsify_adjacency(self, threshold: float) -> int: raise ValueError(f"threshold must be >= 0, got {threshold}") with torch.no_grad(): # masked_fill_ 가 mul_(mask.to(dtype)) 보다 깔끔 + 캐스팅 불필요 (gemini #3301524139) + # .data 명시 — leaf+requires_grad tensor 의 in-place 안전성 (gemini #3301728271 동일 패턴) zero_mask = self.adj.abs() < threshold n_zeroed = int(zero_mask.sum().item()) - self.adj.masked_fill_(zero_mask, 0.0) + self.adj.data.masked_fill_(zero_mask, 0.0) return n_zeroed def extra_repr(self) -> str: diff --git a/tests/neuron/test_graph_channel.py b/tests/neuron/test_graph_channel.py new file mode 100644 index 0000000..95d0117 --- /dev/null +++ b/tests/neuron/test_graph_channel.py @@ -0,0 +1,120 @@ +"""Tests for graphlm.neuron.graph_channel — Phase 10 channel-as-node foundations.""" + +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from graphlm.neuron.graph_channel import ChannelGraphLinear + + +def test_shape_after_init(): + lin = ChannelGraphLinear(8, 16) + assert lin.weight.shape == (16, 8) + assert lin.adj.shape == (16, 8) + assert lin.bias.shape == (16,) + + +def test_full_adj_init_value(): + lin = ChannelGraphLinear(8, 16, adj_init="full") + assert torch.allclose(lin.adj, torch.ones(16, 8)) + + +def test_uniform_small_adj_init_range(): + """adj_init='uniform_small' → uniform[0.05, 0.15] (Phase 2 sweet spot 패턴).""" + torch.manual_seed(0) + lin = ChannelGraphLinear(8, 16, adj_init="uniform_small") + assert (lin.adj >= 0.05).all() and (lin.adj <= 0.15).all() + # 분포가 실제로 spread 되어 있는지 (degenerate 거부) + assert lin.adj.std() > 0.01 + + +@pytest.mark.parametrize("bad_init", ["zero", "zeros"]) +def test_zero_init_rejected(bad_init): + """0-init 옵션 명시적 거부 — Phase 1/7/9 의 vanishing gradient 함정 회피.""" + with pytest.raises(ValueError, match="vanishing gradient"): + ChannelGraphLinear(8, 16, adj_init=bad_init) # type: ignore[arg-type] + + +def test_unknown_adj_init_raises(): + with pytest.raises(ValueError, match="unknown adj_init"): + ChannelGraphLinear(8, 16, adj_init="bogus") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "field,value", + [("in_features", 0), ("in_features", -1), ("out_features", 0), ("out_features", -1)], +) +def test_invalid_features_raises(field, value): + kwargs = {"in_features": 8, "out_features": 16} + kwargs[field] = value + with pytest.raises(ValueError, match=field): + ChannelGraphLinear(**kwargs) + + +def test_function_preservation_equivalent_to_standard_linear(): + """adj=full + 같은 W → standard Linear 와 forward 동일 (atol=1e-5).""" + torch.manual_seed(0) + in_f, out_f = 8, 16 + cg = ChannelGraphLinear(in_f, out_f, adj_init="full") + std = nn.Linear(in_f, out_f, bias=True) + with torch.no_grad(): + std.weight.copy_(cg.weight) + std.bias.copy_(cg.bias) + + x = torch.randn(2, 4, in_f) + y_cg = cg(x) + y_std = std(x) + assert torch.allclose(y_cg, y_std, atol=1e-5), ( + f"function preservation 깨짐: max |diff| = {(y_cg - y_std).abs().max().item()}" + ) + + +def test_weight_and_adj_both_have_gradient(): + lin = ChannelGraphLinear(8, 16, adj_init="uniform_small") + x = torch.randn(4, 8) + 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_adj_sparsity_initial(): + """초기 adj_init='full' 은 모두 1 이므로 (|adj| < 0.05) sparsity = 0.""" + lin = ChannelGraphLinear(8, 16, adj_init="full") + assert lin.adj_sparsity(threshold=0.05) == 0.0 + + +def test_adj_sparsity_after_zero_fill(): + lin = ChannelGraphLinear(8, 16, adj_init="full") + # 일부 entry 를 강제로 0 으로 → sparsity 변화 확인 + with torch.no_grad(): + lin.adj[:8].fill_(0.0) # 절반 (8/16) + assert abs(lin.adj_sparsity(threshold=0.05) - 0.5) < 1e-6 + + +def test_sparsify_adj_zeros_below_threshold(): + """uniform_small (0.05~0.15) 에서 threshold=0.10 으로 sparsify → 평균 절반 zeroed.""" + torch.manual_seed(0) + lin = ChannelGraphLinear(64, 64, adj_init="uniform_small") + n_total = 64 * 64 + n_zeroed = lin.sparsify_adj(threshold=0.10) + # uniform[0.05, 0.15] 의 mid = 0.10 이라 약 절반 (~50%) 이 < 0.10 + assert n_zeroed > n_total * 0.3 + assert n_zeroed < n_total * 0.7 + + +def test_sparsify_adj_negative_threshold_raises(): + lin = ChannelGraphLinear(8, 16) + with pytest.raises(ValueError, match="threshold"): + lin.sparsify_adj(threshold=-0.1) + + +def test_freeze_adjacency(): + lin = ChannelGraphLinear(8, 16) + lin.freeze_adjacency() + assert not lin.adj.requires_grad + assert lin.weight.requires_grad