Skip to content

Method → application → architecture organization + hybrid trainer#49

Open
kmheckel wants to merge 2 commits into
mainfrom
feat/method-app-arch-org
Open

Method → application → architecture organization + hybrid trainer#49
kmheckel wants to merge 2 commits into
mainfrom
feat/method-app-arch-org

Conversation

@kmheckel

@kmheckel kmheckel commented Jul 4, 2026

Copy link
Copy Markdown
Owner

What & why

Spyx had modules but no answer to the first question a new researcher asks: "how do I train this thing?" This PR adds a method-first organization — training methods indexed by the order of information they exploit — and maps it onto applications and architectures, plus two new experimental building blocks it references.

Docs (user-facing)

  • explanation/training-methods.md — the spine. Five families: 0th-order/evolutionary, 1st-order/surrogate gradient, transfer (conversion + QAT), local/bio-inspired (honestly flagged as roadmap — issues Implement Synthetic Gradients #27/neuron models for e-prop implementation #28), and the 0+1 hybrid. Each with when-to-use, the trade-off, and the concrete Spyx entry point.
  • explanation/choosing-an-approach.md — two decision matrices (method × application, method × architecture) and a task → application → architecture → method "start here" flow.
  • Wired first into the Explanation nav; linked from the home page.

spyx.experimental.zoo — reference recipes

Runnable recipes keyed by application, tagged by method × architecture, all on synthetic data (no downloads, sub-second smoke paths):

recipe application method architecture
control-lif-es control / RL evolutionary (evosax Open_ES) LIF-MLP
classification-rsnn classification surrogate (axn.arctan) RSNN
language-s5 language gradient S5 (ssm.S5Diag)

REGISTRY / list_recipes / get; each recipe exposes build / synthetic_batch / demo.

spyx.experimental.hybrid — the 0+1 trainer

The novel method: a surrogate gradient corrected by an antithetic-NES estimate of the true (hard-spike) loss gradient, projected orthogonal to the surrogate — the surrogate supplies the cheap bulk descent direction, ES supplies only the correction in the subspace the surrogate is blind to. It is the complement of Guided-ES (which restricts ES search to the gradient's subspace; here we restrict it to the orthogonal complement).

hybrid_gradient, make_hybrid_train_step, es_gradient, hybrid_diagnostics.

Honest result (research/new/hybrid_evo_surrogate/, three-arm study): in the cheap smoke regime the pure surrogate wins — ‖g_orth‖ ≈ 41 ≫ ‖g_s‖ ≈ 8.6 at λ=0.5, so the high-variance correction dominates. The writeup reports this as a boundary result and proposes a small self-normalising λ (λ·‖g_s‖/‖g_orth‖) with larger K as the fix, not a win it didn't earn.

Verification

  • uv run pytest -m "not network"202 passed, 1 skipped (16 new tests across zoo + hybrid)
  • uv run ruff check / uv run ty check — clean
  • uv run mkdocs build --strict — clean
  • Each stream was built and adversarially verified by an independent agent (strict build re-run, real pytest, and a hand-check of the NES estimator scaling + orthogonality of the projection — ⟨g_orth, g_s⟩ ≈ -3e-7).

Everything new lands under spyx.experimental (unstable-API contract); the stable core is untouched.

🤖 Generated with Claude Code

…d trainer

Organizes Spyx around the question a new researcher actually asks — "how do I
train this?" — by the order of information each method exploits, then maps that
onto applications and architectures.

Docs (user-facing):
- docs/explanation/training-methods.md — the spine: 0th-order/evolutionary,
  1st-order/surrogate, transfer/conversion+QAT, local/bio-inspired (honestly
  flagged roadmap, #27/#28), and the 0+1 hybrid. Each with when-to-use,
  trade-off, and Spyx entry point.
- docs/explanation/choosing-an-approach.md — method×application and
  method×architecture decision matrices + a task→app→arch→method flow.
- Wired into the mkdocs Explanation nav and linked from the home page.

spyx.experimental.zoo — runnable reference recipes keyed by application and
tagged by training method × architecture, all on synthetic data:
- control-lif-es (evolutionary / LIF-MLP, evosax Open_ES)
- classification-rsnn (surrogate / RSNN, axn.arctan)
- language-s5 (gradient / S5, spyx.ssm.S5Diag)
REGISTRY/list_recipes/get; each recipe exposes build/synthetic_batch/demo.

spyx.experimental.hybrid — the 0+1 trainer: a surrogate gradient corrected by
an antithetic-NES estimate of the true (hard-spike) loss gradient, projected
orthogonal to the surrogate (the complement of Guided-ES). hybrid_gradient,
make_hybrid_train_step, es_gradient, hybrid_diagnostics. research/new/
hybrid_evo_surrogate/ runs the three-arm study — honest negative in the smoke
regime (surrogate wins; correction needs a small self-normalising lambda).

Gate: 202 tests pass (16 new), ruff + ty clean, mkdocs build --strict clean.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9fcd8582ef

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

:return: list of per-generation mean rollout costs (lower is better).
"""
import optax # lazy: optax is a core dep but imported here to keep import light
from evosax.algorithms import Open_ES

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Gate the control demo on the evo extra

This import is only available when the optional [evo] extra is installed, but the new zoo tests call every recipe's demo() unconditionally and the CI test job syncs only --extra quant, so control-lif-es will raise ModuleNotFoundError: evosax in the default test environment. Either include the evo extra in the tested dependency set or skip/gate this recipe when evosax is unavailable.

Useful? React with 👍 / 👎.

Follow-up to the honest negative in the initial hybrid study: the raw-lambda
correction can end up much WORSE than plain surrogate descent because the
antithetic-NES estimate has norm ||g_orth|| that easily exceeds ||g_s||, so at a
fixed lambda the noisy ES term dominates the update.

Fix: add normalize=True to hybrid_gradient / hybrid_diagnostics /
make_hybrid_train_step. In this mode lambda is a dimensionless FRACTION of the
surrogate step — the correction is rescaled by lam_eff = lambda*||g_s||/||g_orth||
so ||applied correction|| = lambda*||g_s|| exactly, immune to the ES variance /
sigma scaling. New diagnostics: lam_eff, correction_fraction. New test asserts
the correction has norm lambda*||g_s||, stays orthogonal to g_s, and lambda=0
still recovers pure surrogate.

Study: run.py is now a four-arm comparison (surrogate / es / hybrid-raw /
hybrid-norm); sweep.py sweeps seeds x K x lambda over a harder regime and reports
mean delta(true_loss) with seed spread so a "win" must survive noise.

Honest verdict (README): self-normalising lambda removes the failure mode and
brings hybrid to PARITY with the surrogate (never catastrophically worse) — a
real improvement over the raw version — but it does NOT beat the surrogate on
these easy synthetic tasks (task ceiling, little surrogate bias left to correct).
Parity-plus-safety reported, not a win.

Gate: ruff + ty clean; test_hybrid (6) + test_zoo (11) pass.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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