Feature Interpretation › Code Implementation: Feature Circuit Analysis
Language AI HandbookAll chapters Feature Interpretation › Code Implementation: Feature Circuit AnalysisOne section · Sections in this chapter8 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement a simplified version of feature circuit analysis. We'll work with a toy model that mimics the residual stream structure, using the same synthetic SAE activations we created earlier.
Simulating a Residual Stream with Layer Outputs
import numpy as np
np.random.seed(0)
# Toy model parameters
d_model = 64 # Hidden dimension
n_layers = 4 # Number of transformer layers
n_heads = 4 # Attention heads per layer
n_sae_features = 128 # SAE dictionary size
n_circuit_tokens = 100 # Tokens to analyze
# Simulate SAE encoder weights: n_sae_features x d_model
W_enc = np.random.randn(n_sae_features, d_model)
W_enc = W_enc / np.linalg.norm(W_enc, axis=1, keepdims=True) # Normalize
# Simulate residual stream at each layer for each token
# x_layers[l][t] = residual stream at layer l, token t
x_layers = []
x = np.random.randn(n_circuit_tokens, d_model) * 0.5 # Initial embeddings
x_layers.append(x.copy())
# Simulate attention head outputs: (n_layers, n_heads, n_tokens, d_model)
# Most heads produce small random outputs, but some heads are "specialized"
# Layer 1 head 0 and Layer 2 head 2 produce larger, structured outputs
head_outputs = (
np.random.randn(n_layers, n_heads, n_circuit_tokens, d_model) * 0.15
)
# Make specific heads dominant: they output along the direction of W_enc rows
# This creates structured attribution patterns (some heads much more influential)
signal_direction = W_enc[0] # Direction aligned with first feature
head_outputs[1, 0] += np.outer(
np.random.randn(n_circuit_tokens) * 0.8, signal_direction
)
head_outputs[2, 2] += np.outer(
np.random.randn(n_circuit_tokens) * 0.6, signal_direction * 0.7
)
head_outputs[0, 3] += np.outer(
np.random.randn(n_circuit_tokens) * 0.4, signal_direction * 0.4
)
for l in range(n_layers):
# Sum of attention head outputs + MLP output (simulated)
attn_contrib = head_outputs[l].sum(axis=0)
mlp_contrib = np.random.randn(n_circuit_tokens, d_model) * 0.2
x = x + attn_contrib + mlp_contrib
x_layers.append(x.copy())
# Compute SAE feature activations at the final layer
x_final = x_layers[-1]
pre_activations = x_final @ W_enc.T # n_tokens x n_sae_features
feature_acts_final = np.maximum(
0, pre_activations - 0.5
) # ReLU with threshold 0.5Final layer feature activation stats: Shape: (100, 128) Active (non-zero) entries: 3659 Sparsity: 0.7141 Mean activation (when active): 0.5585
Computing Attribution Scores
For each active feature at the final layer, we compute how much each attention head at each layer contributed to it. The attribution is computed as the mean absolute projection of each head's output onto each feature's encoder direction, averaged across all token positions.
def compute_head_to_feature_attribution(
head_outputs_arr, W_enc_arr, target_layer=-1
):
"""
Compute attribution of each attention head to each SAE feature.
Parameters
----------
head_outputs_arr : (n_layers, n_heads, n_tokens, d_model)
W_enc_arr : (n_sae_features, d_model)
target_layer : int, layer at which SAE features are computed
Returns
-------
attributions : (n_layers, n_heads, n_sae_features)
Mean absolute attribution across tokens
"""
n_l, n_h, n_t, d = head_outputs_arr.shape
attributions = np.zeros((n_l, n_h, W_enc_arr.shape[0]))
for l in range(n_l):
for h in range(n_h):
# head output at this layer, all tokens: (n_tokens, d_model)
head_out = head_outputs_arr[l, h]
# Projection onto each feature direction: (n_tokens, n_sae_features)
proj = head_out @ W_enc_arr.T
# Mean absolute contribution across tokens
attributions[l, h] = np.abs(proj).mean(axis=0)
return attributions
attributions = compute_head_to_feature_attribution(head_outputs, W_enc)
# For a specific feature, find its top contributing heads
target_feature = int(
np.argmax(feature_acts_final.sum(axis=0))
) # Most active feature
feature_attrs = attributions[:, :, target_feature] # n_layers x n_headsAttribution analysis for feature 37: Attribution matrix shape (layers x heads): (4, 4) Top-5 contributing (layer, head) pairs: Layer 1, Head 0: attribution = 0.1378 Layer 2, Head 2: attribution = 0.1367 Layer 2, Head 3: attribution = 0.1353 Layer 3, Head 0: attribution = 0.1264 Layer 1, Head 1: attribution = 0.1237
The attribution analysis reveals that not all heads contribute equally to the target feature. The top-contributing (layer, head) pairs represent the upstream circuit for this feature: these are the computation steps that write information into the residual stream in the direction of the feature's encoder vector. In a real model, this would tell you that, for example, "Layer 3, Head 2 is responsible for writing the capital-city relationship feature, and it does so by attending from the country-name token to the capital-city token."
Attribution Heatmap
Feature Co-activation Graph
Features that frequently activate together on the same tokens are likely functionally related. Co-activation can arise for several reasons: two features may be part of the same circuit and tend to be activated together by the same upstream cause, or they may be tracking different aspects of the same input context, or they may be redundant features covering similar concepts (a common occurrence when the SAE dictionary is larger than the effective dimensionality of the feature space).
We measure co-activation with pointwise mutual information (PMI), which controls for the marginal frequencies of each feature. Two frequent features will have many co-activations by chance alone; PMI normalizes this away by comparing the observed co-activation rate against the rate expected under independence.
# Compute feature co-activation matrix from our synthetic data
# activations is (n_tokens, n_features): already computed above
# Binarize for co-activation counts
binary_acts = (activations > 0).astype(np.float32)
# Co-activation counts: n_features x n_features
# co_activation[i, j] = number of tokens where both i and j are active
co_activation = binary_acts.T @ binary_acts # (n_features, n_features)
# Normalize: compute pointwise mutual information (PMI)
marginal = binary_acts.mean(axis=0) # P(feature_i active)
expected = (
np.outer(marginal, marginal) * n_tokens
) # Expected co-activation under independence
# Avoid division by zero
with np.errstate(divide="ignore", invalid="ignore"):
pmi = np.where(
expected > 0, np.log(co_activation / (expected + 1e-10) + 1e-10), 0
)
# Focus on top features for visualization
top_features_for_graph = sorted_by_freq[:20]
pmi_subset = pmi[np.ix_(top_features_for_graph, top_features_for_graph)]
# Normalize to [0, 1] for display
pmi_subset_norm = (pmi_subset - pmi_subset.min()) / (
pmi_subset.max() - pmi_subset.min() + 1e-10
)
np.fill_diagonal(pmi_subset_norm, 0) # Remove self-edgesCo-activation analysis (top-20 features): PMI matrix shape: (20, 20) PMI range: [-0.169, 2.548] Feature pairs with strong co-activation (normalized PMI > 0.5): 0
The PMI heatmap is a useful first filter for identifying which features are worth investigating together in a circuit analysis. Pairs with high PMI that are also semantically similar (based on their feature labels) are strong candidates for being part of the same functional circuit. Pairs with high PMI but dissimilar labels may indicate an interpretation error in one of the features, or a compound concept that spans both labels.