diff --git a/docs/figures/neuron/phase13/hybrid_adj.png b/docs/figures/neuron/phase13/hybrid_adj.png new file mode 100644 index 0000000..2953490 Binary files /dev/null and b/docs/figures/neuron/phase13/hybrid_adj.png differ diff --git a/docs/figures/neuron/phase13/loss_curves.png b/docs/figures/neuron/phase13/loss_curves.png new file mode 100644 index 0000000..b4b4b9e Binary files /dev/null and b/docs/figures/neuron/phase13/loss_curves.png differ diff --git a/notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb b/notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb new file mode 100644 index 0000000..3d731db --- /dev/null +++ b/notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb @@ -0,0 +1,282 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 12-phase13-hybrid-transformer\n", + "\n", + "**neuron Phase 13** — HybridGraphLinear 를 **실제 Transformer 의 FFN 위치** 에 통합 + **RMSNorm** 도입.\n", + "Phase 8~12 는 모두 MLP-LM baseline 이었고, Phase 13 부터 standard pre-norm Transformer block 위에서 graph hidden layer 의 paradigm 이 동작하는지 검증.\n", + "\n", + "핵심 가설:\n", + "1. **function preservation (Transformer 위)** — `hybrid_full_full` ≈ `plain` (RMSNorm + 표준 FFN)?\n", + "2. **dual routing 우위 재현** — Phase 12 에서 hybrid_full_around_one 이 최저 final_loss 였음. Transformer 에서도 유지?\n", + "3. **full scale-corrected (around_one × around_one)** — outer/inner 둘 다 학습 활성화 init 의 효과?\n", + "4. **RMSNorm 안정성** — LayerNorm 대비 학습 동등 또는 우위 (모든 arch 에서 NaN 없음)?\n", + "\n", + "설계: 4 arch × 2 seed = 8 run, max_steps=1500.\n", + "\n", + "데이터: TinyShakespeare (char-LM, block_size=64)\n", + "시드: [42, 123]\n", + "작성일: 2026-05-26\n", + "연관: Issue [#67](https://github.com/EinSofINTEREST/GraphLM/issues/67) / Phase 12 baseline PR [#66](https://github.com/EinSofINTEREST/GraphLM/pull/66)\n", + "\n", + "Phase 13 의 ``identity`` outer 미지원 → 4 arch 는 plain / hybrid_full_full / hybrid_full_around_one / hybrid_around_one_around_one." + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## 0. 환경 / 의존성" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": "from __future__ import annotations\n\nimport math\nimport statistics\nfrom pathlib import Path\n\nimport matplotlib.pyplot as plt\nimport torch\n\nfrom graphlm.data.tinyshakespeare import (\n CharTokenizer,\n TinyShakespeareDataset,\n load_tinyshakespeare_text,\n)\nfrom graphlm.neuron.hybrid_transformer_demo import (\n HybridGraphTransformerLM,\n HybridTransformerTrainConfig,\n count_parameters,\n train_hybrid_transformer_lm,\n)\nfrom graphlm.utils import safe_perplexity\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"device: {device}\")\nprint(f\"torch: {torch.__version__}\")" + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## 1. Config + 데이터" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "# data\n", + "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", + "# model + train hyperparameters\n", + "HIDDEN_DIM = 128\n", + "N_HEADS = 4\n", + "FFN_DIM = 256 # 2x hidden (작게 — sweep 속도 위해)\n", + "N_LAYERS = 4\n", + "GROUP_SIZE = 16 # hidden_dim/group_size = 8, ffn_dim/group_size = 16\n", + "BLOCK_SIZE = 64\n", + "BATCH_SIZE = 32\n", + "LR = 3e-4\n", + "MAX_STEPS = 1500\n", + "SEEDS = [42, 123]\n", + "ARCHS = [\n", + " \"plain\",\n", + " \"hybrid_full_full\",\n", + " \"hybrid_full_around_one\",\n", + " \"hybrid_around_one_around_one\",\n", + "]\n", + "\n", + "# 모델 파라미터 수 비교 (arch 별)\n", + "print(\"\\n== Parameter count by arch ==\")\n", + "for arch in ARCHS:\n", + " m = HybridGraphTransformerLM(\n", + " vocab_size=vocab_size,\n", + " hidden_dim=HIDDEN_DIM,\n", + " n_heads=N_HEADS,\n", + " ffn_dim=FFN_DIM,\n", + " n_layers=N_LAYERS,\n", + " max_seq_len=BLOCK_SIZE,\n", + " arch=arch,\n", + " group_size=GROUP_SIZE,\n", + " )\n", + " print(f\" {arch:32s} params = {count_parameters(m):,}\")" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## 2. Sweep 실행 (4 arch × 2 seed = 8 run)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": "results = {}\nfor arch in ARCHS:\n for seed in SEEDS:\n key = (arch, seed)\n print(f\"\\n== arch={arch} seed={seed} ==\")\n cfg = HybridTransformerTrainConfig(\n dataset=dataset,\n vocab_size=vocab_size,\n hidden_dim=HIDDEN_DIM,\n n_heads=N_HEADS,\n ffn_dim=FFN_DIM,\n n_layers=N_LAYERS,\n group_size=GROUP_SIZE,\n arch=arch,\n block_size=BLOCK_SIZE,\n batch_size=BATCH_SIZE,\n lr=LR,\n max_steps=MAX_STEPS,\n seed=seed,\n device=device,\n )\n out = train_hybrid_transformer_lm(cfg)\n results[key] = out\n print(\n f\" final_loss = {out['final_loss']:.4f} (perplexity = {safe_perplexity(out['final_loss']):.2f})\"\n )" + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## 3. 결과 표 + 자동 verdict" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": "print(f\"{'arch':32s} {'seed':>6s} {'final_loss':>12s} {'perplexity':>12s}\")\nprint(\"-\" * 70)\nfor (arch, seed), out in results.items():\n fl = out[\"final_loss\"]\n print(f\"{arch:32s} {seed:>6d} {fl:>12.4f} {safe_perplexity(fl):>12.2f}\")\n\n# arch 별 평균 / 표준편차\nprint(\"\\n== Arch-level summary (mean ± σ across seeds) ==\")\nsummary = {}\nfor arch in ARCHS:\n vals = [results[(arch, s)][\"final_loss\"] for s in SEEDS]\n summary[arch] = (statistics.mean(vals), statistics.stdev(vals) if len(vals) > 1 else 0.0)\n m, s = summary[arch]\n print(f\" {arch:32s} {m:.4f} ± {s:.4f} (perplexity ≈ {safe_perplexity(m):.2f})\")\n\n# 자동 verdict\nprint(\"\\n== Verdict ==\")\nplain_loss = summary[\"plain\"][0]\nff_loss = summary[\"hybrid_full_full\"][0]\nfa_loss = summary[\"hybrid_full_around_one\"][0]\naa_loss = summary[\"hybrid_around_one_around_one\"][0]\n\n# 1. function preservation (학습 후 hybrid_full_full ≈ plain — 학습 dynamics 가 동일 sweet spot 유지)\ndiff_ff = abs(ff_loss - plain_loss)\nverdict_1 = \"PASS\" if diff_ff < 0.15 else \"FAIL\"\nprint(f\"1. function preservation: |hybrid_full_full - plain| = {diff_ff:.4f} [{verdict_1}]\")\n\n# 2. scale-corrected 우위 (hybrid_full_around_one 이 plain 보다 우수 또는 동등)\nverdict_2 = \"PASS\" if fa_loss <= plain_loss + 0.05 else \"FAIL\"\ndiff_fa = fa_loss - plain_loss\nprint(f\"2. inner around_one ≤ plain + 0.05: diff = {diff_fa:+.4f} [{verdict_2}]\")\n\n# 3. full scale-corrected (around_one × around_one) 안정성 (유한 + plain 근방)\n# CodeRabbit #3304306127 — isfinite 가 inf/-inf 도 거부 (isnan 만 쓰면 inf 가 PASS 됨)\nverdict_3 = \"PASS\" if (math.isfinite(aa_loss) and aa_loss < plain_loss + 0.5) else \"FAIL\"\ndiff_aa = aa_loss - plain_loss\nprint(f\"3. full around_one stable: diff = {diff_aa:+.4f} [{verdict_3}]\")\n\n# 4. 모두 finite 로 학습 종료 (RMSNorm 안정성, inf/nan 모두 거부)\nall_finite = all(math.isfinite(out[\"final_loss\"]) for out in results.values())\nverdict_4 = \"PASS\" if all_finite else \"FAIL\"\nprint(f\"4. RMSNorm stability (all finite): {all_finite} [{verdict_4}]\")" + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "## 4. Loss curve 시각화" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n", + "colors = {\n", + " \"plain\": \"tab:gray\",\n", + " \"hybrid_full_full\": \"tab:blue\",\n", + " \"hybrid_full_around_one\": \"tab:orange\",\n", + " \"hybrid_around_one_around_one\": \"tab:green\",\n", + "}\n", + "window = 50 # rolling mean for smoothing\n", + "\n", + "for arch in ARCHS:\n", + " losses_per_seed = [results[(arch, s)][\"losses\"] for s in SEEDS]\n", + " # rolling mean\n", + " smoothed = []\n", + " for losses in losses_per_seed:\n", + " smoothed.append(\n", + " [\n", + " sum(losses[max(0, i - window) : i + 1]) / min(i + 1, window)\n", + " for i in range(len(losses))\n", + " ]\n", + " )\n", + " # mean ± σ across seeds\n", + " arr = torch.tensor(smoothed)\n", + " mean = arr.mean(dim=0)\n", + " std = arr.std(dim=0)\n", + " steps = list(range(len(mean)))\n", + " color = colors[arch]\n", + " ax.plot(steps, mean, label=arch, color=color, linewidth=1.5)\n", + " ax.fill_between(steps, mean - std, mean + std, color=color, alpha=0.15)\n", + "\n", + "ax.set_xlabel(\"step\")\n", + "ax.set_ylabel(f\"loss (rolling mean w={window})\")\n", + "ax.set_title(\"Phase 13 — Transformer FFN: 4 arch loss curves (mean ± σ over 2 seeds)\")\n", + "ax.legend(loc=\"upper right\")\n", + "ax.grid(alpha=0.3)\n", + "plt.tight_layout()\n", + "\n", + "out_dir = Path(\"../../runs/notebook-neuron-phase13\")\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. adj_outer / adj_inner heatmap (hybrid arch 만)\n", + "\n", + "각 hybrid arch 의 첫 block 의 fc1 adj_outer 와 adj_inner (block-aggregated) 학습 후 모습." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "hybrid_archs = [a for a in ARCHS if a != \"plain\"]\n", + "n_arch = len(hybrid_archs)\n", + "fig, axes = plt.subplots(n_arch, 2, figsize=(10, 3 * n_arch))\n", + "if n_arch == 1:\n", + " axes = axes.reshape(1, -1)\n", + "\n", + "for row, arch in enumerate(hybrid_archs):\n", + " # seed 42 의 첫 block 의 fc1\n", + " snap = results[(arch, 42)][\"final_adj\"][0][\"fc1\"]\n", + " outer = snap[\"outer\"].numpy() # (G_out, G_in)\n", + " inner = (\n", + " snap[\"inner\"].abs().mean(dim=(-1, -2)).numpy()\n", + " ) # (G_out, G_in) — block-aggregated magnitude\n", + "\n", + " ax_o = axes[row, 0]\n", + " im_o = ax_o.imshow(outer, cmap=\"RdBu_r\", vmin=-2, vmax=2)\n", + " ax_o.set_title(f\"{arch} — adj_outer (block 0, fc1)\")\n", + " ax_o.set_xlabel(\"G_in\")\n", + " ax_o.set_ylabel(\"G_out\")\n", + " plt.colorbar(im_o, ax=ax_o, fraction=0.046)\n", + "\n", + " ax_i = axes[row, 1]\n", + " im_i = ax_i.imshow(inner, cmap=\"viridis\")\n", + " ax_i.set_title(f\"{arch} — |adj_inner| block-mean (block 0, fc1)\")\n", + " ax_i.set_xlabel(\"G_in\")\n", + " ax_i.set_ylabel(\"G_out\")\n", + " plt.colorbar(im_i, ax=ax_i, fraction=0.046)\n", + "\n", + "plt.tight_layout()\n", + "fig.savefig(out_dir / \"hybrid_adj.png\", dpi=150, bbox_inches=\"tight\")\n", + "plt.show()\n", + "print(f\"saved: {out_dir / 'hybrid_adj.png'}\")" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## 6. 결론 / 다음 단계\n", + "\n", + "(셀 출력 보고 사용자가 채울 영역)\n", + "\n", + "- function preservation: hybrid_full_full vs plain — Transformer 위에서도 동등?\n", + "- dual routing 우위: Phase 12 패턴 재현?\n", + "- adj 학습 패턴: outer / inner 가 다른 영역을 cover?\n", + "\n", + "**Phase 14 후보**:\n", + "- attention 의 qkv / out 을 HybridGraphLinear 로 교체 — function-level graph 가 attention 까지 확장\n", + "- Net2Net / LiGO 식 growable transformer block — 학습 중 hidden_dim 또는 n_layers 증가" + ] + } + ], + "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 +} \ No newline at end of file diff --git a/src/graphlm/neuron/__init__.py b/src/graphlm/neuron/__init__.py index dafd3e6..84d0c0b 100644 --- a/src/graphlm/neuron/__init__.py +++ b/src/graphlm/neuron/__init__.py @@ -16,6 +16,13 @@ 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 +from graphlm.neuron.hybrid_transformer import ( + HybridGraphFFN, + HybridGraphTransformerBlock, + PlainTransformerBlock, + make_block, +) +from graphlm.neuron.rms_norm import RMSNorm __all__ = [ "ChannelGraphLinear", @@ -23,11 +30,16 @@ "GrowableEmbedding", "GrowableLayerNorm", "GrowableLinear", + "HybridGraphFFN", "HybridGraphLinear", + "HybridGraphTransformerBlock", "NeuronBlock", "NeuronConfig", "NeuronGrowingDecoder", + "PlainTransformerBlock", + "RMSNorm", "SinusoidalAlpha", "add_attn_function_preserving", "add_attn_smooth_start", + "make_block", ] diff --git a/src/graphlm/neuron/hybrid_transformer.py b/src/graphlm/neuron/hybrid_transformer.py new file mode 100644 index 0000000..4e77b4a --- /dev/null +++ b/src/graphlm/neuron/hybrid_transformer.py @@ -0,0 +1,209 @@ +"""Phase 13 — HybridGraphLinear 의 Transformer FFN 통합 + RMSNorm pre-norm block. + +본 paradigm 의 graph 표현 (Phase 12 HybridGraphLinear) 을 **실제 Transformer 아키텍처** 안의 +FFN 위치에 도입. Phase 8~12 는 MLP-LM baseline 이었고, Phase 13 부터 standard pre-norm +Transformer block 위에서 검증. + +Scope (Phase 13): +- FFN 의 fc1 (hidden → ffn) + fc2 (ffn → hidden) 만 ``HybridGraphLinear`` +- attention (qkv / out) 은 표준 ``nn.Linear`` 유지 (Phase 14 검토) +- norm 은 ``RMSNorm`` (modern Transformer 표준) + +function preservation: +- ``HybridGraphFFN`` 의 adj_outer=full + adj_inner=full + 동일 weight 초기화 → standard FFN 동치 +- ``HybridGraphTransformerBlock`` 의 FFN 만 hybrid, attention/norm 은 표준 → standard pre-norm block 과 동치 +- 검증은 tests/neuron/test_hybrid_transformer.py 참조 +""" + +from __future__ import annotations + +from typing import Literal + +import torch.nn.functional as F +from torch import Tensor, nn + +from graphlm.neuron.backbone import CausalSelfAttention +from graphlm.neuron.graph_hybrid import AdjInnerInit, AdjOuterInit, HybridGraphLinear +from graphlm.neuron.rms_norm import RMSNorm + + +class HybridGraphFFN(nn.Module): + """FFN with two ``HybridGraphLinear`` layers + GELU. + + Args: + hidden_dim: in/out 차원 (Transformer hidden size). + ffn_dim: 중간 확장 차원 (보통 4·hidden_dim). + group_size: ``HybridGraphLinear`` 의 block size. hidden_dim / ffn_dim 모두 배수여야 함. + adj_outer_init: outer adj 초기화 (``"full"`` / ``"uniform_around_one"``). + ``"identity"`` 는 FFN 이 정의상 rectangular (hidden ≠ ffn) 이라 지원하지 않음. + adj_inner_init: inner adj 초기화 (``"full"`` / ``"uniform_around_one"``). + dropout: fc2 출력에 적용되는 dropout 확률 (backbone.FFN 과 동일 위치). + + Forward: + ``y = dropout(fc2(GELU(fc1(x))))`` — function preserving when both adj = full and dropout = 0. + """ + + def __init__( + self, + hidden_dim: int, + ffn_dim: int, + group_size: int, + *, + adj_outer_init: AdjOuterInit = "full", + adj_inner_init: AdjInnerInit = "full", + dropout: float = 0.0, + ): + super().__init__() + if adj_outer_init == "identity": + raise ValueError( + "HybridGraphFFN 은 adj_outer_init='identity' 미지원 — " + "FFN 은 hidden_dim → ffn_dim → hidden_dim 으로 rectangular 라 정방 identity 정의 불가. " + "'full' 또는 'uniform_around_one' 사용." + ) + self.fc1 = HybridGraphLinear( + hidden_dim, + ffn_dim, + group_size=group_size, + adj_outer_init=adj_outer_init, + adj_inner_init=adj_inner_init, + bias=False, + ) + self.fc2 = HybridGraphLinear( + ffn_dim, + hidden_dim, + group_size=group_size, + adj_outer_init=adj_outer_init, + adj_inner_init=adj_inner_init, + bias=False, + ) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: Tensor) -> Tensor: + return self.dropout(self.fc2(F.gelu(self.fc1(x)))) + + +class HybridGraphTransformerBlock(nn.Module): + """Pre-norm Transformer block — RMSNorm + CausalSelfAttention + RMSNorm + HybridGraphFFN. + + Forward: + ``x = x + attn(rms1(x))`` + ``x = x + ffn(rms2(x))`` + """ + + def __init__( + self, + hidden_dim: int, + n_heads: int, + ffn_dim: int, + group_size: int, + *, + adj_outer_init: AdjOuterInit = "full", + adj_inner_init: AdjInnerInit = "full", + dropout: float = 0.0, + ): + super().__init__() + self.rms1 = RMSNorm(hidden_dim) + self.attn = CausalSelfAttention(hidden_dim, n_heads, dropout=dropout) + self.rms2 = RMSNorm(hidden_dim) + # dropout 은 attention 과 FFN 양쪽 모두 적용 — backbone.FFN 의 fc2-뒤-dropout 패턴과 일관 + # (Copilot #3303168649) + self.ffn = HybridGraphFFN( + hidden_dim, + ffn_dim, + group_size=group_size, + adj_outer_init=adj_outer_init, + adj_inner_init=adj_inner_init, + dropout=dropout, + ) + + def forward(self, x: Tensor) -> Tensor: + x = x + self.attn(self.rms1(x)) + return x + self.ffn(self.rms2(x)) + + +Arch = Literal[ + "plain", + "hybrid_full_full", + "hybrid_full_around_one", + "hybrid_around_one_around_one", +] + + +class PlainFFN(nn.Module): + """Standard 2-layer FFN (no bias) — Phase 13 ``"plain"`` baseline. + + backbone.FFN 과 동일한 ``dropout(fc2(GELU(fc1(x))))`` 구조. + """ + + def __init__(self, hidden_dim: int, ffn_dim: int, dropout: float = 0.0): + super().__init__() + self.fc1 = nn.Linear(hidden_dim, ffn_dim, bias=False) + self.fc2 = nn.Linear(ffn_dim, hidden_dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: Tensor) -> Tensor: + return self.dropout(self.fc2(F.gelu(self.fc1(x)))) + + +class PlainTransformerBlock(nn.Module): + """Pre-norm Transformer block with standard FFN — Phase 13 ``"plain"`` baseline. + + HybridGraphTransformerBlock 과 동일한 norm/attention/residual 구조 — FFN 만 표준 nn.Linear. + 공정 비교 위해 RMSNorm 동일 사용. dropout 은 attention + FFN 양쪽 모두 적용. + """ + + def __init__(self, hidden_dim: int, n_heads: int, ffn_dim: int, dropout: float = 0.0): + super().__init__() + self.rms1 = RMSNorm(hidden_dim) + self.attn = CausalSelfAttention(hidden_dim, n_heads, dropout=dropout) + self.rms2 = RMSNorm(hidden_dim) + # Copilot #3303168686 — FFN dropout 도 적용해서 backbone.FFN 일관성 확보 + self.ffn = PlainFFN(hidden_dim, ffn_dim, dropout=dropout) + + def forward(self, x: Tensor) -> Tensor: + x = x + self.attn(self.rms1(x)) + return x + self.ffn(self.rms2(x)) + + +def make_block( + arch: Arch, + hidden_dim: int, + n_heads: int, + ffn_dim: int, + group_size: int, + dropout: float = 0.0, +) -> nn.Module: + """4 가지 arch 중 하나로 Phase 13 Transformer block 생성.""" + if arch == "plain": + return PlainTransformerBlock(hidden_dim, n_heads, ffn_dim, dropout=dropout) + if arch == "hybrid_full_full": + return HybridGraphTransformerBlock( + hidden_dim, + n_heads, + ffn_dim, + group_size=group_size, + adj_outer_init="full", + adj_inner_init="full", + dropout=dropout, + ) + if arch == "hybrid_full_around_one": + return HybridGraphTransformerBlock( + hidden_dim, + n_heads, + ffn_dim, + group_size=group_size, + adj_outer_init="full", + adj_inner_init="uniform_around_one", + dropout=dropout, + ) + if arch == "hybrid_around_one_around_one": + return HybridGraphTransformerBlock( + hidden_dim, + n_heads, + ffn_dim, + group_size=group_size, + adj_outer_init="uniform_around_one", + adj_inner_init="uniform_around_one", + dropout=dropout, + ) + raise ValueError(f"unknown arch: {arch}") diff --git a/src/graphlm/neuron/hybrid_transformer_demo.py b/src/graphlm/neuron/hybrid_transformer_demo.py new file mode 100644 index 0000000..4d3343f --- /dev/null +++ b/src/graphlm/neuron/hybrid_transformer_demo.py @@ -0,0 +1,199 @@ +"""Phase 13 — HybridGraphTransformerLM demo + train helper. + +노트북 분리 규약 준수. Phase 13 노트북 12-phase13-hybrid-transformer.ipynb 에서 import. + +4 arch 비교 (모두 0-init 금지 + magnitude rule 적용): +- ``"plain"`` — RMSNorm + 표준 nn.Linear FFN (baseline) +- ``"hybrid_full_full"`` — outer=full + inner=full (function preserving) +- ``"hybrid_full_around_one"`` — outer=full + inner=uniform_around_one (Phase 11 channel-level) +- ``"hybrid_around_one_around_one"`` — 둘 다 uniform_around_one (fully scale-corrected) + +identity outer 는 Phase 13 에서 미지원 — FFN 의 rectangular 구조상 의미 없음. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass + +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.hybrid_transformer import ( + Arch, + HybridGraphTransformerBlock, + make_block, +) +from graphlm.neuron.rms_norm import RMSNorm +from graphlm.utils import set_seed + + +@dataclass(frozen=True) +class HybridTransformerTrainConfig: + """1 run 학습에 필요한 모든 hyperparameter (CodeRabbit #3303186824 — 14 args 통합). + + 구조: model / data / optim / runtime 4 그룹으로 묶어서 가독성 ↑. + """ + + # data + dataset: TinyShakespeareDataset + vocab_size: int + + # model + hidden_dim: int + n_heads: int + ffn_dim: int + n_layers: int + group_size: int + arch: Arch + dropout: float = 0.0 + + # train + block_size: int = 64 + batch_size: int = 32 + lr: float = 3e-4 + max_steps: int = 1500 + + # runtime + seed: int = 0 + device: str = "cpu" + + +class HybridGraphTransformerLM(nn.Module): + """Small char-LM Transformer with arch-dispatched FFN. + + - token embedding + learned positional embedding + - N × pre-norm block (RMSNorm + attn + RMSNorm + FFN) + - final RMSNorm + LM head (weight not tied for 공정 비교) + """ + + def __init__( + self, + vocab_size: int, + hidden_dim: int, + n_heads: int, + ffn_dim: int, + n_layers: int, + max_seq_len: int, + arch: Arch, + group_size: int, + dropout: float = 0.0, + ): + super().__init__() + self.arch = arch + self.max_seq_len = max_seq_len + self.tok_emb = nn.Embedding(vocab_size, hidden_dim) + self.pos_emb = nn.Embedding(max_seq_len, hidden_dim) + self.blocks = nn.ModuleList( + [ + make_block( + arch, + hidden_dim=hidden_dim, + n_heads=n_heads, + ffn_dim=ffn_dim, + group_size=group_size, + dropout=dropout, + ) + for _ in range(n_layers) + ] + ) + self.final_norm = RMSNorm(hidden_dim) + self.lm_head = nn.Linear(hidden_dim, vocab_size, bias=False) + + def forward(self, x: Tensor) -> Tensor: + _batch, seq_len = x.shape + if seq_len > self.max_seq_len: + raise ValueError(f"seq_len {seq_len} > max_seq_len {self.max_seq_len}") + pos = torch.arange(seq_len, device=x.device) + h = self.tok_emb(x) + self.pos_emb(pos) + for block in self.blocks: + h = block(h) + h = self.final_norm(h) + return self.lm_head(h) + + +def _block_iter(model: HybridGraphTransformerLM) -> Iterator[HybridGraphTransformerBlock]: + """모델의 hybrid block 만 yield (plain 은 skip).""" + for blk in model.blocks: + if isinstance(blk, HybridGraphTransformerBlock): + yield blk + + +def _snapshot_adj(model: HybridGraphTransformerLM) -> list[dict[str, dict[str, Tensor]]] | None: + """hybrid arch 인 경우 각 block 의 FFN adj snapshot (Phase 12 demo 와 동일 hierarchy).""" + if model.arch == "plain": + return None + snapshots: list[dict[str, dict[str, Tensor]]] = [] + for blk in _block_iter(model): + snapshots.append( + { + "fc1": { + "outer": blk.ffn.fc1.adj_outer.detach().cpu().clone(), + "inner": blk.ffn.fc1.adj_inner.detach().cpu().clone(), + }, + "fc2": { + "outer": blk.ffn.fc2.adj_outer.detach().cpu().clone(), + "inner": blk.ffn.fc2.adj_inner.detach().cpu().clone(), + }, + } + ) + return snapshots + + +def train_hybrid_transformer_lm(config: HybridTransformerTrainConfig) -> dict: + """1 run 학습 — Phase 13 sweep unit. + + Returns: ``losses``, ``final_loss`` (last 100 mean), ``final_adj`` + (hybrid arch 인 경우 block 별 fc1/fc2 outer/inner snapshot list). + """ + set_seed(config.seed) + model = HybridGraphTransformerLM( + vocab_size=config.vocab_size, + hidden_dim=config.hidden_dim, + n_heads=config.n_heads, + ffn_dim=config.ffn_dim, + n_layers=config.n_layers, + max_seq_len=config.block_size, + arch=config.arch, + group_size=config.group_size, + dropout=config.dropout, + ).to(config.device) + data_iter = iter_random_batches( + config.dataset, batch_size=config.batch_size, block_size=config.block_size, seed=config.seed + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) + losses: list[float] = [] + model.train() + for _step in range(1, config.max_steps + 1): + x, y = next(data_iter) + x, y = x.to(config.device), y.to(config.device) + optimizer.zero_grad() + logits = model(x) + loss = F.cross_entropy(logits.reshape(-1, config.vocab_size), y.reshape(-1)) + 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 + + return { + "losses": losses, + "final_loss": final_loss, + "final_adj": _snapshot_adj(model), + } + + +def count_parameters(model: nn.Module) -> int: + """전체 학습 가능 파라미터 수 (arch 간 공정 비교용 reporting).""" + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +__all__ = [ + "HybridGraphTransformerLM", + "HybridTransformerTrainConfig", + "count_parameters", + "train_hybrid_transformer_lm", +] diff --git a/src/graphlm/neuron/rms_norm.py b/src/graphlm/neuron/rms_norm.py new file mode 100644 index 0000000..a6d4184 --- /dev/null +++ b/src/graphlm/neuron/rms_norm.py @@ -0,0 +1,61 @@ +"""Phase 13 — RMSNorm (root-mean-square layer normalization). + +Modern Transformer (LLaMA / Mistral / Gemma) 의 표준 norm. nn.LayerNorm 대비: + +- mean centering 생략 → 연산량 ↓ +- bias 파라미터 없음 → 파라미터 수 ↓ (hidden_dim 만큼 절약) +- numerical 안정성 동등 또는 우위 (특히 large hidden_dim) + +수식: ``y = x / RMS(x) * weight``, ``RMS(x) = sqrt(mean(x^2) + eps)`` + +본 paradigm 의 Phase 13 backbone (HybridGraphTransformer) 에서 nn.LayerNorm 대체. +backbone.py 의 기존 LayerNorm 은 Phase 1~12 호환성 위해 그대로 유지. +""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn + + +class RMSNorm(nn.Module): + """Root-mean-square layer normalization. + + Args: + hidden_dim: 정규화 대상 마지막 축의 크기. + eps: RMS 분모 numerical 안정성 (LLaMA 의 1e-6 동일). + + Shape: + - input: ``(..., hidden_dim)`` + - output: same as input + """ + + def __init__(self, hidden_dim: int, eps: float = 1e-6): + super().__init__() + if not isinstance(hidden_dim, int) or hidden_dim < 1: + raise ValueError(f"hidden_dim must be a positive int, got {hidden_dim!r}") + if eps <= 0: + raise ValueError(f"eps must be > 0, got {eps}") + self.hidden_dim = hidden_dim + self.eps = eps + # weight 만 학습 (LayerNorm 의 bias 없음) — function preservation 위해 1.0 으로 시작 + self.weight = nn.Parameter(torch.ones(hidden_dim)) + + def forward(self, x: Tensor) -> Tensor: + if x.shape[-1] != self.hidden_dim: + raise ValueError(f"expected last dim {self.hidden_dim}, got {x.shape[-1]}") + # int 등 비-floating 입력은 silent cast 위험 — nn.LayerNorm 과 동일 정책으로 차단 + # (Copilot #3303168589) + if not x.is_floating_point(): + raise TypeError(f"RMSNorm requires floating-point input, got dtype={x.dtype}") + # float32 로 cast 해서 RMS 계산 (mixed precision 안전성) + x_dtype = x.dtype + x_f = x.float() + rms = torch.rsqrt(x_f.pow(2).mean(dim=-1, keepdim=True) + self.eps) + y = (x_f * rms).to(x_dtype) + # weight 도 x_dtype 로 cast — 그래야 mixed precision (FP16/BF16) 에서 residual + # connection dtype mismatch 회피 (gemini #3303153077) + return y * self.weight.to(x_dtype) + + def extra_repr(self) -> str: + return f"hidden_dim={self.hidden_dim}, eps={self.eps}" diff --git a/src/graphlm/utils/__init__.py b/src/graphlm/utils/__init__.py index 4e93c16..e442a19 100644 --- a/src/graphlm/utils/__init__.py +++ b/src/graphlm/utils/__init__.py @@ -1,7 +1,14 @@ """Utility functions for GraphLM (seed, exceptions, paths, logging helpers, etc.).""" from graphlm.utils.exceptions import FunctionPreservationError, GraphLMError +from graphlm.utils.metrics import safe_perplexity from graphlm.utils.paths import repo_root from graphlm.utils.seed import set_seed -__all__ = ["FunctionPreservationError", "GraphLMError", "repo_root", "set_seed"] +__all__ = [ + "FunctionPreservationError", + "GraphLMError", + "repo_root", + "safe_perplexity", + "set_seed", +] diff --git a/src/graphlm/utils/metrics.py b/src/graphlm/utils/metrics.py new file mode 100644 index 0000000..fd36e86 --- /dev/null +++ b/src/graphlm/utils/metrics.py @@ -0,0 +1,36 @@ +"""Numerical metrics helpers — safe perplexity 등. + +노트북에서 직접 정의하던 helper 들을 ``src/graphlm/`` 로 이전 — 노트북은 analysis flow ++ 시각화에만 집중 (CodeRabbit #3304306120, project rule). +""" + +from __future__ import annotations + +import math + + +def safe_perplexity(loss: float, cap: float = 20.0) -> float: + """``exp(loss)`` overflow 방지. + + 학습 초반 큰 loss 또는 발산 시 ``math.exp(loss)`` 는 ``OverflowError`` 발생. + cap=20 → max perplexity ≈ 4.85e8 (충분히 큰 ceiling, 발산 식별 가능). + + Args: + loss: cross-entropy loss (음수 또는 양수, 무한 가능). + cap: exp 적용 전 loss 의 상한 (default 20.0). + + Returns: + ``math.exp(min(loss, cap))`` — 항상 유한 양수. + + Raises: + ValueError: ``cap`` 이 음수면 (의도적 down-clip 으로 underflow 가능). + + See Also: + - gemini #3303153101 rationale: log-domain loss → exp-domain perplexity 변환 시 + overflow 회피. + """ + if cap < 0: + raise ValueError(f"cap must be non-negative, got {cap}") + if math.isnan(loss): + return math.nan + return math.exp(min(loss, cap)) diff --git a/tests/neuron/test_hybrid_transformer.py b/tests/neuron/test_hybrid_transformer.py new file mode 100644 index 0000000..f366a52 --- /dev/null +++ b/tests/neuron/test_hybrid_transformer.py @@ -0,0 +1,209 @@ +"""Tests for graphlm.neuron.hybrid_transformer — Phase 13 Transformer integration.""" + +from __future__ import annotations + +import pytest +import torch +from torch import nn + +from graphlm.neuron.hybrid_transformer import ( + HybridGraphFFN, + HybridGraphTransformerBlock, + PlainFFN, + PlainTransformerBlock, + make_block, +) + +# ── HybridGraphFFN ─────────────────────────────────────────── + + +def test_ffn_shape(): + ffn = HybridGraphFFN(hidden_dim=32, ffn_dim=64, group_size=8) + x = torch.randn(2, 16, 32) + assert ffn(x).shape == (2, 16, 32) + + +def test_ffn_identity_outer_rejected(): + """FFN 은 rectangular 라 identity outer 미지원.""" + with pytest.raises(ValueError, match="rectangular"): + HybridGraphFFN(hidden_dim=32, ffn_dim=64, group_size=8, adj_outer_init="identity") + + +def test_ffn_function_preservation_full_full(): + """adj_outer=full + adj_inner=full + 같은 weight → standard PlainFFN 동일.""" + torch.manual_seed(0) + hidden, ffn_d, k = 16, 32, 4 + hg_ffn = HybridGraphFFN(hidden, ffn_d, group_size=k) + plain_ffn = PlainFFN(hidden, ffn_d) + + # hg_ffn 의 block weight 를 standard nn.Linear weight 로 변환해서 plain 에 복사 + _copy_hybrid_to_plain(hg_ffn.fc1, plain_ffn.fc1) + _copy_hybrid_to_plain(hg_ffn.fc2, plain_ffn.fc2) + + x = torch.randn(2, 8, hidden) + y_hg = hg_ffn(x) + y_plain = plain_ffn(x) + assert torch.allclose(y_hg, y_plain, atol=1e-5), ( + f"function preservation 깨짐: max |diff| = {(y_hg - y_plain).abs().max().item()}" + ) + + +def test_ffn_all_params_have_gradient(): + ffn = HybridGraphFFN( + 16, + 32, + group_size=4, + adj_outer_init="uniform_around_one", + adj_inner_init="uniform_around_one", + ) + x = torch.randn(2, 16) + ffn(x).sum().backward() + for layer in [ffn.fc1, ffn.fc2]: + for attr in ["weight", "adj_outer", "adj_inner"]: + grad = getattr(layer, attr).grad + assert grad is not None, f"{layer}.{attr}.grad is None" + assert (grad.abs().sum() > 0).item(), f"{layer}.{attr}.grad all zero" + + +# ── HybridGraphTransformerBlock ────────────────────────────── + + +def test_block_shape(): + block = HybridGraphTransformerBlock(hidden_dim=32, n_heads=4, ffn_dim=64, group_size=8) + x = torch.randn(2, 16, 32) + assert block(x).shape == (2, 16, 32) + + +def test_block_function_preservation_against_plain(): + """hybrid block (adj=full/full) + plain block 의 동일 weight 로 forward 동일.""" + torch.manual_seed(0) + hidden, n_heads, ffn_d, k = 16, 4, 32, 4 + hybrid = HybridGraphTransformerBlock(hidden, n_heads, ffn_d, group_size=k) + plain = PlainTransformerBlock(hidden, n_heads, ffn_d) + + # rms / attn 은 standard module 이므로 state_dict copy 가능 + plain.rms1.load_state_dict(hybrid.rms1.state_dict()) + plain.rms2.load_state_dict(hybrid.rms2.state_dict()) + plain.attn.load_state_dict(hybrid.attn.state_dict()) + # FFN 만 block → standard 변환 + _copy_hybrid_to_plain(hybrid.ffn.fc1, plain.ffn.fc1) + _copy_hybrid_to_plain(hybrid.ffn.fc2, plain.ffn.fc2) + + x = torch.randn(2, 8, hidden) + y_hybrid = hybrid(x) + y_plain = plain(x) + assert torch.allclose(y_hybrid, y_plain, atol=1e-5), ( + f"block forward 차이: max |diff| = {(y_hybrid - y_plain).abs().max().item()}" + ) + + +def test_block_gradient_flows_all_params(): + block = HybridGraphTransformerBlock( + hidden_dim=16, + n_heads=4, + ffn_dim=32, + group_size=4, + adj_outer_init="uniform_around_one", + adj_inner_init="uniform_around_one", + ) + x = torch.randn(2, 4, 16) + block(x).sum().backward() + null_grad = [n for n, p in block.named_parameters() if p.grad is None] + assert not null_grad, f"grad 없는 파라미터: {null_grad}" + + +# ── make_block dispatch ────────────────────────────────────── + + +@pytest.mark.parametrize( + "arch", + ["plain", "hybrid_full_full", "hybrid_full_around_one", "hybrid_around_one_around_one"], +) +def test_make_block_all_archs_forward(arch): + block = make_block(arch, hidden_dim=16, n_heads=4, ffn_dim=32, group_size=4) + x = torch.randn(2, 8, 16) + assert block(x).shape == (2, 8, 16) + + +def test_make_block_unknown_raises(): + with pytest.raises(ValueError, match="unknown arch"): + make_block("bogus", 16, 4, 32, 4) # type: ignore[arg-type] + + +def test_make_block_plain_uses_nn_linear(): + block = make_block("plain", 16, 4, 32, 4) + assert isinstance(block, PlainTransformerBlock) + assert isinstance(block.ffn.fc1, nn.Linear) + + +def test_make_block_hybrid_uses_hybrid_ffn(): + block = make_block("hybrid_full_full", 16, 4, 32, 4) + assert isinstance(block, HybridGraphTransformerBlock) + assert isinstance(block.ffn, HybridGraphFFN) + + +# ── dropout consistency (Copilot #3303168649 / #3303168686) ── + + +def test_hybrid_ffn_applies_dropout_when_training(): + """HybridGraphFFN 도 fc2 뒤 dropout — backbone.FFN 패턴 일관.""" + ffn = HybridGraphFFN(16, 32, group_size=4, dropout=0.5) + ffn.train() + torch.manual_seed(0) + x = torch.randn(64, 16) + # dropout 활성 시 stochastic — eval 모드와 다른 output 보장 + y_train = ffn(x) + ffn.eval() + y_eval = ffn(x) + assert not torch.allclose(y_train, y_eval), "dropout 이 train 모드에서 적용되지 않음" + + +def test_plain_ffn_applies_dropout_when_training(): + """PlainFFN 도 dropout 인자 적용 — Hybrid 와 API 일관.""" + ffn = PlainFFN(16, 32, dropout=0.5) + ffn.train() + torch.manual_seed(0) + x = torch.randn(64, 16) + y_train = ffn(x) + ffn.eval() + y_eval = ffn(x) + assert not torch.allclose(y_train, y_eval), "dropout 이 train 모드에서 적용되지 않음" + + +@pytest.mark.parametrize( + "block_cls,kwargs", + [ + ( + PlainTransformerBlock, + {"hidden_dim": 16, "n_heads": 4, "ffn_dim": 32, "dropout": 0.5}, + ), + ( + HybridGraphTransformerBlock, + {"hidden_dim": 16, "n_heads": 4, "ffn_dim": 32, "group_size": 4, "dropout": 0.5}, + ), + ], +) +def test_block_dropout_propagates_to_ffn(block_cls, kwargs): + """block 의 dropout 인자가 attention 만이 아니라 FFN 까지 전달.""" + block = block_cls(**kwargs) + assert block.ffn.dropout.p == 0.5, "FFN 의 dropout p 가 block dropout 과 불일치" + + +# ── helpers ────────────────────────────────────────────────── + + +def _copy_hybrid_to_plain(hg, plain): + """HybridGraphLinear 의 block weight → nn.Linear 표준 weight 형식으로 복사. + + Phase 12 test_function_preservation_full_full_equivalent_to_linear 와 동일 로직. + """ + in_f, out_f = hg.in_features, hg.out_features + k = hg.group_size + W_std = torch.zeros(out_f, in_f) + for go in range(hg.n_groups_out): + for gi in range(hg.n_groups_in): + W_std[go * k : (go + 1) * k, gi * k : (gi + 1) * k] = hg.weight[go, gi].T + with torch.no_grad(): + plain.weight.copy_(W_std) + if plain.bias is not None and hg.bias is not None: + plain.bias.copy_(hg.bias) diff --git a/tests/neuron/test_rms_norm.py b/tests/neuron/test_rms_norm.py new file mode 100644 index 0000000..cf8c548 --- /dev/null +++ b/tests/neuron/test_rms_norm.py @@ -0,0 +1,82 @@ +"""Tests for graphlm.neuron.rms_norm — Phase 13 RMSNorm.""" + +from __future__ import annotations + +import pytest +import torch + +from graphlm.neuron.rms_norm import RMSNorm + + +def test_shape_preserved(): + norm = RMSNorm(16) + x = torch.randn(2, 8, 16) + assert norm(x).shape == x.shape + + +def test_init_weight_is_ones(): + norm = RMSNorm(16) + assert torch.allclose(norm.weight, torch.ones(16)) + + +def test_rms_normalizes_to_unit_rms(): + """초기 weight=1 에서 출력의 RMS 가 1 에 매우 가까워야 함.""" + norm = RMSNorm(64) + x = torch.randn(4, 16, 64) * 5.0 # arbitrary scale + y = norm(x) + rms = y.pow(2).mean(dim=-1).sqrt() + # eps 때문에 정확히 1 은 아니지만 매우 근사 + assert torch.allclose(rms, torch.ones_like(rms), atol=1e-3) + + +def test_weight_is_learnable(): + norm = RMSNorm(16) + x = torch.randn(2, 16) + out = norm(x) + out.sum().backward() + assert norm.weight.grad is not None + assert (norm.weight.grad.abs().sum() > 0).item() + + +def test_zero_dim_raises(): + with pytest.raises(ValueError, match="positive int"): + RMSNorm(0) + + +def test_negative_eps_raises(): + with pytest.raises(ValueError, match="eps must be > 0"): + RMSNorm(16, eps=-1e-6) + + +def test_wrong_last_dim_raises(): + norm = RMSNorm(16) + with pytest.raises(ValueError, match="expected last dim 16"): + norm(torch.randn(2, 8)) + + +def test_scaling_via_weight(): + """weight=2.0 로 setting → 출력도 2배 scale.""" + norm = RMSNorm(16) + with torch.no_grad(): + norm.weight.fill_(2.0) + x = torch.randn(4, 16) + y = norm(x) + rms = y.pow(2).mean(dim=-1).sqrt() + assert torch.allclose(rms, torch.full_like(rms, 2.0), atol=1e-3) + + +def test_non_floating_input_raises(): + """int 등 비-floating 입력은 silent cast 회피 위해 차단 (Copilot #3303168589).""" + norm = RMSNorm(16) + x_int = torch.zeros(4, 16, dtype=torch.long) + with pytest.raises(TypeError, match="floating-point"): + norm(x_int) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_dtype_roundtrip_mixed_precision(dtype): + """출력 dtype 이 입력 dtype 과 일치 (gemini #3303153077). residual connection 안전.""" + norm = RMSNorm(32) + x = torch.randn(4, 32, dtype=dtype) + y = norm(x) + assert y.dtype == dtype, f"expected {dtype}, got {y.dtype}" diff --git a/tests/utils/test_metrics.py b/tests/utils/test_metrics.py new file mode 100644 index 0000000..f8fb59d --- /dev/null +++ b/tests/utils/test_metrics.py @@ -0,0 +1,50 @@ +"""Tests for graphlm.utils.metrics — safe_perplexity etc.""" + +from __future__ import annotations + +import math + +import pytest + +from graphlm.utils import safe_perplexity + + +def test_normal_loss(): + """일반적인 char-LM loss 값 (~2.0) 이 정상 perplexity 반환.""" + assert safe_perplexity(2.0) == pytest.approx(math.exp(2.0)) + + +def test_caps_at_default_20(): + """loss=30 → cap=20 적용 → exp(20).""" + assert safe_perplexity(30.0) == pytest.approx(math.exp(20.0)) + + +def test_below_cap_unchanged(): + """loss < cap 일 때 cap 영향 없음.""" + assert safe_perplexity(5.0, cap=20.0) == pytest.approx(math.exp(5.0)) + + +def test_negative_loss(): + """음수 loss 도 정상 처리 (분류 confidence 높을 때).""" + assert safe_perplexity(-1.0) == pytest.approx(math.exp(-1.0)) + + +def test_custom_cap(): + """cap 인자로 ceiling 조정 가능.""" + assert safe_perplexity(15.0, cap=10.0) == pytest.approx(math.exp(10.0)) + + +def test_nan_input_returns_nan(): + """NaN loss → NaN perplexity (silent overflow 회피).""" + assert math.isnan(safe_perplexity(math.nan)) + + +def test_inf_input_caps(): + """+inf loss → cap 적용 → 유한 값.""" + assert safe_perplexity(math.inf) == pytest.approx(math.exp(20.0)) + + +def test_neg_cap_rejected(): + """음수 cap 거부 — exp(음수 large) underflow 위험.""" + with pytest.raises(ValueError, match="non-negative"): + safe_perplexity(2.0, cap=-1.0)