Skip to content

Add one-head nanoGPT Attention Residuals baseline - #60

Draft
charlesmartin14 wants to merge 10 commits into
mainfrom
agent/nanogpt-attnres-baseline
Draft

Add one-head nanoGPT Attention Residuals baseline#60
charlesmartin14 wants to merge 10 commits into
mainfrom
agent/nanogpt-attnres-baseline

Conversation

@charlesmartin14

Copy link
Copy Markdown
Member

What changed

  • Adds Full Attention Residuals (residual_mode: full_attnres) to the existing one-head nanoGPT model.
  • Implements content-dependent depth routing with a zero-initialized learned pseudo-query and parameter-free RMS-normalized keys before both attention and MLP sublayers.
  • Preserves separate W_Q, W_K, W_V, W_O, W_MLP_IN, and W_MLP_OUT matrices for Muon and WeightWatcher.
  • Keeps AttnRes pseudo-query vectors in the auxiliary AdamW optimizer group rather than the 2-D Muon matrix group.
  • Adds an exact 10-epoch matched AttnRes/Muon config using the same FineWeb-Edu data, context length, width, one block, one head, optimizer schedule, probes, and WW settings as the existing long-Muon baseline.
  • Adds a conservative full-10-epoch cosine Muon schedule as a matched standard-vs-AttnRes pair. This is an experimental long-horizon candidate, not claimed as an empirically optimized schedule.
  • Adds unit tests for forward/backward finiteness, routing weights, matrix separation, optimizer partitioning, and multi-depth readiness.
  • Adds ATTENTION_RESIDUALS.md with controlled run commands and convergence-comparison criteria.

Why

The goal is to test whether learned residual routing accelerates convergence under exactly matched token/data conditions, while retaining the RG/WeightWatcher matrix instrumentation used by the existing nanoGPT baselines. Keeping an exact matched schedule and a separate matched long-cosine pair avoids conflating an architectural gain with a learning-rate-schedule gain.

Validation

The branch diff was checked for scope and the AttnRes equations were cross-checked against the Kimi Team Full-AttnRes formulation. Connector-only repository access in this session did not provide a local executable checkout, so the added pytest suite is included for CI/runtime validation rather than claiming a local test pass.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant