Repository navigation
Add subspace interchange interventions - #1892
Conversation
Add transformer_lens.tools.analysis.interchange: an intervention that sets a base input's activation inside an orthonormal subspace to its value on a source input, at one hook point and optionally one position per example. It comes with difference-of-means, probe and random bases to intervene on and a norm-matched random-subspace placebo, plus unit tests, GPT-2 integration tests and a docs page. Part of TransformerLensOrg#1879.
koriyoshi2041
left a comment
koriyoshi2041
left a comment
There was a problem hiding this comment.
collect_activations rejects a hook that fires more than once, but run_with_intervention does not enforce the same single-fire contract. A model with the same HookPoint invoked twice silently applies the intervention twice (I reproduced hook_calls == 2, with no exception). That makes the source activation ambiguous and can change a later occurrence using an activation collected for the earlier one. Could run_with_intervention count hook invocations and raise after the forward pass, with a symmetric regression test alongside test_a_hook_that_fires_twice_is_refused?
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
The first phase of the DAS issue looks pretty good! I have a couple comments, let me know if you have any questions
| return torch.cat(out) | ||
|
|
||
|
|
||
| def run_with_intervention( |
There was a problem hiding this comment.
run_with_intervention sends each base through one forward pass and unless the caller wraps it in no_grad, the full model's autograd graph persists. Measuring IIA over a realistic number of pairs will run out of memory long before collect_activations would. The cost row in analysis_tools.md also describes batched passes that this function doesn't make. Would it be possible to apply the same bounded-memory solution used for collection to evaluation?
There was a problem hiding this comment.
Agreed. run_with_intervention stays the only differentiable pass that DAS trains through (#1893 calls it once per micro-batch).
I also added intervened_logits, which mirrors collect_activations. Batched by batch_size, under no_grad, with logits returned on the CPU. The docs example and the cost row now use it, and I made sure #1893 and #1894 evaluate with it too.
|
|
||
| def _select(act: torch.Tensor, positions: Optional[torch.Tensor]) -> torch.Tensor: | ||
| if positions is None: | ||
| if act.ndim != 2: |
There was a problem hiding this comment.
With positions set, _select indexes axis 1 whatever the activation's rank. On a 4-D hook such as attn.hook_pattern, that axis is heads, so collect_activations quietly returns a [batch, q, k] slice and the error only surfaces later, somewhere unrelated. Can we enforce the [batch, pos, d_model] contract on this branch the way the positions=None branch checks its shape?
There was a problem hiding this comment.
Done. Anything unexpected raises a ValueError with the shape named. There’s a test with a 4-D pattern-shaped hook.
| def test_positions_select_one_vector_per_example(): | ||
| model = _Toy(seq=True) | ||
| base, source = torch.randn(4, 3, D_MODEL), torch.randn(4, 3, D_MODEL) | ||
| last = torch.full((4,), 2) |
There was a problem hiding this comment.
Every test passes the same position for every row, so per-example indexing is never exercised. A _select that applies the first row's position to all rows passes the whole suite. Lets adjust at least one case to use different positions across rows.
There was a problem hiding this comment.
Good point. I've made the position test now use a different position per row, [0, 2, 1, 2], on two toys. 1 reads only the last position, and one sums every position. That means every ro'ws read and write is checked. Trying to apply the first row’s position to every row now fails it.
run_with_intervention now raises unless the hook fires exactly once in the forward pass, as collect_activations already did. With positions, the activation must be [batch, pos, d_model] and positions must hold one index per example; other shapes raise instead of slicing the wrong axis. Add intervened_logits, which evaluates an intervention in batches without gradients and returns CPU logits, so memory is bounded by batch_size. The position tests now use a different position per row.
|
Good catch, thanks. run_with_intervention now counts hook calls and raises unless the hook fires exactly once in the forward pass, which also catches a hook that does nto fire. The regression test checks both functions, for two calls and for zero, and the docstring is now saying task_logits must make one forward pass. Fixed in b450b36. |
|
Verified at |
|
Looks good, moving on to the next part and merging |
3 participants
Description
Part 1 of 3 for #1879.
Adds
tools.analysis.interchange: an intervention that sets a base input's activation inside an orthonormal subspace to its value on a source input, at one hook point and optionally one position per example. Also adds difference-of-means, probe and random bases, a norm-matched random-subspace placebo, tests and a docs page.Type of change
Checklist: