feat(eval): per-site faithfulness observability (rel/abs Frobenius)

param_decomp/README.md

@@ -27,6 +27,7 @@ root with sibling packages `pretrain/` (the in-house target-LM pretrainer) and
 | `ci_fn.py` | shared-transformer CI fn over ordered site specs; the two leaky-hard squashings (SPEC §4.6, S5/S6) |
 | `checkpoint.py` | orbax sharded save/resume of `TrainState` (adversary sources + moments included, no full-gather on the loop, SPEC S22) |
 | `eval.py` | in-loop eval pass: the six CE/KL masking variants + per-site CI-L0 in one jitted step, logged under the torch `EvalLoop` keys (`eval/ce_kl/*`, `eval/l0/*`) — enabled by the optional `eval:` config block |
+| `per_site_faith_eval.py` | per-site faithfulness observability: a jitted step returning per-site `‖Δ_s‖²_F` (the same per-site deltas `FaithfulnessLoss` reduces away; `Σ_s / Σ_s numel == FaithfulnessLoss`), plus `per_site_faith_scalars` (rank-indexed rel-Frobenius top-k). The LM eval logs `eval/faith/{rel_frob_top*,abs_frob_max,rel_frob_bar_chart}`. Observability only — weight-space error, NOT behavioral sensitivity |
 | `slow_eval.py` | LIBRARY for the in-loop slow (plot) tier (SPEC S28, in-loop only — no offline CLI): the `CIHistograms` / `ComponentActivationDensity` / `CIMeanPerComponent` reductions + renders, the config-gated `PermutedCIPlots` / `IdentityCIError` (off the `(T, C)` position CI), the `UVPlots` figure (`render_uv_figure` / `plot_uv_matrices`, shared by the LM in-loop naive-gather path and the toy `toy_uv_eval` cheap path), and the hidden-acts recon scalars. Torch-free numpy/matplotlib; logged under `slow_eval/figures/*` |
 | `run_state.py` | optimizer + initial-`TrainState` construction from an `ExperimentConfig` (orbax restores onto this reference) |
 | `tools/` | `convert_llama_simple_mlp_checkpoint.py` (torch venv) — one-off `.pt` → safetensors conversion of the pile pretrain checkpoint; `migrate_c49k_checkpoint.py` — one-off remap of the frozen C49k clone's orbax `TrainState` (legacy `components.{Vg..Ud}` `(1,*,*)` + flat `sources.<site>`) onto the current layout (site-keyed `components.vu`, `sources.<state_key>.<site>`) so a fine-tune can `restore_latest` it |

param_decomp/per_site_faith_eval.py

@@ -0,0 +1,63 @@
+"""Per-site faithfulness observability — the per-site Frobenius of the weight deltas.
+
+The global `FaithfulnessLoss` (`losses.faithfulness_loss`) is `Σ_s ‖Δ_s‖² / Σ_s numel`
+over the per-site deltas `Δ_s = W_s − V_s@U_s` — an aggregate that HIDES catastrophically
+broken small-norm sites: a handful of sites can carry nearly all of the unmasked-KL damage
+while the mean sits at ~1e-8 (lore `2026-07-03--unmasked-kl-vs-faith-no-eval-bug-4-culprit-sites`).
+This exposes the per-site breakdown so those sites are visible at eval cadence.
+
+Scope: relative Frobenius `‖Δ_s‖_F / ‖W_s‖_F` measures per-site WEIGHT-space error, which
+is NOT the same as behavioral sensitivity — the worst rel-Frobenius sites need not be the
+ones that dominate KL (in the run above they were behaviorally inert). This is pure
+observability, not a culprit detector; the KL-under-delta-ablation probe is what ranks
+behavioral sensitivity.
+"""
+
+import jax.numpy as jnp
+from jaxtyping import Array, Float
+
+from param_decomp.components import DecompVU
+from param_decomp.jit_util import filter_jit
+from param_decomp.lm import DecomposedModel
+
+
+def make_per_site_faith_step(
+    lm: DecomposedModel, compiler_options: dict[str, bool | int | str] | None = None
+):
+    """Build the `filter_jit`'d `per_site_faith(model, components) -> {site: ‖Δ_s‖²_F}`
+    (fp32 squared Frobenius per site). Reuses `model.weight_deltas` — the SAME per-site
+    `Δ_s = W_s − V_s@U_s` the global `FaithfulnessLoss` reduces — so
+    `Σ_s result[s] / Σ_s numel == FaithfulnessLoss` exactly.
+
+    `model` (the frozen-weight-bearing `DecomposedModel`) is the jit ARG — array leaves
+    traced, never baked (the HLO-baking rule). Each `Δ_s` is materialized then reduced to a
+    scalar, so XLA frees it after its reduction: peak memory is ~one delta, not the whole
+    model. Passing a zeroed `components` yields `‖W_s‖²_F` (`Δ_s = W_s − 0`) — the constant
+    denominator for relative Frobenius.
+    """
+    site_names = lm.site_names  # static, read off the closed-over model (HLO-baking rule)
+
+    def per_site_faith(model: DecomposedModel, components: DecompVU) -> dict[str, Float[Array, ""]]:
+        deltas = model.weight_deltas(components)
+        return {site: (deltas[site].astype(jnp.float32) ** 2).sum() for site in site_names}
+
+    return filter_jit(per_site_faith, compiler_options=compiler_options)
+
+
+PER_SITE_FAITH_TOP_K = 8
+
+
+def per_site_faith_scalars(
+    rel_frob: dict[str, float], abs_frob: dict[str, float], top_k: int
+) -> dict[str, float]:
+    """Stable-keyed scalar summaries of the per-site Frobenius breakdown: the `top_k`
+    LARGEST relative-Frobenius values (rank-indexed `1..k`, so the keys stay fixed as the
+    ranking reshuffles over training), plus the largest absolute `‖Δ_s‖_F`. Which SITE holds
+    each rank rides the bar chart, not these keys."""
+    ranked_rel = sorted(rel_frob.values(), reverse=True)
+    out = {
+        f"eval/faith/rel_frob_top{rank}": ranked_rel[rank - 1]
+        for rank in range(1, min(top_k, len(ranked_rel)) + 1)
+    }
+    out["eval/faith/abs_frob_max"] = max(abs_frob.values())
+    return out

