from __future__ import annotations
import hashlib
import html
import json
import math
import os
import random
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from .catalog import FeatureCatalog
from .config import SETTINGS, Settings
from .interventions import (
InterventionSpec,
joint_residual_delta,
normalized_random_control,
residual_delta,
)
from .metrics import (
contrastive_log_odds,
decoder_cosine_matrix,
joint_direction_norm_ratio,
js_divergence_from_logits,
reconstruction_metrics,
sequence_logprob_summary,
sparse_topk_cosine,
)
from .sae import SAEStore, SparseEncoding
def _dtype_from_name(name: str) -> torch.dtype:
name = name.lower().strip()
if name in {'float16', 'fp16', 'half'}:
return torch.float16
if name in {'bfloat16', 'bf16'}:
return torch.bfloat16
if name in {'float32', 'fp32'}:
return torch.float32
raise ValueError(f'Unsupported dtype: {name}')
def _default_device() -> torch.device:
# ZeroGPU exposes CUDA emulation at module load time. Hugging Face recommends
# placing models on CUDA at module scope so startup transfers can be optimized.
if os.getenv('SPACE_ID'):
return torch.device('cuda')
if torch.cuda.is_available():
return torch.device('cuda')
return torch.device('cpu')
@dataclass
class AnalysisResult:
tokens: list[str]
token_index: int
layer: int
features: SparseEncoding
rows: list[list[object]]
metrics: dict[str, float]
@dataclass
class InterventionResult:
baseline_text: str
modified_text: str
feature_activation: float
delta_activation: float
perturbation_norm: float
js_divergence: float
random_js_divergence: float
random_js_std: float
js_specificity_ratio: float
js_empirical_p: float
random_control_count: int
execution_drift_js: float
execution_drift_mean_logprob: float | None
target_text: str
target_token_count: int
target_tokens: list[str]
baseline_target_prob: float | None
modified_target_prob: float | None
random_target_prob: float | None
baseline_sequence_logprob: float | None
modified_sequence_logprob: float | None
random_sequence_logprob: float | None
sequence_logprob_delta: float | None
random_sequence_logprob_delta: float | None
mean_logprob_delta: float | None
random_mean_logprob_delta: float | None
random_abs_mean_logprob_delta: float | None
random_mean_logprob_std: float | None
target_specificity_ratio: float | None
target_empirical_p: float | None
target_token_rows: list[list[object]]
top_token_rows: list[list[object]]
@dataclass
class LayerSweepResult:
tokens: list[str]
token_index: int
rows: list[list[object]]
@dataclass
class DoseResponseResult:
feature_activation: float
target_tokens: list[str]
execution_drift_mean_logprob: float
execution_drift_js: float
rows: list[list[object]]
@dataclass
class FeatureSetResult:
feature_ids: list[int]
feature_rows: list[list[object]]
perturbation_norm: float
js_divergence: float
random_js_divergence: float
random_js_std: float
js_specificity_ratio: float
js_empirical_p: float
random_control_count: int
execution_drift_mean_logprob: float
execution_drift_js: float
baseline_sequence_logprob: float
modified_sequence_logprob: float
random_sequence_logprob: float
sequence_logprob_delta: float
random_sequence_logprob_delta: float
mean_logprob_delta: float
random_mean_logprob_delta: float
random_abs_mean_logprob_delta: float
random_mean_logprob_std: float
target_specificity_ratio: float
target_empirical_p: float
target_tokens: list[str]
target_token_rows: list[list[object]]
@dataclass
class FeatureSetSweepResult:
target_tokens: list[str]
random_control_count: int
execution_drift_mean_logprob: float
execution_drift_js: float
rows: list[list[object]]
@dataclass
class FeatureInteractionResult:
feature_ids: list[int]
target_tokens: list[str]
rows: list[list[object]]
additive_expected_mean_delta: float
joint_mean_delta: float
interaction_excess_mean_delta: float
normalized_interaction: float
execution_drift_mean_logprob: float
@dataclass
class ConceptContrastResult:
feature_id: int
layer: int
prompts_per_concept: int
rows: list[list[object]]
chart_rows: list[list[object]]
leading_concept: str | None
leading_ratio: float | None
active_prompt_count: int
total_prompt_count: int
@dataclass
class FeatureTraceResult:
feature_id: int
layer: int
tokens: list[str]
rows: list[list[object]]
chart_rows: list[list[object]]
active_token_count: int
token_count: int
max_activation: float
max_token_index: int | None
@dataclass
class FeatureGeometryResult:
feature_ids: list[int]
layer: int
rows: list[list[object]]
chart_rows: list[list[object]]
mean_abs_decoder_cosine: float
max_abs_decoder_cosine: float
joint_ablation_norm: float
independent_norm: float
alignment_ratio: float
@dataclass
class ContrastiveCausalResult:
feature_id: int
layer: int
feature_activation: float
perturbation_norm: float
target_a_tokens: list[str]
target_b_tokens: list[str]
rows: list[list[object]]
baseline_log_odds: float
modified_log_odds: float
delta_log_odds: float
baseline_normalized_preference: float
modified_normalized_preference: float
delta_normalized_preference: float
random_signed_mean_delta: float
random_abs_mean_delta: float
random_delta_std: float
specificity_ratio: float
empirical_p: float
random_control_count: int
@dataclass
class ConceptFeatureDiscoveryResult:
concept: str
layer: int
prompts_per_concept: int
top_n: int
ranking_mode: str
rows: list[list[object]]
chart_rows: list[list[object]]
candidate_ids: list[int]
default_candidate_id: int | None
current_context_available: bool
current_token_index: int | None
displayed_current_active_count: int
split_half_k: int | None
split_half_shared_count: int
split_half_jaccard: float | None
resample_replicates: int
resample_mean_support: float | None
resample_high_support_count: int
@dataclass
class CandidateCausalScreenResult:
feature_ids: list[int]
target_tokens: list[str]
rows: list[list[object]]
chart_rows: list[list[object]]
active_feature_count: int
candidate_count: int
execution_drift_mean_logprob: float
execution_drift_js: float
@dataclass
class CandidateSpecificityResult:
feature_ids: list[int]
target_tokens: list[str]
rows: list[list[object]]
chart_rows: list[list[object]]
active_feature_count: int
candidate_count: int
random_control_count: int
execution_drift_mean_logprob: float
execution_drift_js: float
@dataclass
class CandidateCrossTargetResult:
feature_ids: list[int]
targets: list[str]
rows: list[list[object]]
chart_rows: list[list[object]]
summary_rows: list[list[object]]
pairwise_rows: list[list[object]]
active_feature_count: int
@dataclass
class FeatureCueScanResult:
feature_id: int
layer: int
prompt_stem: str
rows: list[list[object]]
chart_rows: list[list[object]]
active_cue_count: int
cue_count: int
@dataclass
class FeatureCueContextResult:
feature_id: int
layer: int
stems: list[str]
cues: list[str]
rows: list[list[object]]
chart_rows: list[list[object]]
active_condition_count: int
condition_count: int
cue_active_context_counts: dict[str, int]
cue_mean_activations: dict[str, float]
dominant_cue: str | None
dominant_cue_context_count: int
off_dominant_active_count: int
@dataclass
class ParaphraseResult:
tokens_a: list[str]
token_index_a: int
tokens_b: list[str]
token_index_b: int
topk_jaccard: float
sparse_cosine: float
promptwide_jaccard: float
promptwide_cosine: float
shared_top_n: int
top_n: int
rows: list[list[object]]
chart_rows: list[list[object]]
class FeatureLensRuntime:
def __init__(self, settings: Settings = SETTINGS) -> None:
self.settings = settings
self.device = _default_device()
self.model_dtype = torch.float16 if self.device.type == 'cuda' else torch.float32
self.sae_dtype = _dtype_from_name(settings.sae_dtype)
if self.device.type == 'cpu' and self.sae_dtype != torch.float32:
self.sae_dtype = torch.float32
self.model = None
self.tokenizer = None
self.sae_store: SAEStore | None = None
self.catalog = FeatureCatalog()
self.load_error: str | None = None
@property
def ready(self) -> bool:
return self.model is not None and self.tokenizer is not None and self.sae_store is not None
def ensure_ready(self, preload_saes: bool = False) -> None:
if self.ready:
assert self.sae_store is not None
if preload_saes:
self.sae_store.preload()
return
self.tokenizer = AutoTokenizer.from_pretrained(self.settings.model_id)
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.padding_side = 'left'
self.model = AutoModelForCausalLM.from_pretrained(
self.settings.model_id,
torch_dtype=self.model_dtype,
low_cpu_mem_usage=True,
)
self.model.to(self.device)
self.model.eval()
self.sae_store = SAEStore(
repo_id=self.settings.sae_repo_id,
layers=self.settings.layers,
device=self.device,
dtype=self.sae_dtype,
top_k=self.settings.sae_top_k,
)
if preload_saes:
self.sae_store.preload()
self.load_error = None
def token_choices(self, text: str) -> list[tuple[str, int]]:
self.ensure_ready(preload_saes=False)
assert self.tokenizer is not None
ids = self.tokenizer(text, add_special_tokens=True)['input_ids']
tokens = [self.tokenizer.decode([token_id]) for token_id in ids]
return [(f'{idx}: {token!r}', idx) for idx, token in enumerate(tokens)]
def _inputs(self, text: str) -> dict[str, torch.Tensor]:
assert self.tokenizer is not None
batch = self.tokenizer(
text,
return_tensors='pt',
truncation=True,
max_length=self.settings.max_prompt_tokens,
)
return {key: value.to(self.device) for key, value in batch.items()}
def _target_ids(self, target_text: str) -> list[int]:
assert self.tokenizer is not None
ids = self.tokenizer(target_text, add_special_tokens=False)['input_ids']
if not ids:
raise ValueError('Target continuation tokenized to an empty sequence.')
return [int(x) for x in ids]
def _append_target(
self,
prompt_inputs: dict[str, torch.Tensor],
target_ids: Sequence[int],
) -> dict[str, torch.Tensor]:
prompt_ids = prompt_inputs['input_ids']
target = torch.tensor(
list(target_ids),
dtype=prompt_ids.dtype,
device=prompt_ids.device,
).unsqueeze(0)
full_ids = torch.cat([prompt_ids, target], dim=1)
if 'attention_mask' in prompt_inputs:
target_mask = torch.ones(
(prompt_ids.shape[0], len(target_ids)),
dtype=prompt_inputs['attention_mask'].dtype,
device=prompt_ids.device,
)
attention = torch.cat([prompt_inputs['attention_mask'], target_mask], dim=1)
else:
attention = torch.ones_like(full_ids)
return {'input_ids': full_ids, 'attention_mask': attention}
@staticmethod
def _repeat_inputs(inputs: dict[str, torch.Tensor], repeats: int) -> dict[str, torch.Tensor]:
return {key: value.repeat(int(repeats), 1) for key, value in inputs.items()}
@staticmethod
def _hidden_from_output(output):
return output[0] if isinstance(output, tuple) else output
@staticmethod
def _replace_hidden_in_output(output, hidden: torch.Tensor):
if isinstance(output, tuple):
return (hidden, *output[1:])
return hidden
@contextmanager
def _capture_hook(self, layer: int, bucket: dict) -> Iterator[None]:
assert self.model is not None
def hook(_module, _inputs, output):
hidden = self._hidden_from_output(output)
if 'hidden' not in bucket:
bucket['hidden'] = hidden.detach()
handle = self.model.model.layers[int(layer)].register_forward_hook(hook)
try:
yield
finally:
handle.remove()
@contextmanager
def _capture_hooks(self, layers: Sequence[int], buckets: dict[int, dict]) -> Iterator[None]:
assert self.model is not None
handles = []
for layer in layers:
bucket = buckets[int(layer)]
def hook(_module, _inputs, output, *, target=bucket):
hidden = self._hidden_from_output(output)
if 'hidden' not in target:
target['hidden'] = hidden.detach()
handles.append(self.model.model.layers[int(layer)].register_forward_hook(hook))
try:
yield
finally:
for handle in handles:
handle.remove()
@contextmanager
def _delta_hook(self, layer: int, token_index: int, delta: torch.Tensor) -> Iterator[None]:
assert self.model is not None
applied = {'done': False}
def hook(_module, _inputs, output):
if applied['done']:
return output
hidden = self._hidden_from_output(output)
if hidden.ndim != 3:
return output
seq_len = hidden.shape[1]
idx = self._resolve_index(int(token_index), seq_len)
modified = hidden.clone()
modified[:, idx, :] = modified[:, idx, :] + delta.to(hidden.device, hidden.dtype)
applied['done'] = True
return self._replace_hidden_in_output(output, modified)
handle = self.model.model.layers[int(layer)].register_forward_hook(hook)
try:
yield
finally:
handle.remove()
@contextmanager
def _batch_delta_hook(
self,
layer: int,
token_index: int,
deltas: torch.Tensor,
) -> Iterator[None]:
"""Apply one residual delta per batch row in a single model forward."""
assert self.model is not None
if deltas.ndim != 2:
raise ValueError('deltas must have shape [batch, d_model].')
applied = {'done': False}
def hook(_module, _inputs, output):
if applied['done']:
return output
hidden = self._hidden_from_output(output)
if hidden.ndim != 3:
return output
if hidden.shape[0] != deltas.shape[0]:
raise ValueError('Delta batch size does not match model batch size.')
idx = self._resolve_index(int(token_index), hidden.shape[1])
modified = hidden.clone()
modified[:, idx, :] = modified[:, idx, :] + deltas.to(hidden.device, hidden.dtype)
applied['done'] = True
return self._replace_hidden_in_output(output, modified)
handle = self.model.model.layers[int(layer)].register_forward_hook(hook)
try:
yield
finally:
handle.remove()
@staticmethod
def _resolve_index(token_index: int, seq_len: int) -> int:
idx = int(token_index)
if idx < 0:
idx = seq_len + idx
if idx < 0 or idx >= seq_len:
raise IndexError(f'Token index {token_index} outside prompt length {seq_len}.')
return idx
@staticmethod
def _control_seed(text: str, layer: int, key: str, mode: str, coefficient: float) -> int:
payload = f'{text}\0{layer}\0{key}\0{mode}\0{coefficient:.8g}'.encode()
return int.from_bytes(hashlib.sha256(payload).digest()[:4], 'big', signed=False)
@staticmethod
def _random_controls(delta: torch.Tensor, seed: int, count: int) -> list[torch.Tensor]:
if int(count) < 1:
raise ValueError('Random-control ensemble must contain at least one direction.')
return [
normalized_random_control(delta, seed=int(seed) + 104729 * idx)
for idx in range(int(count))
]
@staticmethod
def _random_effect_summary(values: Sequence[float], target_effect: float) -> tuple[float, float, float, float]:
if not values:
raise ValueError('Random-control values must not be empty.')
tensor = torch.tensor([float(x) for x in values], dtype=torch.float64)
signed_mean = float(tensor.mean().item())
abs_mean = float(tensor.abs().mean().item())
std = float(tensor.std(unbiased=False).item())
empirical_p = float(
(1 + int((tensor.abs() >= abs(float(target_effect))).sum().item()))
/ (len(values) + 1)
)
return signed_mean, abs_mean, std, empirical_p
@staticmethod
def _dict_cosine(a: dict[int, float], b: dict[int, float]) -> float:
if not a and not b:
return 1.0
if not a or not b:
return 0.0
dot = sum(float(value) * float(b.get(feature_id, 0.0)) for feature_id, value in a.items())
norm_a = math.sqrt(sum(float(value) ** 2 for value in a.values()))
norm_b = math.sqrt(sum(float(value) ** 2 for value in b.values()))
if norm_a == 0.0 or norm_b == 0.0:
return 0.0
return float(dot / (norm_a * norm_b))
@staticmethod
def _dict_jaccard(a: dict[int, float], b: dict[int, float]) -> float:
set_a = {feature_id for feature_id, value in a.items() if float(value) > 0}
set_b = {feature_id for feature_id, value in b.items() if float(value) > 0}
union = set_a | set_b
return float(len(set_a & set_b) / len(union)) if union else 1.0
@staticmethod
def _max_pool_encoding(encoding: SparseEncoding) -> dict[int, float]:
indices = encoding.indices.detach().cpu()
values = encoding.values.detach().float().cpu()
if indices.ndim == 1:
indices = indices.unsqueeze(0)
values = values.unsqueeze(0)
pooled: dict[int, float] = {}
for row_ids, row_values in zip(indices.tolist(), values.tolist(), strict=True):
for feature_id, value in zip(row_ids, row_values, strict=True):
value = float(value)
if value <= 0:
continue
feature_id = int(feature_id)
pooled[feature_id] = max(pooled.get(feature_id, 0.0), value)
return pooled
@staticmethod
def _encoding_map(encoding: SparseEncoding) -> dict[int, float]:
return {
int(feature_id): float(value)
for feature_id, value in zip(
encoding.indices.detach().cpu().tolist(),
encoding.values.detach().float().cpu().tolist(),
strict=True,
)
if float(value) > 0
}
@torch.inference_mode()
def _analyze_and_pool(
self,
text: str,
layer: int,
token_index: int = -1,
top_n: int = 12,
) -> tuple[AnalysisResult, dict[int, float]]:
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
if int(layer) not in self.settings.layers:
raise ValueError(f'Layer must be one of {self.settings.layers}.')
inputs = self._inputs(text)
bucket: dict = {}
with self._capture_hook(int(layer), bucket):
self.model(**inputs, use_cache=False)
hidden = bucket['hidden'][0]
idx = self._resolve_index(int(token_index), hidden.shape[0])
sae = self.sae_store.get(int(layer))
all_encoding = sae.encode(hidden)
encoding = SparseEncoding(
indices=all_encoding.indices[idx],
values=all_encoding.values[idx],
)
residual = hidden[idx]
reconstruction = sae.decode_sparse(encoding)
metrics = reconstruction_metrics(residual, reconstruction)
metrics['active_features'] = float(encoding.active_count)
values = encoding.values.float().clamp_min(0)
total = float(values.sum().item())
metrics['top5_mass_fraction'] = (
float(values[: min(5, values.numel())].sum().item()) / total if total > 0 else 0.0
)
ids = inputs['input_ids'][0].tolist()
tokens = [self.tokenizer.decode([token_id]) for token_id in ids]
rows: list[list[object]] = []
rank_count = min(int(top_n), encoding.indices.numel())
for rank in range(rank_count):
feature_id = int(encoding.indices[rank].item())
activation = float(encoding.values[rank].item())
rows.append(
[rank + 1, feature_id, activation, self.catalog.hint(int(layer), feature_id)]
)
result = AnalysisResult(
tokens=tokens,
token_index=idx,
layer=int(layer),
features=encoding,
rows=rows,
metrics=metrics,
)
return result, self._max_pool_encoding(all_encoding)
@torch.inference_mode()
def analyze(self, text: str, layer: int, token_index: int = -1, top_n: int = 12) -> AnalysisResult:
result, _ = self._analyze_and_pool(text, layer, token_index, top_n)
return result
@torch.inference_mode()
def layer_sweep(self, text: str, token_index: int = -1) -> LayerSweepResult:
self.ensure_ready(preload_saes=True)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
inputs = self._inputs(text)
buckets = {int(layer): {} for layer in self.settings.layers}
with self._capture_hooks(self.settings.layers, buckets):
self.model(**inputs, use_cache=False)
seq_len = int(inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), seq_len)
rows: list[list[object]] = []
for layer in self.settings.layers:
residual = buckets[int(layer)]['hidden'][0, idx]
sae = self.sae_store.get(int(layer))
encoding = sae.encode(residual)
reconstruction = sae.decode_sparse(encoding)
metrics = reconstruction_metrics(residual, reconstruction)
values = encoding.values.float().clamp_min(0)
positive = values[values > 0]
total = positive.sum()
if positive.numel() <= 1 or float(total.item()) <= 0:
entropy = 0.0
else:
probs = positive / total
entropy = float(
(-(probs * torch.log(probs)).sum() / math.log(positive.numel())).item()
)
top5_fraction = (
float(values[: min(5, values.numel())].sum().item() / total.item())
if float(total.item()) > 0
else 0.0
)
rows.append(
[
int(layer),
float(metrics['cosine']),
float(metrics['nmse']),
int(encoding.active_count),
float(values[0].item()) if values.numel() else 0.0,
top5_fraction,
entropy,
]
)
ids = inputs['input_ids'][0].tolist()
tokens = [self.tokenizer.decode([token_id]) for token_id in ids]
return LayerSweepResult(tokens=tokens, token_index=idx, rows=rows)
def token_html(self, tokens: list[str], selected_index: int) -> str:
chips = []
for idx, token in enumerate(tokens):
safe = html.escape(token if token.strip() else repr(token))
selected = idx == int(selected_index)
cls = 'token selected' if selected else 'token'
chips.append(f'{idx}{safe}')
return '' + ''.join(chips) + '
'
@staticmethod
def _top_token_rows(tokenizer, baseline_logits: torch.Tensor, modified_logits: torch.Tensor, k: int = 8):
p = torch.softmax(baseline_logits.float(), dim=-1)
q = torch.softmax(modified_logits.float(), dim=-1)
union_ids = torch.unique(torch.cat([torch.topk(p, k).indices, torch.topk(q, k).indices]))
rows = []
for token_id in union_ids.tolist():
token = tokenizer.decode([int(token_id)])
bp = float(p[token_id].item())
mp = float(q[token_id].item())
rows.append([repr(token), bp, mp, mp - bp])
rows.sort(key=lambda row: max(row[1], row[2]), reverse=True)
return rows[: min(len(rows), 12)]
def _target_rows(
self,
target_ids: Sequence[int],
baseline_token_logps: Sequence[float],
modified_token_logps: Sequence[float],
random_token_logps: Sequence[float],
) -> list[list[object]]:
assert self.tokenizer is not None
rows = []
for idx, (token_id, bp, mp, rp) in enumerate(
zip(
target_ids,
baseline_token_logps,
modified_token_logps,
random_token_logps,
strict=True,
)
):
rows.append(
[
idx,
repr(self.tokenizer.decode([int(token_id)])),
float(bp),
float(mp),
float(rp),
float(mp - bp),
float(rp - bp),
]
)
return rows
@torch.inference_mode()
def intervene(
self,
text: str,
layer: int,
token_index: int,
feature_id: int,
mode: str,
coefficient: float,
target_text: str = '',
max_new_tokens: int = 24,
) -> InterventionResult:
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
sae = self.sae_store.get(int(layer))
target_ids: list[int] = []
capture: dict = {}
single_baseline_logits: torch.Tensor
single_baseline_mean: float | None = None
if target_text.strip():
target_ids = self._target_ids(target_text)
scoring_inputs = self._append_target(prompt_inputs, target_ids)
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**scoring_inputs, use_cache=False)
single_baseline_logits = single_baseline_out.logits[0]
_, single_baseline_mean, _ = sequence_logprob_summary(
single_baseline_logits,
prompt_length=prompt_len,
target_ids=target_ids,
)
else:
scoring_inputs = prompt_inputs
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**scoring_inputs, use_cache=False)
single_baseline_logits = single_baseline_out.logits[0]
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
original_activation = encoding.activation_for(int(feature_id))
spec = InterventionSpec(mode=mode, coefficient=float(coefficient))
delta = residual_delta(sae.decoder_direction(int(feature_id)), original_activation, spec)
seed = self._control_seed(
text, int(layer), str(int(feature_id)), mode, float(coefficient)
)
random_controls = self._random_controls(
delta, seed=seed, count=self.settings.live_random_controls
)
zero = torch.zeros_like(delta)
all_deltas = torch.stack([zero, delta, *random_controls], dim=0)
repeated = self._repeat_inputs(scoring_inputs, all_deltas.shape[0])
with self._batch_delta_hook(int(layer), idx, all_deltas):
edited = self.model(**repeated, use_cache=False)
baseline_logits = edited.logits[0]
modified_logits = edited.logits[1]
random_logits = [edited.logits[row] for row in range(2, edited.logits.shape[0])]
next_idx = prompt_len - 1
baseline_next_logits = baseline_logits[next_idx]
modified_next_logits = modified_logits[next_idx]
random_next_logits = [row[next_idx] for row in random_logits]
execution_drift_js = js_divergence_from_logits(
single_baseline_logits[next_idx], baseline_next_logits
)
js = js_divergence_from_logits(baseline_next_logits, modified_next_logits)
random_js_values = [
js_divergence_from_logits(baseline_next_logits, logits)
for logits in random_next_logits
]
random_js_mean, random_js_abs_mean, random_js_std, js_empirical_p = (
self._random_effect_summary(random_js_values, js)
)
# JS divergence is non-negative, so signed and absolute means are identical up to numerical noise.
random_js_reference = random_js_abs_mean
js_ratio = abs(js) / max(random_js_reference, 1e-12)
baseline_seq = baseline_mean = modified_seq = modified_mean = None
random_seq_mean = random_mean_signed = random_abs_mean = random_mean_std = None
sequence_delta = random_sequence_delta = mean_delta = specificity = target_p = None
target_rows: list[list[object]] = []
target_tokens: list[str] = []
bp = mp = rp = None
execution_drift_mean = None
if target_ids:
baseline_seq, baseline_mean, baseline_token_logps = sequence_logprob_summary(
baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
modified_seq, modified_mean, modified_token_logps = sequence_logprob_summary(
modified_logits, prompt_length=prompt_len, target_ids=target_ids
)
random_summaries = [
sequence_logprob_summary(logits, prompt_length=prompt_len, target_ids=target_ids)
for logits in random_logits
]
random_seqs = [item[0] for item in random_summaries]
random_means = [item[1] for item in random_summaries]
random_token_matrix = [item[2] for item in random_summaries]
random_token_mean = [
float(sum(row[token_pos] for row in random_token_matrix) / len(random_token_matrix))
for token_pos in range(len(target_ids))
]
random_mean_deltas = [float(value - baseline_mean) for value in random_means]
random_seq_deltas = [float(value - baseline_seq) for value in random_seqs]
mean_delta = float(modified_mean - baseline_mean)
sequence_delta = float(modified_seq - baseline_seq)
random_mean_signed, random_abs_mean, random_mean_std, target_p = self._random_effect_summary(
random_mean_deltas, mean_delta
)
random_sequence_delta = float(sum(random_seq_deltas) / len(random_seq_deltas))
random_seq_mean = float(sum(random_seqs) / len(random_seqs))
specificity = abs(mean_delta) / max(random_abs_mean, 1e-12)
_, single_mean, _ = sequence_logprob_summary(
single_baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
execution_drift_mean = float(baseline_mean - single_mean)
p = torch.softmax(baseline_next_logits.float(), dim=-1)
q = torch.softmax(modified_next_logits.float(), dim=-1)
random_probs = [torch.softmax(logits.float(), dim=-1) for logits in random_next_logits]
first_id = int(target_ids[0])
bp = float(p[first_id].item())
mp = float(q[first_id].item())
rp = float(sum(prob[first_id].item() for prob in random_probs) / len(random_probs))
target_tokens = [self.tokenizer.decode([int(token_id)]) for token_id in target_ids]
target_rows = self._target_rows(
target_ids, baseline_token_logps, modified_token_logps, random_token_mean
)
generation_kwargs = {
'max_new_tokens': min(int(max_new_tokens), self.settings.max_new_tokens),
'do_sample': False,
'return_dict_in_generate': True,
'output_scores': False,
'pad_token_id': self.tokenizer.eos_token_id,
}
baseline_generation = self.model.generate(**prompt_inputs, **generation_kwargs)
with self._delta_hook(int(layer), idx, delta):
modified_generation = self.model.generate(**prompt_inputs, **generation_kwargs)
baseline_ids = baseline_generation.sequences[0, prompt_len:]
modified_ids = modified_generation.sequences[0, prompt_len:]
baseline_text = self.tokenizer.decode(baseline_ids, skip_special_tokens=True)
modified_text = self.tokenizer.decode(modified_ids, skip_special_tokens=True)
return InterventionResult(
baseline_text=baseline_text,
modified_text=modified_text,
feature_activation=float(original_activation),
delta_activation=float(spec.delta_activation(original_activation)),
perturbation_norm=float(torch.linalg.vector_norm(delta.float()).item()),
js_divergence=float(js),
random_js_divergence=float(random_js_reference),
random_js_std=float(random_js_std),
js_specificity_ratio=float(js_ratio),
js_empirical_p=float(js_empirical_p),
random_control_count=len(random_controls),
execution_drift_js=float(execution_drift_js),
execution_drift_mean_logprob=execution_drift_mean,
target_text=target_text,
target_token_count=len(target_ids),
target_tokens=target_tokens,
baseline_target_prob=bp,
modified_target_prob=mp,
random_target_prob=rp,
baseline_sequence_logprob=baseline_seq,
modified_sequence_logprob=modified_seq,
random_sequence_logprob=random_seq_mean,
sequence_logprob_delta=sequence_delta,
random_sequence_logprob_delta=random_sequence_delta,
mean_logprob_delta=mean_delta,
random_mean_logprob_delta=random_mean_signed,
random_abs_mean_logprob_delta=random_abs_mean,
random_mean_logprob_std=random_mean_std,
target_specificity_ratio=specificity,
target_empirical_p=target_p,
target_token_rows=target_rows,
top_token_rows=self._top_token_rows(
self.tokenizer, baseline_next_logits, modified_next_logits, k=8
),
)
@torch.inference_mode()
def dose_response(
self,
text: str,
layer: int,
token_index: int,
feature_id: int,
target_text: str,
multipliers: Sequence[float] = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0),
) -> DoseResponseResult:
if not target_text.strip():
raise ValueError('Dose-response requires a target continuation.')
if not any(abs(float(multiplier) - 1.0) < 1e-12 for multiplier in multipliers):
raise ValueError('Dose-response multipliers must include 1.0 as the zero-edit reference.')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
target_ids = self._target_ids(target_text)
full_inputs = self._append_target(prompt_inputs, target_ids)
sae = self.sae_store.get(int(layer))
capture: dict = {}
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**full_inputs, use_cache=False)
single_baseline_logits = single_baseline_out.logits[0]
_, single_baseline_mean, _ = sequence_logprob_summary(
single_baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
original_activation = encoding.activation_for(int(feature_id))
direction = sae.decoder_direction(int(feature_id))
deltas: list[torch.Tensor] = []
delta_coefficients: list[float] = []
norms: list[float] = []
for multiplier in multipliers:
spec = InterventionSpec('scale', float(multiplier))
delta = residual_delta(direction, original_activation, spec)
deltas.append(delta)
delta_coefficients.append(float(spec.delta_activation(original_activation)))
norms.append(float(torch.linalg.vector_norm(delta.float()).item()))
repeated = self._repeat_inputs(full_inputs, len(deltas))
with self._batch_delta_hook(int(layer), idx, torch.stack(deltas, dim=0)):
outputs = self.model(**repeated, use_cache=False)
reference_idx = next(
idx for idx, multiplier in enumerate(multipliers) if abs(float(multiplier) - 1.0) < 1e-12
)
reference_logits = outputs.logits[reference_idx]
baseline_seq, baseline_mean, _ = sequence_logprob_summary(
reference_logits, prompt_length=prompt_len, target_ids=target_ids
)
baseline_next = reference_logits[prompt_len - 1]
execution_drift_mean = float(baseline_mean - single_baseline_mean)
execution_drift_js = js_divergence_from_logits(
single_baseline_logits[prompt_len - 1], baseline_next
)
rows: list[list[object]] = []
for row_idx, multiplier in enumerate(multipliers):
modified_logits = outputs.logits[row_idx]
modified_seq, modified_mean, _ = sequence_logprob_summary(
modified_logits, prompt_length=prompt_len, target_ids=target_ids
)
rows.append(
[
float(multiplier),
delta_coefficients[row_idx],
norms[row_idx],
float(baseline_mean),
float(modified_mean),
float(modified_mean - baseline_mean),
float(modified_seq - baseline_seq),
float(js_divergence_from_logits(baseline_next, modified_logits[prompt_len - 1])),
]
)
return DoseResponseResult(
feature_activation=float(original_activation),
target_tokens=[self.tokenizer.decode([int(token_id)]) for token_id in target_ids],
execution_drift_mean_logprob=execution_drift_mean,
execution_drift_js=float(execution_drift_js),
rows=rows,
)
@torch.inference_mode()
def intervene_feature_set(
self,
text: str,
layer: int,
token_index: int,
feature_ids: Sequence[int],
mode: str,
coefficient: float,
target_text: str,
) -> FeatureSetResult:
if not target_text.strip():
raise ValueError('Feature-set causal testing requires a target continuation.')
ids = list(dict.fromkeys(int(x) for x in feature_ids))
if not ids:
raise ValueError('Select at least one feature.')
if len(ids) > 12:
raise ValueError('Select at most 12 features for a live feature-set intervention.')
if mode not in {'ablate', 'scale'}:
raise ValueError("Feature-set mode must be 'ablate' or 'scale'.")
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
target_ids = self._target_ids(target_text)
full_inputs = self._append_target(prompt_inputs, target_ids)
sae = self.sae_store.get(int(layer))
capture: dict = {}
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**full_inputs, use_cache=False)
single_baseline_logits = single_baseline_out.logits[0]
_, single_baseline_mean, _ = sequence_logprob_summary(
single_baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
activations = [encoding.activation_for(feature_id) for feature_id in ids]
directions = torch.stack([sae.decoder_direction(feature_id) for feature_id in ids], dim=0)
spec = InterventionSpec(mode, float(coefficient))
delta, coefficient_deltas = joint_residual_delta(directions, activations, spec)
seed = self._control_seed(
text, int(layer), ','.join(str(x) for x in ids), mode, float(coefficient)
)
controls = self._random_controls(
delta, seed=seed, count=self.settings.live_random_controls
)
zero = torch.zeros_like(delta)
all_deltas = torch.stack([zero, delta, *controls], dim=0)
repeated = self._repeat_inputs(full_inputs, all_deltas.shape[0])
with self._batch_delta_hook(int(layer), idx, all_deltas):
outputs = self.model(**repeated, use_cache=False)
baseline_logits = outputs.logits[0]
modified_logits = outputs.logits[1]
random_logits = [outputs.logits[row] for row in range(2, outputs.logits.shape[0])]
baseline_seq, baseline_mean, baseline_tokens = sequence_logprob_summary(
baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
modified_seq, modified_mean, modified_tokens = sequence_logprob_summary(
modified_logits, prompt_length=prompt_len, target_ids=target_ids
)
random_summaries = [
sequence_logprob_summary(logits, prompt_length=prompt_len, target_ids=target_ids)
for logits in random_logits
]
random_seqs = [item[0] for item in random_summaries]
random_means = [item[1] for item in random_summaries]
random_token_matrix = [item[2] for item in random_summaries]
random_token_mean = [
float(sum(row[token_pos] for row in random_token_matrix) / len(random_token_matrix))
for token_pos in range(len(target_ids))
]
mean_delta = float(modified_mean - baseline_mean)
sequence_delta = float(modified_seq - baseline_seq)
random_mean_deltas = [float(value - baseline_mean) for value in random_means]
random_seq_deltas = [float(value - baseline_seq) for value in random_seqs]
random_mean_signed, random_abs_mean, random_mean_std, target_p = self._random_effect_summary(
random_mean_deltas, mean_delta
)
next_idx = prompt_len - 1
js = js_divergence_from_logits(baseline_logits[next_idx], modified_logits[next_idx])
random_js_values = [
js_divergence_from_logits(baseline_logits[next_idx], logits[next_idx])
for logits in random_logits
]
_, random_js_abs_mean, random_js_std, js_p = self._random_effect_summary(
random_js_values, js
)
execution_drift_mean = float(baseline_mean - single_baseline_mean)
execution_drift_js = js_divergence_from_logits(
single_baseline_logits[next_idx], baseline_logits[next_idx]
)
feature_rows = [
[
feature_id,
float(activation),
float(delta_coefficient),
self.catalog.hint(int(layer), feature_id),
]
for feature_id, activation, delta_coefficient in zip(
ids, activations, coefficient_deltas, strict=True
)
]
return FeatureSetResult(
feature_ids=ids,
feature_rows=feature_rows,
perturbation_norm=float(torch.linalg.vector_norm(delta.float()).item()),
js_divergence=float(js),
random_js_divergence=float(random_js_abs_mean),
random_js_std=float(random_js_std),
js_specificity_ratio=float(abs(js) / max(random_js_abs_mean, 1e-12)),
js_empirical_p=float(js_p),
random_control_count=len(controls),
execution_drift_mean_logprob=execution_drift_mean,
execution_drift_js=float(execution_drift_js),
baseline_sequence_logprob=float(baseline_seq),
modified_sequence_logprob=float(modified_seq),
random_sequence_logprob=float(sum(random_seqs) / len(random_seqs)),
sequence_logprob_delta=sequence_delta,
random_sequence_logprob_delta=float(sum(random_seq_deltas) / len(random_seq_deltas)),
mean_logprob_delta=mean_delta,
random_mean_logprob_delta=random_mean_signed,
random_abs_mean_logprob_delta=random_abs_mean,
random_mean_logprob_std=float(random_mean_std),
target_specificity_ratio=float(abs(mean_delta) / max(random_abs_mean, 1e-12)),
target_empirical_p=float(target_p),
target_tokens=[self.tokenizer.decode([int(token_id)]) for token_id in target_ids],
target_token_rows=self._target_rows(
target_ids, baseline_tokens, modified_tokens, random_token_mean
),
)
@torch.inference_mode()
def feature_set_size_sweep(
self,
text: str,
layer: int,
token_index: int,
target_text: str,
sizes: Sequence[int] = (1, 3, 5),
) -> FeatureSetSweepResult:
"""Jointly ablate the strongest k active features and compare to random ensembles."""
if not target_text.strip():
raise ValueError('Feature-set size sweep requires a target continuation.')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
target_ids = self._target_ids(target_text)
full_inputs = self._append_target(prompt_inputs, target_ids)
sae = self.sae_store.get(int(layer))
capture: dict = {}
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**full_inputs, use_cache=False)
single_baseline_logits = single_baseline_out.logits[0]
_, single_baseline_mean, _ = sequence_logprob_summary(
single_baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
active_ids = [
int(feature_id)
for feature_id, value in zip(
encoding.indices.detach().cpu().tolist(),
encoding.values.detach().float().cpu().tolist(),
strict=True,
)
if float(value) > 0
]
valid_sizes = [int(size) for size in sizes if int(size) > 0 and int(size) <= len(active_ids)]
if not valid_sizes:
raise ValueError('Not enough active features for the requested set sizes.')
condition_deltas: list[torch.Tensor] = [torch.zeros_like(residual)]
metadata: list[tuple[int, str, list[int], float]] = []
for size in valid_sizes:
selected = active_ids[:size]
activations = [encoding.activation_for(feature_id) for feature_id in selected]
directions = torch.stack([sae.decoder_direction(feature_id) for feature_id in selected])
delta, _ = joint_residual_delta(
directions, activations, InterventionSpec('ablate', 0.0)
)
norm = float(torch.linalg.vector_norm(delta.float()).item())
condition_deltas.append(delta)
metadata.append((size, 'sae', selected, norm))
seed = self._control_seed(text, int(layer), f'top-{size}', 'ablate_set', 0.0)
controls = self._random_controls(
delta, seed=seed, count=self.settings.live_random_controls
)
for control_idx, control in enumerate(controls):
condition_deltas.append(control)
metadata.append((size, f'random_{control_idx}', selected, norm))
all_deltas = torch.stack(condition_deltas, dim=0)
repeated = self._repeat_inputs(full_inputs, all_deltas.shape[0])
with self._batch_delta_hook(int(layer), idx, all_deltas):
outputs = self.model(**repeated, use_cache=False)
baseline_logits = outputs.logits[0]
baseline_seq, baseline_mean, _ = sequence_logprob_summary(
baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
baseline_next = baseline_logits[prompt_len - 1]
execution_drift_mean = float(baseline_mean - single_baseline_mean)
execution_drift_js = js_divergence_from_logits(
single_baseline_logits[prompt_len - 1], baseline_next
)
grouped: dict[int, dict[str, object]] = {
size: {'features': active_ids[:size], 'norm': None, 'sae': None, 'random': []}
for size in valid_sizes
}
for output_idx, meta in enumerate(metadata, start=1):
size, kind, selected, norm = meta
logits = outputs.logits[output_idx]
seq_logp, mean_logp, _ = sequence_logprob_summary(
logits, prompt_length=prompt_len, target_ids=target_ids
)
item = {
'seq': float(seq_logp),
'mean': float(mean_logp),
'js': float(js_divergence_from_logits(baseline_next, logits[prompt_len - 1])),
}
grouped[size]['norm'] = norm
grouped[size]['features'] = selected
if kind == 'sae':
grouped[size]['sae'] = item
else:
grouped[size]['random'].append(item)
rows: list[list[object]] = []
for size in valid_sizes:
group = grouped[size]
sae_item = group['sae']
random_items = group['random']
assert isinstance(sae_item, dict)
assert isinstance(random_items, list) and random_items
sae_mean_delta = float(sae_item['mean'] - baseline_mean)
sae_seq_delta = float(sae_item['seq'] - baseline_seq)
random_mean_deltas = [float(item['mean'] - baseline_mean) for item in random_items]
random_js_values = [float(item['js']) for item in random_items]
random_signed, random_abs, random_std, target_p = self._random_effect_summary(
random_mean_deltas, sae_mean_delta
)
_, random_js_abs, random_js_std, js_p = self._random_effect_summary(
random_js_values, float(sae_item['js'])
)
rows.append(
[
int(size),
', '.join(str(x) for x in group['features']),
float(group['norm']),
float(baseline_mean),
float(sae_item['mean']),
sae_mean_delta,
random_signed,
random_abs,
random_std,
float(abs(sae_mean_delta) / max(random_abs, 1e-12)),
float(target_p),
sae_seq_delta,
float(sae_item['js']),
random_js_abs,
random_js_std,
float(js_p),
]
)
return FeatureSetSweepResult(
target_tokens=[self.tokenizer.decode([int(token_id)]) for token_id in target_ids],
random_control_count=self.settings.live_random_controls,
execution_drift_mean_logprob=execution_drift_mean,
execution_drift_js=float(execution_drift_js),
rows=rows,
)
@torch.inference_mode()
def feature_interaction_test(
self,
text: str,
layer: int,
token_index: int,
feature_ids: Sequence[int],
target_text: str,
) -> FeatureInteractionResult:
"""Compare individual ablations with their joint ablation to measure non-additivity."""
if not target_text.strip():
raise ValueError('Feature interaction testing requires a target continuation.')
ids = list(dict.fromkeys(int(x) for x in feature_ids))
if len(ids) < 2:
raise ValueError('Select at least two features for the interaction decomposition.')
if len(ids) > 5:
raise ValueError('Select at most five features for the live interaction decomposition.')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
target_ids = self._target_ids(target_text)
full_inputs = self._append_target(prompt_inputs, target_ids)
sae = self.sae_store.get(int(layer))
capture: dict = {}
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**full_inputs, use_cache=False)
single_logits = single_baseline_out.logits[0]
_, single_mean, _ = sequence_logprob_summary(
single_logits, prompt_length=prompt_len, target_ids=target_ids
)
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
activations = [encoding.activation_for(feature_id) for feature_id in ids]
directions = [sae.decoder_direction(feature_id) for feature_id in ids]
individual_deltas = [
residual_delta(direction, activation, InterventionSpec('ablate', 0.0))
for direction, activation in zip(directions, activations, strict=True)
]
joint_delta = torch.stack(individual_deltas, dim=0).sum(dim=0)
all_deltas = torch.stack([torch.zeros_like(joint_delta), *individual_deltas, joint_delta], dim=0)
repeated = self._repeat_inputs(full_inputs, all_deltas.shape[0])
with self._batch_delta_hook(int(layer), idx, all_deltas):
outputs = self.model(**repeated, use_cache=False)
baseline_logits = outputs.logits[0]
baseline_seq, baseline_mean, _ = sequence_logprob_summary(
baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
baseline_next = baseline_logits[prompt_len - 1]
rows: list[list[object]] = []
individual_mean_deltas: list[float] = []
for feature_idx, feature_id in enumerate(ids, start=1):
logits = outputs.logits[feature_idx]
seq_logp, mean_logp, _ = sequence_logprob_summary(
logits, prompt_length=prompt_len, target_ids=target_ids
)
delta_mean = float(mean_logp - baseline_mean)
individual_mean_deltas.append(delta_mean)
rows.append(
[
f'Feature {feature_id}',
str(feature_id),
float(activations[feature_idx - 1]),
float(torch.linalg.vector_norm(individual_deltas[feature_idx - 1].float()).item()),
delta_mean,
float(seq_logp - baseline_seq),
float(js_divergence_from_logits(baseline_next, logits[prompt_len - 1])),
]
)
joint_logits = outputs.logits[len(ids) + 1]
joint_seq, joint_mean, _ = sequence_logprob_summary(
joint_logits, prompt_length=prompt_len, target_ids=target_ids
)
joint_mean_delta = float(joint_mean - baseline_mean)
rows.append(
[
'Joint ablation',
', '.join(str(x) for x in ids),
float(sum(activations)),
float(torch.linalg.vector_norm(joint_delta.float()).item()),
joint_mean_delta,
float(joint_seq - baseline_seq),
float(js_divergence_from_logits(baseline_next, joint_logits[prompt_len - 1])),
]
)
additive_expected = float(sum(individual_mean_deltas))
interaction_excess = float(joint_mean_delta - additive_expected)
scale = max(sum(abs(value) for value in individual_mean_deltas), 1e-12)
return FeatureInteractionResult(
feature_ids=ids,
target_tokens=[self.tokenizer.decode([int(token_id)]) for token_id in target_ids],
rows=rows,
additive_expected_mean_delta=additive_expected,
joint_mean_delta=joint_mean_delta,
interaction_excess_mean_delta=interaction_excess,
normalized_interaction=float(interaction_excess / scale),
execution_drift_mean_logprob=float(baseline_mean - single_mean),
)
@torch.inference_mode()
def compare_paraphrases(
self,
text_a: str,
text_b: str,
layer: int,
token_index_a: int = -1,
token_index_b: int = -1,
top_n: int = 12,
) -> ParaphraseResult:
if not text_a.strip() or not text_b.strip():
raise ValueError('Enter both the original prompt and a paraphrase.')
a, pooled_a = self._analyze_and_pool(text_a, int(layer), int(token_index_a), max(int(top_n), 12))
b, pooled_b = self._analyze_and_pool(text_b, int(layer), int(token_index_b), max(int(top_n), 12))
map_a = self._encoding_map(a.features)
map_b = self._encoding_map(b.features)
set_a = set(map_a)
set_b = set(map_b)
union = set_a | set_b
jaccard = len(set_a & set_b) / len(union) if union else 1.0
cosine = sparse_topk_cosine(
a.features.indices, a.features.values, b.features.indices, b.features.values
)
promptwide_jaccard = self._dict_jaccard(pooled_a, pooled_b)
promptwide_cosine = self._dict_cosine(pooled_a, pooled_b)
top_ids_a = [int(row[1]) for row in a.rows[: int(top_n)]]
top_ids_b = [int(row[1]) for row in b.rows[: int(top_n)]]
top_union = list(dict.fromkeys(top_ids_a + top_ids_b))
shared_top_n = len(set(top_ids_a) & set(top_ids_b))
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
for feature_id in top_union:
va = float(map_a.get(feature_id, 0.0))
vb = float(map_b.get(feature_id, 0.0))
status = 'shared' if va > 0 and vb > 0 else ('original only' if va > 0 else 'paraphrase only')
rows.append(
[feature_id, va, vb, status, self.catalog.hint(int(layer), feature_id)]
)
chart_rows.append([str(feature_id), 'Original', va])
chart_rows.append([str(feature_id), 'Paraphrase', vb])
rows.sort(key=lambda row: max(float(row[1]), float(row[2])), reverse=True)
return ParaphraseResult(
tokens_a=a.tokens,
token_index_a=a.token_index,
tokens_b=b.tokens,
token_index_b=b.token_index,
topk_jaccard=float(jaccard),
sparse_cosine=float(cosine),
promptwide_jaccard=float(promptwide_jaccard),
promptwide_cosine=float(promptwide_cosine),
shared_top_n=int(shared_top_n),
top_n=int(top_n),
rows=rows,
chart_rows=chart_rows,
)
def _contrast_prompt_rows(self, prompts_per_concept: int) -> list[dict]:
data_path = Path(__file__).resolve().parents[1] / 'data' / 'prompts.jsonl'
rows = [
json.loads(line)
for line in data_path.read_text(encoding='utf-8').splitlines()
if line.strip()
]
selected: list[dict] = []
by_concept: dict[str, list[dict]] = {}
for row in rows:
# Use one wording per paraphrase pair so the live contrast is not dominated by near-duplicates.
if int(row.get('variant', 0)) != 0:
continue
by_concept.setdefault(str(row['concept']), []).append(row)
for concept in sorted(by_concept):
selected.extend(by_concept[concept][: int(prompts_per_concept)])
return selected
@torch.inference_mode()
def feature_token_trace(
self,
text: str,
layer: int,
feature_id: int,
) -> FeatureTraceResult:
"""Trace one SAE feature across every non-padding token in a prompt."""
if not text.strip():
raise ValueError('Enter a prompt first.')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
if int(layer) not in self.settings.layers:
raise ValueError(f'Layer must be one of {self.settings.layers}.')
if int(feature_id) < 0 or int(feature_id) >= self.settings.sae_width:
raise ValueError(f'Feature id must be in [0, {self.settings.sae_width - 1}].')
inputs = self._inputs(text)
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**inputs, use_cache=False)
hidden = capture['hidden'][0]
sae = self.sae_store.get(int(layer))
encoding = sae.encode(hidden)
mask = encoding.indices == int(feature_id)
activations = torch.where(mask, encoding.values, torch.zeros_like(encoding.values)).sum(dim=-1)
values = activations.detach().float().cpu().tolist()
ids = inputs['input_ids'][0].tolist()
tokens = [self.tokenizer.decode([int(token_id)]) for token_id in ids]
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
for idx, (token, value) in enumerate(zip(tokens, values, strict=True)):
value = float(value)
rows.append([idx, repr(token), value, bool(value > 0)])
chart_rows.append([f'{idx}: {token if token.strip() else repr(token)}', value])
active = [idx for idx, value in enumerate(values) if float(value) > 0]
if active:
max_idx = max(active, key=lambda idx: float(values[idx]))
max_value = float(values[max_idx])
else:
max_idx = None
max_value = 0.0
return FeatureTraceResult(
feature_id=int(feature_id),
layer=int(layer),
tokens=tokens,
rows=rows,
chart_rows=chart_rows,
active_token_count=len(active),
token_count=len(tokens),
max_activation=max_value,
max_token_index=max_idx,
)
@torch.inference_mode()
def feature_geometry(
self,
text: str,
layer: int,
token_index: int,
feature_ids: Sequence[int],
) -> FeatureGeometryResult:
"""Inspect pairwise SAE decoder geometry and activation-weighted ablation geometry."""
if not text.strip():
raise ValueError('Enter a prompt first.')
ids = list(dict.fromkeys(int(x) for x in feature_ids))
if len(ids) < 2:
raise ValueError('Select at least two distinct features for geometry analysis.')
if len(ids) > 8:
raise ValueError('Geometry analysis supports at most eight features in the live app.')
if any(feature_id < 0 or feature_id >= self.settings.sae_width for feature_id in ids):
raise ValueError(f'Feature ids must be in [0, {self.settings.sae_width - 1}].')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.sae_store is not None
inputs = self._inputs(text)
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**inputs, use_cache=False)
hidden = capture['hidden'][0]
idx = self._resolve_index(int(token_index), hidden.shape[0])
sae = self.sae_store.get(int(layer))
encoding = sae.encode(hidden[idx])
activations = [float(encoding.activation_for(feature_id)) for feature_id in ids]
directions = torch.stack([sae.decoder_direction(feature_id).float() for feature_id in ids])
gram = decoder_cosine_matrix(directions)
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
offdiag: list[float] = []
for i in range(len(ids)):
for j in range(i + 1, len(ids)):
cosine = float(gram[i, j].item())
offdiag.append(abs(cosine))
rows.append([ids[i], ids[j], activations[i], activations[j], cosine])
chart_rows.append([f'{ids[i]} ↔ {ids[j]}', cosine])
individual_deltas = torch.stack(
[-float(activation) * direction for activation, direction in zip(activations, directions, strict=True)]
)
joint_norm, independent_norm, alignment_ratio = joint_direction_norm_ratio(individual_deltas)
return FeatureGeometryResult(
feature_ids=ids,
layer=int(layer),
rows=rows,
chart_rows=chart_rows,
mean_abs_decoder_cosine=float(sum(offdiag) / len(offdiag)) if offdiag else 0.0,
max_abs_decoder_cosine=float(max(offdiag)) if offdiag else 0.0,
joint_ablation_norm=joint_norm,
independent_norm=independent_norm,
alignment_ratio=float(alignment_ratio),
)
@torch.inference_mode()
def contrastive_intervention(
self,
text: str,
layer: int,
token_index: int,
feature_id: int,
mode: str,
coefficient: float,
target_a: str,
target_b: str,
) -> ContrastiveCausalResult:
"""Measure whether an SAE edit shifts preference between two exact continuations."""
if not text.strip():
raise ValueError('Enter a prompt first.')
if not target_a:
raise ValueError('Enter preferred continuation A.')
if not target_b:
raise ValueError('Enter comparison continuation B.')
if target_a == target_b:
raise ValueError('Continuations A and B must be different.')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**prompt_inputs, use_cache=False)
residual = capture['hidden'][0, idx]
sae = self.sae_store.get(int(layer))
encoding = sae.encode(residual)
activation = float(encoding.activation_for(int(feature_id)))
delta = residual_delta(
sae.decoder_direction(int(feature_id)),
activation,
InterventionSpec(mode=mode, coefficient=float(coefficient)),
)
seed = self._control_seed(text, int(layer), str(int(feature_id)), mode, float(coefficient))
controls = self._random_controls(delta, seed=seed, count=self.settings.live_random_controls)
deltas = torch.stack([torch.zeros_like(delta), delta, *controls], dim=0)
def score(target_text: str):
target_ids = self._target_ids(target_text)
scoring_inputs = self._append_target(prompt_inputs, target_ids)
repeated = self._repeat_inputs(scoring_inputs, deltas.shape[0])
with self._batch_delta_hook(int(layer), idx, deltas):
output = self.model(**repeated, use_cache=False)
summaries = [
sequence_logprob_summary(output.logits[row], prompt_length=prompt_len, target_ids=target_ids)
for row in range(output.logits.shape[0])
]
return target_ids, summaries
ids_a, scores_a = score(target_a)
ids_b, scores_b = score(target_b)
base_a_seq, base_a_mean, _ = scores_a[0]
edit_a_seq, edit_a_mean, _ = scores_a[1]
base_b_seq, base_b_mean, _ = scores_b[0]
edit_b_seq, edit_b_mean, _ = scores_b[1]
baseline_log_odds, modified_log_odds, delta_log_odds = contrastive_log_odds(
base_a_seq, edit_a_seq, base_b_seq, edit_b_seq
)
baseline_norm_pref, modified_norm_pref, delta_norm_pref = contrastive_log_odds(
base_a_mean, edit_a_mean, base_b_mean, edit_b_mean
)
random_delta_log_odds: list[float] = []
for row in range(2, len(scores_a)):
random_a_seq = float(scores_a[row][0])
random_b_seq = float(scores_b[row][0])
random_delta_log_odds.append(
float((random_a_seq - random_b_seq) - baseline_log_odds)
)
random_signed, random_abs, random_std, empirical_p = self._random_effect_summary(
random_delta_log_odds, delta_log_odds
)
ratio = abs(delta_log_odds) / max(random_abs, 1e-12)
rows = [
[
'A (preferred)',
target_a,
len(ids_a),
float(base_a_seq),
float(edit_a_seq),
float(edit_a_seq - base_a_seq),
float(base_a_mean),
float(edit_a_mean),
float(edit_a_mean - base_a_mean),
],
[
'B (comparison)',
target_b,
len(ids_b),
float(base_b_seq),
float(edit_b_seq),
float(edit_b_seq - base_b_seq),
float(base_b_mean),
float(edit_b_mean),
float(edit_b_mean - base_b_mean),
],
]
return ContrastiveCausalResult(
feature_id=int(feature_id),
layer=int(layer),
feature_activation=activation,
perturbation_norm=float(torch.linalg.vector_norm(delta.float()).item()),
target_a_tokens=[self.tokenizer.decode([int(x)]) for x in ids_a],
target_b_tokens=[self.tokenizer.decode([int(x)]) for x in ids_b],
rows=rows,
baseline_log_odds=baseline_log_odds,
modified_log_odds=modified_log_odds,
delta_log_odds=delta_log_odds,
baseline_normalized_preference=baseline_norm_pref,
modified_normalized_preference=modified_norm_pref,
delta_normalized_preference=delta_norm_pref,
random_signed_mean_delta=random_signed,
random_abs_mean_delta=random_abs,
random_delta_std=random_std,
specificity_ratio=float(ratio),
empirical_p=float(empirical_p),
random_control_count=len(controls),
)
@torch.inference_mode()
def concept_contrast_scan(
self,
feature_id: int,
layer: int,
prompts_per_concept: int | None = None,
) -> ConceptContrastResult:
"""Measure one SAE feature using prompt-wide max activation on a balanced concept batch."""
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
if int(feature_id) < 0 or int(feature_id) >= self.settings.sae_width:
raise ValueError(f'Feature id must be in [0, {self.settings.sae_width - 1}].')
n = int(prompts_per_concept or self.settings.contrast_prompts_per_concept)
if n < 1 or n > 8:
raise ValueError('Contrast prompts per concept must be between 1 and 8.')
rows = self._contrast_prompt_rows(n)
if not rows:
raise RuntimeError('No controlled contrast prompts are available.')
texts = [str(row['text']) for row in rows]
batch = self.tokenizer(
texts,
return_tensors='pt',
padding=True,
truncation=True,
max_length=self.settings.max_prompt_tokens,
)
batch = {key: value.to(self.device) for key, value in batch.items()}
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**batch, use_cache=False)
residuals = capture['hidden']
sae = self.sae_store.get(int(layer))
encoding = sae.encode(residuals)
feature_mask = encoding.indices == int(feature_id)
token_activations = torch.where(
feature_mask, encoding.values, torch.zeros_like(encoding.values)
).sum(dim=-1)
attention = batch.get('attention_mask', torch.ones_like(batch['input_ids'])).bool()
token_activations = torch.where(attention, token_activations, torch.zeros_like(token_activations))
prompt_activations = token_activations.max(dim=1).values
values = prompt_activations.detach().float().cpu().tolist()
grouped: dict[str, list[float]] = {}
for row, value in zip(rows, values, strict=True):
grouped.setdefault(str(row['concept']), []).append(float(value))
table_rows: list[list[object]] = []
for concept in sorted(grouped):
vals = grouped[concept]
tensor = torch.tensor(vals, dtype=torch.float64)
active = [value for value in vals if value > 0]
positive_mean = float(sum(active) / len(active)) if active else 0.0
table_rows.append(
[
concept,
len(vals),
float(tensor.mean().item()),
float(torch.median(tensor).item()),
float(len(active) / len(vals)),
positive_mean,
float(max(vals) if vals else 0.0),
]
)
table_rows.sort(key=lambda row: (float(row[2]), float(row[4]), float(row[6])), reverse=True)
active_prompt_count = sum(1 for value in values if float(value) > 0)
if table_rows and float(table_rows[0][2]) > 0:
leader: str | None = str(table_rows[0][0])
first = float(table_rows[0][2])
second = float(table_rows[1][2]) if len(table_rows) > 1 else 0.0
ratio: float | None = first / second if second > 0 else None
else:
leader = None
ratio = None
chart_rows = [[str(row[0]), float(row[2])] for row in table_rows]
return ConceptContrastResult(
feature_id=int(feature_id),
layer=int(layer),
prompts_per_concept=n,
rows=table_rows,
chart_rows=chart_rows,
leading_concept=leader,
leading_ratio=ratio,
active_prompt_count=int(active_prompt_count),
total_prompt_count=len(values),
)
@torch.inference_mode()
def concept_feature_discovery(
self,
concept: str,
layer: int,
prompts_per_concept: int | None = None,
top_n: int = 12,
ranking_mode: str = 'balanced_selectivity',
current_text: str | None = None,
current_token_index: int = -1,
) -> ConceptFeatureDiscoveryResult:
"""Find exploratory concept candidates and show whether they are usable in the current Workbench context.
``balanced_selectivity`` downweights globally high-activation features by combining target selectivity,
target coverage, and target activation magnitude. ``raw_mean_difference`` preserves the simpler raw
mean-difference ranking for comparison. ``causal_ready`` further requires activation at the currently
selected Workbench token and ranks those compatible candidates by balanced evidence plus a log-scaled
current-token activation term. If a Workbench prompt is supplied it is appended to the same model batch,
so current-prompt compatibility does not require another forward pass.
"""
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
concept = str(concept).strip()
if not concept:
raise ValueError('Choose a target concept.')
n = int(prompts_per_concept or self.settings.contrast_prompts_per_concept)
if n < 1 or n > 8:
raise ValueError('Discovery prompts per concept must be between 1 and 8.')
top_n = int(top_n)
if top_n < 1 or top_n > 25:
raise ValueError('Number of candidate features must be between 1 and 25.')
if ranking_mode not in {'balanced_selectivity', 'raw_mean_difference', 'causal_ready'}:
raise ValueError(
"ranking_mode must be 'balanced_selectivity', 'raw_mean_difference', or 'causal_ready'."
)
rows = self._contrast_prompt_rows(n)
available = sorted({str(row['concept']) for row in rows})
if concept not in available:
raise ValueError(f'Concept must be one of {available}.')
current_text = str(current_text or '').strip()
current_context_available = bool(current_text)
texts = [str(row['text']) for row in rows]
if current_context_available:
texts.append(current_text)
batch = self.tokenizer(
texts,
return_tensors='pt',
padding=True,
truncation=True,
max_length=self.settings.max_prompt_tokens,
)
batch = {key: value.to(self.device) for key, value in batch.items()}
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**batch, use_cache=False)
sae = self.sae_store.get(int(layer))
encoding = sae.encode(capture['hidden'])
attention = batch.get('attention_mask', torch.ones_like(batch['input_ids'])).bool()
valid = attention.unsqueeze(-1).expand_as(encoding.values)
values = torch.where(valid, encoding.values, torch.zeros_like(encoding.values)).float()
dense = torch.zeros(
(values.shape[0], self.settings.sae_width),
device=values.device,
dtype=torch.float32,
)
dense.scatter_reduce_(
1,
encoding.indices.reshape(values.shape[0], -1),
values.reshape(values.shape[0], -1),
reduce='amax',
include_self=True,
)
controlled_count = len(rows)
controlled_dense = dense[:controlled_count]
target_mask = torch.tensor(
[str(row['concept']) == concept for row in rows],
device=dense.device,
dtype=torch.bool,
)
other_mask = ~target_mask
target = controlled_dense[target_mask]
other = controlled_dense[other_mask]
target_mean = target.mean(dim=0)
other_mean = other.mean(dim=0)
target_rate = (target > 0).float().mean(dim=0)
other_rate = (other > 0).float().mean(dim=0)
mean_diff = target_mean - other_mean
selectivity = mean_diff / (target_mean + other_mean + 1e-8)
# A scale-aware but selectivity-first exploratory score. log1p prevents very large SAE coefficients from
# overwhelming features that are much more exclusive to the target concept.
balanced_score = selectivity.clamp_min(0) * target_rate * torch.log1p(target_mean.clamp_min(0))
current_prompt_max = torch.zeros(self.settings.sae_width, device=dense.device, dtype=torch.float32)
current_token_dense = torch.zeros_like(current_prompt_max)
resolved_current_idx: int | None = None
if current_context_available:
current_row = controlled_count
current_prompt_max = dense[current_row]
valid_positions = torch.nonzero(attention[current_row], as_tuple=False).flatten()
prompt_len = int(valid_positions.numel())
if prompt_len:
resolved_current_idx = self._resolve_index(int(current_token_index), prompt_len)
padded_position = int(valid_positions[resolved_current_idx].item())
token_ids = encoding.indices[current_row, padded_position]
token_values = encoding.values[current_row, padded_position].float().clamp_min(0)
current_token_dense.scatter_reduce_(
0,
token_ids,
token_values,
reduce='amax',
include_self=True,
)
# Stability diagnostics reuse the same controlled activation batch, so they add no model inference.
# Split-half overlap is intentionally simple; deterministic balanced bootstrap support gives a second
# view of how often each displayed feature survives small changes to the live prompt sample.
split_half_k: int | None = None
split_half_shared_count = 0
split_half_jaccard: float | None = None
resample_replicates = 0
resample_rank_lists: list[list[int]] = []
row_concepts = [str(row['concept']) for row in rows]
def _rank_indices(row_indices: list[int]) -> list[int]:
if not row_indices:
return []
index_tensor = torch.tensor(row_indices, device=dense.device, dtype=torch.long)
sub_dense = controlled_dense.index_select(0, index_tensor)
sub_labels = [row_concepts[index] for index in row_indices]
sub_target_mask = torch.tensor(
[label == concept for label in sub_labels],
device=dense.device,
dtype=torch.bool,
)
sub_other_mask = ~sub_target_mask
sub_target = sub_dense[sub_target_mask]
sub_other = sub_dense[sub_other_mask]
if sub_target.shape[0] == 0 or sub_other.shape[0] == 0:
return []
sub_target_mean = sub_target.mean(dim=0)
sub_other_mean = sub_other.mean(dim=0)
sub_target_rate = (sub_target > 0).float().mean(dim=0)
sub_mean_diff = sub_target_mean - sub_other_mean
sub_selectivity = sub_mean_diff / (sub_target_mean + sub_other_mean + 1e-8)
sub_balanced = (
sub_selectivity.clamp_min(0)
* sub_target_rate
* torch.log1p(sub_target_mean.clamp_min(0))
)
sub_eligible = (sub_target_mean > 0) & (sub_mean_diff > 0)
if ranking_mode == 'causal_ready':
sub_eligible = sub_eligible & (current_token_dense > 0)
sub_ranking = sub_balanced * torch.log1p(current_token_dense.clamp_min(0))
elif ranking_mode == 'balanced_selectivity':
sub_ranking = sub_balanced
else:
sub_ranking = sub_mean_diff
sub_idx = torch.nonzero(sub_eligible, as_tuple=False).flatten()
if sub_idx.numel() == 0:
return []
sub_order = torch.argsort(sub_ranking[sub_idx], descending=True)
return [int(value.item()) for value in sub_idx[sub_order[:top_n]]]
if n >= 2:
half = max(1, n // 2)
seen_by_concept: dict[str, int] = {}
indices_a: list[int] = []
indices_b: list[int] = []
for row_idx, row_concept in enumerate(row_concepts):
local_idx = seen_by_concept.get(row_concept, 0)
seen_by_concept[row_concept] = local_idx + 1
if local_idx < half:
indices_a.append(row_idx)
else:
indices_b.append(row_idx)
ids_a = _rank_indices(indices_a)
ids_b = _rank_indices(indices_b)
if ids_a and ids_b:
set_a, set_b = set(ids_a), set(ids_b)
shared = set_a & set_b
union = set_a | set_b
split_half_k = min(len(ids_a), len(ids_b), top_n)
split_half_shared_count = len(shared)
split_half_jaccard = float(len(shared) / len(union)) if union else 1.0
pools: dict[str, list[int]] = {}
for row_idx, row_concept in enumerate(row_concepts):
pools.setdefault(row_concept, []).append(row_idx)
seed = 13013 + int(layer) * 97 + n * 17 + sum(ord(ch) for ch in concept)
rng = random.Random(seed)
resample_replicates = 32
for _ in range(resample_replicates):
sampled_indices: list[int] = []
for row_concept in sorted(pools):
pool = pools[row_concept]
sampled_indices.extend(rng.choice(pool) for _ in range(len(pool)))
ranked = _rank_indices(sampled_indices)
if ranked:
resample_rank_lists.append(ranked)
resample_replicates = len(resample_rank_lists)
eligible = (target_mean > 0) & (mean_diff > 0)
if ranking_mode == 'causal_ready':
if not current_context_available or resolved_current_idx is None:
raise ValueError(
"Causal-ready ranking requires a current Workbench prompt/token. "
"Set the Workbench prompt/layer/token first."
)
eligible = eligible & (current_token_dense > 0)
ranking_values = balanced_score * torch.log1p(current_token_dense.clamp_min(0))
elif ranking_mode == 'balanced_selectivity':
ranking_values = balanced_score
else:
ranking_values = mean_diff
candidate_idx = torch.nonzero(eligible, as_tuple=False).flatten()
if candidate_idx.numel() == 0:
return ConceptFeatureDiscoveryResult(
concept=concept,
layer=int(layer),
prompts_per_concept=n,
top_n=top_n,
ranking_mode=ranking_mode,
rows=[],
chart_rows=[],
candidate_ids=[],
default_candidate_id=None,
current_context_available=current_context_available,
current_token_index=resolved_current_idx,
displayed_current_active_count=0,
split_half_k=split_half_k,
split_half_shared_count=split_half_shared_count,
split_half_jaccard=split_half_jaccard,
resample_replicates=resample_replicates,
resample_mean_support=None,
resample_high_support_count=0,
)
order = torch.argsort(ranking_values[candidate_idx], descending=True)
candidate_idx = candidate_idx[order[:top_n]]
resample_support: dict[int, float] = {}
resample_median_rank: dict[int, float | None] = {}
displayed_supports: list[float] = []
if resample_replicates:
for feature_tensor in candidate_idx:
fid = int(feature_tensor.item())
ranks = [ranked.index(fid) + 1 for ranked in resample_rank_lists if fid in ranked]
support = len(ranks) / resample_replicates
if ranks:
ordered_ranks = sorted(ranks)
midpoint = len(ordered_ranks) // 2
if len(ordered_ranks) % 2:
median_rank = float(ordered_ranks[midpoint])
else:
median_rank = float((ordered_ranks[midpoint - 1] + ordered_ranks[midpoint]) / 2)
else:
median_rank = None
resample_support[fid] = float(support)
resample_median_rank[fid] = median_rank
displayed_supports.append(float(support))
resample_mean_support = (
float(sum(displayed_supports) / len(displayed_supports)) if displayed_supports else None
)
resample_high_support_count = sum(support >= 0.75 for support in displayed_supports)
table_rows: list[list[object]] = []
chart_rows: list[list[object]] = []
default_candidate_id: int | None = None
for rank, feature_tensor in enumerate(candidate_idx, start=1):
fid = int(feature_tensor.item())
current_max = float(current_prompt_max[fid].item()) if current_context_available else 0.0
current_token = float(current_token_dense[fid].item()) if current_context_available else 0.0
if default_candidate_id is None and current_token > 0:
default_candidate_id = fid
score = float(ranking_values[fid].item())
row = [
rank,
fid,
score,
float(target_mean[fid].item()),
float(other_mean[fid].item()),
float(mean_diff[fid].item()),
float(selectivity[fid].item()),
float(target_rate[fid].item()),
float(other_rate[fid].item()),
current_max,
current_token,
bool(current_token > 0),
resample_support.get(fid) if resample_replicates else None,
resample_median_rank.get(fid) if resample_replicates else None,
]
table_rows.append(row)
chart_rows.append([str(fid), score])
if default_candidate_id is None and candidate_idx.numel():
default_candidate_id = int(candidate_idx[0].item())
return ConceptFeatureDiscoveryResult(
concept=concept,
layer=int(layer),
prompts_per_concept=n,
top_n=top_n,
ranking_mode=ranking_mode,
rows=table_rows,
chart_rows=chart_rows,
candidate_ids=[int(x.item()) for x in candidate_idx],
default_candidate_id=default_candidate_id,
current_context_available=current_context_available,
current_token_index=resolved_current_idx,
displayed_current_active_count=sum(bool(row[11]) for row in table_rows),
split_half_k=split_half_k,
split_half_shared_count=split_half_shared_count,
split_half_jaccard=split_half_jaccard,
resample_replicates=resample_replicates,
resample_mean_support=resample_mean_support,
resample_high_support_count=resample_high_support_count,
)
@torch.inference_mode()
def candidate_causal_screen(
self,
text: str,
layer: int,
token_index: int,
feature_ids: Sequence[int],
target_text: str,
) -> CandidateCausalScreenResult:
"""Cheaply triage several candidate features with one batched ablation screen.
This deliberately omits random controls. Its purpose is to rank candidates before
spending a full live causal test (with the random-control ensemble) on one or two
promising features. All feature ablations share the same batched zero-edit reference.
"""
if not text.strip():
raise ValueError('Enter and inspect a Workbench prompt first.')
if not target_text.strip():
raise ValueError('Enter a target continuation for candidate causal screening.')
ids = list(dict.fromkeys(int(x) for x in feature_ids))
if not ids:
raise ValueError('Select at least one candidate feature to screen.')
if len(ids) > 8:
raise ValueError('Candidate causal screening supports at most eight features per run.')
if any(feature_id < 0 or feature_id >= self.settings.sae_width for feature_id in ids):
raise ValueError(f'Feature ids must be in [0, {self.settings.sae_width - 1}].')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
target_ids = self._target_ids(target_text)
full_inputs = self._append_target(prompt_inputs, target_ids)
sae = self.sae_store.get(int(layer))
capture: dict = {}
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**full_inputs, use_cache=False)
single_logits = single_baseline_out.logits[0]
_, single_mean, _ = sequence_logprob_summary(
single_logits, prompt_length=prompt_len, target_ids=target_ids
)
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
activations = [float(encoding.activation_for(feature_id)) for feature_id in ids]
deltas = [
residual_delta(
sae.decoder_direction(feature_id),
activation,
InterventionSpec('ablate', 0.0),
)
for feature_id, activation in zip(ids, activations, strict=True)
]
all_deltas = torch.stack([torch.zeros_like(deltas[0]), *deltas], dim=0)
repeated = self._repeat_inputs(full_inputs, all_deltas.shape[0])
with self._batch_delta_hook(int(layer), idx, all_deltas):
outputs = self.model(**repeated, use_cache=False)
baseline_logits = outputs.logits[0]
baseline_seq, baseline_mean, _ = sequence_logprob_summary(
baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
baseline_next = baseline_logits[prompt_len - 1]
execution_drift_mean = float(baseline_mean - single_mean)
execution_drift_js = js_divergence_from_logits(
single_logits[prompt_len - 1], baseline_next
)
scored: list[dict[str, object]] = []
for row_idx, (feature_id, activation, delta) in enumerate(
zip(ids, activations, deltas, strict=True), start=1
):
logits = outputs.logits[row_idx]
seq_logp, mean_logp, _ = sequence_logprob_summary(
logits, prompt_length=prompt_len, target_ids=target_ids
)
mean_delta = float(mean_logp - baseline_mean)
seq_delta = float(seq_logp - baseline_seq)
js = float(js_divergence_from_logits(baseline_next, logits[prompt_len - 1]))
norm = float(torch.linalg.vector_norm(delta.float()).item())
scored.append(
{
'feature_id': int(feature_id),
'activation': float(activation),
'active': bool(activation > 0),
'norm': norm,
'mean_delta': mean_delta,
'seq_delta': seq_delta,
'js': js,
}
)
scored.sort(
key=lambda item: (abs(float(item['mean_delta'])), float(item['js'])),
reverse=True,
)
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
for rank, item in enumerate(scored, start=1):
feature_id = int(item['feature_id'])
mean_delta = float(item['mean_delta'])
rows.append(
[
rank,
feature_id,
float(item['activation']),
bool(item['active']),
float(item['norm']),
mean_delta,
float(item['seq_delta']),
float(item['js']),
]
)
chart_rows.append([str(feature_id), mean_delta])
return CandidateCausalScreenResult(
feature_ids=[int(item['feature_id']) for item in scored],
target_tokens=[self.tokenizer.decode([int(token_id)]) for token_id in target_ids],
rows=rows,
chart_rows=chart_rows,
active_feature_count=sum(bool(item['active']) for item in scored),
candidate_count=len(scored),
execution_drift_mean_logprob=execution_drift_mean,
execution_drift_js=float(execution_drift_js),
)
@torch.inference_mode()
def candidate_specificity_screen(
self,
text: str,
layer: int,
token_index: int,
feature_ids: Sequence[int],
target_text: str,
) -> CandidateSpecificityResult:
"""Compare a small candidate set against per-feature norm-matched random ensembles.
This is the controlled follow-up to :meth:`candidate_causal_screen`. Each SAE
ablation gets its own deterministic random-control ensemble with the same L2 norm.
All targeted and random conditions share one batched zero-edit reference so the
comparison is both GPU-efficient and execution-context consistent.
"""
if not text.strip():
raise ValueError('Enter and inspect a Workbench prompt first.')
if not target_text.strip():
raise ValueError('Enter a target continuation for controlled candidate comparison.')
ids = list(dict.fromkeys(int(x) for x in feature_ids))
if not ids:
raise ValueError('Select at least one candidate feature for controlled comparison.')
if len(ids) > 3:
raise ValueError('Controlled candidate comparison supports at most three features per run.')
if any(feature_id < 0 or feature_id >= self.settings.sae_width for feature_id in ids):
raise ValueError(f'Feature ids must be in [0, {self.settings.sae_width - 1}].')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
target_ids = self._target_ids(target_text)
full_inputs = self._append_target(prompt_inputs, target_ids)
sae = self.sae_store.get(int(layer))
capture: dict = {}
with self._capture_hook(int(layer), capture):
single_baseline_out = self.model(**full_inputs, use_cache=False)
single_logits = single_baseline_out.logits[0]
_, single_mean, _ = sequence_logprob_summary(
single_logits, prompt_length=prompt_len, target_ids=target_ids
)
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
activations = [float(encoding.activation_for(feature_id)) for feature_id in ids]
zero = torch.zeros_like(residual)
all_deltas: list[torch.Tensor] = [zero]
metadata: list[tuple[int, str, float, float]] = []
for feature_id, activation in zip(ids, activations, strict=True):
delta = residual_delta(
sae.decoder_direction(feature_id),
activation,
InterventionSpec('ablate', 0.0),
)
norm = float(torch.linalg.vector_norm(delta.float()).item())
all_deltas.append(delta)
metadata.append((feature_id, 'sae', activation, norm))
seed = self._control_seed(text, int(layer), str(feature_id), 'candidate_specificity', 0.0)
controls = self._random_controls(
delta, seed=seed, count=self.settings.live_random_controls
)
for control_idx, control in enumerate(controls):
all_deltas.append(control)
metadata.append((feature_id, f'random_{control_idx}', activation, norm))
delta_batch = torch.stack(all_deltas, dim=0)
repeated = self._repeat_inputs(full_inputs, delta_batch.shape[0])
with self._batch_delta_hook(int(layer), idx, delta_batch):
outputs = self.model(**repeated, use_cache=False)
baseline_logits = outputs.logits[0]
baseline_seq, baseline_mean, _ = sequence_logprob_summary(
baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
baseline_next = baseline_logits[prompt_len - 1]
execution_drift_mean = float(baseline_mean - single_mean)
execution_drift_js = js_divergence_from_logits(
single_logits[prompt_len - 1], baseline_next
)
grouped: dict[int, dict[str, object]] = {
feature_id: {
'activation': activation,
'norm': 0.0,
'sae': None,
'random': [],
}
for feature_id, activation in zip(ids, activations, strict=True)
}
for output_idx, meta in enumerate(metadata, start=1):
feature_id, kind, activation, norm = meta
logits = outputs.logits[output_idx]
seq_logp, mean_logp, _ = sequence_logprob_summary(
logits, prompt_length=prompt_len, target_ids=target_ids
)
item = {
'seq_delta': float(seq_logp - baseline_seq),
'mean_delta': float(mean_logp - baseline_mean),
'js': float(js_divergence_from_logits(baseline_next, logits[prompt_len - 1])),
}
grouped[feature_id]['activation'] = float(activation)
grouped[feature_id]['norm'] = float(norm)
if kind == 'sae':
grouped[feature_id]['sae'] = item
else:
random_items = grouped[feature_id]['random']
assert isinstance(random_items, list)
random_items.append(item)
scored: list[dict[str, object]] = []
for feature_id in ids:
group = grouped[feature_id]
sae_item = group['sae']
random_items = group['random']
assert isinstance(sae_item, dict)
assert isinstance(random_items, list) and random_items
mean_delta = float(sae_item['mean_delta'])
js = float(sae_item['js'])
random_mean_deltas = [float(item['mean_delta']) for item in random_items]
random_js_values = [float(item['js']) for item in random_items]
random_signed, random_abs, random_std, target_p = self._random_effect_summary(
random_mean_deltas, mean_delta
)
_, random_js_abs, random_js_std, js_p = self._random_effect_summary(
random_js_values, js
)
scored.append(
{
'feature_id': int(feature_id),
'activation': float(group['activation']),
'active': bool(float(group['activation']) > 0),
'norm': float(group['norm']),
'mean_delta': mean_delta,
'seq_delta': float(sae_item['seq_delta']),
'random_signed': float(random_signed),
'random_abs': float(random_abs),
'random_std': float(random_std),
'target_ratio': float(abs(mean_delta) / max(random_abs, 1e-12)),
'target_p': float(target_p),
'js': js,
'random_js': float(random_js_abs),
'random_js_std': float(random_js_std),
'js_ratio': float(js / max(random_js_abs, 1e-12)),
'js_p': float(js_p),
}
)
scored.sort(
key=lambda item: (float(item['target_ratio']), abs(float(item['mean_delta']))),
reverse=True,
)
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
for rank, item in enumerate(scored, start=1):
feature_id = int(item['feature_id'])
rows.append(
[
rank,
feature_id,
float(item['activation']),
bool(item['active']),
float(item['norm']),
float(item['mean_delta']),
float(item['random_signed']),
float(item['random_abs']),
float(item['random_std']),
float(item['target_ratio']),
float(item['target_p']),
float(item['seq_delta']),
float(item['js']),
float(item['random_js']),
float(item['random_js_std']),
float(item['js_ratio']),
float(item['js_p']),
]
)
chart_rows.extend(
[
[str(feature_id), 'Target specificity', float(item['target_ratio'])],
[str(feature_id), 'JS specificity', float(item['js_ratio'])],
]
)
return CandidateSpecificityResult(
feature_ids=[int(item['feature_id']) for item in scored],
target_tokens=[self.tokenizer.decode([int(token_id)]) for token_id in target_ids],
rows=rows,
chart_rows=chart_rows,
active_feature_count=sum(bool(item['active']) for item in scored),
candidate_count=len(scored),
random_control_count=self.settings.live_random_controls,
execution_drift_mean_logprob=execution_drift_mean,
execution_drift_js=float(execution_drift_js),
)
@torch.inference_mode()
def candidate_cross_target_profile(
self,
text: str,
layer: int,
token_index: int,
feature_ids: Sequence[int],
targets: Sequence[str],
) -> CandidateCrossTargetResult:
"""Profile native candidate ablations across several exact target continuations.
This is a screening diagnostic, not a random-controlled causal claim. The residual
representation and native feature deltas are captured once from the Workbench prompt.
Each target is then evaluated with one small batched forward containing the zero-edit
reference plus every selected candidate ablation.
"""
if not text.strip():
raise ValueError('Enter and inspect a Workbench prompt first.')
ids = list(dict.fromkeys(int(x) for x in feature_ids))
if not ids:
raise ValueError('Select at least one candidate feature for cross-target profiling.')
if len(ids) > 3:
raise ValueError('Cross-target profiling supports at most three features per run.')
if any(feature_id < 0 or feature_id >= self.settings.sae_width for feature_id in ids):
raise ValueError(f'Feature ids must be in [0, {self.settings.sae_width - 1}].')
target_list: list[str] = []
seen_targets: set[str] = set()
for raw in targets:
target = str(raw)
if not target.strip() or target in seen_targets:
continue
seen_targets.add(target)
target_list.append(target)
if len(target_list) < 2:
raise ValueError('Enter at least two distinct target continuations.')
if len(target_list) > 5:
raise ValueError('Cross-target profiling supports at most five target continuations.')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompt_inputs = self._inputs(text)
prompt_len = int(prompt_inputs['input_ids'].shape[1])
idx = self._resolve_index(int(token_index), prompt_len)
sae = self.sae_store.get(int(layer))
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**prompt_inputs, use_cache=False)
residual = capture['hidden'][0, idx]
encoding = sae.encode(residual)
activations = [float(encoding.activation_for(feature_id)) for feature_id in ids]
deltas = [
residual_delta(
sae.decoder_direction(feature_id),
activation,
InterventionSpec('ablate', 0.0),
)
for feature_id, activation in zip(ids, activations, strict=True)
]
delta_batch = torch.stack([torch.zeros_like(residual), *deltas], dim=0)
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
by_feature: dict[int, list[tuple[str, float, float]]] = {feature_id: [] for feature_id in ids}
for target_text in target_list:
target_ids = self._target_ids(target_text)
full_inputs = self._append_target(prompt_inputs, target_ids)
repeated = self._repeat_inputs(full_inputs, delta_batch.shape[0])
with self._batch_delta_hook(int(layer), idx, delta_batch):
outputs = self.model(**repeated, use_cache=False)
baseline_logits = outputs.logits[0]
baseline_seq, baseline_mean, _ = sequence_logprob_summary(
baseline_logits, prompt_length=prompt_len, target_ids=target_ids
)
baseline_next = baseline_logits[prompt_len - 1]
token_count = len(target_ids)
for output_idx, (feature_id, activation, delta) in enumerate(
zip(ids, activations, deltas, strict=True), start=1
):
logits = outputs.logits[output_idx]
seq_logp, mean_logp, _ = sequence_logprob_summary(
logits, prompt_length=prompt_len, target_ids=target_ids
)
mean_delta = float(mean_logp - baseline_mean)
seq_delta = float(seq_logp - baseline_seq)
js = float(js_divergence_from_logits(baseline_next, logits[prompt_len - 1]))
rows.append(
[
int(feature_id),
target_text,
int(token_count),
float(activation),
float(torch.linalg.vector_norm(delta.float()).item()),
mean_delta,
seq_delta,
js,
]
)
chart_rows.append([target_text, str(feature_id), mean_delta])
by_feature[feature_id].append((target_text, mean_delta, js))
summary_rows: list[list[object]] = []
pairwise_rows: list[list[object]] = []
for feature_id in ids:
items = by_feature[feature_id]
strongest = max(items, key=lambda item: abs(item[1]))
strongest_abs = abs(float(strongest[1]))
other_abs = [abs(float(item[1])) for item in items if item is not strongest]
mean_other = float(sum(other_abs) / len(other_abs)) if other_abs else 0.0
profile_ratio = float(strongest_abs / max(mean_other, 1e-12))
deltas = [float(item[1]) for item in items]
abs_values = [abs(delta) for delta in deltas]
total_abs = float(sum(abs_values))
if total_abs > 0 and len(abs_values) > 1:
proportions = [value / total_abs for value in abs_values if value > 0]
normalized_entropy = float(
-sum(value * math.log(value) for value in proportions) / math.log(len(abs_values))
)
else:
normalized_entropy = 0.0
effect_concentration = float(1.0 - normalized_entropy)
signed_bias = float(sum(deltas) / total_abs) if total_abs > 0 else 0.0
signs = {1 if delta > 0 else -1 if delta < 0 else 0 for delta in deltas}
nonzero_signs = {sign for sign in signs if sign != 0}
sign_consistency = 'same sign' if len(nonzero_signs) <= 1 else 'mixed signs'
if effect_concentration >= 0.25:
profile_pattern = (
'target-concentrated / mixed-sign'
if sign_consistency == 'mixed signs'
else 'target-concentrated / same-sign'
)
elif signed_bias <= -0.8:
profile_pattern = 'broad same-sign suppression'
elif signed_bias >= 0.8:
profile_pattern = 'broad same-sign enhancement'
else:
profile_pattern = 'broad mixed-sign'
summary_rows.append(
[
int(feature_id),
str(strongest[0]),
float(strongest[1]),
strongest_abs,
mean_other,
profile_ratio,
sign_consistency,
normalized_entropy,
effect_concentration,
signed_bias,
profile_pattern,
max(float(item[2]) for item in items),
]
)
for left_idx in range(len(items)):
for right_idx in range(left_idx + 1, len(items)):
target_a, delta_a, _ = items[left_idx]
target_b, delta_b, _ = items[right_idx]
preference_shift = float(delta_a - delta_b)
direction = (
f'toward {target_a}'
if preference_shift > 0
else f'toward {target_b}' if preference_shift < 0 else 'no shift'
)
pairwise_rows.append(
[
int(feature_id),
str(target_a),
str(target_b),
preference_shift,
abs(preference_shift),
direction,
]
)
summary_rows.sort(key=lambda row: float(row[3]), reverse=True)
pairwise_rows.sort(key=lambda row: float(row[4]), reverse=True)
return CandidateCrossTargetResult(
feature_ids=ids,
targets=target_list,
rows=rows,
chart_rows=chart_rows,
summary_rows=summary_rows,
pairwise_rows=pairwise_rows,
active_feature_count=sum(activation > 0 for activation in activations),
)
@staticmethod
def _cue_prompt(stem: str, cue: str) -> str:
stem = stem.rstrip()
cue = cue.strip()
if not cue:
return stem
if cue[0] in ':;,.!?=)]}':
return stem + cue
return stem + ' ' + cue
@torch.inference_mode()
def feature_cue_scan(
self,
feature_id: int,
layer: int,
prompt_stem: str,
cues: Sequence[str],
) -> FeatureCueScanResult:
"""Measure one feature at the final token after appending controlled completion cues."""
if not prompt_stem.strip():
raise ValueError('Enter a prompt stem.')
cue_list: list[str] = []
seen: set[str] = set()
for raw in cues:
cue = str(raw).strip()
if not cue or cue in seen:
continue
seen.add(cue)
cue_list.append(cue)
if not cue_list:
raise ValueError('Enter at least one cue.')
if len(cue_list) > 12:
raise ValueError('Cue scan supports at most 12 cues per run.')
if int(feature_id) < 0 or int(feature_id) >= self.settings.sae_width:
raise ValueError(f'Feature id must be in [0, {self.settings.sae_width - 1}].')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
prompts = [self._cue_prompt(prompt_stem, cue) for cue in cue_list]
batch = self.tokenizer(
prompts,
return_tensors='pt',
padding=True,
truncation=True,
max_length=self.settings.max_prompt_tokens,
)
batch = {key: value.to(self.device) for key, value in batch.items()}
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**batch, use_cache=False)
sae = self.sae_store.get(int(layer))
encoding = sae.encode(capture['hidden'])
attention = batch.get('attention_mask', torch.ones_like(batch['input_ids'])).bool()
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
active_count = 0
for row_idx, (cue, full_prompt) in enumerate(zip(cue_list, prompts, strict=True)):
valid_positions = torch.nonzero(attention[row_idx], as_tuple=False).flatten()
final_pos = int(valid_positions[-1].item())
indices = encoding.indices[row_idx, final_pos]
values = encoding.values[row_idx, final_pos]
mask = indices == int(feature_id)
activation = float(values[mask][0].item()) if bool(mask.any()) else 0.0
active = activation > 0
active_count += int(active)
token_id = int(batch['input_ids'][row_idx, final_pos].item())
final_token = self.tokenizer.decode([token_id])
rows.append([cue, full_prompt, repr(final_token), activation, active])
chart_rows.append([cue, activation])
return FeatureCueScanResult(
feature_id=int(feature_id),
layer=int(layer),
prompt_stem=prompt_stem,
rows=rows,
chart_rows=chart_rows,
active_cue_count=active_count,
cue_count=len(cue_list),
)
@torch.inference_mode()
def feature_cue_context_scan(
self,
feature_id: int,
layer: int,
stems: Sequence[str],
cues: Sequence[str],
) -> FeatureCueContextResult:
"""Cross completion cues with multiple prompt stems in one batch.
This distinguishes a cue-specific response (for example a feature that responds to ``is`` everywhere)
from a context-sensitive completion-boundary response.
"""
stem_list: list[str] = []
seen_stems: set[str] = set()
for raw in stems:
stem = str(raw).strip()
if not stem or stem in seen_stems:
continue
seen_stems.add(stem)
stem_list.append(stem)
cue_list: list[str] = []
seen_cues: set[str] = set()
for raw in cues:
cue = str(raw).strip()
if not cue or cue in seen_cues:
continue
seen_cues.add(cue)
cue_list.append(cue)
if not stem_list:
raise ValueError('Enter at least one prompt stem.')
if not cue_list:
raise ValueError('Enter at least one completion cue.')
if len(stem_list) > 8:
raise ValueError('Cue-context scan supports at most 8 prompt stems.')
if len(cue_list) > 8:
raise ValueError('Cue-context scan supports at most 8 completion cues.')
if len(stem_list) * len(cue_list) > 40:
raise ValueError('Cue-context scan supports at most 40 stem × cue conditions per run.')
if int(feature_id) < 0 or int(feature_id) >= self.settings.sae_width:
raise ValueError(f'Feature id must be in [0, {self.settings.sae_width - 1}].')
self.ensure_ready(preload_saes=False)
assert self.model is not None and self.tokenizer is not None and self.sae_store is not None
conditions: list[tuple[str, str, str]] = []
for stem in stem_list:
for cue in cue_list:
conditions.append((stem, cue, self._cue_prompt(stem, cue)))
prompts = [full for _, _, full in conditions]
batch = self.tokenizer(
prompts,
return_tensors='pt',
padding=True,
truncation=True,
max_length=self.settings.max_prompt_tokens,
)
batch = {key: value.to(self.device) for key, value in batch.items()}
capture: dict = {}
with self._capture_hook(int(layer), capture):
self.model(**batch, use_cache=False)
sae = self.sae_store.get(int(layer))
encoding = sae.encode(capture['hidden'])
attention = batch.get('attention_mask', torch.ones_like(batch['input_ids'])).bool()
rows: list[list[object]] = []
chart_rows: list[list[object]] = []
active_count = 0
cue_active_context_counts = {cue: 0 for cue in cue_list}
cue_activation_values = {cue: [] for cue in cue_list}
for row_idx, (stem, cue, full_prompt) in enumerate(conditions):
valid_positions = torch.nonzero(attention[row_idx], as_tuple=False).flatten()
final_pos = int(valid_positions[-1].item())
indices = encoding.indices[row_idx, final_pos]
values = encoding.values[row_idx, final_pos]
mask = indices == int(feature_id)
activation = float(values[mask][0].item()) if bool(mask.any()) else 0.0
active = activation > 0
active_count += int(active)
cue_active_context_counts[cue] += int(active)
cue_activation_values[cue].append(activation)
token_id = int(batch['input_ids'][row_idx, final_pos].item())
final_token = self.tokenizer.decode([token_id])
short_stem = stem if len(stem) <= 42 else stem[:39] + '…'
rows.append([stem, cue, full_prompt, repr(final_token), activation, active])
chart_rows.append([short_stem, cue, activation])
cue_mean_activations = {
cue: float(sum(vals) / len(vals)) if vals else 0.0
for cue, vals in cue_activation_values.items()
}
dominant_cue = max(
cue_list,
key=lambda cue: (cue_active_context_counts[cue], cue_mean_activations[cue]),
) if cue_list else None
dominant_count = cue_active_context_counts.get(dominant_cue, 0) if dominant_cue else 0
off_dominant_active_count = sum(
count for cue, count in cue_active_context_counts.items() if cue != dominant_cue
)
return FeatureCueContextResult(
feature_id=int(feature_id),
layer=int(layer),
stems=stem_list,
cues=cue_list,
rows=rows,
chart_rows=chart_rows,
active_condition_count=active_count,
condition_count=len(conditions),
cue_active_context_counts=cue_active_context_counts,
cue_mean_activations=cue_mean_activations,
dominant_cue=dominant_cue,
dominant_cue_context_count=dominant_count,
off_dominant_active_count=off_dominant_active_count,
)
RUNTIME = FeatureLensRuntime()
if SETTINGS.eager_load:
try:
RUNTIME.ensure_ready(preload_saes=True)
except Exception as exc: # Keep the UI alive and surface a useful error on first call.
RUNTIME.load_error = f'{type(exc).__name__}: {exc}'
print(f'FeatureLens eager load failed: {RUNTIME.load_error}')