Add kurtosis-profile validation test for JacobianLens (#1539 Tier-1) - #1616
Merged
jlarson4 merged 2 commits intoAug 7, 2026
Merged
Conversation
…sOrg#1539 Tier-1) Asserts the workspace-band signature as relative structure rather than absolute levels, per the cross-family measurement in TransformerLensOrg#1539: band rise vs the model's own early-third baseline, lens-specificity vs the logit-lens control through the identical code path, a gpt2-small negative control, and final-layer identity-transport agreement between arms.
SanjidMzi
force-pushed
the
test/jlens-kurtosis-profile
branch
from
August 6, 2026 21:46
0a4a86e to
8b7e61c
Compare
Collaborator
|
Looks good @SanjidMzi, thanks for taking this on |
jlarson4
added a commit
that referenced
this pull request
Aug 7, 2026
* Add Jacobian lens fitting guide (#1544) * Quantify Jacobian Lens causal swap success (#1545) Co-authored-by: Dreamer431 <113128214+Dreamer431@users.noreply.github.com> * test(integration): add oracle parity test for JacobianLens (#1539 Tier-1) (#1543) * test(integration): add oracle parity test for JacobianLens (#1539 Tier-1) Compares TransformerBridge JacobianLens.readout() against the reference anthropics/jacobian-lens oracle (pinned to 581d398) on google/gemma-2-2b-it across 75 layer x prompt cells (5 prompts x 15 sampled layers). Pass criteria per the #1539 spec (matching #1505 spike numbers): - Worst-case top-8 token overlap >= 7/8 in every cell - Spearman rank-correlation >= 0.95 on the top-64 logit union per cell The oracle is installed at test time via pip from the pinned commit so the threshold is reproducible independent of upstream drift. Reuses the bridge's original_model + tokenizer to avoid a second model copy in RAM. * style: apply black formatting (line-length=100) * fix: remove unused Dict, Tuple typing imports (pycln) * style: fix black formatting for py310 target (double blank lines) * test: use pytest.importorskip for oracle dep; add oracle-parity CI workflow Replace subprocess pip-install fixture with pytest.importorskip so the test skips gracefully in standard uv venvs (no pip present) and does not mutate the developer environment with no cleanup. Add .github/workflows/oracle-parity.yml — a dedicated workflow that reads ORACLE_COMMIT from the test file (single source of truth) and installs the oracle out-of-band before running the @pytest.mark.slow suite. Triggers on workflow_dispatch and on pushes that touch the test or workflow file, keeping oracle runs opt-in for PR checks. * Visual Encoders (ViT, DeiT) Support Rollout (#1546) * Add ViTArchitectureAdapter to supported architectures * Add ViTArchitectureAdapter to architecture factory * Add ViT and DeiT models to model registry * Add new model descriptions for Vision Transformers and Wav2Vec2 * Add ViTArchitectureAdapter for vision models Implement ViT/DeiT architecture adapter for model bridging. * Create vision_embedings.py * Add VisionClassifierHeadBridge for CLS token classification Implement VisionClassifierHeadBridge to handle CLS token slicing for classification. * Add visual model configuration to transformer bridge * Clarify pixel_values usage for multimodal and vision models Updated documentation for pixel_values parameter to clarify its use with vision models. * Update bridge.py * Update bridge.py * Rename vision_embedings.py to vision_embeddings.py * Define vision model and classification architectures Added vision model architectures and classification heads. * Add support for vision architectures in transformers * Refactor VisionClassifierHeadBridge to use pooled output Updated the VisionClassifierHeadBridge to directly use an already-pooled CLS token instead of slicing from the sequence output. Adjusted the forward method to reflect this change and improved error handling for the original component. * Update vit.py * Add unit tests for ViTArchitectureAdapter This file contains unit tests for the ViTArchitectureAdapter, covering component mapping, configuration flags, weight conversions, and model preparation methods. * Create test_vit_adapter.py * Update transformers.py * Update vit.py * Update vit.py * Update vit.py * Fix type hint for get_remote_component method * Fix type hint for get_remote_component method * Change import of torch to torch.nn in vit.py * Update vit.py * Re-add dummy 'mlp' attribute injection for ViTLayer Reintroduce a patch_layers function to inject a dummy 'mlp' attribute into ViTLayer blocks for MLPBridge compatibility. * Refactor ViTLayer handling by removing patch_layers Removed the patch_layers function and its call, which injected a dummy 'mlp' attribute into ViTLayer blocks. Updated comments for clarity regarding the MLPBridge container. * Add dummy 'mlp' attribute to ViTLayer blocks Inject a dummy 'mlp' attribute into ViTLayer blocks to satisfy hasattr check for TransformerLens. * Update vit.py * Remove TestViTConfigNCtx and related test case Removed deprecated TestViTConfigNCtx class and its test case for n_ctx. * Enhance ViTLayer with MLP wrapper and fix forward method Added a non-circular MLP wrapper to ViTLayer blocks and fixed tuple-chaining bug in forward method. * Refactor ViT layer forward pass handling Refactor forward pass handling for ViT layers to safely unpack tuple outputs and ensure compatibility with the model's internal structure. * Refactor ViTLayer forward pass handling Refactor forward pass handling for ViTLayer to fix tuple-chaining bug and ensure compatibility with HF model outputs. * Fix tuple handling in ViTLayer forward method Modified the forward method to handle tuple inputs and outputs for ViTLayer, ensuring compatibility with Tensor expectations. * Reorder model prefix checks for better clarity * Update vit.py * Detect model class name in prepare_model method Added detection for model class name in prepare_model method. * Simplify prefix determination for ViT models Refactor model prefix detection logic for ViT and DeiT models. * Implement fixture for distilled DeiT model testing Added a fixture to load the distilled DeiT model for testing. * Update DeiT bridge tests for bare model handling Refactor tests for DeiT bridge to accommodate bare model behavior and update assertions accordingly. * Set architecture in Hugging Face model configuration * Support DeiTLayer in patch_layers function * sort * Replace direct attribute assignment with setattr * black fix * fix formatting after merge * Update vit.py * Update ViT adapter test paths for consistency * Remove redundant test for n_ctx in prepare_loading Removed test for prepare_loading not affecting n_ctx. * black sorted * black reorder * Refactor vit_bridge and vit_bare_bridge fixtures * temp support up to transformers 5.8.0 * support transformers 5.13.0 * format fixed. Unit test all passed. Intergration test all passed. should be good to go * Update vit.py * Clarify tokenizer support in ViTArchitectureAdapter Added comment to clarify the lack of tokenizer support for vision models. * Remove head_dim assignment from hf_config Removed unused head_dim assignment from hf_config. * Update vit.py * Update bridge.py * Add VisionEmbeddingsBridge and VisionClassifierHeadBridge * Update test_vit_adapter.py * Update test_vit_adapter.py * Update test_vit_adapter.py * Update test_vit_adapter.py * Improve compatibility mode error and output handling Updated error message for clarity and added handling for last_hidden_state in output. * Refactor test to check output type and shape Update test to verify that the forward method returns a tensor instead of a raw HF output object. Adjust assertions to match the expected behavior after changes in bridge.py. * Update bridge.py * Update test_vit_adapter.py * Remove obsolete tests from TestViTPrepareLoading Removed deprecated tests for prepare_loading() in TestViTPrepareLoading. * Update test_vit_adapter.py * Update test_vit_adapter.py * formatted --------- Co-authored-by: Jonah Larson <jonahalarson@comcast.net> * Add Starcoder2 architecture adapter (#1533) Co-authored-by: jlarson4 <jonahalarson@comcast.net> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> * Docs: Jlens Qwen3.5-4b Demo (#1547) * Run experiments with Qwen-3.5 architecture support * Run notebook cells * Add lfm2 tiny integration test (#1552) * Add lfm2 tiny integration test * Fix formatting * Add AST (Audio Spectrogram Transformer) Adapter (#1484) * chore: save WIP on V3 transformerbridge migration re add ast import and add to factory :wq y wq :wq * feat(ast): migrate AST adapter to V3 TransformerBridge and component_mapping * refactor(ast): resolve PR feedback for docstrings, prefix-awareness, unit tests * test(ast): split parity to integration folder, add load_weights boot test, fix audio classification load path, and add to ARCHITECTURE_DESCRIPTIONS * fix(ast): union audio classification sets, specific boot test assertions and two comment typo fixes * Verificaiton for lapa (#1556) * ViT and AST model verification (#1582) * verified a few models for ViT and AST * improved vision testing for ViT models * fix(bridge): return W_in/W_out/W_gate in TL orientation for nn.Linear-backed models (#1558) * feat: respect prepend_bos and add return_input_tokens flag * fix torch orientation bug * review changes * review changes * pipeline fix * deprecate remaining hooked entry points (#1592) * deprecate remaining hooked entry points * test: account for hooked transformer warning in notebook * fix: address deprecation warning review feedback * fix: correct encoder deprecation warning stacklevel * support loading fit checkpoints in JacobianLens.load() (#1574) * support loading fit checkpoints in JacobianLens.load() * fix: black formatting and update conflicting test for checkpoint load * address jlarson4 review: preserve target_layer, add reference fixture, fix tuned-lens note - Remove "target_layer" from _FIT_RESERVED_KEYS so it survives checkpoint conversion and validate_model() can refuse non-final-target lenses - Add test_load_checkpoint_mirrors_fit_payload_schema: fixture matches the exact keys fit() produces so format drift causes a test failure - Fix tuned-lens note: it is the Jacobian artifact format that has no bias slot, not the tuned-lens format; tuned-lens translators are affine (weight + bias) * address jlarson4 review: read n_done key in _from_checkpoint_payload Real checkpoint writers (reference package) store the prompt count as n_done, not n_prompts. Prefer n_done with n_prompts as fallback so genuine checkpoints are not rejected with n_prompts=0. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NiwNUm3YFj9yAuSBuGDnd8 * address jlarson4 review: align checkpoint loader with reference write_checkpoint() schema - _from_checkpoint_payload: read n_done first (real checkpoints use n_done, not n_prompts) - _from_checkpoint_payload: infer d_model from jacobian_sum matrix shape (real checkpoints have no d_model key) - _from_checkpoint_payload: harvest top-level target_layer into metadata (reference format stores it at top level, not nested) - _from_checkpoint_payload: guard empty jacobian_sum with a clear ValueError before attempting shape derivation - load() docstring: update Fit checkpoint schema to reflect the real 6-key reference format - tests: replace test_load_checkpoint_with_zero_n_prompts_raises with two tests - tests: rewrite test_load_checkpoint_mirrors_fit_payload_schema to use verbatim 6-key reference payload - tests: add test_load_checkpoint_harvests_flat_provenance_and_strips_fit_keys - docs: update schema table — replace n_prompts/d_model with real 6-key format - docs: note d_model inferred from matrix shape - docs: document target_layer as deliberate exception to fit-key stripping --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> * Drop optional unused aliases on hybrid architectures (#1579) * Drop optional unused aliases * feat: assign fallbacks on pruned * Improve unittests * Add kurtosis-profile validation test for JacobianLens (#1539 Tier-1) (#1616) * Add kurtosis-profile validation test for JacobianLens (#1539 Tier-1) Asserts the workspace-band signature as relative structure rather than absolute levels, per the cross-family measurement in #1539: band rise vs the model's own early-third baseline, lens-specificity vs the logit-lens control through the identical code path, a gpt2-small negative control, and final-layer identity-transport agreement between arms. * Fix gpt oss olmo3 parity (#1621) * Resolution for issue 1619 * Updated for 1620 * Add verification script * Fixed 1619 on HF * Verification script repair * fixing per-layer olmo * cleanup --------- Co-authored-by: abhi <abhinavbellapu@berkeley.edu> Co-authored-by: emerardd <113128214+emerardd@users.noreply.github.com> Co-authored-by: Dreamer431 <113128214+Dreamer431@users.noreply.github.com> Co-authored-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: Jiankun Wei <72998341+david-wei-01001@users.noreply.github.com> Co-authored-by: SanjidMzi <56235075+SanjidMzi@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Kyle Yin Xu <123780557+KYinXu@users.noreply.github.com> Co-authored-by: Syed Adil Ahmed <tensorcruncher@gmail.com> Co-authored-by: Dylan <159935143+dylanberens@users.noreply.github.com> Co-authored-by: Md.Sadiq <mohammadsadiq4950@gmail.com> Co-authored-by: msaule <sau24006@byui.edu> Co-authored-by: Priyanka Bajaj <42418272+priyanka25aug@users.noreply.github.com>
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.
Description
Adds the Tier-1 kurtosis-profile check from #1539. The §4 signature is asserted as relative structure rather than absolute levels, following the cross-family measurement posted in that thread: absolute early-third kurtosis ranges from 0.2 to ~140 across published artifacts, so a fixed "≈0 early" assert only holds for one family.
Four checks, on two pinned artifacts:
Prompts come from the wikitext-103 validation split; the published lenses were fitted on train, so the readout corpus is held out. Lens artifacts are pinned to the same neuronpedia revision as
test_jacobian_lens_oracle_parity.py. Marked slow: downloads Qwen3.5-0.8B (~1.6 GB), runs on CPU in a few minutes. QUARANTINES.md row added.Part of #1539 (Tier 1, kurtosis profile).
Type of change
Checklist: