Skip to content

v0.12.0 — sae-moe-forge

Latest

Choose a tag to compare

@jascal jascal released this 05 Jun 02:20

Routed mixture-of-experts forge. forge_to_moe / ForgedMoE project a polygram-compressed SAE basis into a routed MoE whose per-token decode cost scales as k_experts / n_experts of the flat SAE. Inference-only, zero new parameters in v1.

Promotes the runtime-MoE play that bio-sae's fine-tune ceiling sweep re-opened — the residual sharp-feature (cov95) tax distillation couldn't close is a routed-expert target ("lossless at fixed runtime cost, not fixed param count").

Highlights

  • forge_to_moe(basis, expert_dictionary=None, …) — explicit-ExpertDictionary or auto-cluster-from-checkpoint.
  • ForgedMoE (nn.Module, buffers only): forward(track_load=…), route, expert_load, faithfulness_report, coherence_diagnostic, self-contained save_pretrained/load_pretrained.
  • SubDictionaryExpertSet (deterministic W_dec slices; vectorised decode) + PolygramHeuristicRouter (bit-for-bit parity with polygram routing).
  • Torch-free import saeforge preserved via lazy exports.

Acceptance bands (all green)

  • Ak = n_experts collapses to flat SAE (MSE/coord ≤ 1e-5; measured 0.0)
  • B — counted decode-cost ratio within 2/E ± 0.05
  • C-strict — clusterable basis routed-vs-flat ≤ 0.5× flat-vs-host (0.111 at d=768); C-advisory on isotropic bases
  • D — config + reloaded-module forward byte-identical

Docs: docs/moe-forge.md. Spec: openspec/specs/sae-moe-forge/. Full CHANGELOG entry under [0.12.0].

CI note: the 3.11 check on the source PR was red solely due to a transient HuggingFace 429 rate-limit on gpt2 weight downloads (unrelated to this change); 3.12 and the full local suite were green.