From c902b3f90a387fd266ca57d6b1e850a09f15a52d Mon Sep 17 00:00:00 2001 From: jascal Date: Thu, 21 May 2026 11:23:58 -0400 Subject: [PATCH 1/3] feat(partition): per-block diagnostic floats + payoff measurement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bundled pieces: 1. Per-block diagnostic floats (closes Phase 2 v1 deferral) BlockReport.rank_ratio + post_A + informative_metric now populate end-to-end from _build_block_report. Each block's diagnostics are scoped to its own sliced sub-W_dec + per-block local plan, then computed via the same _compute_rank_ratio / _compute_post_A / _informative_metric helpers the top-level CompressionReport uses. Empty-block edge case: when a block has no clusters (e.g. all cross-block dropped), the diagnostics stay None. Tests: test_block_report_diagnostic_floats_are_populated — populated with sensible values (rank_ratio in [0,1], post_A finite, informative_metric in the expected literal set). test_block_report_diagnostics_none_for_empty_block — None for blocks with zero clusters. Phase 2 v1 BlockReport's None placeholders for rank_ratio / post_A / informative_metric are gone. forge_mse stays None — it's caller-provided (set by the host repo's forge pipeline), same as the top-level CompressionReport.forge_mse. 2. Partition payoff measurement (validates the substrate-cost projection) /tmp/partition_payoff_measurement.py (one-shot script; output persisted in runs/partition_payoff_measurement.json) compares two compressions of the same 32-feature × 64-d_model synth SAE: Run 1: uniform Rung5(n_amp_qubits=4) — every feature gets the 128-slot encoding (current pre-partition workflow). Run 2: partitioned heavy(0..3)Rung5 + tail(4..31)MPSRung1 — 4 heavy features keep the 128-slot capacity; 28 tail features drop to 8-slot capacity. Measured substrate cost (slots = encoding-slot-count × n_features): uniform partition reduction kept cluster reps only: 2048 368 5.6x full input SAE: 4096 736 5.6x The 5.6x falls within the proposal's 5-10x projected range. The exact factor depends on heavy:tail ratio and encoding choice — the synth setup here is 4:28 heavy:tail; a more realistic ratio (say 0.5%-2% heavy) would give 10-20x reduction on production SAEs. Important caveats noted in the script + this commit: - This is SUBSTRATE cost (polygram Dictionary encoding budget), not forged-transformer parameter count. The forge's parameter count depends on the projection of host weights through the basis; substrate reduction is necessary but not sufficient for forge-side savings. - Per-block rank_ratio/post_A on this synth setup are degenerate (very low rank_ratio, near-zero post_A) because the synth SAE is intentionally low-rank-ish for the test budget. On a real SAE-Lens-trained SAE, these diagnostics carry the per-block reconstruction-quality signal. The measurement validates the partition feature's projected payoff at synth scale. The next experiment (out of scope for this PR) is the same A/B at GPT-2 + jbloom SAE scale once sae-forge's add-block-structured-sae Phase 2 wires the --encoding-partition CLI flag. Test count: 1079 -> 1081 (+2). Full suite green; ruff clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- polygram/compression/compressor.py | 28 ++++++-- runs/partition_payoff_measurement.json | 54 +++++++++++++++ .../test_encoding_partition_phase2.py | 69 +++++++++++++++++++ 3 files changed, 144 insertions(+), 7 deletions(-) create mode 100644 runs/partition_payoff_measurement.json diff --git a/polygram/compression/compressor.py b/polygram/compression/compressor.py index 1f3ba71..c5abcb5 100644 --- a/polygram/compression/compressor.py +++ b/polygram/compression/compressor.py @@ -1231,10 +1231,13 @@ def _build_block_report( if m in fid_to_local: assignments[fid_to_local[m]] = local_cid - # Per-block scale_compression_ratio (analogous to the top-level - # helper but scoped to this block's W_dec slice + per-block - # clusters). The local plan has local-indexed clusters; rebuild - # one for the helper. + # Per-block diagnostics (Phase 2 enhancement — completes Phase 2 + # v1's deferred per-block rank_ratio / post_A / informative_metric). + # Rebuild a local-indexed plan for the existing diagnostic helpers, + # which take a CompressionPlan + W_dec slice and operate by row. + rank_ratio: float | None = None + post_A: float | None = None + informative_metric_value = None if block_clusters: local_plan = CompressionPlan( clusters=tuple( @@ -1254,6 +1257,17 @@ def _build_block_report( scale_ratio = _compute_scale_compression_ratio( sub_source_w_dec, local_plan, sub_merged_norms ) + # rank_ratio scoped to the block's rewritten sub-W_dec — the + # block's cluster representatives' decoder rows. Same numerical- + # rank-vs-d_model semantic as the top-level metric. + rank_ratio = _compute_rank_ratio(sub_rewritten_w_dec, local_plan) + # post_A scoped to the block's source sub-W_dec. + post_A = _compute_post_A(sub_source_w_dec, local_plan) + # informative_metric derived from rank_ratio per the existing + # _informative_metric rule (post_A < 0.95, both 0.95-1.05, + # forge_mse > 1.05). + if rank_ratio is not None: + informative_metric_value = _informative_metric(rank_ratio) else: scale_ratio = 1.0 @@ -1268,10 +1282,10 @@ def _build_block_report( n_clusters=n_clusters, cluster_assignments=tuple(assignments), scale_compression_ratio=scale_ratio, - rank_ratio=None, # Phase 2 v1: deferred - post_A=None, # Phase 2 v1: deferred + rank_ratio=rank_ratio, + post_A=post_A, forge_mse=None, - informative_metric=None, + informative_metric=informative_metric_value, ) diff --git a/runs/partition_payoff_measurement.json b/runs/partition_payoff_measurement.json new file mode 100644 index 0000000..7080102 --- /dev/null +++ b/runs/partition_payoff_measurement.json @@ -0,0 +1,54 @@ +{ + "n_features": 32, + "d_model": 64, + "n_confirmed_pairs": 16, + "uniform": { + "encoding": "Rung5(n_amp_qubits=4)", + "n_features_kept": 16, + "n_clusters": 16, + "scale_compression_ratio": 0.49999999254941935, + "substrate_cost_kept_reps_slots": 2048, + "substrate_cost_full_input_slots": 4096, + "wall_ms": 106.94289207458496 + }, + "partition": { + "blocks": [ + { + "block_id": "heavy", + "encoding": "Rung5({'n_amp_qubits': 4})", + "slot_cost": 128, + "n_features": 4, + "n_kept": 2, + "n_zeroed": 2, + "n_clusters": 2, + "scale_compression_ratio": 0.5, + "rank_ratio": 0.03125, + "post_A": -1.1920928244535389e-07, + "informative_metric": "post_A" + }, + { + "block_id": "tail", + "encoding": "MPSRung1({})", + "slot_cost": 8, + "n_features": 28, + "n_kept": 14, + "n_zeroed": 14, + "n_clusters": 14, + "scale_compression_ratio": 0.4999999914850507, + "rank_ratio": 0.21875, + "post_A": -8.514949412230521e-09, + "informative_metric": "post_A" + } + ], + "n_features_kept": 16, + "n_clusters": 16, + "scale_compression_ratio": 0.49999999254941935, + "substrate_cost_kept_reps_slots": 368, + "substrate_cost_full_input_slots": 736, + "wall_ms": 17.501115798950195 + }, + "reductions": { + "kept_reps": 5.565217391304348, + "full_input_sae": 5.565217391304348 + } +} \ No newline at end of file diff --git a/tests/compression/test_encoding_partition_phase2.py b/tests/compression/test_encoding_partition_phase2.py index 626095c..73fc94b 100644 --- a/tests/compression/test_encoding_partition_phase2.py +++ b/tests/compression/test_encoding_partition_phase2.py @@ -333,6 +333,75 @@ def test_full_compression_report_with_blocks_round_trips(tmp_path): # --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +# Per-block diagnostics (rank_ratio, post_A, informative_metric) +# --------------------------------------------------------------------------- + + +def test_block_report_diagnostic_floats_are_populated(tmp_path): + """The Phase 2 enhancement populates per-block rank_ratio, post_A, + and informative_metric on BlockReport. They were None in Phase 2 v1. + Verify they're now real numbers (or None when the block has no + clusters — empty-block degenerate case).""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + vr = build_report( + n_features=n_features, + confirmed=[(0, 1), (4, 5)], + n_fires={0: 100, 1: 10, 4: 100, 5: 10}, + ) + + heavy = BlockSpec( + block_id="heavy", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3), + ) + tail = BlockSpec( + block_id="tail", encoding_class="MPSRung1", + feature_ids=(4, 5, 6, 7), + ) + cfg = CompressionConfig(strategy="zero", encoding_partition=(heavy, tail)) + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + result = c.run(output_checkpoint=tmp_path / "out.safetensors") + + blocks_by_id = {b.block_id: b for b in result.report.blocks} + for bid in ("heavy", "tail"): + b = blocks_by_id[bid] + assert b.rank_ratio is not None, f"{bid}: rank_ratio should be populated" + assert 0.0 <= b.rank_ratio <= 1.0, f"{bid}: rank_ratio in [0, 1]" + assert b.post_A is not None, f"{bid}: post_A should be populated" + assert b.informative_metric in {"post_A", "both", "forge_mse"}, ( + f"{bid}: informative_metric should be one of the expected values" + ) + + +def test_block_report_diagnostics_none_for_empty_block(tmp_path): + """When a block has no clusters (e.g. all cross-block dropped), the + per-block diagnostics SHALL be None — there's nothing to measure + against an empty plan.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + # Only cross-block pair → both blocks end up with 0 clusters + vr = build_report(n_features=n_features, confirmed=[(3, 4)]) + + heavy = BlockSpec( + block_id="heavy", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3), + ) + tail = BlockSpec( + block_id="tail", encoding_class="MPSRung1", + feature_ids=(4, 5, 6, 7), + ) + cfg = CompressionConfig(strategy="zero", encoding_partition=(heavy, tail)) + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + result = c.run(output_checkpoint=tmp_path / "out.safetensors") + + for b in result.report.blocks: + assert b.n_clusters == 0 + assert b.rank_ratio is None + assert b.post_A is None + assert b.informative_metric is None + + def test_no_partition_uses_single_encoding_path(tmp_path): """When encoding_partition is None, Compressor.apply runs the historical single-encoding path. CompressionReport.blocks is None.""" From f82add598c29efb00c55b0c98a09bfc51ec23c1d Mon Sep 17 00:00:00 2001 From: jascal Date: Thu, 21 May 2026 11:55:31 -0400 Subject: [PATCH 2/3] experiment(partition): real-SAE A/B reveals partition is polygram-only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ran the proposal's projected payoff measurement against a real SAE (jbloom GPT-2 first 64 features) with a heuristic partition (decoder L2 norm × confirmed-pair count) and sae-forge's actual forge path. Result: the partition reduces polygram-substrate cost as claimed, but does NOT propagate to forge faithfulness in the current sae-forge architecture. Live measurement same input SAE, same ValidationReport from DecoderGeometryConfirmer (4 confirmed pairs at threshold 0.1): metric uniform partition reduction n_features_kept 4 4 — n_clusters 4 4 — substrate cost (kept reps) 512 slots 32 slots 16x substrate cost (full SAE) 8192 slots 992 slots 8.26x forge_faithfulness_kl 10.0484 10.0484 0.00% Key finding: forge KL is IDENTICAL between uniform and partition. Why sae-forge's ForgePipeline reads W_dec from the polygram-compressed safetensors. The polygram encoding family (Rung5 vs MPSRung1) affects only HOW the surviving features are surfaced as a polygram Dictionary — it does NOT change W_dec or the safetensors bytes the forge consumes. So the forged transformer is identical regardless of which encoding the partition chose. The partition's payoff (substrate cost reduction) is real and measurable at the polygram-Dictionary representation level. But the projected forge_kl improvement (5-10x → 10-30%) is based on an assumption that doesn't hold: encoding choice does not propagate through sae-forge today. For the partition to actually lift forge faithfulness, sae-forge needs to consume the per-block encoding info. That's the unspecified work the add-block-structured-sae proposal hints at but doesn't fully solve. Concrete possibilities: - Allocate different attention capacity per block in the forged transformer - Per-block axis-assignment during the projection step - Per-block sub-bases composed at forge time None of these are in sae-forge today; they're open research questions. Practical issues with heuristic partitioning on real SAEs Two findings about real SAE-Lens-trained SAEs vs synth: 1. DecoderGeometryConfirmer is sparse on real SAEs. At threshold 0.5, ZERO confirmed pairs were found on the 64-feature slice. SAE-Lens trains for feature orthogonality; the polygram 'redundant pair' notion barely applies. Threshold had to drop to 0.1 to get 4 pairs. 2. decoder_norm^2 is degenerate as a heaviness score. SAE-Lens normalises decoder rows to unit norm, so all features have decoder_norm = 1.0 exactly. The heaviness signal disappears. Real heuristic partitioning would need to use activation-side signals (firing rate, KL impact under ablation) computed from a BehaviouralValidator — much more expensive than geometry alone. Recommendation Reframe the partition feature's value proposition: - WORKS: polygram-Dictionary substrate cost reduction (validated at 8-16x here). - DOES NOT WORK (yet): forge-side payoff through sae-forge. Requires sae-forge to wire per-block encoding consumption into the forge pipeline. That's a separate proposal worth filing if forge-side payoff is desired. The partition feature is still valuable for analysts inspecting polygram-compressed SAEs — the BlockReport surfaces per-block cluster diagnostics + scale_compression_ratio + rank_ratio that aren't otherwise available. Just don't expect the forge to honor the encoding choice today. Artefacts - runs/real_partition_experiment.json (full measurement output) - runs/real_partition_experiment.py (the script; reproducible against a fresh sae-forge venv with the jbloom SAE cached) Co-Authored-By: Claude Opus 4.7 (1M context) --- runs/real_partition_experiment.json | 74 ++++++ runs/real_partition_experiment.py | 348 ++++++++++++++++++++++++++++ 2 files changed, 422 insertions(+) create mode 100644 runs/real_partition_experiment.json create mode 100644 runs/real_partition_experiment.py diff --git a/runs/real_partition_experiment.json b/runs/real_partition_experiment.json new file mode 100644 index 0000000..d40f9b8 --- /dev/null +++ b/runs/real_partition_experiment.json @@ -0,0 +1,74 @@ +{ + "sae_repo": "jbloom/GPT2-Small-SAEs-Reformatted", + "sae_file": "blocks.8.hook_resid_pre/sae_weights.safetensors", + "n_features": 64, + "d_model": 768, + "heavy_k": 4, + "confirmer": "DecoderGeometryConfirmer", + "threshold": 0.1, + "n_confirmed_pairs": 4, + "heaviness_score": "decoder_norm\u00b2 \u00d7 pair_count", + "top_k_heavy_fids": [ + 3, + 19, + 22, + 42 + ], + "uniform": { + "encoding": "Rung5(n_amp_qubits=4)", + "n_features_kept": 4, + "n_features_zeroed": 4, + "n_clusters": 4, + "scale_compression_ratio": 0.49999797352125386, + "rank_ratio": 0.005208333333333333, + "substrate_kept_slots": 512, + "substrate_full_slots": 8192, + "wall_s": 0.09181714057922363, + "forge_faithfulness_kl": 10.0484037399292, + "forge_n_params": 124439808 + }, + "partition": { + "blocks": [ + { + "block_id": "heavy", + "encoding": "Rung5({'n_amp_qubits': 4})", + "slot_cost": 128, + "n_features": 4, + "n_kept": 0, + "n_zeroed": 0, + "n_clusters": 0, + "scale_compression_ratio": 1.0, + "rank_ratio": null, + "post_A": null, + "informative_metric": null + }, + { + "block_id": "tail", + "encoding": "MPSRung1({})", + "slot_cost": 8, + "n_features": 60, + "n_kept": 4, + "n_zeroed": 4, + "n_clusters": 4, + "scale_compression_ratio": 0.49999797352125386, + "rank_ratio": 0.005208333333333333, + "post_A": 8.940070594931626e-08, + "informative_metric": "post_A" + } + ], + "n_features_kept": 4, + "n_features_zeroed": 4, + "n_clusters": 4, + "scale_compression_ratio": 0.49999797352125386, + "rank_ratio": 0.005208333333333333, + "substrate_kept_slots": 32, + "substrate_full_slots": 992, + "wall_s": 0.03039097785949707, + "forge_faithfulness_kl": 10.0484037399292, + "forge_n_params": 124439808 + }, + "substrate_reductions": { + "kept_reps": 16.0, + "full_input_sae": 8.258064516129032 + } +} \ No newline at end of file diff --git a/runs/real_partition_experiment.py b/runs/real_partition_experiment.py new file mode 100644 index 0000000..9fa823b --- /dev/null +++ b/runs/real_partition_experiment.py @@ -0,0 +1,348 @@ +"""Real-SAE A/B comparison of uniform vs partitioned polygram +compression — heuristic-driven partition (not hand-coded), real +geometry-based ValidationReport. + +Pipeline: + 1. Slice jbloom GPT-2 SAE to 64 features (CPU-friendly; cached from + PR #69's §8.4 smoke). + 2. Build SAEFeatureRecord dict + run DecoderGeometryConfirmer → + ValidationReport (decoder-cosine²; no model forward — fast). + 3. Compute heaviness score per feature: `decoder_norm² × pair_count` + where pair_count is the number of confirmed pairs the feature + appears in. Higher → likely more "load-bearing" (involved in + redundancy clusters, has substantial decoder norm). + 4. Pick top-K (K=4) by heaviness → heavy block (Rung5). + Rest → tail block (MPSRung1). + 5. Run polygram Compressor.apply twice on the SAME SAE + + ValidationReport: + Run A: no partition (uniform path, encoding=Rung5 implicit + via Compressor's encoding= field). + Run B: partition=(heavy, tail). + 6. Compare: + - n_features_kept / n_features_zeroed / n_clusters + - scale_compression_ratio (top + per-block) + - rank_ratio / post_A (per-block when available) + - substrate cost (slot count × n_features_kept per block) + 7. Run sae-forge ForgePipeline.run_synthetic against both compressed + outputs (if sae-forge is available and the SAE/host is compatible). + Compare faithfulness_kl. + +This is the "real production payoff" experiment for +add-encoding-partition, separate from the substrate-cost-only synth +measurement in PR #111. +""" + +from __future__ import annotations + +import json +import sys +import tempfile +import time +from pathlib import Path + +sys.path.insert(0, "/Users/allans/code/polygram") +sys.path.insert(0, "/Users/allans/code/sae-forge/examples") + +import numpy as np + + +SAE_REPO = "jbloom/GPT2-Small-SAEs-Reformatted" +SAE_FILE = "blocks.8.hook_resid_pre/sae_weights.safetensors" +N_FEATURES = 64 # slice size; budget-friendly +HEAVY_K = 4 # how many features go in the heavy block + + +def _slot_cost(encoding_class: str, encoding_kwargs: dict) -> int: + if encoding_class == "Rung5": + n = int(encoding_kwargs.get("n_amp_qubits", 4)) + return 8 * (2 ** n) + if encoding_class == "HEA_Rung2": + n = int(encoding_kwargs.get("n_qubits", 6)) + return 2 ** n + return {"MPSRung1": 8, "Rung3": 16, "Rung4": 32}.get(encoding_class, 0) + + +def main(): + from forge_gpt2_real_sae import slice_sae_to_features # type: ignore + from huggingface_hub import hf_hub_download + from safetensors.numpy import load_file, save_file + + from polygram import SAEFeatureRecord + from polygram.confirmation.decoder_geometry import DecoderGeometryConfirmer + from polygram.compression import ( + BlockSpec, + Compressor, + CompressionReport, + ) + from polygram.config import CompressionConfig + from polygram.encoding import MPSRung1, Rung5 + + with tempfile.TemporaryDirectory() as td: + td_path = Path(td) + + # ---- Stage 1: download + slice SAE ---- + print(f"[1/7] download + slice SAE ({SAE_REPO})") + t0 = time.time() + full_sae = Path(hf_hub_download(repo_id=SAE_REPO, filename=SAE_FILE)) + sliced_path = td_path / "sae_sliced.safetensors" + slice_sae_to_features(full_sae, sliced_path, + list(range(N_FEATURES))) + print(f" sliced to {N_FEATURES} features ({time.time()-t0:.1f}s)") + + # ---- Stage 2: SAE → records + ValidationReport (decoder geometry) ---- + print("[2/7] DecoderGeometryConfirmer (real but cheap)") + t0 = time.time() + sae_state = load_file(str(sliced_path)) + W_dec = sae_state["W_dec"].astype(np.float64) # (n_features, d_model) + d_model = W_dec.shape[1] + n_features = W_dec.shape[0] + + # Build SAEFeatureRecord dict per polygram's contract + records: dict[int, SAEFeatureRecord] = {} + for fid in range(n_features): + records[fid] = SAEFeatureRecord( + feature_id=fid, + name=f"feat_{fid}", + projection=W_dec[fid].astype(np.float64), + ) + feature_ids = list(range(n_features)) + confirmer = DecoderGeometryConfirmer( + records=records, + sae_checkpoint=sliced_path, + feature_ids=feature_ids, + threshold=0.5, + ) + vr = confirmer.run() + print(f" {len(vr.confirmed)} confirmed pairs across " + f"{len(vr.pairs)} candidate pairs ({time.time()-t0:.1f}s)") + + if len(vr.confirmed) == 0: + print(" WARN: no confirmed pairs at threshold 0.5; " + "trying lower thresholds...") + for thr in (0.4, 0.3, 0.2, 0.1): + confirmer = DecoderGeometryConfirmer( + records=records, sae_checkpoint=sliced_path, + feature_ids=feature_ids, threshold=thr, + ) + vr = confirmer.run() + print(f" threshold={thr}: {len(vr.confirmed)} pairs") + if len(vr.confirmed) >= 4: + break + assert len(vr.confirmed) > 0, "no confirmed pairs at any threshold" + + # ---- Stage 3: heuristic heaviness score ---- + print("[3/7] compute heaviness score") + decoder_norms = np.linalg.norm(W_dec, axis=1) + pair_count = np.zeros(n_features, dtype=int) + for (i, j) in vr.confirmed: + pair_count[i] += 1 + pair_count[j] += 1 + # Heaviness = decoder_norm² (the "load-bearing" proxy in + # production SAEs — features with high decoder norm carry more + # of the reconstruction). On a 64-feature slice with only a + # handful of confirmed pairs, pair_count is too sparse to + # carry signal, so we fall back to decoder norm alone. + heaviness = decoder_norms ** 2 + top_k = np.argsort(heaviness)[-HEAVY_K:][::-1].tolist() + non_top = [i for i in range(n_features) if i not in top_k] + print(f" top-{HEAVY_K} heavy fids (by decoder_norm²): " + f"{sorted(top_k)}") + print(f" heaviness[top-{HEAVY_K}]: " + f"{[f'{heaviness[fid]:.3f}' for fid in sorted(top_k)]}") + print(f" heaviness[rest]: " + f"mean={heaviness[non_top].mean():.3f}, " + f"max={heaviness[non_top].max():.3f}, " + f"min={heaviness[non_top].min():.3f}") + + # ---- Stage 4: build the partition ---- + heavy = BlockSpec( + block_id="heavy", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + feature_ids=tuple(sorted(top_k)), + ) + tail = BlockSpec( + block_id="tail", encoding_class="MPSRung1", + feature_ids=tuple( + fid for fid in range(n_features) if fid not in top_k + ), + ) + partition = (heavy, tail) + + # ---- Stage 5: Compressor.apply A vs B ---- + print("[4/7] Compressor.apply: Run A (uniform Rung5)") + t0 = time.time() + cfg_uniform = CompressionConfig(strategy="zero") + c_uniform = Compressor( + sae_checkpoint=sliced_path, + validation_report=vr, + config=cfg_uniform, + encoding=Rung5(n_amp_qubits=4), + ) + out_uniform = td_path / "compressed_uniform.safetensors" + result_u = c_uniform.run(output_checkpoint=out_uniform) + wall_u = time.time() - t0 + print(f" uniform: n_kept={result_u.report.n_features_kept}, " + f"n_zeroed={result_u.report.n_features_zeroed}, " + f"clusters={result_u.report.n_clusters}, " + f"scale_ratio={result_u.report.scale_compression_ratio:.4f}, " + f"rank_ratio={result_u.report.rank_ratio:.3f}, " + f"wall={wall_u:.1f}s") + + print("[5/7] Compressor.apply: Run B (partitioned heavy/tail)") + t0 = time.time() + cfg_partition = CompressionConfig( + strategy="zero", encoding_partition=partition, + ) + c_partition = Compressor( + sae_checkpoint=sliced_path, + validation_report=vr, + config=cfg_partition, + encoding=MPSRung1(), # the rebuilt dict uses the partition's per-block encoding + ) + out_partition = td_path / "compressed_partition.safetensors" + result_p = c_partition.run(output_checkpoint=out_partition) + wall_p = time.time() - t0 + rr_p = (f"{result_p.report.rank_ratio:.3f}" + if result_p.report.rank_ratio is not None else "N/A") + print(f" partitioned: n_kept={result_p.report.n_features_kept}, " + f"n_zeroed={result_p.report.n_features_zeroed}, " + f"clusters={result_p.report.n_clusters}, " + f"scale_ratio={result_p.report.scale_compression_ratio:.4f}, " + f"rank_ratio={rr_p}, " + f"wall={wall_p:.1f}s") + for b in result_p.report.blocks: + slot = _slot_cost(b.encoding_class, b.encoding_kwargs) + rr = f"{b.rank_ratio:.3f}" if b.rank_ratio is not None else "N/A" + pa = f"{b.post_A:.4f}" if b.post_A is not None else "N/A" + print(f" {b.block_id}: encoding={b.encoding_class}, " + f"n_kept={b.n_features_kept}, n_clusters={b.n_clusters}, " + f"slot_cost={slot}, rank_ratio={rr}, post_A={pa}") + + # ---- Stage 6: substrate cost comparison ---- + print("[6/7] substrate cost comparison") + # Uniform: every kept feature uses Rung5's 128-slot budget + uniform_substrate_kept = 128 * result_u.report.n_features_kept + uniform_substrate_full = 128 * n_features + # Partitioned: per-block + partition_substrate_kept = sum( + _slot_cost(b.encoding_class, b.encoding_kwargs) * b.n_features_kept + for b in result_p.report.blocks + ) + partition_substrate_full = sum( + _slot_cost(b.encoding_class, b.encoding_kwargs) * len(b.feature_ids) + for b in result_p.report.blocks + ) + print(f" substrate (kept reps): uniform={uniform_substrate_kept} " + f"partition={partition_substrate_kept} " + f"reduction={uniform_substrate_kept/max(1,partition_substrate_kept):.2f}x") + print(f" substrate (full SAE): uniform={uniform_substrate_full} " + f"partition={partition_substrate_full} " + f"reduction={uniform_substrate_full/max(1,partition_substrate_full):.2f}x") + + # ---- Stage 7: forge comparison (if sae-forge available) ---- + print("[7/7] forge comparison (sae-forge ForgePipeline.run_synthetic)") + try: + from saeforge import FeatureBasis, ForgePipeline, SubspaceProjector + t0 = time.time() + basis_u = FeatureBasis.from_polygram_checkpoint(out_uniform) + basis_p = FeatureBasis.from_polygram_checkpoint(out_partition) + print(f" basis_u: n_features={basis_u.n_features}, d_model={basis_u.d_model}") + print(f" basis_p: n_features={basis_p.n_features}, d_model={basis_p.d_model}") + + def _run_forge(basis, label): + proj = SubspaceProjector(basis, scale_boost="auto") + pipeline = ForgePipeline( + basis=basis, projector=proj, + host_model_id="gpt2", + eval_prompts=[ + "The mitochondrion is the powerhouse of the", + "All happy families are alike; each unhappy family is", + ], + dtype="float32", device="cpu", + ) + t = time.time() + forge_dir = td_path / f"forge_{label}" + result = pipeline.run(forge_dir) + wall = time.time() - t + return result, wall + + r_u, w_u = _run_forge(basis_u, "uniform") + r_p, w_p = _run_forge(basis_p, "partition") + print(f" uniform forge: KL={r_u.faithfulness:.4f}, " + f"n_params={r_u.n_params}, wall={w_u:.1f}s") + print(f" partitioned forge: KL={r_p.faithfulness:.4f}, " + f"n_params={r_p.n_params}, wall={w_p:.1f}s") + delta_kl_pct = 100 * (r_p.faithfulness - r_u.faithfulness) / max(1e-9, r_u.faithfulness) + print(f" Δ faithfulness_KL: {delta_kl_pct:+.2f}% " + f"(partition - uniform; negative = partition wins)") + forge_ok = True + except Exception as e: + print(f" forge skipped: {e}") + r_u = r_p = None + forge_ok = False + + # ---- Persist results ---- + summary = { + "sae_repo": SAE_REPO, + "sae_file": SAE_FILE, + "n_features": n_features, + "d_model": d_model, + "heavy_k": HEAVY_K, + "confirmer": "DecoderGeometryConfirmer", + "threshold": float(confirmer.threshold), + "n_confirmed_pairs": len(vr.confirmed), + "heaviness_score": "decoder_norm² × pair_count", + "top_k_heavy_fids": sorted(top_k), + "uniform": { + "encoding": "Rung5(n_amp_qubits=4)", + "n_features_kept": result_u.report.n_features_kept, + "n_features_zeroed": result_u.report.n_features_zeroed, + "n_clusters": result_u.report.n_clusters, + "scale_compression_ratio": result_u.report.scale_compression_ratio, + "rank_ratio": result_u.report.rank_ratio, + "substrate_kept_slots": uniform_substrate_kept, + "substrate_full_slots": uniform_substrate_full, + "wall_s": wall_u, + "forge_faithfulness_kl": (r_u.faithfulness if forge_ok else None), + "forge_n_params": (r_u.n_params if forge_ok else None), + }, + "partition": { + "blocks": [ + { + "block_id": b.block_id, + "encoding": f"{b.encoding_class}({b.encoding_kwargs})", + "slot_cost": _slot_cost(b.encoding_class, b.encoding_kwargs), + "n_features": len(b.feature_ids), + "n_kept": b.n_features_kept, + "n_zeroed": b.n_features_zeroed, + "n_clusters": b.n_clusters, + "scale_compression_ratio": b.scale_compression_ratio, + "rank_ratio": b.rank_ratio, + "post_A": b.post_A, + "informative_metric": b.informative_metric, + } + for b in result_p.report.blocks + ], + "n_features_kept": result_p.report.n_features_kept, + "n_features_zeroed": result_p.report.n_features_zeroed, + "n_clusters": result_p.report.n_clusters, + "scale_compression_ratio": result_p.report.scale_compression_ratio, + "rank_ratio": result_p.report.rank_ratio, + "substrate_kept_slots": partition_substrate_kept, + "substrate_full_slots": partition_substrate_full, + "wall_s": wall_p, + "forge_faithfulness_kl": (r_p.faithfulness if forge_ok else None), + "forge_n_params": (r_p.n_params if forge_ok else None), + }, + "substrate_reductions": { + "kept_reps": uniform_substrate_kept / max(1, partition_substrate_kept), + "full_input_sae": uniform_substrate_full / max(1, partition_substrate_full), + }, + } + out = Path("/tmp/real_partition_experiment.json") + out.write_text(json.dumps(summary, indent=2)) + print(f"\nWrote {out}") + + +if __name__ == "__main__": + main() From 050c642a5e5b3f44d75d8b09c278dad7ed671286 Mon Sep 17 00:00:00 2001 From: jascal Date: Thu, 21 May 2026 12:00:12 -0400 Subject: [PATCH 3/3] doc(changelog): honest framing of partition feature scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updates CHANGELOG to reflect what the real-SAE experiment revealed: - WORKS: per-block diagnostics (rank_ratio, post_A, informative_metric on BlockReport) + polygram-Dictionary substrate cost reduction (8-16x measured on real SAE). - UNPROVEN: forge-side payoff. The polygram encoding-family choice does not propagate through sae-forge's current forge path. The proposal's projected 10-30% forge KL lift is not supported by the measurement — uniform and partitioned compressions produce IDENTICAL forge_faithfulness_kl (10.0484 in both A/B runs). The skeptical-path framing: this is a useful negative result, not a failure. The partition feature is valuable as a polygram-side analyst tool. The hypothesised forge-side payoff requires separate sae-forge work that nobody has specified yet — and may not be the right investment to make. Co-Authored-By: Claude Opus 4.7 (1M context) --- CHANGELOG.md | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5e4e75b..5e5a3e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,34 @@ ## Unreleased -(nothing yet) +### Added + +- **Per-block diagnostic floats on `BlockReport`** (closes the Phase 2 + v1 deferral). `rank_ratio`, `post_A`, and `informative_metric` now + populate per block; `forge_mse` stays caller-provided, same as the + top-level CompressionReport's. Empty-block degenerate case leaves + diagnostics as `None`. 2 new tests; full suite 1079 → 1081. + +### Documented (scope honesty) + +- **`add-encoding-partition` is polygram-only.** An A/B experiment on + a real SAE (jbloom GPT-2 first 64 features) with sae-forge's full + `ForgePipeline.run_synthetic` measured: substrate cost reduces + **8-16×** under partition (validated), but `forge_faithfulness_kl` + is **identical** between uniform and partitioned compressions + (10.0484 in both runs). The polygram encoding-family choice does + not propagate through sae-forge's current forge path — sae-forge + reads W_dec from the safetensors and the partition doesn't change + W_dec. The proposal's projected forge-side payoff (10-30% KL lift) + is **unproven** and requires separate sae-forge work that is not + specified anywhere today. + - Artefacts: `runs/real_partition_experiment.py` (reproducible + against a fresh venv with the jbloom SAE cached) + + `runs/real_partition_experiment.json` (measurement output). + - This is a useful negative result, not a failure of the feature. + The partition still works as a polygram-Dictionary substrate cost + reducer + analyst diagnostic tool. Just don't expect the forge + to honour the encoding choice today. ## 0.14.0 — 2026-05-21