[flyingdoubleG]

Description

JacobianLens.ablation_hooks is documented as projecting the given tokens' directions out of the residual stream. With several tokens, however, it removes their unit lens vectors one at a time:

for unit in local_units:
    result = result - (result @ unit)[..., None] * unit

This only zeroes the projections when the vectors are orthogonal. Lens vectors of related tokens are not. Removing v̂₂ puts back a component -(h·v̂₂)(v̂₂·v̂₁) along v̂₁, and so on down the list, so the "ablated" activation still has a large projection onto the directions it was meant to remove. The paper ablates whole token sets (e.g. 176 ethics-related lens vectors) as one intervention.

Measured on GPT-2 small with the published neuronpedia/jacobian-lens lens, on one prompt, at blocks.{layer}.hook_out:

layer	token set	mean / max |cos| between lens vectors	max |h·v̂| before	after current ablation	after joint projection
4	France, Paris, French, Lyon, Europe	0.72 / 0.85	168.6	38.2	3e-5
4	one … ten	0.64 / 0.95	127.5	77.9	6e-5
8	France … Europe	0.76 / 0.86	273.0	49.0	2e-4
8	one … ten	0.74 / 0.95	235.6	170.9	2e-4

With the ten number tokens at layer 8, about 73% of the largest projection is still there after the current ablation.

The fix projects onto the orthogonal complement of the span, h <- h - Q Qᵀ h. Q is an orthonormal basis of the unit lens vectors, built with the existing orthonormal_subspace from projection_kernel (rank-revealing SVD, so repeated or linearly dependent tokens are handled). For a single token this is the same h - (h·v̂)v̂ as before. The docstring is updated to state the guarantee.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

test_ablation_hooks_remove_every_direction_of_non_orthogonal_tokens uses the toy fitted lens with three tokens whose lens vectors are not orthogonal. It checks that the ablated activation has zero projection onto every vector, and that only their span changes. It fails on dev (max residual projection 0.34) and passes with the change.

Local runs (Windows, CPU, torch 2.9.1 and the locked 2.11.0, with --jaxtyping-packages=transformer_lens,beartype.beartype): tests/unit/tools gives 1113 passed and 4 failed. The same 4 fail on dev without this change, while dev CI is green: test_faithfulness_recovery_is_graded_not_binary (recovery comes out ≈ -1e-8), and three test_sparse_probing tests that pin L-BFGS stop reasons and refinement counts. They look platform-sensitive, and none of them touches the code changed here. isort --check-only passes; mypy reports no errors in jacobian_lens.py.

[jlarson4]

Looks good, thanks for putting this together. Merging now