Open In Colab

Circuit Tracer Analysis Backend Demo#

This notebook demonstrates the Interpretune analysis ops pipeline for circuit-tracer analysis, running the full semantic concept intervention flow as a single composite pipeline:

  • Concept Direction — Compute a directional concept vector (e.g. “Capitals − States”)

  • Attribution Graph — Generate a circuit-level attribution graph from the concept direction

  • Node Influence — Score graph nodes by their causal influence on the concept

  • Top Features — Extract the most influential transcoder features

  • Feature Intervention — Amplify those features and measure how predictions shift

The five ops above are composed into the intervention_from_concept pipeline, which chains them automatically and returns results in a single AnalysisBatch.

Tip: For individual op dispatching and the DISPATCHER API, see the op_collection_example notebook.

The backend parameter selects the circuit-tracer backend: NNsight (default) or TransformerLens — both are validated by the parameterized notebook tests. (The TransformerLens backend uses the legacy HookedTransformer path; circuit-tracer does not yet support TransformerBridge — see the tracking notes in docs/circuit_tracer_backend_support.md.)

Hide code cell content

# @title Imports { display-mode: "form" }
from pprint import pformat

import interpretune as it
from it_examples import _ACTIVE_PATCHES  # noqa: F401
from it_examples.example_module_registry import MODULE_EXAMPLE_REGISTRY
from it_examples.utils.example_helpers import required_os_env
from it_examples.utils.nb_ui_utils import (
    display_target_gap,
    display_top_features_comparison,
    display_topk_token_predictions,
    resolve_feature_explanations,
)
from interpretune import ITSession, ITSessionConfig
from interpretune.analysis.ops.base import AnalysisBatch
from interpretune.config import AnalysisCfg, init_analysis_cfgs

Notebook Parameters#

This cell contains parameters that can be injected by papermill during parameterized test runs.

dashboard_mode selects where feature-dashboard links point: "public" uses neuronpedia.org, "local" uses a local Neuronpedia dev webapp (local_webapp_url) — see the Neuronpedia localhost guide.

# Parameters - These will be injected by papermill during parameterized test runs
backend = "nnsight"  # Options: "transformerlens", "nnsight"
core_log_dir = None  # Directory to save analysis logs (if None, a temp directory will be created)
intervention_scale_factor = 10.0  # Multiplier for feature intervention amplitudes
dashboard_mode = "public"  # Options: "public" (neuronpedia.org links), "local" (local Neuronpedia dev webapp)
local_webapp_url = "http://localhost:3000"  # Feature-dashboard base URL used when dashboard_mode="local"

Hide code cell content

# @title Environment Setup { display-mode: "form" }
env_path: str | None = None  # set to '/full/path/to/.env' to override
os_env_reqs = None
assert required_os_env(env_path=env_path, env_reqs=os_env_reqs)

Session Configuration#

We load a registered example module configuration for Gemma-2-2b with the RTE task, then configure it for the selected backend. The "gemma" transcoder set uses Gemma Scope transcoders, matching the upstream attribution_targets_demo.

Note: the first MODULE_EXAMPLE_REGISTRY access hydrates every example registry entry. Per-entry config-normalization feedback (categorized as ITInstantiationFeedbackWarning) is suppressed during that bulk hydration so this cell only surfaces messages relevant to the requested configuration; directly instantiating a config still shows its feedback.

# Load the demo configuration from the example module registry
base_itdm_cfg, base_it_cfg, dm_cls, m_cls = MODULE_EXAMPLE_REGISTRY.get("gemma2.rte_demo.circuit_tracer")

# Configure backend
base_it_cfg.circuit_tracer_cfg.backend = backend
print(f"Backend: {backend}")

# Optionally override core_log_dir
if core_log_dir:
    base_it_cfg.core_log_dir = core_log_dir

# Use Gemma Scope transcoders (matches upstream attribution_targets_demo)
base_it_cfg.circuit_tracer_cfg.transcoder_set = "gemma"

# Configure intervention settings
base_it_cfg.circuit_tracer_cfg.intervention_value_source = "top_feature_activation_values"
base_it_cfg.circuit_tracer_cfg.intervention_scale_factor = intervention_scale_factor

print(pformat(base_it_cfg.circuit_tracer_cfg))
Backend: nnsight
CircuitTracerConfig(backend='nnsight',
                    model_name=None,
                    transcoder_set='gemma',
                    dtype=torch.bfloat16,
                    max_n_logits=10,
                    desired_logit_prob=0.95,
                    batch_size=256,
                    max_feature_nodes=8192,
                    offload='cpu',
                    lazy_encoder=None,
                    lazy_decoder=True,
                    verbose=True,
                    default_node_threshold=0.8,
                    default_edge_threshold=0.98,
                    save_graphs=True,
                    graph_output_dir=None,
                    analysis_target_tokens=['▁Dallas', '▁Austin'],
                    target_token_ids=None,
                    use_neuronpedia=False,
                    intervention_scale_factor=10.0,
                    intervention_max_influence_norm_scale=False,
                    intervention_sign_aware_scale=True,
                    intervention_value=None,
                    intervention_value_source='top_feature_activation_values',
                    intervention_constrained_layers=None,
                    intervention_freeze_attention=None,
                    intervention_apply_activation_function=None,
                    intervention_sparse=False,
                    intervention_return_activations=False,
                    nnsight_remote=False,
                    ndif_api_key=None)
# Configure the session with the appropriate adapter composition
if backend == "nnsight":
    adapter_ctx = (it.Adapter.core, it.Adapter.nnsight, it.Adapter.circuit_tracer)
else:
    # the TL circuit-tracer backend needs the transformer_lens adapter in the composition
    # (it provides the replacement-model init path; see docs/circuit_tracer_backend_support.md)
    adapter_ctx = (it.Adapter.core, it.Adapter.transformer_lens, it.Adapter.circuit_tracer)

session_cfg = ITSessionConfig(
    adapter_ctx=adapter_ctx,
    datamodule_cfg=base_itdm_cfg,
    module_cfg=base_it_cfg,
    datamodule_cls=dm_cls,
    module_cls=m_cls,
)

it_session = ITSession(session_cfg)

# Initialize session (loads model, sets up hooks, etc.)
it.it_init(**it_session)

# Set up analysis config on module
module = it_session.module
tokenizer = module.replacement_model.tokenizer

graph_op = it.compute_attribution_graph
module.analysis_cfg = AnalysisCfg(target_op=graph_op, ignore_manual=True, save_tokens=False)
init_analysis_cfgs(module, [module.analysis_cfg])

print(f"Module type: {type(module).__name__}")
print("Session initialized successfully!")
[INFO] interpretune.utils.logging: Loading ReplacementModel with backend: nnsight
INFO:interpretune.utils.logging:Loading ReplacementModel with backend: nnsight
[INFO] interpretune.utils.logging: NNsight ReplacementModel initialized for Circuit Tracer
INFO:interpretune.utils.logging:NNsight ReplacementModel initialized for Circuit Tracer
[INFO] interpretune.utils.logging: Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute 'quantization_config'
INFO:interpretune.utils.logging:Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute 'quantization_config'
[INFO] interpretune.utils.logging: Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute '_pre_quantization_dtype'
INFO:interpretune.utils.logging:Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma2Config' object has no attribute '_pre_quantization_dtype'
[INFO] interpretune.utils.logging: Preparing data: InterpretunableDataModule
INFO:interpretune.utils.logging:Preparing data: InterpretunableDataModule
[INFO] interpretune.utils.logging: The following columns  don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns  don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: The following columns  don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns  don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: The following columns  don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns  don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: Setting up datamodule: InterpretunableDataModule
INFO:interpretune.utils.logging:Setting up datamodule: InterpretunableDataModule
[INFO] interpretune.utils.logging: Setting up model: InterpretunableModule
INFO:interpretune.utils.logging:Setting up model: InterpretunableModule
[INFO] interpretune.utils.logging: initializing optimizers and schedulers: InterpretunableModule
INFO:interpretune.utils.logging:initializing optimizers and schedulers: InterpretunableModule
[INFO] interpretune.utils.logging: Input gradient requirements handled by circuit tracer internally.
INFO:interpretune.utils.logging:Input gradient requirements handled by circuit tracer internally.
Module type: InterpretunableModule
Session initialized successfully!

Analysis Pipeline#

We use the intervention_from_concept composite op to chain all five analysis operations in a single pipeline call. Virtual logit target IDs (from concept-direction attribution targets) are automatically resolved to real vocabulary token IDs.

Step

Op

Purpose

Pipeline

concept_direction

Compute a directional concept vector via paired rejection

Pipeline

compute_attribution_graph

Generate a circuit-level attribution graph for the prompt

Pipeline

graph_node_influence

Score graph nodes by causal influence on the concept

Pipeline

extract_top_features

Rank and extract the most influential transcoder features

Pipeline

feature_intervention_forward

Amplify top features and measure prediction shifts

We define the concept groups, prompt, and key tokens, then run the pipeline.

# Concept groups (SentencePiece tokens with leading ▁)
capitals = ["▁Austin", "▁Sacramento", "▁Olympia", "▁Atlanta"]
states = ["▁Texas", "▁California", "▁Washington", "▁Georgia"]
concept_label = "Concept: Capitals − States"

# Analysis prompt
prompt = "Fact: the capital of the state containing Dallas is"

# Key tokens for result comparison (matching upstream attribution_targets_demo)
austin_id = tokenizer.encode("▁Austin", add_special_tokens=False)[-1]
dallas_id = tokenizer.encode("▁Dallas", add_special_tokens=False)[-1]
key_tokens = [("Austin", austin_id), ("Dallas", dallas_id)]

print(f"Prompt: '{prompt}'")
print(f"Key tokens: Austin={austin_id}, Dallas={dallas_id}")
Prompt: 'Fact: the capital of the state containing Dallas is'
Key tokens: Austin=22605, Dallas=26865
# Run the full intervention_from_concept pipeline (all 5 ops in one call)
full_pipeline = it.intervention_from_concept
print(f"Full pipeline: {' → '.join(op.name for op in full_pipeline.composition)}")

results = full_pipeline(
    module,
    AnalysisBatch(
        concept_group_a=capitals,
        concept_group_b=states,
        concept_label=concept_label,
        concept_direction_mode="paired_rejection",
        prompts=[prompt],
    ),
    None,
    0,
    top_n=10,
    intervention_scale_factor=intervention_scale_factor,
)
print("Pipeline complete!")
Full pipeline: concept_direction → compute_attribution_graph → graph_node_influence → extract_top_features → feature_intervention_forward
Pipeline complete!
Phase 0: Precomputing activations and vectors
Precomputation completed in 1.82s
Found 9179 active features
Phase 1: Running forward pass
Forward pass completed in 0.54s
Phase 2: Building input vectors
Using 1 custom attribution targets with total weight 0.1072
Will include 8192 of 9179 feature nodes
Input vectors built in 0.81s
Phase 3: Computing logit attributions
1 logit attribution(s) completed in 0.20s
Phase 4: Computing feature attributions
Feature influence computation:   0%|          | 0/8192 [00:00<?, ?it/s]
Feature influence computation:   3%|▎         | 256/8192 [00:00<00:05, 1554.34it/s]
Feature influence computation:   9%|▉         | 768/8192 [00:00<00:03, 2225.51it/s]
Feature influence computation:  16%|█▌        | 1280/8192 [00:00<00:02, 2441.49it/s]
Feature influence computation:  22%|██▏       | 1792/8192 [00:00<00:02, 2737.41it/s]
Feature influence computation:  28%|██▊       | 2304/8192 [00:00<00:01, 2952.55it/s]
Feature influence computation:  34%|███▍      | 2816/8192 [00:01<00:01, 2982.66it/s]
Feature influence computation:  41%|████      | 3328/8192 [00:01<00:01, 2896.20it/s]
Feature influence computation:  47%|████▋     | 3840/8192 [00:01<00:01, 2864.07it/s]
Feature influence computation:  53%|█████▎    | 4352/8192 [00:01<00:01, 2844.00it/s]
Feature influence computation:  59%|█████▉    | 4864/8192 [00:01<00:01, 2896.82it/s]
Feature influence computation:  66%|██████▌   | 5376/8192 [00:01<00:01, 2740.72it/s]
Feature influence computation:  72%|███████▏  | 5888/8192 [00:02<00:00, 2772.83it/s]
Feature influence computation:  78%|███████▊  | 6400/8192 [00:02<00:00, 2632.40it/s]
Feature influence computation:  84%|████████▍ | 6912/8192 [00:02<00:00, 2683.25it/s]
Feature influence computation:  91%|█████████ | 7424/8192 [00:02<00:00, 2536.09it/s]
Feature influence computation:  97%|█████████▋| 7936/8192 [00:02<00:00, 2623.62it/s]
Feature influence computation: 100%|██████████| 8192/8192 [00:03<00:00, 2699.63it/s]

