[FEAT#55] neuron Phase 8 — structural axis foundations (GrowableLinear/LayerNorm/Embedding + AdamW state 확장)

[juhy0987]

연관 이슈

Closes #55

배경

Phase 4~7 의 모든 gating function 표현력 확장 (per_channel → positional → sinusoidal smooth-start) 이 char-LM 에서 동일 final_loss (~1.78) 로 수렴 → gating axis 의 한계 명확. paradigm 정의 (training-time dynamic parameter count) 에 가장 가까운 axis = structural (nn.Parameter 정적 shape 초월).

구현 내용

1. src/graphlm/neuron/growable.py (신규 모듈)

PyTorch 기반 동적 expansion foundations. TensorFlow / JAX 의존성 없음.

• GrowableLinear — nn.Linear 호환 + expand_out(delta) / expand_in(delta)
  • init="zero" → function preservation (새 dim 의 forward 기여 = 0)
  • init="normal" → small random
  • optimizer=adamw 인자 시 AdamW state (m, v) 동기 확장
• GrowableLayerNorm — expand(delta), 새 dim 의 weight=1/bias=0 (identity-style)
• GrowableEmbedding — expand_dim(delta) (vocab 고정, embedding_dim 만)
• _replace_param_with_optimizer_state helper — Parameter 교체 + AdamW state 확장
  • param_groups 자동 치환 (기존 dim m/v 보존, 새 dim zero, step counter 보존)

bert2BERT (Chen et al. ACL 2022), LiGO (Wang et al. ICLR 2023), MSG (Yuan et al. NeurIPS 2023) 의 expansion 메커니즘과 동일 패러다임.

2. tests/neuron/test_growable.py — 9 신규 테스트

• shape after expand (out, in, dim)
• function preservation — zero-init expand 직후 forward 불변 (atol=1e-6)
• AdamW state 보존 — 기존 m/v 동일, 새 dim = 0, step counter 보존
• 반복 expansion 안정성 (3회 expansion 후 정상 학습)
• LayerNorm identity-preserving expand
• invalid init mode → ValueError

3. notebooks/11-function-level/18-phase8-growable-foundations.ipynb

Growable MLP-LM demo (Transformer 아닌 단순 MLP — foundation 검증 우선).

• 학습 1500 step, 매 500/1000 step 에 hidden_dim 128 → 192 → 256 expansion
• 2 × 4 sweep: seed × (state_preserve, init_mode) = 8 run
  • state_preserve=True/False (AdamW state 보존 vs reset)
  • init_mode="zero"/"normal" (function preservation vs random)
• §6 expansion spike 측정 + final_loss config 별 verdict
• §7 function preservation sanity check (학습 안 한 model 로 zero/normal init 직접 비교)
• §8 시각화: config 별 loss curve (mean ± σ) + hidden width 진화 그래프

검증할 가설

1. AdamW state 보존 시 expansion 직후 loss spike 최소화?
2. zero-init function preservation 의 실질적 효과?
3. 반복 expansion 안정성?
4. state reset vs 보존의 final_loss 차이?

CI / 머지 게이트 점검

• make fmt 통과
• make lint 통과
• make test 통과 — 64 passed (기존 55 + 신규 9)

변경 영향 범위 + 위험도

• 신규 모듈 growable.py — 기존 backbone.py / growth.py / 모든 Phase 1~7 코드 완전 격리
• 신규 테스트 + 노트북 — 기존 영향 없음
• 위험도: 매우 낮음 (additive, isolated)

롤백 계획

• growable.py + test_growable.py + 노트북 단일 commit 으로 격리되어 git revert 한 번에 복구

Phase 9+ 계획

Phase	내용
9	본 foundation 을 활용한 GrowableTransformer block 구현 (NeuronGrowingDecoder 와 병렬 architecture)
10	attention head 수 동적 증가 (head_dim 고정, n_heads 증가 → hidden_dim 자동 확장)
11	layer 자체 동적 추가 (depth growth, 기존 graphlm.growth.net2deeper 와 통합 검토)

PyTorch 기반 유지 명시

본 PR 및 향후 paradigm 진행 모두 torch.nn / torch.optim 만 사용. TensorFlow / JAX 도입하지 않음.

Summary by CodeRabbit

• New Features

  • Added growable neural components that can widen at runtime with function-preserving init options and synchronized optimizer state updates.
  • Added a growable MLP demo, training utilities, and an experimental Phase 8 notebook that runs scheduled hidden-dimension expansions and visualizes loss/width trajectories.

• Tests

  • Expanded test suite covering shape updates, initialization behavior, optimizer-state preservation, gradient handling, repeated expansions, and invalid-input checks.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ae580bcc-a677-4eb8-9160-ca81b0c393f4

📥 Commits

Reviewing files that changed from the base of the PR and between 7d06f7f and 9941f51.

📒 Files selected for processing (4)

• notebooks/11-function-level/18-phase8-growable-foundations.ipynb
• src/graphlm/neuron/growable.py
• src/graphlm/neuron/growable_demo.py
• tests/neuron/test_growable.py

---

📝 Walkthrough

Walkthrough

Adds growable neural components (Linear, LayerNorm, Embedding) with function-preserving expansion and optimizer-state-aware parameter replacement, a Growable MLP demo and training loop, notebook orchestration for ablation/visualization, public exports, and comprehensive unit tests.

Changes

Phase 8: Growable Neural Module Foundations

Layer / File(s)	Summary
Type, validation, tensor & init helpers src/graphlm/neuron/growable.py	Module docstring and helpers: InitMode type, _expand_tensor_with_zeros, _init_new_block, and _validate_delta.
Parameter replacement & optimizer state migration src/graphlm/neuron/growable.py	_replace_param_with_optimizer_state that swaps a Parameter for a larger one, expands .grad if present, migrates/expands optimizer.state entries (tensor-valued moments) and updates param_groups.
GrowableLinear implementation src/graphlm/neuron/growable.py	GrowableLinear with forward, expand_out, and expand_in that grow weight/bias, initialize new blocks ("zero"/"normal"), migrate optimizer state, and update feature sizes.
GrowableLayerNorm & GrowableEmbedding src/graphlm/neuron/growable.py	GrowableLayerNorm.expand uses identity-style tails (weight=1, bias=0); GrowableEmbedding.expand_dim appends embedding dims; both migrate optimizer state and update shapes.
Public module exports src/graphlm/neuron/__init__.py	Adds GrowableLinear, GrowableLayerNorm, and GrowableEmbedding to __all__.
Demo model & training src/graphlm/neuron/growable_demo.py	Adds GrowableMLPLM, expand_hidden, make_ngram_iter, and train_growable_mlp that run scheduled expansions, record losses/widths, and optionally preserve optimizer state.
Unit tests tests/neuron/test_growable.py	Pytest suite covering shape checks, zero-init function preservation, AdamW state migration, gradient preservation during expansion, repeated expansion stability, and invalid-argument errors for all growable types.
Notebook orchestration notebooks/11-function-level/18-phase8-growable-foundations.ipynb	Experiment sweep over (state_preserve × init_mode) on TinyShakespeare, runs training, reports loss-spike deltas and final-loss statistics, performs untrained function-preservation checks, and plots smoothed loss curves + hidden-width evolution.

Sequence Diagram(s)

sequenceDiagram
  participant ComponentA
  participant ComponentB
  ComponentA->>ComponentB: observable interaction

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~65 minutes

Possibly related PRs

• EinSofINTEREST/GraphLM#56: Contains the same growable module implementation, tests, and Phase 8 notebook experiment that align closely with these changes.

Poem

🐰 I nibbled code and planted rows,

New dims unfurl where curiosity grows,
AdamW keeps its careful pace,
Old outputs steady in their place,
A tiny rabbit cheers the growable space 🌱

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: introducing Phase 8 foundations with growable neural components (GrowableLinear, GrowableLayerNorm, GrowableEmbedding) and AdamW optimizer state expansion.
Linked Issues check	✅ Passed	All primary coding objectives from issue #55 are met: GrowableLinear/LayerNorm/Embedding implemented with shape expansion, AdamW state preservation helper included, zero/normal init modes supported, comprehensive tests validate function preservation and state semantics, and demo notebook demonstrates scheduled expansions with ablations.
Out of Scope Changes check	✅ Passed	All changes directly support the Phase 8 foundations objective: growable module implementations, optimizer state expansion, comprehensive tests, and demo notebook with scheduled expansions. No out-of-scope modifications detected.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feature/#55/neuron-phase8-growable-foundations

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces the foundations for "Phase 8" structural axis changes, implementing growable neural network modules (GrowableLinear, GrowableLayerNorm, and GrowableEmbedding) that allow dynamic parameter expansion during training while preserving function outputs and AdamW optimizer states. It also includes a verification notebook and comprehensive unit tests. The review feedback highlights a critical issue where parameter gradients (.grad) are not preserved during parameter replacement, which could lead to lost gradients if expansion occurs between backward() and step(). A code suggestion is provided to safely expand and copy the gradient tensor.

[Copilot]

Pull request overview

본 PR은 neuron Phase 8의 structural axis foundation으로, 학습 중 nn.Parameter의 정적 shape 제약을 우회하기 위해 모듈 파라미터 교체 + AdamW optimizer state(m/v) 동기 확장을 지원하는 growable 레이어들을 추가합니다. 이를 통해 hidden width 등 구조적 차원의 training-time expansion 실험을 위한 기반을 graphlm.neuron 패키지에 도입합니다.

Changes:

• GrowableLinear/GrowableLayerNorm/GrowableEmbedding 및 optimizer state 확장 helper 신규 구현
• shape/function-preserving/AdamW state 보존/반복 확장 안정성에 대한 pytest 9건 추가
• Phase 8 데모 노트북(MLP-LM sweep + 시각화) 추가 및 graphlm.neuron 공개 API로 export

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

File	Description
src/graphlm/neuron/growable.py	growable 레이어 + AdamW state 확장 기반 구현
tests/neuron/test_growable.py	expand 동작/함수 보존/optimizer state 보존 등 단위 테스트 추가
src/graphlm/neuron/init.py	growable 레이어를 패키지 공개 API로 export
notebooks/11-function-level/18-phase8-growable-foundations.ipynb	Phase 8 foundation 검증용 실험/시각화 노트북 추가

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

tests/neuron/test_growable.py (1)

131-147: ⚡ Quick win

Rename this test or assert the weaker contract explicitly.

The body only verifies tail initialization and output shape. Since expanding LayerNorm changes the statistics seen by the original channels, identity_preserving_expand reads as a stronger guarantee than this test actually covers.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/neuron/test_growable.py` around lines 131 - 147, The test named
test_growable_layernorm_identity_preserving_expand claims a stronger
identity-preserving contract than it verifies; update it to reflect the actual
weaker check by either renaming the test to something like
test_growable_layernorm_tail_init_and_shape or by changing the assertions to
explicitly state the weaker contract (e.g., assert only that ln.expand(n)
produces correct ln.weight/ln.bias initialization for new dims and that output
shape matches), referencing the GrowableLayerNorm instance ln and its expand
method ln.expand to locate the code to update.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@notebooks/11-function-level/18-phase8-growable-foundations.ipynb`:
- Around line 179-192: The docstring for expand_hidden and the §7 notebook text
overstate “function preservation”: update the expand_hidden docstring (method
expand_hidden) to explicitly state that with init="zero" logits remain
numerically close for the tested inputs (not exactly identical) because
self.ln.expand (GrowableLayerNorm) renormalizes channels; likewise edit the §7
explanatory text to say "`zero`-init keeps forward/logits numerically close
(perturbs less than normal-init) for the tested inputs" and reference the
measured max_diff_zero vs max_diff_normal; if you need a true internal-channel
identity check, note the alternative of skipping ln.expand (or removing
LayerNorm) in the test instead of claiming exact preservation.

In `@src/graphlm/neuron/growable.py`:
- Around line 124-131: Add an explicit guard that validates delta > 0 at the
start of public expansion entry points to fail fast before touching parameters
or optimizer state: in the expand_out method (and the other public expand_*
methods in this module), check if delta <= 0 and raise a ValueError with a clear
message (e.g., "delta must be > 0"), returning/raising immediately so no
parameter swaps or optimizer changes occur; place the check as the first
statements in the functions (before any parameter/optimizer handling or shape
manipulation).
- Around line 203-217: GrowableLayerNorm.expand currently increases
normalized_shape which changes LayerNorm's mean/variance and breaks function
preservation for existing channels; instead, keep the LayerNorm's operative
normalization size unchanged and treat the new tail as inactive until explicitly
activated: modify GrowableLayerNorm to track an active_normalized_shape (or
active_channels) separate from total capacity, leave self.normalized_shape (or
the value used in torch.nn.functional.layer_norm) set to the original size,
append new weight/bias tails via _replace_param_with_optimizer_state (as done)
but do NOT increase the normalization size in expand; add logic in the forward
method (or wherever normalization stats are computed) to include tail channels
only when active_normalized_shape == total capacity (or when activated), and
expose an activate(delta) or set_active(new_active) method to increase
active_normalized_shape and then update normalization size to include those
channels when they are turned on.

---

Nitpick comments:
In `@tests/neuron/test_growable.py`:
- Around line 131-147: The test named
test_growable_layernorm_identity_preserving_expand claims a stronger
identity-preserving contract than it verifies; update it to reflect the actual
weaker check by either renaming the test to something like
test_growable_layernorm_tail_init_and_shape or by changing the assertions to
explicitly state the weaker contract (e.g., assert only that ln.expand(n)
produces correct ln.weight/ln.bias initialization for new dims and that output
shape matches), referencing the GrowableLayerNorm instance ln and its expand
method ln.expand to locate the code to update.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 82028003-9d4c-4114-b67e-1017b36e804f

📥 Commits

Reviewing files that changed from the base of the PR and between 797ec66 and 7d06f7f.

📒 Files selected for processing (4)

• notebooks/11-function-level/18-phase8-growable-foundations.ipynb
• src/graphlm/neuron/__init__.py
• src/graphlm/neuron/growable.py
• tests/neuron/test_growable.py