Method → application → architecture organization + hybrid trainer#49
Method → application → architecture organization + hybrid trainer#49kmheckel wants to merge 2 commits into
Conversation
…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>
There was a problem hiding this comment.
💡 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 |
There was a problem hiding this comment.
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>
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 atask → application → architecture → method"start here" flow.spyx.experimental.zoo— reference recipesRunnable recipes keyed by application, tagged by method × architecture, all on synthetic data (no downloads, sub-second smoke paths):
control-lif-esOpen_ES)classification-rsnnaxn.arctan)language-s5ssm.S5Diag)REGISTRY/list_recipes/get; each recipe exposesbuild/synthetic_batch/demo.spyx.experimental.hybrid— the 0+1 trainerThe 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.6at λ=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— cleanuv run mkdocs build --strict— clean⟨g_orth, g_s⟩ ≈ -3e-7).Everything new lands under
spyx.experimental(unstable-API contract); the stable core is untouched.🤖 Generated with Claude Code