Feature attributions completed in 3.04s
Attribution completed in 8.88s

Results: Concept Direction#

The concept direction vector captures the “capital-ness” concept via paired rejection, projecting out state-specific components from each capital embedding.

The reported direction norm should read 1.0000: concept_direction returns a unit-normalized vector, so the print is a cheap sanity check — a 0/nan norm flags a degenerate concept pair (groups canceling under paired rejection), and any non-unit value flags an aggregation/normalization regression. Because the direction is unit-norm, the applied perturbation magnitude in direct add-mode steering is exactly direction_scale_factor, and the attribution/node-influence rankings below are unaffected by direction scale.

print(f"Concept direction shape: {results.concept_direction.shape}")
print(f"Concept label: {results.concept_label}")
direction_norm = float(results.concept_direction.norm())
print(f"Direction norm: {direction_norm:.4f}")
assert abs(direction_norm - 1.0) < 1e-3, "concept_direction should return a unit-normalized vector"
Concept direction shape: torch.Size([2304])
Concept label: Concept: Capitals − States
Direction norm: 1.0000

Results: Top Features#

The top transcoder features ranked by node influence scores. Each feature is a (layer, position, feature_index) tuple.

For interactive node influence visualization, see the attribution_analysis notebook.

# Build feature tuples and scores for display
features = [tuple(f.tolist()) for f in results.top_feature_ids]
scores = results.top_feature_scores.tolist()

# Feature-dashboard links: public neuronpedia.org or a local Neuronpedia dev webapp
neuronpedia_base_url = "https://www.neuronpedia.org" if dashboard_mode == "public" else local_webapp_url

# Best-effort explanation text from the feature API (public or local webapp); unmapped features
# simply render an empty Explanation cell
feature_explanations = resolve_feature_explanations(
    model_id="gemma-2-2b",
    source_set="gemmascope-transcoder-16k",
    feature_tuples=[(f[0], f[-1]) for f in features],
    base_url=neuronpedia_base_url,
)

