[Closed: Duplicate of #1596] feat(jacobian_lens): J-space sparse decomposition - #1597
Closed
janmenjayap wants to merge 16 commits into
Closed
[Closed: Duplicate of #1596] feat(jacobian_lens): J-space sparse decomposition#1597janmenjayap wants to merge 16 commits into
janmenjayap wants to merge 16 commits into
Conversation
Co-authored-by: Dreamer431 <113128214+Dreamer431@users.noreply.github.com>
…erLensOrg#1539 Tier-1) (TransformerLensOrg#1543) * test(integration): add oracle parity test for JacobianLens (TransformerLensOrg#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 TransformerLensOrg#1539 spec (matching TransformerLensOrg#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.
* 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>
Co-authored-by: jlarson4 <jonahalarson@comcast.net> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
* Run experiments with Qwen-3.5 architecture support * Run notebook cells
* Add lfm2 tiny integration test * Fix formatting
) * 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
* verified a few models for ViT and AST * improved vision testing for ViT models
- Add `get_sparse_decomposition` to decompose an activation into a k-sparse nonnegative combination of J-lens vectors (Gurnee et al., 2026). - Support `nonnegative_orthogonal_matching_pursuit` (default, exact NNLS re-solve) and `gradient_pursuit` algorithms. - Return both the nonnegative coordinates and the orthogonal-projection J-space component. - Distinguish the projection from the coefficient reconstruction; the projection residual matches `swap_hooks`. - Keep the implementation model-free by operating directly on the raw dictionary tensor. - Add tests covering both algorithms, exact-resolve NNLS correctness, a brute-force optimum oracle, and input validation. Part of TransformerLensOrg#1539 (Tier 2).
- Add `JacobianLens.lens_vector_dictionary(model, layer)` returning the `[d_vocab, d_model]` dictionary whose rows are the J-lens vectors `v_t = J[layer]^T W_U[:, t]`. - Cache the dictionary per (layer, device) and release it in `clear_device_cache`, so a sparse decomposition can reuse it; document its vocabulary-sized memory cost. - Add tests asserting the dictionary matches `lens_vectors` over every token, is cached and invalidated by `clear_device_cache`, and rejects an unfitted layer.
- Add `JacobianLens.decompose(model, activation_or_prompt, layer, *, position, k, algorithm)` decomposing either a raw activation vector or the `blocks.{layer}.hook_out` activation at a prompt position, validating inputs before building the dictionary.
- Build and cache the layer dictionary via `lens_vector_dictionary` and solve with `get_sparse_decomposition`.
- Export `JSpaceDecomposition` and `get_sparse_decomposition` from `transformer_lens.tools.analysis`.
- Add end-to-end tests for the raw-activation and prompt paths, the algorithm passthrough, and the input-validation error paths.
- Add a GPT-2 integration test (regular CI): `decompose` on a real `blocks.6.hook_out` activation returns k nonnegative atoms, the non-J-space residual is orthogonal to every selected J-lens vector, and the J-space component plus residual recover the activation. - Add a slow gemma-2-2b-it integration test validating `decompose` on the published lens artifact: support size, nonnegative coordinates, in-vocabulary token ids, and component-plus-residual reconstruction. - Document J-space sparse decomposition in `jacobian_lens_fitting.md`: the `decompose` API, local coordinates versus the orthogonal-projection J-space component, and the paper's variance facts with closed-model caveats.
Add a References section to the decomposition module docstring: Gurnee et al. (2026) for the J-space method, Pati et al. (1993) for the greedy orthogonal-matching-pursuit selection, Blumensath & Davies (2008) for the gradient-pursuit update, and Lawson & Hanson (1974) for the active-set nonnegative least-squares re-solve.
janmenjayap
changed the title
janmenjayap
changed the title
janmenjayap
changed the title
janmenjayap
changed the title
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
10 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.
Closing as a duplicate of PR.