Skip to content

Refine initialization, objectives, and evaluation - #1001

Merged
danbraunai-goodfire merged 1 commit into
mainfrom
bridge/task-public-release-2026-08-31
Sep 4, 2026
Merged

Refine initialization, objectives, and evaluation#1001
danbraunai-goodfire merged 1 commit into
mainfrom
bridge/task-public-release-2026-08-31

Conversation

@claude-spd1

@claude-spd1 claude-spd1 commented Sep 4, 2026 •
edited by danbraunai-goodfire
Loading

Copy link
Copy Markdown
Collaborator

Description

  • Generalizes neuron-aligned initialization across supported targets.
  • Consolidates importance minimality around the smooth-L0 ImportanceMinimalityLoss and updates canonical configs and guidance.
  • Routes evaluation through shared context, including correct query/KV-head usage accounting, and removes the standalone hidden-activation evaluator.
  • Makes canonical train/eval data sources explicitly disjoint and reorganizes the test suite around the public package structure.

Related Issue

N/A.

Motivation and Context

These changes keep initialization, minimality, and evaluation behavior consistent as the library grows, while making the supported configuration and test surfaces easier to follow.

How Has This Been Tested?

  • uv lock --check
  • uv run ruff format --check
  • uv run ruff check
  • uv run basedpyright
  • Python 3.12 import smoke
  • llama-golden and multi-device test shards
  • Full core/targets shard on the exact PR commit: 618 passed, 61 skipped, 5 expected xfails, 0 failures/errors

Does this PR introduce a breaking change?

Yes. Legacy Lp importance-minimality configuration is replaced by the smooth-L0 ImportanceMinimalityLoss fields, and the standalone hidden-activation evaluator is removed in favor of shared evaluation context. Existing callers using those retired paths must migrate.

Diff breakdown

  • library/runtime: +1,555 -2,214
  • tests/fixtures: +1,105 -476
  • docs: +16 -15
  • project/config: +22 -13