param_decomp/tests/test_per_site_faith_eval.py

@@ -0,0 +1,72 @@
+"""CPU tests for per-site faithfulness observability.
+
+The load-bearing invariant: the per-site `‖Δ_s‖²_F` this step emits are the SAME per-site
+deltas the global `FaithfulnessLoss` reduces, so `Σ_s result / Σ_s numel == FaithfulnessLoss`
+exactly. Also checks the zeroed-vu `‖W_s‖²_F` denominator and the rank-indexed scalar helper.
+"""
+
+import jax
+import jax.numpy as jnp
+
+from param_decomp.components import init_decomp_vu
+from param_decomp.losses import faithfulness_loss
+from param_decomp.per_site_faith_eval import make_per_site_faith_step, per_site_faith_scalars
+from param_decomp.targets.glu_transformer import glu_site_specs, mlp_family_site_cs
+from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm
+
+
+def _tiny_lm_and_vu():
+    cfg = _tiny_cfg()
+    sites = glu_site_specs(cfg, mlp_family_site_cs(4, 5, 8))
+    lm = _tiny_decomposed_lm(cfg, sites, jax.random.PRNGKey(0))
+    vu = init_decomp_vu(sites, jax.random.PRNGKey(1))
+    return lm, vu
+
+
+def test_per_site_delta_sq_matches_weight_deltas():
+    lm, vu = _tiny_lm_and_vu()
+    delta_sq = make_per_site_faith_step(lm)(lm, vu)
+    deltas = lm.weight_deltas(vu)
+    assert set(delta_sq) == set(deltas)
+    for site, delta in deltas.items():
+        expected = (delta.astype(jnp.float32) ** 2).sum()
+        assert jnp.allclose(delta_sq[site], expected, rtol=1e-6)
+
+
+def test_sum_matches_global_faithfulness_loss():
+    lm, vu = _tiny_lm_and_vu()
+    delta_sq = make_per_site_faith_step(lm)(lm, vu)
+    deltas = lm.weight_deltas(vu)
+    total_numel = sum(d.size for d in deltas.values())
+    reconstructed = sum(float(v) for v in delta_sq.values()) / total_numel
+    assert jnp.allclose(reconstructed, faithfulness_loss(deltas), rtol=1e-6)
+
+
+def test_zeroed_vu_gives_frozen_weight_norms():
+    lm, vu = _tiny_lm_and_vu()
+    step = make_per_site_faith_step(lm)
+    zero_vu = jax.tree.map(jnp.zeros_like, vu)
+    w_sq = step(lm, zero_vu)
+    for site, w in lm.weight_deltas(zero_vu).items():  # Δ = W − 0 = W
+        assert jnp.allclose(w_sq[site], (w.astype(jnp.float32) ** 2).sum(), rtol=1e-6)
+        assert float(w_sq[site]) > 0.0
+
+
+def test_per_site_faith_scalars_ranks_and_keys():
+    rel_frob = {"a": 0.1, "b": 0.9, "c": 0.3, "d": 0.5}
+    abs_frob = {"a": 2.0, "b": 9.0, "c": 3.0, "d": 5.0}
+    out = per_site_faith_scalars(rel_frob, abs_frob, top_k=3)
+    assert out["eval/faith/rel_frob_top1"] == 0.9
+    assert out["eval/faith/rel_frob_top2"] == 0.5
+    assert out["eval/faith/rel_frob_top3"] == 0.3
+    assert "eval/faith/rel_frob_top4" not in out
+    assert out["eval/faith/abs_frob_max"] == 9.0
+
+
+def test_per_site_faith_scalars_top_k_exceeds_site_count():
+    rel_frob = {"a": 0.2, "b": 0.7}
+    abs_frob = {"a": 1.0, "b": 4.0}
+    out = per_site_faith_scalars(rel_frob, abs_frob, top_k=8)
+    assert out["eval/faith/rel_frob_top1"] == 0.7
+    assert out["eval/faith/rel_frob_top2"] == 0.2
+    assert "eval/faith/rel_frob_top3" not in out

