`JacobianLens.swap_hooks` re-reads live coordinates, causing odd/even cancellation across layer bands

[koriyoshi2041]

Describe the bug

JacobianLens.swap_hooks re-reads lens coordinates from the live activation at every hooked layer:

coords = selected.float() @ local_pinv.T
delta = alpha * ((coords[..., [1, 0]] - coords) @ local_basis.T)

A coordinate exchange is an involution. Across a band whose in-plane coordinates are approximately preserved, the second hook therefore undoes the first, the third redoes it, and so on. This makes the intervention alternate with band-width parity; at alpha=2, the antisymmetric coordinate is multiplied by -3 per application and can diverge rapidly.

This differs from the paper's clamp control, which uses clean-pass coordinate values. The released hook instead exchanges whatever coordinates are present at each layer.

Across three fitted lenses, the odd-minus-even paired margin effect was:

• Qwen3.5-9B: +2.13 nat, 95% CI [1.26, 3.07]
• Qwen3.5-4B: +1.87 nat, 95% CI [1.09, 2.57]
• Qwen3-4B (dense attention, independently fitted): +3.99 nat, 95% CI [1.66, 6.49] (+3.54 after fp32 refitting)

On Qwen3.5-9B, band widths 1/2/3/4/13 gave flip rates .19/.00/.19/.02/.10, while a clean-coordinate clamp rose to .41. At width 13, alpha=2 used about 3e10 times the intervention energy of alpha=1.

Code example

With any model and fitted lens, compare consecutive band widths:

def margin(layers):
    hooks = lens.swap_hooks(model, source_token, target_token, layers=layers)
    with model.hooks(fwd_hooks=hooks):
        logits = model(tokens)[0, -1].float().log_softmax(-1)
    return float(logits[target_answer] - logits[source_answer])

print(margin([L]))
print(margin([L, L + 1]))      # approximately returns toward clean
print(margin([L, L + 1, L + 2]))

The full experiment and artifacts are public at https://github.com/koriyoshi2041/rlens-probe-swap.

Suggested fix

Keep the current live-coordinate exchange for backward compatibility, but expose a clamp form that fixes each layer's target to the exchanged coordinates from a clean run:

h <- h + V (c_target - V⁺h)

The existing method should also document that repeated live swaps may cancel across layers, and direct multi-layer users to the clamp form.

System info

• TransformerLens main at c03d51037e32ae5e34b3e23d63f63d0e6e16bd1 (the same implementation remains on current main and in v3.8.1)
• Source installation
• Linux, CUDA, Python 3.11
• Models: Qwen3.5-9B, Qwen3.5-4B, Qwen3-4B via TransformerBridge

Checklist

• I checked for similar open and closed issues and PRs in the repository.