feat(eval): per-site faithfulness observability (rel/abs Frobenius) - #993
Open
claude-spd1 wants to merge 2 commits into
Open
feat(eval): per-site faithfulness observability (rel/abs Frobenius)#993claude-spd1 wants to merge 2 commits into
claude-spd1 wants to merge 2 commits into
Conversation
The global FaithfulnessLoss (Σ_s‖Δ_s‖²/Σ_s numel) hides catastrophically
broken small-norm sites (lore 2026-07-03--unmasked-kl-vs-faith-no-eval-bug-4-culprit-sites:
4 of 224 sites carried ~all of the 6.6 unmasked KL while faith sat at ~1e-8).
Add per_site_faith_eval.make_per_site_faith_step: a jitted (model, components) ->
{site: ‖Δ_s‖²_F} reusing model.weight_deltas — the SAME per-site deltas the global
FaithfulnessLoss reduces, so Σ_s/Σ_s numel == FaithfulnessLoss exactly. ‖W_s‖²_F (the
constant relative denominator) comes from the same step on a zeroed vu (Δ=W−0), cached
once in the LM eval_fn closure.
The LM in-loop eval logs, at eval cadence (always-on, cheap — one faith-loss-sized
materialization per eval, which already runs every train step):
- eval/faith/rel_frob_top{1..8} — 8 largest rel-Frobenius, RANK-indexed (stable line keys)
- eval/faith/abs_frob_max
- eval/faith/rel_frob_bar_chart — all-site, site-labeled (rank 0)
Observability only — relative Frobenius is per-site WEIGHT-space error, NOT behavioral
sensitivity (the culprit sites above were behaviorally inert on rel-Frobenius). The
KL-under-delta-ablation probe is the diagnostic that ranks behavioral sensitivity;
deferred pending a compute/design decision.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…transformer split) llama_site_specs moved/renamed to glu_site_specs and mlp_family_site_cs moved to targets.glu_transformer since this branch's original base — the test's imports were stale after rebasing onto current feature/jax. Crew-Address: agent/6pmp
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Reopens #934 (closed unmerged in a bulk stale-PR sweep on 2026-07-24, no comments/review) rebased onto current
feature/jaxtip. Same content, no behavior changes from the original.Per-site faithfulness observability: turns the opaque aggregate
FaithfulnessLossinto per-site signal.param_decomp/per_site_faith_eval.py:make_per_site_faith_step(lm, co)→ jitted(model, components) -> {site: ‖Δ_s‖²_F}, reusingmodel.weight_deltas— the SAME deltas the globalFaithfulnessLossreduces (Σ_s / Σ_s numel == FaithfulnessLoss, asserted in tests).‖W_s‖²_Fvia the same step on a zeroedvu(Δ=W−0), cached once per run._make_lm_eval_fn). Always-on at eval cadence: logseval/faith/rel_frob_top{1..8}(rank-indexed, stable line keys; top1==max),eval/faith/abs_frob_max, andeval/faith/rel_frob_bar_chart(rank-0, all-site, site-labeled — avoids a 224-key wandb explosion).make checkclean.Context
Follow-up from the investigation that found global per-element-MSE faith hid 4 sites carrying ~99% of unmasked KL. This PR is genuine weight-space observability but — per that investigation — does NOT by itself rank the behavioral culprits (worst rel-Frob sites were inert). It's a foundation; the actual culprit-ranking probe (a per-site delta-ablation/KL-sensitivity measure) is still pending a go/no-go from Oli — see bridge task #55 for that open decision.
Testing
ruff check/ruff format --checkclean on the touched files. Could not execute the JAX test suite in this environment (no AVX support on this node — jaxlib refuses to load); relying on CI (the original #934 run was CI-green on the identical diff pre-rebase; this rebase only resolved a two-line merge conflict against_make_arithmetic_evalinrun.py, no logic change).Crew-Address: agent/6pmp