param_decomp_lab/experiments/lm/run.py

@@ -70,6 +70,11 @@
 )
 from param_decomp.lm import DecomposedModel
 from param_decomp.log import setup_logger
+from param_decomp.per_site_faith_eval import (
+    PER_SITE_FAITH_TOP_K,
+    make_per_site_faith_step,
+    per_site_faith_scalars,
+)
 from param_decomp.run import (
     BackgroundRenderer,
     install_sigterm_flag,
@@ -402,8 +407,14 @@ def _make_lm_eval_fn(
     position_ci_step = make_position_ci_step(lm, co) if want_position_ci else None
 
     arithmetic_eval = _make_arithmetic_eval(eval, built.target, lm, mesh, n_proc, is_main, co)
+    site_names = lm.site_names
+    per_site_faith_step = make_per_site_faith_step(lm, co)
+    # ‖W_s‖²_F is constant over training — computed once (lazily, from the first eval's own
+    # components pytree so shape/sharding are exact) via the same step on a zeroed vu.
+    frozen_weight_sq: dict[str, float] | None = None
 
     def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
+        nonlocal frozen_weight_sq
         eval_pass_index = now_step // eval.every
         # uniform-average of per-batch scalars; mean-safe vs torch's accumulate-then-
         # compute() ONLY because every emitted key is a per-batch reduction that torch also
@@ -428,6 +439,19 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
         eval_record: dict[str, LogValue] = {
             f"eval/{k}": float(v) / eval.n_steps for k, v in metric_sums.items()
         }
+
+        # Per-site faithfulness observability: the per-site ‖Δ_s‖²_F the global
+        # FaithfulnessLoss reduces away. Collective (an all-reduce per site), so it runs in
+        # lockstep on every rank; the wandb bar chart is rank-0-only below.
+        delta_sq = {s: float(v) for s, v in per_site_faith_step(lm, state.components).items()}
+        if frozen_weight_sq is None:
+            zero_vu = jax.tree.map(jnp.zeros_like, state.components)
+            frozen_weight_sq = {s: float(v) for s, v in per_site_faith_step(lm, zero_vu).items()}
+            assert all(w > 0.0 for w in frozen_weight_sq.values()), "zero-norm frozen site"
+        abs_frob = {s: delta_sq[s] ** 0.5 for s in site_names}
+        rel_frob = {s: abs_frob[s] / frozen_weight_sq[s] ** 0.5 for s in site_names}
+        eval_record |= per_site_faith_scalars(rel_frob, abs_frob, PER_SITE_FAITH_TOP_K)
+
         for class_name, attn_step in attn_steps.items():
             # token-weighted (Σ sum_kl / Σ n), NOT the uniform per-batch average above — KL
             # is summed over distributions, divided by their count.
@@ -515,11 +539,27 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
                 "l0",
                 title=f"L0_{eval.l0_ci_alive_threshold}",
             )
+            # All-site per-site relative Frobenius, site-labeled — the rich per-site view the
+            # rank-indexed rel_frob_top* scalars can't carry (their keys are stable but drop
+            # site identity). Descending so the worst sites lead.
+            eval_record["eval/faith/rel_frob_bar_chart"] = wandb.plot.bar(
+                wandb.Table(
+                    columns=["site", "rel_frob"],
+                    data=[
+                        [s, v]
+                        for s, v in sorted(rel_frob.items(), key=lambda kv: kv[1], reverse=True)
+                    ],
+                ),
+                "site",
+                "rel_frob",
+                title="per-site ‖Δ‖_F / ‖W‖_F",
+            )
         if is_main:
             headline = {
                 k: eval_record[f"eval/{k}"]
                 for k in ("ce_kl/kl_ci_masked", "ce_kl/ce_difference_ci_masked")
             }
+            headline["faith/rel_frob_top1"] = eval_record["eval/faith/rel_frob_top1"]
             print(f"[eval @ {now_step}] {headline}", flush=True)
         return eval_record