[Proposal] Dynamic hooks for J-space coordinate patching

[janmenjayap]

---

Proposal

Add JacobianLens.coordinate_patch_hooks(...), a forward-hook variant of the offline
coordinate_patch(...) primitive shipped in #1741, so an anchored J-space coordinate edit can run
live inside model.run_with_hooks(...) / model.generate(...) instead of only on one
already-captured activation.

Suggested labels: enhancement, complexity-moderate, TransformerBridge
cc @jlarson4 (following the review lineage from #1739 / #1741)

---

Motivation

This follows the Jacobian-Lens proposal (#1505), the sparse-decomposition follow-up (#1539), the
merged decomposition implementation (#1596), and the merged offline coordinate-patching primitive
(#1739 / PR #1741).

JacobianLens.coordinate_patch(...) edits one already-captured activation offline and returns a
report-only CoordinatePatch; it installs no forward hooks. #1739 deferred a hooked variant to a
separate follow-up for a concrete reason: each hooked position would otherwise require its own
vocabulary-scale sparse decomposition (get_sparse_decomposition scans the full lens dictionary)
fired live during the forward pass — "do not hide the cost behind a convenient API."

The gap that leaves: researchers cannot study how a patched coordinate propagates through the rest
of a forward pass, and cannot patch during model.generate(...) or any multi-token pass. They must
manually re-capture, patch, and re-inject activations outside the hook system. This issue tracks
closing that gap by wrapping the existing primitive as a forward hook, while making the
per-position cost explicit rather than silent.

---

Proposed semantics

Given a forward-pass residual chunk activations of shape [batch, num_positions, d_model] at a
requested layer, for every (batch_idx, position) pair the hook body:

1. Slices that pair's activation x.
2. Reuses a validated JSpaceDecomposition from decomposition_cache[(layer, batch_idx, position)]
   if present; otherwise runs get_sparse_decomposition(x, dictionary, k, ...) once and stores it.
3. Calls solve_coordinate_patch(x, dictionary, source_idx, target_idx, decomposition=..., mode=..., alpha=...) — the exact PR1 primitive — to anchor-edit only the source/target coordinates while
   holding every other recovered coordinate and the residual fixed.
4. Writes that pair's patched activation back into the requested slice; untouched positions pass
   through unchanged.

A source concept active in one pair never affects another pair: each pair gets its own independent
decomposition and edit, matching the anchored, per-pair contract of the underlying primitive.

---

API sketch

Model-free core (Commit 1, in jacobian_lens_coordinate_patch.py):

def solve_coordinate_patch_positions(
    activations: torch.Tensor,                # [batch, num_positions, d_model]
    dictionary: torch.Tensor,
    position_labels: Sequence[int],           # real sequence position per column; keys the cache
    source_idx: int,
    target_idx: int,
    *,
    layer: int,
    decomposition_cache: Optional[MutableMapping[Tuple[int, int, int], JSpaceDecomposition]] = None,
    k: int = DEFAULT_K,
    mode: str = "substitute",
    alpha: float = 1.0,
    algorithm: str = "nonnegative_orthogonal_matching_pursuit",
) -> Tuple[torch.Tensor, Dict[Tuple[int, int, int], CoordinatePatch]]:

Loops solve_coordinate_patch(...) over every (batch_idx, position) pair — never once per
position shared across the batch — and returns (patched, patches) where patches maps
(layer, batch_idx, position) to that pair's full CoordinatePatch.

Bridge wrapper (Commit 2b, in jacobian_lens.py):

def coordinate_patch_hooks(
    self,
    model: Any,
    source_token: TokenInput,
    target_token: TokenInput,
    layers: Sequence[int],
    *,
    positions: Sequence[int],                 # required — no "defaults to all"
    decomposition_cache: Optional[MutableMapping[Tuple[int, int, int], JSpaceDecomposition]] = None,
    k: int = DEFAULT_K,
    mode: str = "substitute",
    alpha: float = 1.0,
    algorithm: str = "nonnegative_orthogonal_matching_pursuit",
) -> List[Tuple[str, Any]]:

Resolves token inputs and layers exactly like coordinate_patch(...), then builds one forward hook
per layer at _resid_post_hook_name(layer) = blocks.{layer}.hook_out, following the same eager
dictionary-build pattern swap_hooks already uses. Returns [(hook_name, fn), ...] for
model.hooks(fwd_hooks=...) / run_with_hooks(...).

Two departures from swap_hooks, both deliberate:

1. positions is required, not defaulted to all. swap_hooks defaults to all positions because
   its transform is a cheap closed-form pseudoinverse over a fixed 2-vector basis; coordinate
   patching's transform is a full sparse solve per position, so a silent full-sequence default would
   be a cost trap the codebase elsewhere warns against explicitly.
2. Optional caller-owned decomposition_cache keyed (layer, batch_idx, position). When a
   caller re-runs the same activation repeatedly (e.g. an interactive causal-tracing loop holding the
   prompt fixed while varying alpha or mode), a cache hit skips the vocabulary-scale scan and
   reuses the strict-compatibility path already validated in solve_coordinate_patch; a miss solves
   and stores. Purely a performance path — a cache hit and a fresh solve produce an identical patch.

---

Correctness contract

• alpha=0 is an exact no-op, position by position.
• Only the requested (layer, position) entries change; every other position in the forward pass is
  bit-identical to an un-hooked run.
• Oracle parity: the patched activation at a given position matches an offline coordinate_patch(...)
  call on that position's pre-hook activation.
• Each (batch_idx, position) pair is decomposed and edited independently — batch rows do not share
  a support.
• If source_token is inactive at any pair touched by a hook firing, the whole forward pass
  raises (fail-fast; see decision 2), rather than silently patching a subset.
• positions must be non-empty; source_token and target_token must resolve to distinct ids.
• Inputs are not mutated; the output preserves the decomposition module's float32 compute convention
  and device (it does not promise the activation's original storage dtype).
• The conditioning and near-parallel warnings from solve_coordinate_patch still fire from inside
  the hook, per pair, exactly as they would offline — never swallowed or re-wrapped.

---

Design decisions (pinned; maintainer sign-off requested on 2 and 3)

1. Batch handling — settled by construction. Loop (batch_idx, position) pairs explicitly and
   solve once per pair, never once per position shared across the batch: two batch rows can have
   different active supports. Required correctness, not a style choice.
2. Fail-fast on any inactive source (needs sign-off). If source_idx is inactive at any one
   (batch, position) pair touched by a hook firing, solve_coordinate_patch's existing
   ValueError: "... is not in the decomposition's active support" propagates out of the whole
   forward pass. There is no try/except around the per-pair loop, so nothing is a partial write.
   The alternative — skip that pair with a warning and copy the row through unedited — is a two-line
   change if preferred. Fail-fast is recommended: a silently partial intervention is a worse failure
   mode than a loud error. But it changes ergonomics for batched prompts where not every prompt has
   the source concept active, so it warrants explicit sign-off.
3. decomposition_cache is caller-owned (needs sign-off). A plain dict (any MutableMapping) the
   caller passes in and can inspect, not a JacobianLens attribute — matching CoordinatePatch's
   report-only, no-retained-state design.
4. Warning timing. One UserWarning per call at hook-construction time, naming
   len(layers) x len(positions) — not a per-position warning (that would spam), and not an attempt
   to report the exact number of live (non-cache-hit) solves, which is only known at forward time.
   This is the plan's own call, flagged for review.

---

Relationship to #1746 / #1747

Many thanks to the authors of #1746 and #1747 for the careful write-up there — it's a useful
reference for this section. #1746 reports that swap_hooks re-reads live lens coordinates at
every hooked layer, which makes a multi-layer swap band involutive (odd/even cancellation); #1747
proposes a separate clean-coordinate clamp for J-lens swaps to address that. coordinate_patch_hooks
shares the structural property of re-decomposing the live activation per layer, so we'd expect a
multi-layer install to compound across the band in a similar way — but since a substitute patch is
not an involution, it should not exhibit the exact cancellation #1746 describes. We'd welcome a
clean-coordinate clamp variant of coordinate patching as a follow-up if the multi-layer concern turns
out to be material in practice; we've left it out of scope for this issue. To be clear, this work
does not fix or depend on #1746 / #1747 — we just want to flag that both areas touch
jacobian_lens.py, the __init__.py exports, and jacobian_lens_fitting.md, so whichever lands
second will need a small mechanical rebase

---

Proposed repository changes

• New model-free core solve_coordinate_patch_positions(...) in
  transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py.
• Additive JacobianLens.coordinate_patch_hooks(...) in
  transformer_lens/tools/analysis/jacobian_lens.py (no change to swap_hooks / coordinate_patch).
• Export solve_coordinate_patch_positions from transformer_lens/tools/analysis/__init__.py.
• Extract the shared _ToyBridge test fixture into tests/unit/tools/conftest.py (mechanical, no
  behavior change) so a new hook-test file can reuse it.
• New model-free unit tests, a dedicated wrapper test file
  (tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py, not folded into the already-large
  test_jacobian_lens.py), and one cached-model integration invariant test.
• Short docs in jacobian_lens_fitting.md: cost model, required positions, cache semantics,
  fail-fast behavior.

---

Non-goals

• Arbitrary multi-slot permutations (already out of scope for the primitive).
• Any behavioral / causal-swap benchmark (separate follow-up, PR3).
• HookedTransformer support (separate follow-up, PR4).
• A clean-coordinate clamp variant of coordinate patching (candidate future PR; see above).
• Claims that a successful intervention proves unique causal mediation.

---

Validation

Unit tests use exact planted dictionaries with analytic expected patches; integration tests assert
algebraic invariants only (alpha=0 no-op end-to-end, unedited positions bit-identical through a
real run_with_hooks pass) — not behavioral success rates or token-flip claims, per PR1's own policy
against behavior-dependent assertions. Cache-hit and fresh-solve paths are checked to produce
identical patches. Full formatting, typing, unit, and integration surfaces run before submission.

---

Checklist

• Confirm no competing open issue/PR for dynamic coordinate-patch hooks
• Maintainer sign-off on design decisions 2 (fail-fast vs. skip-with-warning) and 3
  (caller-owned decomposition_cache).

---