Repository navigation
Conversation
…on tensors (#1774) * feat(sparse_probing): add leakage-safe core k-sparse probe fit Add fit_sparse_probe: a stratified train/test split is computed before any learned statistic, features are selected by train-only mean difference, and a CPU-float64 LBFGS logistic fit reports explicit objective/gradient-norm convergence diagnostics instead of assuming success. Report held-out accuracy, precision, recall, and F1 so callers can judge decodability without the result implying causal model use, neuron monosemanticity, or superposition. Cover exact score/index selection, leakage isolation, deterministic ties/RNG, optimizer convergence and failure, constant columns, and invalid-input rejection. Sweep, controls, exports, and docs land in a follow-up commit. * feat(sparse_probing): add k-sweep, null controls, exports, and docs Add sweep_sparse_probe: fits a shared-split k-grid and reports raw random-coordinate and shuffled-training-label control distributions alongside each k, so callers can judge coordinate concentration without the result implying an automatic significance test. Export fit_sparse_probe, sweep_sparse_probe, and the result dataclasses from transformer_lens.tools.analysis, and add the sparse_probing guide to the docs toctree, documenting the leakage-safe contract and a run_with_cache composition example. * test(sparse_probing): cover MPS-to-CPU transfer before float64 fit fit_sparse_probe accepts feature matrices on any device, including MPS, where float64 is unsupported. Add a regression test that fits a probe over MPS-resident features and asserts the selected train/test tensors land on CPU with float64 dtype instead of raising when the float64 cast is attempted while still on the Metal device. * fix(sparse_probing): accept jaxtyping>=0.3 type-check exception and dtype policy Rebasing onto origin/dev pulled in jaxtyping>=0.3 (#1732), which re-raises type-check violations as jaxtyping.TypeCheckError instead of letting BeartypeCallHintParamViolation propagate, and now classifies float8_e4m3fn as a valid Float dtype at the annotation level. Use the project's tests/typecheck_errors.TYPECHECK_ERRORS convention for the annotation-level cases, and move the float8_e4m3fn case to the explicit ValueError dtype-rejection table, where the function's own dtype guard now catches it. * fix(sparse_probing): make gradient acceptance scale-relative * fix(sparse_probing): stamp the realized k on each sweep result * test(sparse_probing): pin metrics to the held-out rows and control labels * docs(sparse_probing): correct result-field list and add grouped-split caveat * test(sparse_probing): cover the single-example-per-class rejection * fix(sparse_probing): key gradient acceptance to feature magnitude, not initial correlation The fit accepted when the final gradient inf-norm stayed under gradient_tolerance * max(1, initial gradient inf-norm), where the initial gradient at the zero start measures label-feature correlation rather than activation scale. A real probe support got a generous bound while random-subset and label-shuffle controls, uncorrelated by construction, were pinned to the bare absolute floor and could reject a genuinely converged fit at activation scale, discarding the whole sweep. Key the bound to the magnitude of the selected training columns instead: gradient_tolerance * max(1, k * max abs(selected training features)), computed from the same matrix used for the main fit and both control arms. LBFGS's own tolerance_grad is left unchanged. Docstrings and the guide describe the new bound. The large-scale guard test now scales activations into the thousands so the achieved gradient lands strictly between the absolute floor and the relative bound, and asserts both sides independently of the production formula. A new test drives the controls through sweep_sparse_probe at the same scale to cover the case the old bound broke. * fix(sparse_probing): cap LBFGS evaluations explicitly and report the stop reason torch.optim.LBFGS was constructed without max_eval, silently falling back to torch's undocumented max_iter * 5 // 4 default. Combined with a tolerance_grad that activation-scale fits may never reach, a fit can run to either cap and return a truncated, off-optimum solve that still clears the scale-relative acceptance check, with nothing on the result letting a caller tell it apart from a fit that actually met its gradient tolerance. Set max_eval explicitly as a named local and derive a stop_reason from the realized iteration count, evaluation count, and gradient norm, in priority order so a met tolerance outranks a simultaneously hit cap and max_iter outranks max_eval. Thread stop_reason through _FitOutcome, _fit_result, and SparseProbeResult, and document the new field in the guide. Add a test that distinguishes a converged (tolerance_grad) fit from iteration-capped (max_iter) and evaluation-capped (max_eval) truncations. * fix(sparse_probing): reject gradient_tolerance at the top of its range _validate_inputs accepted gradient_tolerance == 1, which is passed straight to LBFGS as tolerance_grad. Since the acceptance threshold has a max(1.0, ...) floor, a tolerance_grad of 1.0 lets LBFGS stop before the first step whenever the initial gradient is already <= 1, returning an all-zero-coefficient untrained probe as a successful fit. Reject the whole closed top of the range: the valid interval is (0, 1), not (0, 1]. Update the validation message and the guide doc to match, and add the newly-rejected boundary value to the invalid-input parametrization. * test(sparse_probing): assert class_weight changes the fit against a weighted Newton reference The previous test fit a balanced fixture in both class_weight="balanced" and class_weight=None modes and only checked metadata round-tripping. On a balanced fixture every class weight is 1.0 in both modes, so the two fits are numerically identical: a bug that ignores class_weight=None and always balances, or one that forces all-ones weights, both leave that test green. Fit an imbalanced (~1:4) fixture in each mode at a shared seed so the split and selected support are identical, then assert the coefficient vectors differ and each matches an independent Newton reference computed with the matching weighting. Generalize the reference helpers to take the weighting mode explicitly instead of hardcoding the balanced formula. * test(sparse_probing): guard the standardized held-out recompute against leakage The held-out recompute ran only under the default preprocess="none", where preprocess_mean/preprocess_scale are the identity, so it never standardized anything and could not catch a leakage bug on the preprocess="standardize" path, such as standardizing held-out rows with their own mean and std instead of the returned train-only statistics. Parametrize the test over preprocess in ("none", "standardize") and apply the returned train-only preprocess_mean/preprocess_scale to the raw held-out rows before scoring. Under "none" this is a no-op; under "standardize" it pins the recompute to the train statistics the module reports, so a held-out-own-stats leakage bug would diverge from result.metrics and fail the equality asserts.
* Fixing issue with model_structure/compatibility_mode docs lists of aliases * Additional doc cleanups for hook semantics
* fix PositionEmbeddingsAttentionBridge * pipeline fix
* docs: align analysis tools guide with v4 capabilities * docs: clarify analysis tool architecture limits
* fix removed names from import * pipeline fix
…l patch gate (#1775) * fix(svd_circuits): correct the OV output-direction docstring and add a cross-check test HeadSVD's docstring had U and V swapped for the OV map: for a factored map A @ B, U's columns live in A's input space and V's columns live in B's output space, so for OV (A = W_V_h, B = W_O_h) it is V, not U, whose columns are the residual-stream directions this head writes into and should be projected through W_U for a vocab or logit readout. The swap never raised a shape error because both spaces are d_model-dimensional, so it would have silently used the wrong basis in every downstream readout, projection, or patch built on top of it. Adds a regression test that cross-checks HeadSVD.OV.V against the already-shipped SVDInterpreter (which projects its OV singular vectors through the unembedding the other way), on a tiny no-download TransformerBridge, so the assertion cannot pass under either labeling by construction. * feat(svd_circuits): add OV vocab readout and logit signature vocab_readout projects a head's top-k OV output directions (HeadSVD.V, per the corrected convention) through the unembedding, gated on TransformerBridge compatibility mode so LayerNorm stays folded into W_U. logit_signature reconstructs one direction's rank-1 OV output and reads its signed logit effect on given tokens; it requires the direction to be isolated first, since a rotation-ambiguous or null direction's signature is not attributable to it alone. * feat(svd_circuits): add project_activations firing coefficients Add ActivationProjection and project_activations to recover per-position firing coefficients by projecting a head's actual OV output (hook_result) onto its V basis. Summing the coefficients against V reconstructs the head's output because V's columns are orthonormal and define the basis the head writes in. Restore the model's prior use_attn_result setting after the forward pass so read-only analysis does not leave a configuration side effect. * feat(svd_circuits): add patch_along_directions causal gate patch_along_directions causally validates a claimed OV subfunction by reconstructing a head's hook_result onto a chosen span of HeadSVD.V directions (keep) or its complement (ablate), then comparing a caller's metric before and after against an equally-sized random-subspace baseline. gated is true only when the requested subspace moves the metric by more than the random baseline does, so a claim cannot be called causal merely because patching moved the metric at all. Before touching the model, the retained set is checked against every degenerate block reported by HeadSVD.degenerate_blocks(): a block must be kept whole or dropped whole, since attributing an effect to part of a rotation-ambiguous or null block would let a caller hand-pick around the same guard require_isolated already enforces per direction. The forward passes run under torch.no_grad(), matching the other analysis tools that repeat a forward pass per call, and the model's prior use_attn_result setting is restored in a finally block so this read-only check leaves no configuration side effect. * feat(svd_circuits): export public API and validate Bridge compatibility mode Adds decompose_head, project_activations, patch_along_directions, vocab_readout, logit_signature, and their supporting types (HeadSVD, HeadDecomposition, RankReportRow, ActivationProjection, LogitSignature, PatchResult, DegenerateDirectionError) to transformer_lens.tools.analysis's imports, __all__, and tool-listing docstring, in the same ASCII-sorted order the module already uses for its other tools. No compatibility-mode logic changes here: vocab_readout and logit_signature already gate on TransformerBridge compatibility mode, and patch_along_directions already restores use_attn_result. This commit only widens the public surface, so the causal gate ships alongside the readouts rather than either landing without the other. * test(svd_circuits): add GPT-2 small IOI patch integration Exercise vocab_readout and patch_along_directions against GPT-2 small's layer 9 head 9 (name-mover head) on an IOI-style prompt with a Mary/John logit-diff metric, checking that the readout and causal gate agree on a real Bridge with compatibility mode enabled. * docs(svd_circuits): drop HookedTransformer references from the module HookedTransformer is being removed, so the four consumer docstrings (vocab_readout, logit_signature, project_activations, patch_along_directions) and _validate_bridge_compatibility no longer list it as an accepted model or cite its LayerNorm folding as the reason the compatibility-mode guard fires only for TransformerBridge. The contract is now stated in terms of TransformerBridge alone: vocab_readout and logit_signature require compatibility mode so that W_U carries the folded final LayerNorm weights. Docstring-only, no behavior change. The guard still returns early for a non-TransformerBridge model; tightening that early-return to fail loudly is a behavior change, not a documentation one, so it is left out of this commit. * fix(svd_circuits): bind each decomposition to its compatibility-mode state enable_compatibility_mode folds ln1 into W_V and centres W_O and W_U, so a HeadSVD taken before the call describes a different OV map than the model then computes. The existing guard inspected only the live model, so re-running a consumer with a decomposition captured under the other state returned a silently wrong result, and project_activations and patch_along_directions had no guard at all. Record the model's compatibility-mode state on HeadSVD at decomposition time and add an orthogonal guard that refuses a mismatch in all four consumers, pointing the caller at re-running decompose_head. * fix(svd_circuits): project W_U and V in a common dtype * fix(svd_circuits): validate direction indices against the map rank * fix(svd_circuits): place the subspace projector on the activation's device * fix(svd_circuits): draw the causal baseline in-span, average it, and gate by mode * test(svd_circuits): pin the OV write basis and per-head isolation against mutation * fix(svd_circuits): filter project_activations cache to one hook and reject a batched tensor pre-forward * fix(svd_circuits): threshold the gate on the mean control magnitude Averaging the signed control deltas and then taking abs() understates the typical control effect and shrinks as n_baseline grows wherever the controls mix sign, so ablate gates pass too easily and keep too rarely. Make baseline_delta_metric the mean of the per-draw delta magnitudes, a non-negative typical control effect, and compare abs(delta_metric) against it directly. Update the PatchResult and patch_along_directions docstrings to state the field is a magnitude, and add a test on a mixed-sign stub that pins the threshold to the magnitude mean and guards against a regression to abs(signed mean). * fix(svd_circuits): gate a full-rank retained set like the empty set A full-rank keep reconstructs the head onto span(V) itself (V V^T), and every width-rank in-span control also spans span(V), so the kept and control projectors coincide and their deltas tie by construction, just as the empty set does on the zero subspace. Extend the empty-set guard to reject a full-rank retained set unless an explicit threshold is given, reword the message and comment to cover both degenerate cases, and update the docstring. The integration full-keep call and a new unit case pass a threshold accordingly. * fix(svd_circuits): check the folded-LN state, not the compatibility flag * docs(svd_circuits): document NotImplementedError from use_attn_result Both project_activations and patch_along_directions enable use_attn_result via set_use_attn_result(True) to read a head's per-head output. On an attention adapter with no per-head result fork point that call raises NotImplementedError, but neither Raises block listed it. Document it in both. * docs(svd_circuits): correct the in-span control geometry note * refactor(svd_circuits): default HeadSVD.compatibility_mode to False * fix(svd_circuits): accept a 1-D prompt tensor in project_activations Requiring exactly two dimensions rejected a 1-D [pos] token tensor, which the previous revision accepted and which patch_along_directions still accepts. Reject only a tensor whose leading dimension is larger than one, so a bare [pos] prompt flows through to the forward while the post-cache batch guard still backstops a batched string list. * test(svd_circuits): pin the control baseline to its Haar expectation The span-aligned stub's readout equals its head output, so every control delta shares a sign and no gate test on it can tell the magnitude mean from the signed mean's magnitude; the same geometry lets a full-residual-stream control and a single draw pass every existing test. Add a many-draw assertion against the analytic Haar expectation E[baseline_delta] = pos * ||h||^2 * (1 - w/rank), whose 0.5 band separates the in-span value 3.0 from the full-stream value 9.0, and a second test that threading one generator through single-draw calls reproduces the multi-draw average bit for bit. * test(svd_circuits): run the causal-direction gate on both heads The sibling-head test called _make_subspace_hook directly and every patch_along_directions test used head 0, so the control-draw call inside patch_along_directions could have been hardcoded to head 0 without any unit test failing. Parametrize test_patch_along_directions_discriminates_causal_direction over both heads and add a wrong-head assertion: a copy mislabeled via dataclasses.replace(ov, head=other) keeps the same V but points at the sibling slice, so a head-faithful baseline draw yields a different baseline_delta_metric while a draw hardcoded to head 0 would tie the two. The strong-vs-weak 5x margin holds for both heads unchanged. * style(svd_circuits): wrap long call sites to satisfy black * fix(svd_circuits): require fold support in the W_U readout guard Thirty-five adapters declare no fold support, so enable_compatibility_mode() leaves the final LayerNorm in place while the bridge still reports processed weights. The readout guard read only _weights_processed and _fold_ln_requested, so on those adapters vocab_readout projected through an unfolded W_U. Compute the folded-LayerNorm predicate in one helper so the readout guard and the staleness guard cannot drift apart. * fix(svd_circuits): bind a decomposition to the folded-LN state, not the flag * fix(svd_circuits): restore the rank check in the prompt guard * fix(svd_circuits): raise the control-baseline default and document its seed sensitivity * test(svd_circuits): cover the states that separate the flag from the fold * fix(svd_circuits): require eval mode in the two forward-pass tools
* Fix GatedDeltaNet hooked forward on transformers >= 5.15 transformers 5.15 stopped setting causal_conv1d_fn and chunk_gated_delta_rule on the GatedDeltaNet instance (they are now module-level functions wrapped by use_kernelized_func), so the bridge's hooked forward raised AttributeError for Qwen3Next / Qwen3.5 / Qwen3.5-MoE linear-attention layers. Look the attributes up with getattr and fall back to the torch conv path and the module-level torch_chunk_gated_delta_rule when they are absent. * Resolve GatedDeltaNet conv from the HF module and name the transformers range on failure
* fix(vllm): respect left-padded attention masks * fix(vllm): restore masked prompt coordinates
… first (#1805) * Fix circular import when the architecture adapter factory is imported first Importing transformer_lens.factories.architecture_adapter_factory before transformer_lens.model_bridge fails: the factory imports model_bridge, whose __init__ imports model_bridge.sources, which imports names from the still partially initialized factory module. Import the factory module in model_bridge.sources and look the names up at call time instead. * Update the inspect vLLM dtype test's factory patch target
* fix(bridge): honor 4D boolean mask polarity in causal loss * fix(bridge): preserve MPT mask polarity and broaden loss coverage
…1811) Compatibility mode masks attention scores with -inf, so the leading pad queries of a left-padded row have no visible key and softmax to NaN. The NaN was scrubbed from the forward pattern, but softmax's backward still multiplied by the saved NaN output, and the additive padding mask passed it through to Q/K, so every position of the padded row got non-finite gradients. Add AttentionBridge._masked_softmax: in compatibility mode it zeroes fully masked rows before softmax and zeroes their pattern afterwards, keeping both passes finite while still producing HookedTransformer's all-zero pattern. Outside compatibility mode it is plain softmax. Use it in _softmax_dropout_pattern and at the two direct softmax sites (PositionEmbeddingsAttentionBridge and the LLaDA attention bridge). Fixes #1809
* feat(analysis): causal-swap benchmark corpus and answer metrics * feat(analysis): baseline-capability filter * feat(analysis): norm-matched control-token selection * feat(analysis): coordinate-patch causal-swap trial runner * feat(analysis): bootstrap CI and frozen artifact schema * docs(notebook): causal-swap benchmark for coordinate_patch_hooks Add the generation entry point (module main()), a cached-GPT-2 integration smoke test, the frozen GPT-2-small benchmark artifact, and a notebook that loads that artifact only and never calls the model. Establish demos/data/ as the directory for notebook-adjacent frozen data files. Register the new notebook in the nbval CI matrix, the docs make_docs copy list, and the docs index, and add a Causal-swap benchmark subsection to jacobian_lens_fitting.md with the interpretation caveats. Export run_causal_swap_benchmark, BenchmarkCorpus, FunctionSpec, bootstrap_success_rate_ci, and load_artifact from the analysis package. The artifact is regenerable with: uv run python -m transformer_lens.tools.analysis.jacobian_lens_causal_swap_benchmark * test(make_docs): expect coordinate-patch benchmark demo in copied notebooks * fix(benchmark): select control tokens on CPU and match edit displacement * fix(benchmark): derive trial skips from an up-front support check * feat(benchmark): sample each trial's control arm across seeds * fix(benchmark): report exact Clopper-Pearson intervals with trial counts * feat(benchmark): record executed layers and fingerprint the full artifact * test(benchmark): pin condition equivalence and a real executing trial * feat(benchmark): add a CLI entry point for artifact generation * docs(demo): regenerate the artifact and require a checkout
* Removed remaining references to HookedTransformer * Fragile MPS test fix
) * Keep Dream/LLaDA2 rope shim out of the shared ROPE_INIT_FUNCTIONS The Dream and LLaDA2-MoE adapters restored transformers v4's "default" rope init by inserting it into the global ROPE_INIT_FUNCTIONS dict. From transformers 5.17, PreTrainedModel._init_weights builds its rope lookup as {"default": module.compute_default_rope_parameters, **ROPE_INIT_FUNCTIONS}, so the global entry overrides every native model's own default rope init. After loading either model, constructing e.g. a LlamaForCausalLM in the same process raises AttributeError: 'LlamaConfig' object has no attribute 'rope_theta'. Patch only the remote modeling modules instead, as the Ouro adapter already does, via a shared restore_default_rope_init helper. * Forward the load revision to the Dream/LLaDA2 rope shim boot(..., revision=...) imports a separate copy of the remote modeling file; force-importing only the default revision left that copy without the 'default' rope entry (KeyError: 'default').
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
6 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
A series of bug fixes and smaller features for v4
Type of change
Checklist: