Skip to content

Add CircuitLens feature-circuit attribution - #131

Open
Xmaster6y wants to merge 1 commit into
mainfrom
codex/issue-125-circuitlens
Open

Add CircuitLens feature-circuit attribution#131
Xmaster6y wants to merge 1 commit into
mainfrom
codex/issue-125-circuitlens

Conversation

@Xmaster6y

Copy link
Copy Markdown
Owner

Summary

  • add model-agnostic CircuitLens artifacts and tensor decompositions for upstream transcoder features, frozen-pattern attention head/source-token contributions, and downstream output logits
  • run attribution through public Target and HookSession activation/gradient capture APIs, including graph-retaining activation capture and restoration of existing parameter gradients
  • make local-Jacobian, frozen-attention, and frozen-nonlinearity conventions explicit and serialize them with contributor scores
  • keep transformer-lens and circuit-tracer in an optional circuit-lens integration extra rather than TDHook core dependencies
  • add an analytically checkable toy circuit covering all three attribution paths

Validation

  • uv run pre-commit run --all-files
  • uv run pytest tests -q (645 passed, 1 skipped)
  • uv build

Closes #125

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Aug 31, 2026

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review Completed 2026-08-31T20:27:29.795513Z 7fa3842 PR opened
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@codecov

codecov Bot commented Aug 31, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 82.72727% with 38 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.56%. Comparing base (c403cd2) to head (7fa3842).

Files with missing lines Patch % Lines
src/tdhook/attribution/circuit_lens.py 81.90% 38 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #131      +/-   ##
==========================================
- Coverage   98.36%   97.56%   -0.81%     
==========================================
  Files          53       54       +1     
  Lines        4091     4305     +214     
==========================================
+ Hits         4024     4200     +176     
- Misses         67      105      +38     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

@cubic-dev-ai cubic-dev-ai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 issues found across 6 files

Confidence score: 3/5

  • In src/tdhook/attribution/circuit_lens.py, mismatched pattern and values head selections can combine data from different heads while reporting the pattern index, producing incorrect attribution results; require matching head selections before positional decomposition.
  • In src/tdhook/attribution/circuit_lens.py, target_position=-1 attributes the final query but returns an invalid negative token index, which can mislead downstream consumers; reject negative positions consistently with AttentionSite validation.

You’re at about 99% of the monthly reviewed-line limit. You may want to disable incremental reviews to conserve quota. Reviews will continue until that limit is exceeded. If you need help avoiding interruptions, please contact contact@cubic.dev.

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="src/tdhook/attribution/circuit_lens.py">

<violation number="1" location="src/tdhook/attribution/circuit_lens.py:202">
P3: When callers use `target_position=-1`, the helper attributes the final query while returning an invalid negative token index. Reject negative positions here, matching `AttentionSite` validation.</violation>

<violation number="2" location="src/tdhook/attribution/circuit_lens.py:316">
P2: When `pattern` and `values` select different head indices, this pairs one head’s pattern with another head’s values and reports the pattern index. Require matching head selections before positional decomposition.</violation>
</file>

Reply with feedback, questions, or to request a fix.

Re-trigger cubic

_captured(gradient, "attention output gradient"),
layer=site.layer,
target_position=site.target_position,
head_indices=site.pattern.indices,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: When pattern and values select different head indices, this pairs one head’s pattern with another head’s values and reports the pattern index. Require matching head selections before positional decomposition.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/tdhook/attribution/circuit_lens.py, line 316:

<comment>When `pattern` and `values` select different head indices, this pairs one head’s pattern with another head’s values and reports the pattern index. Require matching head selections before positional decomposition.</comment>

<file context>
@@ -0,0 +1,439 @@
+                    _captured(gradient, "attention output gradient"),
+                    layer=site.layer,
+                    target_position=site.target_position,
+                    head_indices=site.pattern.indices,
+                )
+            )
</file context>

raise ValueError("output_weight head and head_dim axes do not match values")
if output_gradient.shape != (queries, output_weight.shape[2]):
raise ValueError("output_gradient must have shape [query, model]")
if target_position >= queries:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3: When callers use target_position=-1, the helper attributes the final query while returning an invalid negative token index. Reject negative positions here, matching AttentionSite validation.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/tdhook/attribution/circuit_lens.py, line 202:

<comment>When callers use `target_position=-1`, the helper attributes the final query while returning an invalid negative token index. Reject negative positions here, matching `AttentionSite` validation.</comment>

<file context>
@@ -0,0 +1,439 @@
+        raise ValueError("output_weight head and head_dim axes do not match values")
+    if output_gradient.shape != (queries, output_weight.shape[2]):
+        raise ValueError("output_gradient must have shape [query, model]")
+    if target_position >= queries:
+        raise IndexError("target_position is outside the attention query axis")
+    if head_indices is None:
</file context>

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7fa3842d8d

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +299 to +300
_captured(activations, "upstream activation"),
_captured(gradients, "upstream gradient"),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Match repeated feature captures by invocation

When an upstream module is invoked more than once and multiple invocations contribute to the target, activation captures are recorded in forward order while gradient captures are recorded in reverse order. Reading .value for both therefore multiplies the last invocation's activation by the first invocation's gradient, silently corrupting feature scores. Pair CapturedTarget.values in the appropriate order or reject sites reached multiple times.

Useful? React with 👍 / 👎.

output = model(*model_args, **kwargs)
live_target = _site_scalar(_captured(target_activation, "target feature"), target_feature)
retain_graph = bool(logit_indices)
live_target.backward(retain_graph=retain_graph)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve caller-owned input gradients

When a tensor in model_args is a leaf with requires_grad=True, this backward pass—and each later logit backward pass—accumulates attribution gradients into its existing .grad. Only parameter gradients are saved and restored, so invoking this helper during training or input optimization silently contaminates caller-owned input gradients even though model gradients are restored.

Useful? React with 👍 / 👎.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add CircuitLens feature-circuit attribution

1 participant