Skip to content

Releases: jascal/sae-forge

v0.12.0 — sae-moe-forge

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.

v0.5.0 — GroundTruthTarget

Choose a tag to compare

@jascal jascal released this 18 May 05:09
3a81174

Install

Since sae-forge is not yet on PyPI, install from this tag:

pip install git+https://github.com/jascal/sae-forge.git@v0.5.0

Headline

add-gt-alignment-target — third built-in FaithfulnessTarget, motivated by jascal/sm-sae's production GroundTruthAlignment scorer. Family defaults (KL for LM hosts, cosine for whisper) are byte-identical to v0.4.0; GT-alignment is opt-in only via ForgePipeline(faithfulness=GroundTruthTarget(labels=L)).

The minor-version bump (vs. a patch) reflects the new scipy>=1.10 runtime dependency — technically a breaking change for callers with strict pins, even though the default surface is additive.

Added

saeforge.eval.targets.GroundTruthTarget (also re-exported as saeforge.eval.GroundTruthTarget). Scores forged residual-stream activations against an (N, M) binary label matrix via per-feature × per-label AUC — the right gate when your eval fixture carries known per-sample categories (synthetic mixtures, BERT-probe-derived datasets, concept-bottleneck suites).

  • Supported pool strategies: "mean" / "max" / "last".
  • Default hidden_extractor covers the six bundled LM-shape families (gpt2 / llama / gemma2 / qwen2 / qwen3 / qwen3_moe) via duck typing; Whisper / exotic forges supply their own.
  • AUC uses scipy.stats.rankdata(method="average") for bit-equal parity with sklearn.metrics.roc_auc_score on tie-heavy fixtures (no sklearn dep).

Demo: examples/forge_with_gt_alignment.py (mixture-of-gaussians, ~20s on CPU).

Quick example

import numpy as np
from saeforge import ForgePipeline
from saeforge.eval import GroundTruthTarget

labels = np.eye(3, dtype=np.float32)[cluster_ids]  # (N, 3) one-hot

pipeline = ForgePipeline(
    basis=basis,
    projector=projector,
    faithfulness=GroundTruthTarget(labels=labels, pool="mean"),
)
result = pipeline.run(output_dir)
print(result.faithfulness, result.faithfulness_target_name)

PRs in this release

  • #51 — openspec proposal
  • #52 — implementation (GroundTruthTarget + scipy AUC + tests + docs)
  • #53 — openspec archive; promote spec onto canonical faithfulness-target/spec.md
  • #54 — version bump + CHANGELOG promotion

What is NOT in this release

  • Additional scorers (Pearson, Spearman, monosemanticity, probe accuracy).
  • Multi-label-hierarchy comparison.
  • Whisper-encoder default hidden_extractor support.
  • requires_host=False protocol opt-out.
  • Family default routing for GT-alignment (still opt-in only).
  • CLI flag (--faithfulness-target gt_alignment). Label matrices have no reasonable CLI representation.