@danbraunai-goodfire
danbraunai-goodfire merged commit 964d8ba into main Sep 4, 2026
5 checks passed
Antovigo added a commit to Antovigo/spd that referenced this pull request Sep 5, 2026
@Antovigo @claude
Merge upstream/main (goodfire-ai#1001) — one init, aligned on every s…
8b38226 
…ite kind

Upstream 964d8ba "Refine initialization, objectives, and evaluation" (goodfire-ai#1001)
into feature/dual_obj_jax. 38 conflicts; the local dual objective, tPD delta
pinning, AB grid, two-stream eval namespace and frozen-stack split all survive.

Resolutions worth naming:

* INITS. goodfire-ai#1001 introduced `ComponentInitializer`, a target-aware V/U seeding seam,
  and its own data-free `neuron_aligned` init. This branch's `pd.weight_init`
  enum is retired wholesale: `default`/`coupled`/`zero_u` are gone (none beat the
  alternatives) and `neuron_aligned_targeted` is re-expressed as a
  `ComponentInitializer`. The config surface moves from `pd.weight_init` to
  `decomposition.sites.initialization` + `.neuron_ranks`, so the arms live where
  the sites they seed are authored.

* NEURON-ALIGNED TARGETED now aligns EVERY decomposed site, attention included —
  no `zero_u` fallback, so `validate_neuron_alignment` demands an entry per site.
  The harvest produces five rankings per block instead of one: mlp neurons by
  write energy, q and k channels by their own projection-output energy (attention
  mixes across positions, never across channels, so q_j and k_j meet only inside
  their head's dot product and rank independently), o channels by
  E[z^2]*||W_o[:,j]||^2, and v channels by that summed over the GQA group — v
  channel (g,d) IS attention channel (h,d) for every h in group g. Statistic
  renamed `unit_energy`; the read side refuses the old `write_energy` artifacts
  rather than misaligning silently. The per-coordinate factorization is now
  `glu_transformer.selected_unit_factors`, shared with the data-free init.

* IMP-MIN. goodfire-ai#1001 deleted the L_p penalty and gave its NAME to smooth-L0, so
  `SmoothL0ImportanceMinimalityLoss` -> `ImportanceMinimalityLoss` everywhere.
  The local `normalize_at_one` rescale survives as a defaulted keyword. The tPD
  TMS goldens are re-recorded: that fixture's `pnorm=1.0` became `gamma=1.0`, a
  different per-value penalty, and the merged step now reproduces upstream's own
  goldens to 12 significant digits under this branch's non-target key spelling.

* EVAL. Upstream's shared per-batch context (one clean forward + CI envelope per
  batch, folded by `BatchedOperation`) replaces per-operation forwards, and this
  branch's per-stream/per-role fan-out rides on top: `LMBatchContext` carries its
  `Stream`, holds the whole `AnyCI` envelope, and each operation folds only the
  stream it measures and reads its own head via `ci_for(role)`. Delta pinning
  (SPEC T4) is preserved on all three non-target recon families and now has a
  BINDING-level guard, not only kernel-level ones — an unpinned non-target probe
  measures a delta->0 attack no component-side training can defend, which reads
  as a decomposition failure that is really a measurement bug.

* Also adopted: GQA-aware nonlinearity locality (QueryHeads/KVHeads, soft USE
  counts), the padded-CI attention API, and the whole-tree test relocation.
  Dropped with upstream: `hidden_acts_eval.py` and both hidden-acts eval metrics.

1317 tests pass. The one failure, `test_no_checkpointing::test_sigterm_exits_
without_saving`, is its own 300s deadline expiring under 16-way xdist contention;
it passes in 41s alone. ruff + basedpyright clean.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Antovigo added a commit to Antovigo/spd that referenced this pull request Sep 6, 2026
@Antovigo @claude
docs: init robustness on the current neuron-aligned run
677b73b 
Replaces p-222d379b (addsub-L18-18, pre-goodfire-ai#1001 pd.weight_init schema) with
p-6540dfdd (addsub-L18-23, post-goodfire-ai#1001), which is the current aligned init and
covers every site kind including attention.

- general text: aligned is ~22% easier to attack at 100 steps (0.0127 vs 0.0104),
  more than twice either run's own adversary-start spread - the older run read 12%
- task distribution: aligned is ~10% HARDER to attack (0.0047 vs 0.0052), a new
  direction the older pairing could not resolve (its spread reached 24%)
- both inits still saturate by ~40 steps, so the nonlinearity penalty's runaway
  remains specific to that penalty

The two runs straddle goodfire-ai#1001 and no single build opens both; verified before
comparing that the probe computation is identical across the merge (core recon/
adversary/masking/recon byte-identical, eval-side change a pure refactor) and
that probe inputs match, so the curves are comparable. Recorded in the caveats.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011EqWYtK2A5PugNTcQx2iAE
Antovigo added a commit to Antovigo/spd that referenced this pull request Sep 11, 2026
@Antovigo @claude
feat(init): port zero_u onto the ComponentInitializer seam
6ec3026 
Antoine wants to test `zero_u` again on the merged branch. The arm was retired with
`pd.weight_init` in the goodfire-ai#1001 merge (8b38226); this brings it back in the shape goodfire-ai#1001
introduced — a plain `ComponentInitializer`, the same seam `random` and the two aligned
inits already use — rather than reviving the enum branch in `run_state`.

`zero_u` is the target-coupled seed's `V` with `U` zeroed: the component sum is exactly
zero at init and the delta carries all of `W`, so a subcomponent acquires norm only as the
reconstruction losses demand it instead of starting with `W`-scale weight a mask adversary
could switch on. `V` still feeds the CI nets a live signal, and `U` has a nonzero gradient
from step 0 (`V`'s is zero until `U` moves off zero).

THE PROPERTY THAT MAKES THE TEST WORTH RUNNING, and the one most likely to rot silently:
the per-site keys are split exactly as the retired arm split them —
`split(key, len(sites))` indexed by site position — and `V` is `_coupled_site_vu`'s `V`.
So a run here reproduces the pre-goodfire-ai#1001 zero_u `V` bit for bit at one seed, which is what
makes p-88665048 a control rather than a different experiment.
`test_v_is_the_coupled_seed_with_the_pre_1001_key_discipline` pins it.

Shape of the port:
- core/components.py restores `zero_component_stacks` and `_coupled_site_vu` (both deleted
  by the merge) and adds `init_component_stacks_zero_u`. `coupled` itself is NOT restored:
  zero_u is built directly as coupled-V-plus-zero-U, so there is no dead public API for an
  arm nobody can select.
- core/init_placed.py adds `zero_u_component_initializer` plus `_site_weights_in_graph`,
  which reads each site's frozen `W` INSIDE the init graph through the zero-stacks
  `weight_deltas` identity, so no full-precision copy crosses the jit.
- Domain-neutral by construction: `W` is reached only via `weight_deltas`, which every
  `DecomposedModel` implements, so the toys can select it as readily as the LM.
- Config: `zero_u` joins `ComponentInitialization`, dispatches in
  `load_run.component_initializer_for`, and the capacity validator short-circuits on
  `random | zero_u` — the aligned arms' `C <= coordinate count` bound does not apply.

1007 tests pass (9 new); ruff and basedpyright clean.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@danbraunai-goodfire danbraunai-goodfire danbraunai-goodfire approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@claude-spd1 @danbraunai-goodfire