display_top_features_comparison(
    {"Top Features (by node influence)": features},
    {"Top Features (by node influence)": scores},
    neuronpedia_model="gemma-2-2b",
    neuronpedia_base_url=neuronpedia_base_url,
    feature_explanations=feature_explanations,
)
Top Features (by node influence)
#NodeScoreExplanation
1(21, 10, 5943)0.0039a mix of location names, political words, and parts of code
2(24, 10, 6394)0.0021place names and words describing geographic locality
3(23, 10, 12237)0.0019locations
4(0, 2, 16200)0.0014code syntax elements
5(24, 10, 5999)0.0013language related to institutions, negative situations, the internet, and programming languages
6(20, 10, 15589)0.0012references to geographic locations, especially in addresses
7(19, 10, 2695)0.0011locations in North America
8(24, 10, 6044)0.0010locations and legal case identifiers
9(22, 10, 4999)0.0010references to metropolitan areas and travel between locations
10(18, 10, 6101)9.49e-04words, phrases, and names related to governments and political entities

Results: Feature Intervention#

Amplifying the top features by the configured scale factor and measuring how predictions shift. If the pipeline correctly identifies “capital” features, amplifying them should increase the probability of Austin relative to Dallas.

pre_logits = results.pre_intervention_logits.float().cpu()
post_logits = results.post_intervention_logits.float().cpu()

# Rich pre/post comparison with key tokens
display_topk_token_predictions(
    prompt,
    pre_logits,
    post_logits,
    tokenizer,
    k=5,
    key_tokens=key_tokens,
)

# Consolidated pre/post probabilities, logits, and the Austin - Dallas gap in one table
pre_gap, post_gap = display_target_gap(
    pre_logits,
    post_logits,
    ("Austin", austin_id),
    ("Dallas", dallas_id),
    title="Feature intervention — Austin vs Dallas",
)

# Sanity gate (both backends): amplifying "capital" features must widen the Austin-Dallas gap
assert post_gap > pre_gap, "feature intervention should widen the Austin-Dallas logit gap"
Input Sentence:
Fact: the capital of the state containing Dallas is
Original Top 5 Tokens
Token Probability Distribution
▁Austin42.045%
42.0%
▁not5.690%
5.7%
▁the5.690%
5.7%
▁Texas5.022%
5.0%
▁Fort3.911%
3.9%
New Top 5 Tokens
Token Probability Distribution
▁Austin64.504%
64.5%
▁San11.209%
11.2%
Austin2.834%
2.8%
▁Fort2.501%
2.5%
▁Irving1.948%
1.9%
Key Tokens
Token Original New Change
Austin42.0452%64.5045%
+0.2246
Dallas3.0457%0.2987%
0.0275
Feature intervention — Austin vs Dallas
Token Pre prob Post prob Pre logit Post logit Δ
Austin42.045%64.504%26.125026.1250+0.0000
Dallas3.046%0.299%23.500020.7500-2.7500
Gap (Austin − Dallas)+2.6250+5.3750+2.7500

Summary#

The intervention_from_concept pipeline ran all five analysis ops in a single call:

Op

Result

concept_direction

Computed a “Capitals − States” concept vector via paired rejection

compute_attribution_graph

Generated a circuit-level attribution graph for the prompt

graph_node_influence

Scored graph nodes by causal influence

extract_top_features

Ranked and extracted the top-10 most influential features

feature_intervention_forward

Amplified top features and measured prediction shifts

Virtual logit target IDs from the attribution graph are automatically resolved to real vocabulary token IDs, so no manual override is needed.

Key result: Amplifying the identified “capital” features shifts the model’s prediction towards the correct answer (Austin), validating that the circuit-tracer pipeline identifies causally relevant features.

What’s Next#

  • Concept-Direction Steering Demo: See ct_concept_steering_demo.ipynb for store- and embed-based, sign-aware multi-feature concept steering with public or locally served feature dashboards

  • RTE research direction: RTE-focused concept-direction research continues in interpretune#220

  • Analysis Injection: See attribution_analysis.ipynb for node influence visualization via injection hooks

  • Op Dispatching: See op_collection_example.ipynb for the DISPATCHER API and individual op usage

  • AnalysisRunner: For batch analysis workflows, use AnalysisRunner with AnalysisCfg objects

  • Basic CT Tutorial: See circuit_tracer_adapter_example_basic.ipynb for the foundational adapter tutorial