#!/usr/bin/env python3 """Write-to-use tracing helpers for recurrent-state SAE features.""" from __future__ import annotations import math import time from collections import defaultdict from collections.abc import Sequence from typing import Any import numpy as np import torch from experiments.steering.causal_clamp import compute_activation_matrix, reconstruct_with_feature_updates from experiments.analysis.probe_features import compute_text_properties PRIMARY_LAYOUT_PROPERTIES = ("last_sent_boundary_pos",) SECONDARY_LAYOUT_PROPERTIES = ("newline_density", "paragraph_count") LAYOUT_PROPERTY_SET = set(PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES) EVENT_TYPES = ("sentence_boundary", "paragraph_boundary", "non_event_control") def _encode_states( sae, sae_type: str, state_batch: torch.Tensor, ) -> torch.Tensor: if sae_type == "flat": return sae.encode(state_batch.reshape(state_batch.shape[0], -1)) return sae.encode(state_batch) def _safe_spearman(x: np.ndarray, y: np.ndarray) -> float: if x.size != y.size or x.size < 4: return 0.0 if float(np.std(x)) < 1e-8 or float(np.std(y)) < 1e-8: return 0.0 x_rank = np.argsort(np.argsort(x)) y_rank = np.argsort(np.argsort(y)) xr = x_rank.astype(np.float64) yr = y_rank.astype(np.float64) xr -= xr.mean() yr -= yr.mean() denom = math.sqrt(float(np.dot(xr, xr) * np.dot(yr, yr))) if denom < 1e-12: return 0.0 return float(np.dot(xr, yr) / denom) def _best_layout_property( act_values: np.ndarray, property_arrays: dict[str, np.ndarray], ) -> tuple[str, float]: best_property = "" best_rho = 0.0 for prop in PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES: rho = _safe_spearman(act_values, property_arrays[prop]) if abs(rho) > abs(best_rho): best_property = prop best_rho = rho return best_property, best_rho def decode_token_piece( tokenizer, token_id: int, cache: dict[int, str], ) -> str: if token_id not in cache: cache[token_id] = tokenizer.decode( [token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False, ) return cache[token_id] def classify_prefix_event(prefix_text: str) -> str: if prefix_text.endswith("\n\n") or prefix_text.endswith("\n"): return "paragraph_boundary" stripped = prefix_text.rstrip() if stripped and stripped[-1] in ".?!": return "sentence_boundary" return "non_event_control" def select_cross_corpus_trace_features( sae, sae_type: str, openwebtext_states: torch.Tensor, openwebtext_texts: Sequence[str], ultrachat_states: torch.Tensor, ultrachat_texts: Sequence[str], min_rho: float = 0.20, max_features: int = 3, batch_size: int = 512, ) -> list[dict[str, Any]]: """Choose stable H4 tracing features with one checkpoint across both corpora.""" n_owt = min(len(openwebtext_texts), int(openwebtext_states.shape[0])) n_uc = min(len(ultrachat_texts), int(ultrachat_states.shape[0])) if n_owt == 0 or n_uc == 0: raise ValueError("Need non-empty states/texts for both corpora") act_owt = compute_activation_matrix(sae, openwebtext_states[:n_owt], batch_size=batch_size) act_uc = compute_activation_matrix(sae, ultrachat_states[:n_uc], batch_size=batch_size) alive_owt = (act_owt > 0).any(axis=0) alive_uc = (act_uc > 0).any(axis=0) shared_alive = np.nonzero(alive_owt & alive_uc)[0].astype(int).tolist() if not shared_alive: raise ValueError("No shared alive features across corpora") owt_props = [compute_text_properties(text) for text in openwebtext_texts[:n_owt]] uc_props = [compute_text_properties(text) for text in ultrachat_texts[:n_uc]] owt_arrays = { prop: np.array([row[prop] for row in owt_props], dtype=np.float64) for prop in PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES } uc_arrays = { prop: np.array([row[prop] for row in uc_props], dtype=np.float64) for prop in PRIMARY_LAYOUT_PROPERTIES + SECONDARY_LAYOUT_PROPERTIES } strict: list[dict[str, Any]] = [] relaxed: list[dict[str, Any]] = [] for feature_idx in shared_alive: owt_best_prop, owt_best_rho = _best_layout_property(act_owt[:, feature_idx], owt_arrays) uc_best_prop, uc_best_rho = _best_layout_property(act_uc[:, feature_idx], uc_arrays) if not owt_best_prop or not uc_best_prop: continue if abs(owt_best_rho) < min_rho or abs(uc_best_rho) < min_rho: continue entry = { "feature_idx": int(feature_idx), "openwebtext_property": owt_best_prop, "openwebtext_rho": float(owt_best_rho), "ultrachat_property": uc_best_prop, "ultrachat_rho": float(uc_best_rho), "score": float(min(abs(owt_best_rho), abs(uc_best_rho))), "mean_score": float((abs(owt_best_rho) + abs(uc_best_rho)) / 2.0), } strict_ok = ( (owt_best_prop == "last_sent_boundary_pos" and uc_best_prop in LAYOUT_PROPERTY_SET) or (uc_best_prop == "last_sent_boundary_pos" and owt_best_prop in LAYOUT_PROPERTY_SET) ) if strict_ok: strict.append(entry) if owt_best_prop in LAYOUT_PROPERTY_SET and uc_best_prop in LAYOUT_PROPERTY_SET: relaxed.append(entry) def _sort_key(item: dict[str, Any]) -> tuple[float, float, int]: primary_hits = int(item["openwebtext_property"] == "last_sent_boundary_pos") + int( item["ultrachat_property"] == "last_sent_boundary_pos" ) return (primary_hits, item["score"], item["mean_score"]) strict.sort(key=_sort_key, reverse=True) relaxed.sort(key=_sort_key, reverse=True) selected = strict[:max_features] if len(selected) < max_features: seen = {item["feature_idx"] for item in selected} for item in relaxed: if item["feature_idx"] in seen: continue selected.append(item) seen.add(item["feature_idx"]) if len(selected) >= max_features: break if not selected: raise ValueError("No cross-corpus layout features met the tracing threshold") return selected def trace_feature_trajectories( model, tokenizer, corpus_batches: list[torch.Tensor], layer_idx: int, head_idx: int, sae, sae_type: str, feature_specs: list[dict[str, Any]], prompt_len: int = 512, max_offset: int = 64, device: str = "cuda", ) -> dict[str, Any]: """Capture per-token coefficients and event-aligned summaries for selected features.""" t0 = time.perf_counter() feature_indices = [int(item["feature_idx"]) for item in feature_specs] feature_to_slot = {feature_idx: slot for slot, feature_idx in enumerate(feature_indices)} piece_cache: dict[int, str] = {} prompt_records: list[dict[str, Any]] = [] aligned_values = { feature_idx: { event: {"coeff_curves": [], "delta_curves": [], "normalized_curves": []} for event in EVENT_TYPES } for feature_idx in feature_indices } print( "[factor-trace] " f"layer={layer_idx} head={head_idx} features={feature_indices} " f"batches={len(corpus_batches)} prompt_len={prompt_len}", flush=True, ) prompt_index = 0 for batch_idx_global, batch in enumerate(corpus_batches, start=1): if batch.numel() == 0: continue token_batch = batch[:, :prompt_len].to(device) batch_size, seq_len = token_batch.shape if seq_len < 2: continue coeffs = np.zeros((batch_size, seq_len, len(feature_indices)), dtype=np.float32) event_labels = [["non_event_control"] * seq_len for _ in range(batch_size)] prefixes = [""] * batch_size cache = None for pos in range(seq_len): step_ids = token_batch[:, pos : pos + 1] if cache is None: out = model(input_ids=step_ids, use_cache=True) else: out = model(input_ids=step_ids, past_key_values=cache, use_cache=True) cache = out.past_key_values state_batch = cache.layers[layer_idx].recurrent_states[:, head_idx].float() act_batch = _encode_states(sae, sae_type, state_batch) coeffs[:, pos, :] = act_batch[:, feature_indices].detach().cpu().numpy() for batch_idx in range(batch_size): token_id = int(token_batch[batch_idx, pos].item()) prefixes[batch_idx] += decode_token_piece(tokenizer, token_id, piece_cache) event_labels[batch_idx][pos] = classify_prefix_event(prefixes[batch_idx]) deltas = np.empty_like(coeffs) deltas[:, 0, :] = coeffs[:, 0, :] deltas[:, 1:, :] = coeffs[:, 1:, :] - coeffs[:, :-1, :] rng = np.random.default_rng(0) offsets = list(range(max_offset + 1)) for batch_idx in range(batch_size): labels = event_labels[batch_idx] prompt_record = { "prompt_index": prompt_index, "token_ids": [int(x) for x in token_batch[batch_idx].detach().cpu().tolist()], "event_labels": labels, "feature_coefficients": { str(feature_idx): coeffs[batch_idx, :, slot].tolist() for feature_idx, slot in feature_to_slot.items() }, "feature_deltas": { str(feature_idx): deltas[batch_idx, :, slot].tolist() for feature_idx, slot in feature_to_slot.items() }, } prompt_records.append(prompt_record) eligible_controls = [pos for pos, label in enumerate(labels) if label == "non_event_control" and pos + max_offset < seq_len] event_positions = [ pos for pos, label in enumerate(labels) if label in ("sentence_boundary", "paragraph_boundary") and pos + max_offset < seq_len ] if event_positions and eligible_controls: control_sample = rng.choice( np.array(eligible_controls), size=min(len(event_positions), len(eligible_controls)), replace=False, ).tolist() else: control_sample = [] for feature_idx, slot in feature_to_slot.items(): for event in ("sentence_boundary", "paragraph_boundary"): positions = [pos for pos in event_positions if labels[pos] == event] for pos in positions: coeff_curve = coeffs[batch_idx, pos : pos + max_offset + 1, slot] delta_curve = deltas[batch_idx, pos : pos + max_offset + 1, slot] denom = max(abs(float(coeff_curve[0])), 1e-8) aligned_values[feature_idx][event]["coeff_curves"].append(coeff_curve.astype(np.float64)) aligned_values[feature_idx][event]["delta_curves"].append(delta_curve.astype(np.float64)) aligned_values[feature_idx][event]["normalized_curves"].append( (coeff_curve / denom).astype(np.float64) ) for pos in control_sample: coeff_curve = coeffs[batch_idx, pos : pos + max_offset + 1, slot] delta_curve = deltas[batch_idx, pos : pos + max_offset + 1, slot] denom = max(abs(float(coeff_curve[0])), 1e-8) aligned_values[feature_idx]["non_event_control"]["coeff_curves"].append(coeff_curve.astype(np.float64)) aligned_values[feature_idx]["non_event_control"]["delta_curves"].append(delta_curve.astype(np.float64)) aligned_values[feature_idx]["non_event_control"]["normalized_curves"].append( (coeff_curve / denom).astype(np.float64) ) prompt_index += 1 if prompt_index % 16 == 0: print( "[factor-trace] " f"prompts={prompt_index} batch={batch_idx_global}/{len(corpus_batches)} " f"elapsed={time.perf_counter() - t0:.1f}s", flush=True, ) print( "[factor-trace] " f"finished_batch={batch_idx_global}/{len(corpus_batches)} " f"prompts={prompt_index} elapsed={time.perf_counter() - t0:.1f}s", flush=True, ) aligned_summary: dict[str, dict[str, Any]] = {} for feature_idx in feature_indices: feature_summary: dict[str, Any] = {} for event in EVENT_TYPES: coeff_curves = aligned_values[feature_idx][event]["coeff_curves"] delta_curves = aligned_values[feature_idx][event]["delta_curves"] norm_curves = aligned_values[feature_idx][event]["normalized_curves"] if coeff_curves: coeff_stack = np.stack(coeff_curves, axis=0) delta_stack = np.stack(delta_curves, axis=0) norm_stack = np.stack(norm_curves, axis=0) feature_summary[event] = { "count": int(coeff_stack.shape[0]), "offsets": offsets, "mean_coeff_curve": coeff_stack.mean(axis=0).tolist(), "mean_delta_curve": delta_stack.mean(axis=0).tolist(), "median_normalized_curve": np.median(norm_stack, axis=0).tolist(), "mean_event_jump": float(delta_stack[:, 0].mean()), } else: feature_summary[event] = { "count": 0, "offsets": offsets, "mean_coeff_curve": [0.0] * len(offsets), "mean_delta_curve": [0.0] * len(offsets), "median_normalized_curve": [0.0] * len(offsets), "mean_event_jump": 0.0, } control_jump = feature_summary["non_event_control"]["mean_event_jump"] for event in ("sentence_boundary", "paragraph_boundary"): feature_summary[event]["jump_minus_control"] = ( feature_summary[event]["mean_event_jump"] - control_jump ) aligned_summary[str(feature_idx)] = feature_summary return { "prompt_len": prompt_len, "max_offset": max_offset, "selected_features": feature_specs, "prompt_records": prompt_records, "event_aligned": aligned_summary, } def select_trace_events( trace_result: dict[str, Any], target_feature_idx: int, preferred_events: Sequence[str] = ("sentence_boundary", "paragraph_boundary"), offsets: Sequence[int] = (16, 32, 64), min_retention: float = 0.5, ) -> list[dict[str, Any]]: """Choose write/use event pairs from traced trajectories.""" candidates: list[dict[str, Any]] = [] key = str(target_feature_idx) for prompt in trace_result.get("prompt_records", []): labels = prompt["event_labels"] coeffs = prompt["feature_coefficients"][key] deltas = prompt["feature_deltas"][key] best: dict[str, Any] | None = None for pos, label in enumerate(labels): if label not in preferred_events: continue if pos >= len(coeffs): continue coeff = float(coeffs[pos]) jump = float(deltas[pos]) if jump <= 0 or coeff <= 0: continue use_offset = None for offset in offsets: if pos + offset >= len(coeffs): continue if float(coeffs[pos + offset]) >= min_retention * coeff: use_offset = int(offset) break if use_offset is None: continue entry = { "prompt_index": int(prompt["prompt_index"]), "write_pos": int(pos), "event_type": label, "write_coeff": coeff, "write_jump": jump, "use_offset": use_offset, "use_pos": int(pos + use_offset), "baseline_use_coeff": float(coeffs[pos + use_offset]), } if best is None or entry["write_jump"] > best["write_jump"]: best = entry if best is not None: candidates.append(best) candidates.sort(key=lambda item: item["write_jump"], reverse=True) return candidates def _context_tokens( prompt_record: dict[str, Any], pos: int, window: int = 24, ) -> set[int]: token_ids = prompt_record["token_ids"] start = max(0, pos - window) end = min(len(token_ids), pos + 1) return {int(tok) for tok in token_ids[start:end]} def _context_jaccard( prompt_a: dict[str, Any], pos_a: int, prompt_b: dict[str, Any], pos_b: int, window: int = 24, ) -> float: ctx_a = _context_tokens(prompt_a, pos_a, window=window) ctx_b = _context_tokens(prompt_b, pos_b, window=window) if not ctx_a and not ctx_b: return 1.0 if not ctx_a or not ctx_b: return 0.0 return float(len(ctx_a & ctx_b) / max(len(ctx_a | ctx_b), 1)) def _nearest_non_event_position( prompt_record: dict[str, Any], write_pos: int, use_pos: int, min_gap: int = 8, ) -> int | None: labels = prompt_record["event_labels"] candidates = [] for pos, label in enumerate(labels): if label != "non_event_control": continue if pos >= use_pos: continue if abs(pos - write_pos) < min_gap: continue candidates.append((abs(pos - write_pos), pos)) if not candidates: return None candidates.sort() return int(candidates[0][1]) def select_transplant_pairs( trace_result: dict[str, Any], target_feature_idx: int, max_pairs: int = 16, preferred_events: Sequence[str] = ("paragraph_boundary", "sentence_boundary"), preferred_offsets: Sequence[int] = (16,), min_write_gain: float = 1.10, context_window: int = 24, min_context_similarity: float = 0.10, ) -> list[dict[str, Any]]: """Pair low-write recipients with stronger matched donors from the trace set.""" prompt_map = { int(prompt["prompt_index"]): prompt for prompt in trace_result.get("prompt_records", []) } candidates = [ item for item in select_trace_events(trace_result, target_feature_idx) if item["event_type"] in preferred_events ] if preferred_offsets: filtered = [item for item in candidates if int(item["use_offset"]) in preferred_offsets] if len(filtered) >= 2: candidates = filtered if len(candidates) < 2: candidates = select_trace_events(trace_result, target_feature_idx) if len(candidates) < 2: return [] donors = sorted(candidates, key=lambda item: (item["write_jump"], item["write_coeff"]), reverse=True) recipients = sorted(candidates, key=lambda item: (item["write_jump"], item["write_coeff"])) donor_pool = donors[: max(4, min(len(donors), max_pairs))] pairs: list[dict[str, Any]] = [] used_recipient_prompts: set[int] = set() for recipient in recipients: if len(pairs) >= max_pairs: break recipient_prompt = prompt_map.get(int(recipient["prompt_index"])) if recipient_prompt is None: continue if int(recipient["prompt_index"]) in used_recipient_prompts: continue best_match: tuple[tuple[int, int, float, float], dict[str, Any]] | None = None for donor in donor_pool: donor_prompt_index = int(donor["prompt_index"]) recipient_prompt_index = int(recipient["prompt_index"]) if donor_prompt_index == recipient_prompt_index: continue if float(donor["write_coeff"]) <= float(recipient["write_coeff"]) * min_write_gain: continue donor_prompt = prompt_map.get(donor_prompt_index) if donor_prompt is None: continue same_event = int(donor["event_type"] == recipient["event_type"]) same_offset = int(donor["use_offset"] == recipient["use_offset"]) context_sim = _context_jaccard( donor_prompt, int(donor["write_pos"]), recipient_prompt, int(recipient["write_pos"]), window=context_window, ) if context_sim < min_context_similarity: continue score = ( same_event, same_offset, context_sim, float(donor["write_jump"]) - float(recipient["write_jump"]), ) candidate = { "donor_prompt_index": donor_prompt_index, "donor_write_pos": int(donor["write_pos"]), "donor_use_pos": int(donor["use_pos"]), "donor_use_offset": int(donor["use_offset"]), "donor_event_type": donor["event_type"], "donor_write_coeff": float(donor["write_coeff"]), "donor_write_jump": float(donor["write_jump"]), "recipient_prompt_index": recipient_prompt_index, "recipient_write_pos": int(recipient["write_pos"]), "recipient_use_pos": int(recipient["use_pos"]), "recipient_use_offset": int(recipient["use_offset"]), "recipient_event_type": recipient["event_type"], "recipient_write_coeff": float(recipient["write_coeff"]), "recipient_write_jump": float(recipient["write_jump"]), "context_similarity": float(context_sim), } if best_match is None or score > best_match[0]: best_match = (score, candidate) if best_match is None: continue pair = best_match[1] wrong_time_pos = _nearest_non_event_position( recipient_prompt, write_pos=pair["recipient_write_pos"], use_pos=pair["recipient_use_pos"], ) pair["wrong_time_pos"] = wrong_time_pos pairs.append(pair) used_recipient_prompts.add(int(recipient["prompt_index"])) return pairs def choose_similar_frequency_feature( target_feature_idx: int, frequencies: np.ndarray, exclude: set[int], ) -> int: target_freq = float(frequencies[target_feature_idx]) best_idx = -1 best_gap = float("inf") for idx, freq in enumerate(frequencies): if idx in exclude: continue gap = abs(float(freq) - target_freq) if gap < best_gap: best_gap = gap best_idx = idx if best_idx < 0: raise ValueError("Could not find a similar-frequency control feature") return int(best_idx) def build_use_site_benchmark( model, tokenizer, corpus_sequences: list[torch.Tensor], trace_result: dict[str, Any], layer_idx: int, head_idx: int, sae, sae_type: str, target_feature_indices: Sequence[int], max_rows: int = 64, device: str = "cuda", ) -> dict[str, Any]: """Mine high-confidence downstream use sites for the signed H4 boundary code.""" if len(target_feature_indices) != 2: raise ValueError("build_use_site_benchmark expects the signed pair [62, 105]") pos_feat, neg_feat = [int(x) for x in target_feature_indices] prompt_map = { int(prompt["prompt_index"]): prompt for prompt in trace_result.get("prompt_records", []) } sequence_map = { int(prompt["prompt_index"]): seq for prompt, seq in zip(trace_result.get("prompt_records", []), corpus_sequences) } readout_token_ids = _readout_class_token_ids(tokenizer) candidate_pairs = select_transplant_pairs( trace_result=trace_result, target_feature_idx=pos_feat, max_pairs=max(max_rows * 8, 128), preferred_events=("paragraph_boundary",), preferred_offsets=(8, 16, 32), min_write_gain=1.20, min_context_similarity=0.20, ) if not candidate_pairs: raise ValueError("No candidate paragraph-boundary pairs for use-site benchmark") def _signed_coeff(prompt_record: dict[str, Any], pos: int) -> float: pos_val = float(prompt_record["feature_coefficients"][str(pos_feat)][pos]) neg_val = float(prompt_record["feature_coefficients"][str(neg_feat)][pos]) return pos_val - neg_val def _signed_delta(prompt_record: dict[str, Any], pos: int) -> tuple[float, float]: pos_delta = float(prompt_record["feature_deltas"][str(pos_feat)][pos]) neg_delta = float(prompt_record["feature_deltas"][str(neg_feat)][pos]) return pos_delta, neg_delta all_rows: list[dict[str, Any]] = [] accepted_rows: list[dict[str, Any]] = [] drop_counts: dict[str, int] = defaultdict(int) t0 = time.perf_counter() print( "[factor-trace-benchmark] " f"target={[pos_feat, neg_feat]} candidate_pairs={len(candidate_pairs)} max_rows={max_rows}", flush=True, ) for pair_idx, pair in enumerate(candidate_pairs, start=1): recipient_prompt_index = int(pair["recipient_prompt_index"]) donor_prompt_index = int(pair["donor_prompt_index"]) if recipient_prompt_index not in prompt_map or recipient_prompt_index not in sequence_map: continue recipient_prompt = prompt_map[recipient_prompt_index] donor_prompt = prompt_map[donor_prompt_index] use_pos = int(pair["recipient_use_pos"]) seq = sequence_map[recipient_prompt_index].to(device) seq_end = min(int(seq.shape[1]), use_pos + 10) if seq_end <= use_pos + 1: continue seq = seq[:, :seq_end] baseline = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=[pos_feat, neg_feat], ) donor_write_signed = _signed_coeff(donor_prompt, int(pair["donor_write_pos"])) recipient_write_signed = _signed_coeff(recipient_prompt, int(pair["recipient_write_pos"])) donor_delta_pos, donor_delta_neg = _signed_delta(donor_prompt, int(pair["donor_write_pos"])) signed_write_gap = float(donor_write_signed - recipient_write_signed) start = max(int(pair["recipient_write_pos"]) + 1, use_pos - 2) end = min(use_pos + 8, int(seq.shape[1]) - 2) for decision_pos in range(start, end + 1): actual_next_token = int(seq[0, decision_pos + 1].item()) readout_class, readout_text = _classify_readout_token(tokenizer, actual_next_token) row = { "donor_prompt_index": donor_prompt_index, "recipient_prompt_index": recipient_prompt_index, "write_pos": int(pair["recipient_write_pos"]), "use_pos": use_pos, "use_offset": int(pair["recipient_use_offset"]), "decision_pos": int(decision_pos), "intervention_pos": int(max(0, decision_pos - 1)), "wrong_time_pos": _nearest_non_event_position( recipient_prompt, write_pos=int(max(0, decision_pos - 1)), use_pos=int(decision_pos), min_gap=4, ), "context_similarity": float(pair["context_similarity"]), "signed_write_gap": signed_write_gap, "donor_delta_62": float(donor_delta_pos), "donor_delta_105": float(donor_delta_neg), "actual_next_token": actual_next_token, "readout_text": readout_text, "readout_class": readout_class, "accepted": False, "drop_reason": "", } if readout_class is None: row["drop_reason"] = "readout_class_excluded" drop_counts[row["drop_reason"]] += 1 all_rows.append(row) continue baseline_logits = baseline["logits_by_pos"][decision_pos] baseline_prob = _token_prob_from_logits(baseline_logits, actual_next_token) row["baseline_actual_token_prob"] = baseline_prob row["baseline_headroom"] = float(baseline_prob * (1.0 - baseline_prob)) row["baseline_readout_class_mass"] = _mass_from_logits( baseline_logits, readout_token_ids.get(readout_class, [actual_next_token]) or [actual_next_token], ) signed_use_coeff = float( baseline["feature_coefficients"][str(pos_feat)][decision_pos] - baseline["feature_coefficients"][str(neg_feat)][decision_pos] ) row["signed_use_coeff"] = signed_use_coeff if abs(signed_write_gap) < 0.02: row["drop_reason"] = "weak_signed_write_gap" elif abs(signed_use_coeff) < 0.01: row["drop_reason"] = "weak_signed_use_coeff" elif baseline_prob < 0.05 or baseline_prob > 0.80: row["drop_reason"] = "low_token_headroom" elif row["wrong_time_pos"] is None: row["drop_reason"] = "no_wrong_time" else: row["accepted"] = True row["drop_reason"] = "accepted" row["score"] = ( row["baseline_headroom"], abs(signed_use_coeff), float(pair["context_similarity"]), abs(signed_write_gap), ) accepted_rows.append(row) print( f"[factor-trace-benchmark] accept {len(accepted_rows)} " f"pair={pair_idx}/{len(candidate_pairs)} " f"recipient={recipient_prompt_index}@{decision_pos} " f"class={readout_class} tok={readout_text!r} " f"p={baseline_prob:.3f} headroom={row['baseline_headroom']:.3f} " f"signed_gap={signed_write_gap:+.4f} signed_use={signed_use_coeff:+.4f}", flush=True, ) drop_counts[row["drop_reason"]] += 1 all_rows.append(row) if pair_idx % 5 == 0 or pair_idx == len(candidate_pairs): print( f"[factor-trace-benchmark] progress pair={pair_idx}/{len(candidate_pairs)} " f"rows={len(all_rows)} accepted={len(accepted_rows)} " f"drops={dict(drop_counts)} " f"elapsed={time.perf_counter() - t0:.1f}s", flush=True, ) accepted_rows.sort(key=lambda r: r.get("score", (0, 0, 0, 0)), reverse=True) selected_rows = accepted_rows[:max_rows] print( "[factor-trace-benchmark] " f"accepted={len(accepted_rows)} selected={len(selected_rows)} " f"elapsed={time.perf_counter() - t0:.1f}s", flush=True, ) return { "target_feature_indices": [pos_feat, neg_feat], "n_candidate_pairs": len(candidate_pairs), "n_rows_total": len(all_rows), "n_rows_accepted": len(accepted_rows), "drop_counts": dict(drop_counts), "accepted_rows": selected_rows, "all_rows": all_rows, "total_time_s": float(time.perf_counter() - t0), } def sample_matched_random_features( target_feature_idx: int, frequencies: np.ndarray, alive_features: list[int], exclude: set[int], n_random: int = 4, seed: int = 0, ) -> list[int]: rng = np.random.default_rng(seed) target_freq = float(frequencies[target_feature_idx]) candidates = [idx for idx in alive_features if idx not in exclude] if not candidates: raise ValueError("No alive features available for random controls") candidates.sort(key=lambda idx: abs(float(frequencies[idx]) - target_freq)) pool = candidates[: max(n_random * 8, 32)] if len(pool) <= n_random: return [int(idx) for idx in pool[:n_random]] picked = rng.choice(np.array(pool), size=n_random, replace=False) return [int(idx) for idx in picked.tolist()] def boundary_token_ids(tokenizer) -> list[int]: token_ids: set[int] = set() for text in [".", "?", "!", "\n", ":"]: ids = tokenizer.encode(text, add_special_tokens=False) if len(ids) == 1: token_ids.add(int(ids[0])) return sorted(token_ids) def boundary_token_groups(tokenizer) -> dict[str, list[int]]: groups = { "newline": set(), "sentence_punct": set(), "colon": set(), } for text, key in [("\n", "newline"), (".", "sentence_punct"), ("?", "sentence_punct"), ("!", "sentence_punct"), (":", "colon")]: ids = tokenizer.encode(text, add_special_tokens=False) if len(ids) == 1: groups[key].add(int(ids[0])) all_boundary = set().union(*groups.values()) return { "boundary": sorted(all_boundary), "newline": sorted(groups["newline"]), "sentence_punct": sorted(groups["sentence_punct"]), "colon": sorted(groups["colon"]), } @torch.no_grad() def teacher_forced_feature_rollout( model, input_ids: torch.Tensor, layer_idx: int, head_idx: int, sae, sae_type: str, tracked_features: Sequence[int], intervention_updates: dict[int, float] | None = None, intervention_feature: int | None = None, intervention_pos: int | None = None, capture_hidden_state_positions: Sequence[int] | None = None, intervention_state_delta: torch.Tensor | None = None, ) -> dict[str, Any]: """Run one sequence token-by-token and optionally ablate one feature at one step.""" seq = input_ids if seq.ndim != 2 or seq.shape[0] != 1: raise ValueError("teacher_forced_feature_rollout expects shape (1, seq_len)") tracked = [int(idx) for idx in tracked_features] cache = None coeff_records = {str(idx): [] for idx in tracked} logits_by_pos: dict[int, torch.Tensor] = {} hidden_state_positions = {int(pos) for pos in (capture_hidden_state_positions or [])} hidden_state_by_pos: dict[int, torch.Tensor] = {} for pos in range(int(seq.shape[1])): step_ids = seq[:, pos : pos + 1] if cache is None: out = model( input_ids=step_ids, use_cache=True, output_hidden_states=bool(hidden_state_positions), ) else: out = model( input_ids=step_ids, past_key_values=cache, use_cache=True, output_hidden_states=bool(hidden_state_positions), ) cache = out.past_key_values state = cache.layers[layer_idx].recurrent_states[:, head_idx].float() acts = _encode_states(sae, sae_type, state) for idx in tracked: coeff_records[str(idx)].append(float(acts[0, idx].item())) logits_by_pos[pos] = out.logits[:, -1, :].float().detach().cpu() if pos in hidden_state_positions and out.hidden_states is not None: hidden_state_by_pos[pos] = out.hidden_states[-1][:, -1, :].float().detach().cpu() if intervention_pos == pos and ( intervention_updates is not None or intervention_feature is not None or intervention_state_delta is not None ): layer_cache = cache.layers[layer_idx] original_head = layer_cache.recurrent_states[0, head_idx].float() if intervention_state_delta is not None: edited = original_head + intervention_state_delta.to( device=original_head.device, dtype=original_head.dtype, ) else: ablate_list = [int(intervention_feature)] if intervention_feature is not None else None edited = reconstruct_with_feature_updates( sae=sae, state_head=original_head, sae_type=sae_type, set_updates=intervention_updates, ablate_features=ablate_list, ) layer_cache.recurrent_states[0, head_idx] = edited.to(layer_cache.recurrent_states.dtype) return { "feature_coefficients": coeff_records, "logits_by_pos": logits_by_pos, "hidden_state_by_pos": hidden_state_by_pos, } def _kl_divergence_from_logits( baseline_logits: torch.Tensor, edited_logits: torch.Tensor, ) -> float: baseline_log_probs = torch.log_softmax(baseline_logits, dim=-1) edited_log_probs = torch.log_softmax(edited_logits, dim=-1) baseline_probs = baseline_log_probs.exp() kl = torch.sum(baseline_probs * (baseline_log_probs - edited_log_probs), dim=-1) return float(kl.item()) def _boundary_mass_shift( baseline_logits: torch.Tensor, edited_logits: torch.Tensor, token_ids: Sequence[int], ) -> float: if not token_ids: return 0.0 baseline_probs = torch.softmax(baseline_logits, dim=-1) edited_probs = torch.softmax(edited_logits, dim=-1) baseline_mass = baseline_probs[:, list(token_ids)].sum(dim=-1) edited_mass = edited_probs[:, list(token_ids)].sum(dim=-1) return float((edited_mass - baseline_mass).item()) def _mass_from_logits( logits: torch.Tensor, token_ids: Sequence[int], ) -> float: if not token_ids: return 0.0 probs = torch.softmax(logits, dim=-1) return float(probs[:, list(token_ids)].sum(dim=-1).item()) def _token_logit_shift( baseline_logits: torch.Tensor, edited_logits: torch.Tensor, token_id: int | None, ) -> float: if token_id is None or token_id < 0: return 0.0 return float((edited_logits.float()[:, token_id] - baseline_logits.float()[:, token_id]).item()) def _decoded_boundary_group(tokenizer, token_id: int | None) -> tuple[str, str]: if token_id is None or token_id < 0: return "other", "" raw = tokenizer.decode([int(token_id)]) if raw in {"\n", "\n\n"}: return "newline", raw stripped = raw.strip() if stripped in {".", "?", "!"} and len(stripped) == 1: return "sentence_punct", stripped if stripped == ":": return "colon", stripped return "other", raw def _token_prob_from_logits(logits: torch.Tensor, token_id: int | None) -> float: if token_id is None or token_id < 0: return 0.0 probs = torch.softmax(logits.float(), dim=-1) return float(probs[:, int(token_id)].item()) def _token_prob_shift( baseline_logits: torch.Tensor, edited_logits: torch.Tensor, token_id: int | None, ) -> float: if token_id is None or token_id < 0: return 0.0 base = torch.softmax(baseline_logits.float(), dim=-1)[:, int(token_id)] edit = torch.softmax(edited_logits.float(), dim=-1)[:, int(token_id)] return float((edit - base).item()) def _classify_readout_token(tokenizer, token_id: int | None) -> tuple[str | None, str]: if token_id is None or token_id < 0: return None, "" raw = tokenizer.decode( [int(token_id)], skip_special_tokens=False, clean_up_tokenization_spaces=False, ) if raw in {"\n", "\n\n"}: return "paragraph_break", raw stripped = raw.strip() if stripped in {".", "?", "!"} and len(stripped) == 1: return "sentence_reset", stripped if stripped in {"The", "However", "In", "This"}: return "discourse_opener", stripped if stripped and stripped[0].isupper() and stripped.replace("'", "").isalpha(): return "sentence_capitalized", stripped return None, raw def _readout_class_token_ids(tokenizer) -> dict[str, list[int]]: buckets = { "paragraph_break": set(), "sentence_reset": set(), "discourse_opener": set(), } for text, bucket in [ ("\n", "paragraph_break"), ("\n\n", "paragraph_break"), (".", "sentence_reset"), ("?", "sentence_reset"), ("!", "sentence_reset"), ("The", "discourse_opener"), ("However", "discourse_opener"), ("In", "discourse_opener"), ("This", "discourse_opener"), ]: ids = tokenizer.encode(text, add_special_tokens=False) if len(ids) == 1: buckets[bucket].add(int(ids[0])) return {k: sorted(v) for k, v in buckets.items()} def _preview_eval_site_from_prompt( tokenizer, prompt_record: dict[str, Any], anchor_use_pos: int, back_window: int = 4, forward_window: int = 8, ) -> dict[str, Any]: token_ids = prompt_record["token_ids"] best: tuple[tuple[int, int], dict[str, Any]] | None = None start = max(0, anchor_use_pos - back_window) end = min(len(token_ids) - 1, anchor_use_pos + forward_window) for pos in range(start, end + 1): if pos + 1 >= len(token_ids): continue actual_next_token = int(token_ids[pos + 1]) actual_group, actual_text = _decoded_boundary_group(tokenizer, actual_next_token) if actual_group == "other": continue # Prefer the closest true boundary token, with a slight bias toward later use-sites. score = (-abs(pos - anchor_use_pos), int(pos >= anchor_use_pos)) candidate = { "eval_pos": int(pos), "eval_offset_from_anchor": int(pos - anchor_use_pos), "eval_actual_next_token": actual_next_token, "eval_actual_group": actual_group, "eval_actual_text": actual_text, } if best is None or score > best[0]: best = (score, candidate) return best[1] if best is not None else {} def _select_eval_site( baseline_logits_by_pos: Sequence[torch.Tensor], input_ids: torch.Tensor, anchor_use_pos: int, token_groups: dict[str, list[int]], tokenizer, back_window: int = 4, forward_window: int = 8, preferred_groups: Sequence[str] = ("newline", "sentence_punct", "colon"), ) -> dict[str, Any]: seq_len = int(input_ids.shape[1]) best: tuple[tuple[float, float], dict[str, Any]] | None = None start = max(0, anchor_use_pos - back_window) end = min(len(baseline_logits_by_pos) - 1, anchor_use_pos + forward_window) preferred_rank = {group: len(preferred_groups) - i for i, group in enumerate(preferred_groups)} for pos in range(start, end + 1): if pos + 1 >= seq_len: continue actual_next_token = int(input_ids[0, pos + 1].item()) actual_group, actual_text = _decoded_boundary_group(tokenizer, actual_next_token) if actual_group == "other": continue logits = baseline_logits_by_pos[pos] boundary_mass = _mass_from_logits(logits, token_groups["boundary"]) actual_group_mass = ( _mass_from_logits(logits, token_groups[actual_group]) if actual_group in token_groups else 0.0 ) score = ( preferred_rank.get(actual_group, 0), actual_group_mass, boundary_mass, -abs(pos - anchor_use_pos), ) candidate = { "eval_pos": int(pos), "eval_offset_from_anchor": int(pos - anchor_use_pos), "eval_actual_next_token": actual_next_token, "eval_actual_group": actual_group, "eval_actual_text": actual_text, "baseline_boundary_mass": boundary_mass, "baseline_actual_group_mass": actual_group_mass, } if best is None or score > best[0]: best = (score, candidate) if best is None: return {} return best[1] def run_factor_trace_intervention( model, tokenizer, corpus_sequences: list[torch.Tensor], trace_result: dict[str, Any], layer_idx: int, head_idx: int, sae, sae_type: str, states: torch.Tensor, target_feature_idx: int, n_prompts: int = 16, n_random: int = 4, device: str = "cuda", ) -> dict[str, Any]: """Measure write-to-use effects for one traced feature.""" candidates = select_trace_events(trace_result, target_feature_idx) if not candidates: raise ValueError(f"No valid trace events found for feature {target_feature_idx}") act_matrix = compute_activation_matrix(sae, states) frequencies = (act_matrix > 0).mean(axis=0) alive_features = np.nonzero((act_matrix > 0).any(axis=0))[0].astype(int).tolist() selected = candidates[:n_prompts] exclude = {int(target_feature_idx)} similar_feature = choose_similar_frequency_feature(target_feature_idx, frequencies, exclude) exclude.add(similar_feature) random_features = sample_matched_random_features( target_feature_idx=target_feature_idx, frequencies=frequencies, alive_features=alive_features, exclude=exclude, n_random=n_random, seed=0, ) control_features = [similar_feature] + random_features tracked_features = [target_feature_idx] + control_features boundary_ids = boundary_token_ids(tokenizer) t0 = time.perf_counter() print( "[factor-trace-intervention] target=" f"{target_feature_idx} similar_control={similar_feature} " f"random_controls={random_features} candidates={len(candidates)} " f"evaluating={min(len(candidates), n_prompts)}", flush=True, ) per_prompt = [] target_kls: list[float] = [] control_kls: list[float] = [] target_boundary_shifts: list[float] = [] control_boundary_shifts: list[float] = [] target_attenuations: list[float] = [] control_attenuations: list[float] = [] sequence_map = { int(prompt["prompt_index"]): seq for prompt, seq in zip(trace_result.get("prompt_records", []), corpus_sequences) } for i, candidate in enumerate(selected, start=1): prompt_index = int(candidate["prompt_index"]) if prompt_index not in sequence_map: print( f"[factor-trace-intervention] skip {i}/{len(selected)} " f"prompt={prompt_index} missing_sequence", flush=True, ) continue seq = sequence_map[prompt_index][:, : candidate["use_pos"] + 1].to(device) prompt_t0 = time.perf_counter() print( f"[factor-trace-intervention] prompt {i}/{len(selected)} " f"prompt={prompt_index} event={candidate['event_type']} " f"write={candidate['write_pos']} use={candidate['use_pos']} " f"offset={candidate['use_offset']} seq_len={seq.shape[1]}", flush=True, ) baseline = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_feature=None, intervention_pos=None, ) target_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_feature=target_feature_idx, intervention_pos=int(candidate["write_pos"]), ) baseline_logits = baseline["logits_by_pos"][candidate["use_pos"]] target_logits = target_rollout["logits_by_pos"][candidate["use_pos"]] target_kl = _kl_divergence_from_logits(baseline_logits, target_logits) target_boundary = _boundary_mass_shift(baseline_logits, target_logits, boundary_ids) baseline_use_coeff = baseline["feature_coefficients"][str(target_feature_idx)][candidate["use_pos"]] target_use_coeff = target_rollout["feature_coefficients"][str(target_feature_idx)][candidate["use_pos"]] target_attenuation = float(baseline_use_coeff - target_use_coeff) control_rows = [] prompt_control_kls = [] prompt_control_boundary = [] prompt_control_attenuation = [] for control_feature in control_features: control_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_feature=int(control_feature), intervention_pos=int(candidate["write_pos"]), ) control_logits = control_rollout["logits_by_pos"][candidate["use_pos"]] control_kl = _kl_divergence_from_logits(baseline_logits, control_logits) control_boundary = _boundary_mass_shift(baseline_logits, control_logits, boundary_ids) control_use_coeff = control_rollout["feature_coefficients"][str(target_feature_idx)][candidate["use_pos"]] control_attenuation = float(baseline_use_coeff - control_use_coeff) prompt_control_kls.append(control_kl) prompt_control_boundary.append(control_boundary) prompt_control_attenuation.append(control_attenuation) control_rows.append({ "feature_idx": int(control_feature), "use_time_kl": control_kl, "boundary_mass_shift": control_boundary, "target_feature_attenuation": control_attenuation, }) target_kls.append(target_kl) control_kls.append(float(np.mean(prompt_control_kls)) if prompt_control_kls else 0.0) target_boundary_shifts.append(target_boundary) control_boundary_shifts.append(float(np.mean(prompt_control_boundary)) if prompt_control_boundary else 0.0) target_attenuations.append(target_attenuation) control_attenuations.append(float(np.mean(prompt_control_attenuation)) if prompt_control_attenuation else 0.0) per_prompt.append({ **candidate, "baseline_use_coeff": float(baseline_use_coeff), "target_use_coeff": float(target_use_coeff), "target_use_time_kl": target_kl, "target_boundary_mass_shift": target_boundary, "target_feature_attenuation": target_attenuation, "controls": control_rows, }) running_target = float(np.mean(target_kls)) if target_kls else 0.0 running_control = float(np.mean(control_kls)) if control_kls else 0.0 ratio = ( running_target / running_control if abs(running_control) > 1e-8 else float("inf") ) print( f"[factor-trace-intervention] done {i}/{len(selected)} " f"target_kl={target_kl:.6f} control_kl={control_kls[-1]:.6f} " f"running_ratio={ratio:.2f}x " f"target_att={target_attenuation:.6f} control_att={control_attenuations[-1]:.6f} " f"elapsed={time.perf_counter() - prompt_t0:.1f}s", flush=True, ) target_kl_mean = float(np.mean(target_kls)) if target_kls else 0.0 control_kl_mean = float(np.mean(control_kls)) if control_kls else 0.0 target_boundary_mean = float(np.mean(target_boundary_shifts)) if target_boundary_shifts else 0.0 control_boundary_mean = float(np.mean(control_boundary_shifts)) if control_boundary_shifts else 0.0 target_att_mean = float(np.mean(target_attenuations)) if target_attenuations else 0.0 control_att_mean = float(np.mean(control_attenuations)) if control_attenuations else 0.0 total_elapsed = time.perf_counter() - t0 print( "[factor-trace-intervention] summary " f"prompts={len(per_prompt)} target_kl_mean={target_kl_mean:.6f} " f"control_kl_mean={control_kl_mean:.6f} " f"ratio={(target_kl_mean / control_kl_mean if abs(control_kl_mean) > 1e-8 else float('inf')):.2f}x " f"target_att_mean={target_att_mean:.6f} control_att_mean={control_att_mean:.6f} " f"elapsed={total_elapsed:.1f}s", flush=True, ) return { "target_feature_idx": int(target_feature_idx), "similar_frequency_control": int(similar_feature), "random_controls": [int(idx) for idx in random_features], "n_prompts_evaluated": len(per_prompt), "total_time_s": total_elapsed, "per_prompt": per_prompt, "summary": { "target_use_time_kl_mean": target_kl_mean, "control_use_time_kl_mean": control_kl_mean, "target_vs_random_kl_ratio": ( target_kl_mean / control_kl_mean if abs(control_kl_mean) > 1e-8 else float("inf") ), "target_boundary_mass_shift_mean": target_boundary_mean, "control_boundary_mass_shift_mean": control_boundary_mean, "target_feature_attenuation_mean": target_att_mean, "control_feature_attenuation_mean": control_att_mean, }, } def run_factor_trace_transplant( model, tokenizer, corpus_sequences: list[torch.Tensor], trace_result: dict[str, Any], layer_idx: int, head_idx: int, sae, sae_type: str, states: torch.Tensor, target_feature_indices: Sequence[int] | int, n_pairs: int = 16, device: str = "cuda", ) -> dict[str, Any]: """Transplant a donor write value into matched recipient prompts and score local readouts.""" if isinstance(target_feature_indices, (int, np.integer)): target_features = [int(target_feature_indices)] else: target_features = [int(idx) for idx in target_feature_indices] if not target_features: raise ValueError("Need at least one target feature for transplant") primary_target_feature = int(target_features[0]) prompt_map = { int(prompt["prompt_index"]): prompt for prompt in trace_result.get("prompt_records", []) } sequence_map = { int(prompt["prompt_index"]): seq for prompt, seq in zip(trace_result.get("prompt_records", []), corpus_sequences) } candidate_pairs = select_transplant_pairs( trace_result=trace_result, target_feature_idx=primary_target_feature, max_pairs=max(n_pairs * 12, 64), preferred_events=("paragraph_boundary",), preferred_offsets=(16,), min_write_gain=1.20, min_context_similarity=0.20, ) if not candidate_pairs: raise ValueError(f"No transplant pairs found for features {target_features}") act_matrix = compute_activation_matrix(sae, states) frequencies = (act_matrix > 0).mean(axis=0) similar_features: list[int] = [] excluded = set(target_features) for feat in target_features: similar_feat = choose_similar_frequency_feature( target_feature_idx=int(feat), frequencies=frequencies, exclude=excluded | set(similar_features), ) similar_features.append(int(similar_feat)) token_groups = boundary_token_groups(tokenizer) t0 = time.perf_counter() per_pair = [] target_boundary = [] wrong_feature_boundary = [] wrong_time_boundary = [] target_newline = [] wrong_feature_newline = [] wrong_time_newline = [] target_sentence = [] wrong_feature_sentence = [] wrong_time_sentence = [] target_kl = [] wrong_feature_kl = [] wrong_time_kl = [] target_coeff_delta = [] wrong_feature_coeff_delta = [] wrong_time_coeff_delta = [] print( "[factor-trace-transplant] " f"target={target_features} similar_control={similar_features} " f"candidate_pairs={len(candidate_pairs)}", flush=True, ) baseline_candidates: list[dict[str, Any]] = [] tracked_features = [*target_features, *similar_features] for pair in candidate_pairs: recipient_prompt_index = int(pair["recipient_prompt_index"]) donor_prompt_index = int(pair["donor_prompt_index"]) donor_prompt = prompt_map[donor_prompt_index] recipient_full = sequence_map[recipient_prompt_index].to(device) donor_coeffs = { int(feat): float(donor_prompt["feature_coefficients"][str(int(feat))][pair["donor_write_pos"]]) for feat in target_features } use_pos = int(pair["recipient_use_pos"]) eval_forward_window = 8 seq_end = min(int(recipient_full.shape[1]), use_pos + eval_forward_window + 2) if seq_end <= use_pos + 1: continue seq = recipient_full[:, :seq_end] actual_next_token = int(recipient_full[0, use_pos + 1].item()) baseline = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, ) eval_meta = _select_eval_site( baseline_logits_by_pos=baseline["logits_by_pos"], input_ids=seq, anchor_use_pos=use_pos, token_groups=token_groups, tokenizer=tokenizer, preferred_groups=("newline", "sentence_punct", "colon"), ) if not eval_meta: continue eval_pos = int(eval_meta["eval_pos"]) baseline_logits = baseline["logits_by_pos"][eval_pos] eval_meta["baseline_boundary_mass"] = _mass_from_logits(baseline_logits, token_groups["boundary"]) eval_meta["baseline_actual_group_mass"] = _mass_from_logits( baseline_logits, token_groups[eval_meta["eval_actual_group"]], ) eval_meta["baseline_actual_token_prob"] = _token_prob_from_logits( baseline_logits, int(eval_meta["eval_actual_next_token"]), ) if ( eval_meta["baseline_actual_token_prob"] < 0.03 or eval_meta["baseline_actual_token_prob"] > 0.75 ): continue target_updates = {} wrong_feature_updates = {} sign_flip_updates = {} for feat, donor_coeff in donor_coeffs.items(): recipient_write_coeff = float( baseline["feature_coefficients"][str(int(feat))][int(pair["recipient_write_pos"])] ) donor_prev_coeff = 0.0 if int(pair["donor_write_pos"]) > 0: donor_prev_coeff = float( donor_prompt["feature_coefficients"][str(int(feat))][int(pair["donor_write_pos"]) - 1] ) donor_delta = float(donor_coeff - donor_prev_coeff) target_updates[int(feat)] = float(recipient_write_coeff + donor_delta) sign_flip_updates[int(feat)] = float(recipient_write_coeff - donor_delta) for feat, ctrl in zip(target_features, similar_features): ctrl_write_coeff = float( baseline["feature_coefficients"][str(int(ctrl))][int(pair["recipient_write_pos"])] ) donor_prev_coeff = 0.0 if int(pair["donor_write_pos"]) > 0: donor_prev_coeff = float( donor_prompt["feature_coefficients"][str(int(feat))][int(pair["donor_write_pos"]) - 1] ) donor_delta = float(donor_coeffs[int(feat)] - donor_prev_coeff) wrong_feature_updates[int(ctrl)] = float(ctrl_write_coeff + donor_delta) headroom = eval_meta["baseline_actual_token_prob"] * (1.0 - eval_meta["baseline_actual_token_prob"]) donor_gap = float(pair["donor_write_coeff"]) - float(pair["recipient_write_coeff"]) baseline_score = ( headroom, float(eval_meta["baseline_actual_group_mass"]), float(pair["context_similarity"]), donor_gap, ) baseline_candidates.append( { "pair": pair, "seq": seq, "actual_next_token": actual_next_token, "baseline": baseline, "eval_meta": eval_meta, "use_pos": use_pos, "target_updates": {int(k): float(v) for k, v in target_updates.items()}, "wrong_feature_updates": {int(k): float(v) for k, v in wrong_feature_updates.items()}, "sign_flip_updates": {int(k): float(v) for k, v in sign_flip_updates.items()}, "score": baseline_score, } ) baseline_candidates.sort(key=lambda item: item["score"], reverse=True) if not baseline_candidates: raise ValueError("No transplant candidates survived exact-token headroom filtering") print( "[factor-trace-transplant] " f"baseline-screened={len(baseline_candidates)} selected={min(len(baseline_candidates), n_pairs)}", flush=True, ) for item in baseline_candidates[:n_pairs]: pair = item["pair"] seq = item["seq"] actual_next_token = int(item["actual_next_token"]) baseline = item["baseline"] eval_meta = dict(item["eval_meta"]) use_pos = int(item["use_pos"]) eval_pos = int(eval_meta["eval_pos"]) target_updates = {int(k): float(v) for k, v in item["target_updates"].items()} wrong_feature_updates = {int(k): float(v) for k, v in item["wrong_feature_updates"].items()} sign_flip_updates = {int(k): float(v) for k, v in item["sign_flip_updates"].items()} wrong_time_pos = pair.get("wrong_time_pos") prompt_t0 = time.perf_counter() recipient_prompt_index = int(pair["recipient_prompt_index"]) donor_prompt_index = int(pair["donor_prompt_index"]) print( f"[factor-trace-transplant] pair {len(per_pair)+1}/{min(len(baseline_candidates), n_pairs)} " f"donor={donor_prompt_index}@{pair['donor_write_pos']} " f"recipient={recipient_prompt_index}@{pair['recipient_write_pos']} " f"use={pair['recipient_use_pos']} offset={pair['recipient_use_offset']} " f"ctx={pair['context_similarity']:.2f} " f"tok={eval_meta['eval_actual_text']!r} p={eval_meta['baseline_actual_token_prob']:.3f}", flush=True, ) target_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=target_updates, intervention_pos=int(pair["recipient_write_pos"]), ) wrong_feature_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=wrong_feature_updates, intervention_pos=int(pair["recipient_write_pos"]), ) sign_flip_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=sign_flip_updates, intervention_pos=int(pair["recipient_write_pos"]), ) target_logits = target_rollout["logits_by_pos"][eval_pos] wrong_feature_logits = wrong_feature_rollout["logits_by_pos"][eval_pos] sign_flip_logits = sign_flip_rollout["logits_by_pos"][eval_pos] baseline_use_coeff = float( sum(float(baseline["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features) ) target_use_coeff = float( sum(float(target_rollout["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features) ) wrong_feature_use_coeff = float( sum(float(wrong_feature_rollout["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features) ) sign_flip_use_coeff = float( sum(float(sign_flip_rollout["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features) ) target_boundary_shift = _boundary_mass_shift(baseline_logits, target_logits, token_groups["boundary"]) wrong_feature_boundary_shift = _boundary_mass_shift( baseline_logits, wrong_feature_logits, token_groups["boundary"] ) sign_flip_boundary_shift = _boundary_mass_shift( baseline_logits, sign_flip_logits, token_groups["boundary"] ) target_newline_shift = _boundary_mass_shift(baseline_logits, target_logits, token_groups["newline"]) wrong_feature_newline_shift = _boundary_mass_shift( baseline_logits, wrong_feature_logits, token_groups["newline"] ) sign_flip_newline_shift = _boundary_mass_shift( baseline_logits, sign_flip_logits, token_groups["newline"] ) target_sentence_shift = _boundary_mass_shift( baseline_logits, target_logits, token_groups["sentence_punct"] ) wrong_feature_sentence_shift = _boundary_mass_shift( baseline_logits, wrong_feature_logits, token_groups["sentence_punct"] ) sign_flip_sentence_shift = _boundary_mass_shift( baseline_logits, sign_flip_logits, token_groups["sentence_punct"] ) wrong_time_boundary_shift = 0.0 wrong_time_newline_shift = 0.0 wrong_time_sentence_shift = 0.0 wrong_time_kl_value = 0.0 wrong_time_coeff_value = baseline_use_coeff wrong_time_actual_logit_shift = 0.0 sign_flip_kl_value = _kl_divergence_from_logits(baseline_logits, sign_flip_logits) sign_flip_actual_logit_shift = _token_logit_shift( baseline_logits, sign_flip_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ) if wrong_time_pos is not None: wrong_time_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=target_updates, intervention_pos=int(wrong_time_pos), ) wrong_time_logits = wrong_time_rollout["logits_by_pos"][eval_pos] wrong_time_boundary_shift = _boundary_mass_shift( baseline_logits, wrong_time_logits, token_groups["boundary"] ) wrong_time_newline_shift = _boundary_mass_shift( baseline_logits, wrong_time_logits, token_groups["newline"] ) wrong_time_sentence_shift = _boundary_mass_shift( baseline_logits, wrong_time_logits, token_groups["sentence_punct"] ) wrong_time_kl_value = _kl_divergence_from_logits(baseline_logits, wrong_time_logits) wrong_time_actual_logit_shift = _token_logit_shift( baseline_logits, wrong_time_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ) wrong_time_coeff_value = float( sum( float(wrong_time_rollout["feature_coefficients"][str(int(feat))][use_pos]) for feat in target_features ) ) row = { **pair, **eval_meta, "target_feature_indices": [int(feat) for feat in target_features], "similar_frequency_controls": [int(feat) for feat in similar_features], "primary_target_feature_idx": primary_target_feature, "donor_values": {str(k): float(v) for k, v in donor_coeffs.items()}, "actual_next_token": int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, "actual_next_token_logit_shift_target": _token_logit_shift( baseline_logits, target_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ), "actual_next_token_prob_shift_target": _token_prob_shift( baseline_logits, target_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ), "actual_next_token_logit_shift_wrong_feature": _token_logit_shift( baseline_logits, wrong_feature_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ), "actual_next_token_prob_shift_wrong_feature": _token_prob_shift( baseline_logits, wrong_feature_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ), "actual_next_token_logit_shift_sign_flip": sign_flip_actual_logit_shift, "actual_next_token_prob_shift_sign_flip": _token_prob_shift( baseline_logits, sign_flip_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ), "actual_next_token_logit_shift_wrong_time": wrong_time_actual_logit_shift, "actual_next_token_prob_shift_wrong_time": _token_prob_shift( baseline_logits, wrong_time_logits, int(eval_meta["eval_actual_next_token"]) if eval_meta["eval_actual_next_token"] is not None else actual_next_token, ) if wrong_time_pos is not None else 0.0, "target_use_time_kl": _kl_divergence_from_logits(baseline_logits, target_logits), "wrong_feature_use_time_kl": _kl_divergence_from_logits(baseline_logits, wrong_feature_logits), "sign_flip_use_time_kl": sign_flip_kl_value, "wrong_time_use_time_kl": wrong_time_kl_value, "target_boundary_mass_shift": target_boundary_shift, "wrong_feature_boundary_mass_shift": wrong_feature_boundary_shift, "sign_flip_boundary_mass_shift": sign_flip_boundary_shift, "wrong_time_boundary_mass_shift": wrong_time_boundary_shift, "target_newline_mass_shift": target_newline_shift, "wrong_feature_newline_mass_shift": wrong_feature_newline_shift, "sign_flip_newline_mass_shift": sign_flip_newline_shift, "wrong_time_newline_mass_shift": wrong_time_newline_shift, "target_sentence_mass_shift": target_sentence_shift, "wrong_feature_sentence_mass_shift": wrong_feature_sentence_shift, "sign_flip_sentence_mass_shift": sign_flip_sentence_shift, "wrong_time_sentence_mass_shift": wrong_time_sentence_shift, "baseline_use_coeff": baseline_use_coeff, "target_use_coeff": target_use_coeff, "wrong_feature_use_coeff": wrong_feature_use_coeff, "sign_flip_use_coeff": sign_flip_use_coeff, "wrong_time_use_coeff": wrong_time_coeff_value, "target_use_coeff_delta": float(target_use_coeff - baseline_use_coeff), "wrong_feature_use_coeff_delta": float(wrong_feature_use_coeff - baseline_use_coeff), "sign_flip_use_coeff_delta": float(sign_flip_use_coeff - baseline_use_coeff), "wrong_time_use_coeff_delta": float(wrong_time_coeff_value - baseline_use_coeff), } per_pair.append(row) target_boundary.append(row["target_boundary_mass_shift"]) wrong_feature_boundary.append(row["wrong_feature_boundary_mass_shift"]) wrong_time_boundary.append(row["wrong_time_boundary_mass_shift"]) target_newline.append(row["target_newline_mass_shift"]) wrong_feature_newline.append(row["wrong_feature_newline_mass_shift"]) wrong_time_newline.append(row["wrong_time_newline_mass_shift"]) target_sentence.append(row["target_sentence_mass_shift"]) wrong_feature_sentence.append(row["wrong_feature_sentence_mass_shift"]) wrong_time_sentence.append(row["wrong_time_sentence_mass_shift"]) target_kl.append(row["target_use_time_kl"]) wrong_feature_kl.append(row["wrong_feature_use_time_kl"]) wrong_time_kl.append(row["wrong_time_use_time_kl"]) target_coeff_delta.append(row["target_use_coeff_delta"]) wrong_feature_coeff_delta.append(row["wrong_feature_use_coeff_delta"]) wrong_time_coeff_delta.append(row["wrong_time_use_coeff_delta"]) print( f"[factor-trace-transplant] done {len(per_pair)}/{n_pairs} " f"boundary={row['target_boundary_mass_shift']:+.6f} " f"wrong_feature={row['wrong_feature_boundary_mass_shift']:+.6f} " f"wrong_time={row['wrong_time_boundary_mass_shift']:+.6f} " f"coeff_delta={row['target_use_coeff_delta']:+.6f} " f"elapsed={time.perf_counter() - prompt_t0:.1f}s", flush=True, ) def _mean(values: list[float]) -> float: return float(np.mean(values)) if values else 0.0 summary = { "target_boundary_mass_shift_mean": _mean(target_boundary), "wrong_feature_boundary_mass_shift_mean": _mean(wrong_feature_boundary), "wrong_time_boundary_mass_shift_mean": _mean(wrong_time_boundary), "target_newline_mass_shift_mean": _mean(target_newline), "wrong_feature_newline_mass_shift_mean": _mean(wrong_feature_newline), "wrong_time_newline_mass_shift_mean": _mean(wrong_time_newline), "target_sentence_mass_shift_mean": _mean(target_sentence), "wrong_feature_sentence_mass_shift_mean": _mean(wrong_feature_sentence), "wrong_time_sentence_mass_shift_mean": _mean(wrong_time_sentence), "target_use_time_kl_mean": _mean(target_kl), "wrong_feature_use_time_kl_mean": _mean(wrong_feature_kl), "wrong_time_use_time_kl_mean": _mean(wrong_time_kl), "target_use_coeff_delta_mean": _mean(target_coeff_delta), "wrong_feature_use_coeff_delta_mean": _mean(wrong_feature_coeff_delta), "wrong_time_use_coeff_delta_mean": _mean(wrong_time_coeff_delta), "target_actual_token_logit_shift_mean": _mean( [float(row["actual_next_token_logit_shift_target"]) for row in per_pair] ), "wrong_feature_actual_token_logit_shift_mean": _mean( [float(row["actual_next_token_logit_shift_wrong_feature"]) for row in per_pair] ), "sign_flip_actual_token_logit_shift_mean": _mean( [float(row["actual_next_token_logit_shift_sign_flip"]) for row in per_pair] ), "target_actual_token_prob_shift_mean": _mean( [float(row["actual_next_token_prob_shift_target"]) for row in per_pair] ), "wrong_feature_actual_token_prob_shift_mean": _mean( [float(row["actual_next_token_prob_shift_wrong_feature"]) for row in per_pair] ), "sign_flip_actual_token_prob_shift_mean": _mean( [float(row["actual_next_token_prob_shift_sign_flip"]) for row in per_pair] ), "wrong_time_actual_token_prob_shift_mean": _mean( [float(row["actual_next_token_prob_shift_wrong_time"]) for row in per_pair] ), "wrong_time_actual_token_logit_shift_mean": _mean( [float(row["actual_next_token_logit_shift_wrong_time"]) for row in per_pair] ), "target_minus_sign_flip_actual_token_logit_shift_mean": _mean( [ float(row["actual_next_token_logit_shift_target"]) - float(row["actual_next_token_logit_shift_sign_flip"]) for row in per_pair ] ), "target_minus_sign_flip_actual_token_prob_shift_mean": _mean( [ float(row["actual_next_token_prob_shift_target"]) - float(row["actual_next_token_prob_shift_sign_flip"]) for row in per_pair ] ), "target_beats_wrong_feature_fraction": float( np.mean([a > b for a, b in zip(target_boundary, wrong_feature_boundary)]) ) if target_boundary else 0.0, "target_beats_sign_flip_fraction": float( np.mean([row["actual_next_token_logit_shift_target"] > row["actual_next_token_logit_shift_sign_flip"] for row in per_pair]) ) if per_pair else 0.0, "target_prob_beats_sign_flip_fraction": float( np.mean([row["actual_next_token_prob_shift_target"] > row["actual_next_token_prob_shift_sign_flip"] for row in per_pair]) ) if per_pair else 0.0, "target_beats_wrong_time_fraction": float( np.mean([a > b for a, b in zip(target_boundary, wrong_time_boundary)]) ) if target_boundary else 0.0, "target_actual_token_beats_wrong_feature_fraction": float( np.mean([row["actual_next_token_logit_shift_target"] > row["actual_next_token_logit_shift_wrong_feature"] for row in per_pair]) ) if per_pair else 0.0, "target_actual_token_beats_wrong_time_fraction": float( np.mean([row["actual_next_token_logit_shift_target"] > row["actual_next_token_logit_shift_wrong_time"] for row in per_pair]) ) if per_pair else 0.0, "target_coeff_beats_sign_flip_fraction": float( np.mean([row["target_use_coeff_delta"] > row["sign_flip_use_coeff_delta"] for row in per_pair]) ) if per_pair else 0.0, "target_minus_sign_flip_use_coeff_delta_mean": _mean( [float(row["target_use_coeff_delta"]) - float(row["sign_flip_use_coeff_delta"]) for row in per_pair] ), } summary["boundary_ratio_vs_wrong_feature"] = ( summary["target_boundary_mass_shift_mean"] / summary["wrong_feature_boundary_mass_shift_mean"] if abs(summary["wrong_feature_boundary_mass_shift_mean"]) > 1e-8 else float("inf") ) summary["boundary_ratio_vs_wrong_time"] = ( summary["target_boundary_mass_shift_mean"] / summary["wrong_time_boundary_mass_shift_mean"] if abs(summary["wrong_time_boundary_mass_shift_mean"]) > 1e-8 else float("inf") ) return { "target_feature_idx": primary_target_feature, "target_feature_indices": [int(feat) for feat in target_features], "similar_frequency_control": int(similar_features[0]), "similar_frequency_controls": [int(feat) for feat in similar_features], "n_pairs_evaluated": len(per_pair), "total_time_s": float(time.perf_counter() - t0), "per_pair": per_pair, "summary": summary, } def run_use_site_signed_causal( model, tokenizer, corpus_sequences: list[torch.Tensor], benchmark_result: dict[str, Any], layer_idx: int, head_idx: int, sae, sae_type: str, states: torch.Tensor, n_rows: int = 32, device: str = "cuda", ) -> dict[str, Any]: """Edit the signed H4 direction at the benchmarked use site and score local decisions.""" target_features = [int(x) for x in benchmark_result.get("target_feature_indices", [])] if target_features != [62, 105]: raise ValueError(f"Expected signed pair [62, 105], got {target_features}") accepted_rows = list(benchmark_result.get("accepted_rows", [])) if not accepted_rows: raise ValueError("No accepted benchmark rows for use-site causal stage") act_matrix = compute_activation_matrix(sae, states) frequencies = (act_matrix > 0).mean(axis=0) wrong_feature_pair = [ choose_similar_frequency_feature(target_features[0], frequencies, exclude=set(target_features)), choose_similar_frequency_feature( target_features[1], frequencies, exclude=set(target_features) | {int(choose_similar_frequency_feature(target_features[0], frequencies, exclude=set(target_features)))}, ), ] wrong_feature_pair = [int(x) for x in wrong_feature_pair] readout_token_ids = _readout_class_token_ids(tokenizer) t0 = time.perf_counter() sequence_map = {idx: seq.to(device) for idx, seq in enumerate(corpus_sequences)} tracked_features = [*target_features, *wrong_feature_pair] accepted_rows.sort(key=lambda r: r.get("score", (0, 0, 0, 0)), reverse=True) rows = accepted_rows[:n_rows] per_row = [] print( "[factor-trace-use-site] " f"target={target_features} reverse_sign_control=True wrong_time=True rows={len(rows)}", flush=True, ) for i, row in enumerate(rows, start=1): recipient_idx = int(row["recipient_prompt_index"]) seq_full = sequence_map[recipient_idx] decision_pos = int(row["decision_pos"]) intervention_pos = int(row["intervention_pos"]) wrong_time_pos = row.get("wrong_time_pos") actual_next_token = int(row["actual_next_token"]) seq_end = min(int(seq_full.shape[1]), decision_pos + 3) if seq_end <= decision_pos + 1: continue seq = seq_full[:, :seq_end] prompt_t0 = time.perf_counter() baseline = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, ) target_updates = {} reverse_updates = {} wrong_feature_updates = {} for feat, delta in zip(target_features, [float(row["donor_delta_62"]), float(row["donor_delta_105"])]): base_coeff = float(baseline["feature_coefficients"][str(int(feat))][intervention_pos]) target_updates[int(feat)] = float(base_coeff + delta) reverse_updates[int(feat)] = float(base_coeff - delta) for ctrl, delta in zip(wrong_feature_pair, [float(row["donor_delta_62"]), float(row["donor_delta_105"])]): base_ctrl = float(baseline["feature_coefficients"][str(int(ctrl))][intervention_pos]) wrong_feature_updates[int(ctrl)] = float(base_ctrl + delta) target_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=target_updates, intervention_pos=intervention_pos, ) reverse_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=reverse_updates, intervention_pos=intervention_pos, ) wrong_feature_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=wrong_feature_updates, intervention_pos=intervention_pos, ) wrong_time_rollout = None if wrong_time_pos is not None: wrong_time_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_updates=target_updates, intervention_pos=int(wrong_time_pos), ) baseline_logits = baseline["logits_by_pos"][decision_pos] target_logits = target_rollout["logits_by_pos"][decision_pos] reverse_logits = reverse_rollout["logits_by_pos"][decision_pos] wrong_feature_logits = wrong_feature_rollout["logits_by_pos"][decision_pos] wrong_time_logits = wrong_time_rollout["logits_by_pos"][decision_pos] if wrong_time_rollout is not None else baseline_logits readout_class = str(row["readout_class"]) class_ids = readout_token_ids.get(readout_class, [actual_next_token]) or [actual_next_token] next_positions = [pos for pos in [decision_pos, min(decision_pos + 1, int(seq.shape[1]) - 1)] if pos in baseline["logits_by_pos"]] def _signed_coeff_change(rollout: dict[str, Any]) -> float: diffs = [] for pos in next_positions: base = float(baseline["feature_coefficients"]["62"][pos] - baseline["feature_coefficients"]["105"][pos]) edit = float(rollout["feature_coefficients"]["62"][pos] - rollout["feature_coefficients"]["105"][pos]) diffs.append(edit - base) return float(np.mean(diffs)) if diffs else 0.0 row_out = { **row, "actual_token_logit_shift_target": _token_logit_shift(baseline_logits, target_logits, actual_next_token), "actual_token_logit_shift_reverse": _token_logit_shift(baseline_logits, reverse_logits, actual_next_token), "actual_token_logit_shift_wrong_feature": _token_logit_shift(baseline_logits, wrong_feature_logits, actual_next_token), "actual_token_logit_shift_wrong_time": _token_logit_shift(baseline_logits, wrong_time_logits, actual_next_token), "actual_token_prob_shift_target": _token_prob_shift(baseline_logits, target_logits, actual_next_token), "actual_token_prob_shift_reverse": _token_prob_shift(baseline_logits, reverse_logits, actual_next_token), "actual_token_prob_shift_wrong_feature": _token_prob_shift(baseline_logits, wrong_feature_logits, actual_next_token), "actual_token_prob_shift_wrong_time": _token_prob_shift(baseline_logits, wrong_time_logits, actual_next_token), "readout_class_prob_shift_target": _boundary_mass_shift(baseline_logits, target_logits, class_ids), "readout_class_prob_shift_reverse": _boundary_mass_shift(baseline_logits, reverse_logits, class_ids), "readout_class_prob_shift_wrong_feature": _boundary_mass_shift(baseline_logits, wrong_feature_logits, class_ids), "readout_class_prob_shift_wrong_time": _boundary_mass_shift(baseline_logits, wrong_time_logits, class_ids), "signed_direction_coeff_shift_target": _signed_coeff_change(target_rollout), "signed_direction_coeff_shift_reverse": _signed_coeff_change(reverse_rollout), "signed_direction_coeff_shift_wrong_feature": _signed_coeff_change(wrong_feature_rollout), "signed_direction_coeff_shift_wrong_time": _signed_coeff_change(wrong_time_rollout) if wrong_time_rollout is not None else 0.0, "local_kl_target": _kl_divergence_from_logits(baseline_logits, target_logits), "local_kl_reverse": _kl_divergence_from_logits(baseline_logits, reverse_logits), "local_kl_wrong_time": _kl_divergence_from_logits(baseline_logits, wrong_time_logits) if wrong_time_rollout is not None else 0.0, } per_row.append(row_out) print( f"[factor-trace-use-site] row {len(per_row)}/{len(rows)} " f"tok={row['readout_text']!r} p={row['baseline_actual_token_prob']:.3f} " f"target_prob={row_out['actual_token_prob_shift_target']:+.4f} " f"reverse_prob={row_out['actual_token_prob_shift_reverse']:+.4f} " f"wrong_time_prob={row_out['actual_token_prob_shift_wrong_time']:+.4f} " f"signed_shift={row_out['signed_direction_coeff_shift_target']:+.4f} " f"elapsed={time.perf_counter() - prompt_t0:.1f}s", flush=True, ) def _mean(xs: list[float]) -> float: return float(np.mean(xs)) if xs else 0.0 summary = { "target_actual_token_prob_shift_mean": _mean([r["actual_token_prob_shift_target"] for r in per_row]), "reverse_actual_token_prob_shift_mean": _mean([r["actual_token_prob_shift_reverse"] for r in per_row]), "wrong_time_actual_token_prob_shift_mean": _mean([r["actual_token_prob_shift_wrong_time"] for r in per_row]), "target_minus_reverse_actual_token_prob_shift_mean": _mean( [r["actual_token_prob_shift_target"] - r["actual_token_prob_shift_reverse"] for r in per_row] ), "target_actual_token_logit_shift_mean": _mean([r["actual_token_logit_shift_target"] for r in per_row]), "reverse_actual_token_logit_shift_mean": _mean([r["actual_token_logit_shift_reverse"] for r in per_row]), "target_minus_reverse_actual_token_logit_shift_mean": _mean( [r["actual_token_logit_shift_target"] - r["actual_token_logit_shift_reverse"] for r in per_row] ), "target_readout_class_prob_shift_mean": _mean([r["readout_class_prob_shift_target"] for r in per_row]), "reverse_readout_class_prob_shift_mean": _mean([r["readout_class_prob_shift_reverse"] for r in per_row]), "target_signed_direction_coeff_shift_mean": _mean([r["signed_direction_coeff_shift_target"] for r in per_row]), "reverse_signed_direction_coeff_shift_mean": _mean([r["signed_direction_coeff_shift_reverse"] for r in per_row]), "wrong_time_signed_direction_coeff_shift_mean": _mean([r["signed_direction_coeff_shift_wrong_time"] for r in per_row]), "target_beats_reverse_prob_fraction": float(np.mean([ r["actual_token_prob_shift_target"] > r["actual_token_prob_shift_reverse"] for r in per_row ])) if per_row else 0.0, "target_beats_wrong_time_prob_fraction": float(np.mean([ r["actual_token_prob_shift_target"] > r["actual_token_prob_shift_wrong_time"] for r in per_row ])) if per_row else 0.0, "target_localized_coeff_beats_wrong_time_fraction": float(np.mean([ r["signed_direction_coeff_shift_target"] > r["signed_direction_coeff_shift_wrong_time"] for r in per_row ])) if per_row else 0.0, } summary["promotion_gate_pass"] = bool( summary["target_beats_reverse_prob_fraction"] >= 0.65 and summary["target_minus_reverse_actual_token_prob_shift_mean"] > 0.0 and summary["target_beats_wrong_time_prob_fraction"] >= 0.60 and summary["target_signed_direction_coeff_shift_mean"] > summary["wrong_time_signed_direction_coeff_shift_mean"] ) return { "target_feature_indices": target_features, "wrong_feature_pair": wrong_feature_pair, "n_rows_evaluated": len(per_row), "per_row": per_row, "summary": summary, "total_time_s": float(time.perf_counter() - t0), } def run_use_site_readout_map( model, tokenizer, corpus_sequences: list[torch.Tensor], benchmark_result: dict[str, Any], layer_idx: int, head_idx: int, sae, sae_type: str, n_rows: int = 32, device: str = "cuda", ) -> dict[str, Any]: """Map localized downstream readouts for the signed H4 boundary direction.""" del device # kept for interface symmetry target_features = [int(x) for x in benchmark_result.get("target_feature_indices", [])] if target_features != [62, 105]: raise ValueError(f"Expected signed pair [62, 105], got {target_features}") accepted_rows = list(benchmark_result.get("accepted_rows", [])) if not accepted_rows: raise ValueError("No accepted benchmark rows for readout-map stage") t0 = time.perf_counter() sequence_map = {idx: seq for idx, seq in enumerate(corpus_sequences)} token_class_ids = _readout_class_token_ids(tokenizer) accepted_rows.sort(key=lambda r: r.get("score", (0, 0, 0, 0)), reverse=True) rows = accepted_rows[:n_rows] per_row = [] hidden_sum: np.ndarray | None = None hidden_abs_sum: np.ndarray | None = None hidden_pos_counts: np.ndarray | None = None logit_sum: np.ndarray | None = None logit_abs_sum: np.ndarray | None = None logit_pos_counts: np.ndarray | None = None print( "[factor-trace-readout-map] " f"target={target_features} reverse_sign_control=True rows={len(rows)}", flush=True, ) for i, row in enumerate(rows, start=1): recipient_idx = int(row["recipient_prompt_index"]) seq_full = sequence_map[recipient_idx] decision_pos = int(row["decision_pos"]) intervention_pos = int(row["intervention_pos"]) wrong_time_pos = row.get("wrong_time_pos") actual_next_token = int(row["actual_next_token"]) seq_end = min(int(seq_full.shape[1]), decision_pos + 3) if seq_end <= decision_pos + 1: continue seq = seq_full[:, :seq_end] prompt_t0 = time.perf_counter() capture_positions = [decision_pos, min(decision_pos + 1, seq_end - 1)] baseline = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=target_features, capture_hidden_state_positions=capture_positions, ) target_updates = {} reverse_updates = {} for feat, delta in zip(target_features, [float(row["donor_delta_62"]), float(row["donor_delta_105"])]): base_coeff = float(baseline["feature_coefficients"][str(int(feat))][intervention_pos]) target_updates[int(feat)] = float(base_coeff + delta) reverse_updates[int(feat)] = float(base_coeff - delta) target_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=target_features, intervention_updates=target_updates, intervention_pos=intervention_pos, capture_hidden_state_positions=capture_positions, ) reverse_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=target_features, intervention_updates=reverse_updates, intervention_pos=intervention_pos, capture_hidden_state_positions=capture_positions, ) wrong_time_rollout = None if wrong_time_pos is not None: wrong_time_rollout = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=target_features, intervention_updates=target_updates, intervention_pos=int(wrong_time_pos), capture_hidden_state_positions=capture_positions, ) baseline_logits = baseline["logits_by_pos"][decision_pos] target_logits = target_rollout["logits_by_pos"][decision_pos] reverse_logits = reverse_rollout["logits_by_pos"][decision_pos] wrong_time_logits = wrong_time_rollout["logits_by_pos"][decision_pos] if wrong_time_rollout is not None else baseline_logits readout_class = str(row["readout_class"]) class_ids = token_class_ids.get(readout_class, [actual_next_token]) or [actual_next_token] hidden_delta_vecs = [] for pos in capture_positions: base_hidden = baseline["hidden_state_by_pos"].get(pos) target_hidden = target_rollout["hidden_state_by_pos"].get(pos) reverse_hidden = reverse_rollout["hidden_state_by_pos"].get(pos) if base_hidden is None or target_hidden is None or reverse_hidden is None: continue delta_vec = (target_hidden - reverse_hidden).squeeze(0).numpy().astype(np.float32) hidden_delta_vecs.append(delta_vec) if hidden_delta_vecs: hidden_delta = np.mean(np.stack(hidden_delta_vecs, axis=0), axis=0) if hidden_sum is None: hidden_sum = np.zeros_like(hidden_delta) hidden_abs_sum = np.zeros_like(hidden_delta) hidden_pos_counts = np.zeros_like(hidden_delta) hidden_sum += hidden_delta hidden_abs_sum += np.abs(hidden_delta) hidden_pos_counts += (hidden_delta > 0).astype(np.float32) else: hidden_delta = np.zeros(0, dtype=np.float32) logit_delta = (target_logits - reverse_logits).squeeze(0).numpy().astype(np.float32) if logit_sum is None: logit_sum = np.zeros_like(logit_delta) logit_abs_sum = np.zeros_like(logit_delta) logit_pos_counts = np.zeros_like(logit_delta) logit_sum += logit_delta logit_abs_sum += np.abs(logit_delta) logit_pos_counts += (logit_delta > 0).astype(np.float32) per_row.append( { **row, "actual_token_logit_delta_target_minus_reverse": _token_logit_shift(reverse_logits, target_logits, actual_next_token), "actual_token_prob_delta_target_minus_reverse": _token_prob_shift(reverse_logits, target_logits, actual_next_token), "readout_class_prob_delta_target_minus_reverse": _boundary_mass_shift(reverse_logits, target_logits, class_ids), "actual_token_logit_delta_target_minus_wrong_time": _token_logit_shift(wrong_time_logits, target_logits, actual_next_token), "actual_token_prob_delta_target_minus_wrong_time": _token_prob_shift(wrong_time_logits, target_logits, actual_next_token), "hidden_state_delta_l2_target_minus_reverse": float(np.linalg.norm(hidden_delta)) if hidden_delta.size else 0.0, "elapsed_s": float(time.perf_counter() - prompt_t0), } ) print( f"[factor-trace-readout-map] row {len(per_row)}/{len(rows)} " f"tok={row['readout_text']!r} class={readout_class} " f"delta_prob={per_row[-1]['actual_token_prob_delta_target_minus_reverse']:+.4f} " f"delta_logit={per_row[-1]['actual_token_logit_delta_target_minus_reverse']:+.4f} " f"hidden_l2={per_row[-1]['hidden_state_delta_l2_target_minus_reverse']:.4f} " f"elapsed={per_row[-1]['elapsed_s']:.1f}s", flush=True, ) n_eval = max(len(per_row), 1) def _rank_dims(sum_arr: np.ndarray | None, abs_arr: np.ndarray | None, pos_arr: np.ndarray | None, top_k: int = 16) -> list[dict[str, float]]: if sum_arr is None or abs_arr is None or pos_arr is None: return [] mean = sum_arr / n_eval mean_abs = abs_arr / n_eval pos_frac = pos_arr / n_eval sign_consistency = np.abs(2.0 * pos_frac - 1.0) score = np.abs(mean) * sign_consistency order = np.argsort(-score)[:top_k] ranked = [] for idx in order: ranked.append( { "index": int(idx), "mean_delta": float(mean[idx]), "mean_abs_delta": float(mean_abs[idx]), "sign_consistency": float(sign_consistency[idx]), "positive_fraction": float(pos_frac[idx]), "score": float(score[idx]), } ) return ranked def _rank_tokens(sum_arr: np.ndarray | None, abs_arr: np.ndarray | None, pos_arr: np.ndarray | None, top_k: int = 16) -> dict[str, list[dict[str, float | int | str]]]: if sum_arr is None or abs_arr is None or pos_arr is None: return {"positive": [], "negative": []} mean = sum_arr / n_eval mean_abs = abs_arr / n_eval pos_frac = pos_arr / n_eval pos_order = np.argsort(-mean)[:top_k] neg_order = np.argsort(mean)[:top_k] ranked: dict[str, list[dict[str, float | int | str]]] = {"positive": [], "negative": []} for label, order in [("positive", pos_order), ("negative", neg_order)]: for idx in order: token_text = tokenizer.decode( [int(idx)], skip_special_tokens=False, clean_up_tokenization_spaces=False, ) ranked[label].append( { "token_id": int(idx), "token_text": token_text, "mean_delta": float(mean[idx]), "mean_abs_delta": float(mean_abs[idx]), "positive_fraction": float(pos_frac[idx]), } ) return ranked summary = { "actual_token_prob_delta_target_minus_reverse_mean": float(np.mean([r["actual_token_prob_delta_target_minus_reverse"] for r in per_row])) if per_row else 0.0, "actual_token_logit_delta_target_minus_reverse_mean": float(np.mean([r["actual_token_logit_delta_target_minus_reverse"] for r in per_row])) if per_row else 0.0, "readout_class_prob_delta_target_minus_reverse_mean": float(np.mean([r["readout_class_prob_delta_target_minus_reverse"] for r in per_row])) if per_row else 0.0, "actual_token_prob_delta_target_minus_wrong_time_mean": float(np.mean([r["actual_token_prob_delta_target_minus_wrong_time"] for r in per_row])) if per_row else 0.0, "stable_hidden_dims": _rank_dims(hidden_sum, hidden_abs_sum, hidden_pos_counts, top_k=16), "top_logit_tokens": _rank_tokens(logit_sum, logit_abs_sum, logit_pos_counts, top_k=12), } return { "target_feature_indices": target_features, "n_rows_evaluated": len(per_row), "per_row": per_row, "summary": summary, "total_time_s": float(time.perf_counter() - t0), } def _period_token_id(tokenizer) -> int: ids = tokenizer.encode(".", add_special_tokens=False) if len(ids) != 1: raise ValueError(f"Expected '.' to map to a single token, got ids={ids}") return int(ids[0]) def _decoder_direction_from_weights( sae, weights: np.ndarray, ) -> tuple[torch.Tensor, list[dict[str, float | int]]]: if hasattr(sae, "V_dec") and hasattr(sae, "W_dec"): v = sae.V_dec.detach().float().cpu() w = sae.W_dec.detach().float().cpu() elif hasattr(sae, "V") and hasattr(sae, "W"): v = sae.V.detach().float().cpu() w = sae.W.detach().float().cpu() else: raise ValueError("SAE does not expose rank-1 decoder factors") if v.ndim == 2: v = v[:, None, :] if w.ndim == 2: w = w[:, None, :] weight_t = torch.from_numpy(weights.astype(np.float32)) direction = torch.einsum("i,irk,irv->kv", weight_t, v, w) fro = float(direction.norm().item()) if fro > 0: direction = direction / fro top = np.argsort(-np.abs(weights))[:16] top_weights = [ {"feature_idx": int(idx), "weight": float(weights[idx])} for idx in top if abs(float(weights[idx])) > 1e-6 ] return direction, top_weights def collect_exact_token_direction_dataset( model, corpus_batches: list[torch.Tensor], layer_idx: int, head_idx: int, sae, sae_type: str, target_token_id: int, prompt_len: int = 512, negative_keep_prob: float = 0.03, max_negative_multiplier: int = 4, device: str = "cuda", ) -> dict[str, Any]: """Collect SAE activations for exact '.' vs sampled non-'.' next-token decisions.""" rng = np.random.default_rng(0) xs: list[np.ndarray] = [] ys: list[int] = [] meta: list[dict[str, int]] = [] positives = 0 negatives = 0 prompt_index = 0 t0 = time.perf_counter() print( "[factor-trace-dot-direction] " f"collect_dataset batches={len(corpus_batches)} prompt_len={prompt_len} target_token={target_token_id}", flush=True, ) for batch_idx, batch in enumerate(corpus_batches, start=1): token_batch = batch[:, :prompt_len].to(device) batch_size, seq_len = token_batch.shape if seq_len < 2: continue cache = None for pos in range(seq_len - 1): step_ids = token_batch[:, pos : pos + 1] if cache is None: out = model(input_ids=step_ids, use_cache=True) else: out = model(input_ids=step_ids, past_key_values=cache, use_cache=True) cache = out.past_key_values state_batch = cache.layers[layer_idx].recurrent_states[:, head_idx].float() acts = _encode_states(sae, sae_type, state_batch).detach().cpu().numpy().astype(np.float32) next_tokens = token_batch[:, pos + 1].detach().cpu().numpy() for row_idx in range(batch_size): label = int(next_tokens[row_idx] == target_token_id) keep = label == 1 or rng.random() < negative_keep_prob if not keep: continue xs.append(acts[row_idx]) ys.append(label) meta.append({"prompt_index": prompt_index + row_idx, "decision_pos": int(pos)}) if label == 1: positives += 1 else: negatives += 1 prompt_index += batch_size if batch_idx % 4 == 0 or batch_idx == len(corpus_batches): print( "[factor-trace-dot-direction] " f"dataset_progress batch={batch_idx}/{len(corpus_batches)} " f"positives={positives} negatives={negatives} elapsed={time.perf_counter() - t0:.1f}s", flush=True, ) if positives == 0: raise ValueError("No positive '.' examples collected for direction training") x = np.stack(xs, axis=0).astype(np.float32) y = np.array(ys, dtype=np.int64) meta_arr = np.array(meta, dtype=object) pos_idx = np.where(y == 1)[0] neg_idx = np.where(y == 0)[0] max_neg = max(len(pos_idx) * max_negative_multiplier, len(pos_idx)) if len(neg_idx) > max_neg: neg_idx = rng.choice(neg_idx, size=max_neg, replace=False) selected = np.concatenate([pos_idx, neg_idx]) rng.shuffle(selected) return { "x": x[selected], "y": y[selected], "meta": meta_arr[selected].tolist(), "n_positive": int(len(pos_idx)), "n_negative_sampled": int(len(neg_idx)), "total_time_s": float(time.perf_counter() - t0), } def run_dot_direction_use_site( model, tokenizer, corpus_batches: list[torch.Tensor], corpus_sequences: list[torch.Tensor], benchmark_result: dict[str, Any], use_site_result: dict[str, Any], layer_idx: int, head_idx: int, sae, sae_type: str, prompt_len: int = 512, direction_prompt_count: int = 256, n_rows: int = 5, alphas: Sequence[float] = (0.5, 1.0, 2.0, 4.0), min_accuracy: float = 0.65, device: str = "cuda", ) -> dict[str, Any]: """Train a '.' logistic direction in H4 feature space and test it on exact '.' use-sites.""" from sklearn.linear_model import LogisticRegression # type: ignore[import-untyped] from sklearn.model_selection import train_test_split # type: ignore[import-untyped] target_token_id = _period_token_id(tokenizer) dataset = collect_exact_token_direction_dataset( model=model, corpus_batches=corpus_batches, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, target_token_id=target_token_id, prompt_len=prompt_len, device=device, ) x = dataset["x"] y = dataset["y"] x_train, x_test, y_train, y_test = train_test_split( x, y, test_size=0.2, random_state=0, stratify=y, ) clf = LogisticRegression( max_iter=2000, class_weight="balanced", solver="liblinear", random_state=0, ) clf.fit(x_train, y_train) test_accuracy = float(clf.score(x_test, y_test)) weights = clf.coef_[0].astype(np.float32) direction, top_weights = _decoder_direction_from_weights(sae, weights) print( "[factor-trace-dot-direction] " f"logistic accuracy={test_accuracy:.3f} positives={dataset['n_positive']} " f"negatives={dataset['n_negative_sampled']}", flush=True, ) print( "[factor-trace-dot-direction] " f"top_weights={top_weights[:8]}", flush=True, ) result: dict[str, Any] = { "target_feature_indices": [62, 105], "target_token_id": target_token_id, "target_token_text": ".", "logistic": { "test_accuracy": test_accuracy, "n_positive": dataset["n_positive"], "n_negative_sampled": dataset["n_negative_sampled"], "top_weights": top_weights, }, } if test_accuracy < min_accuracy: result["status"] = "stopped_low_accuracy" result["rows"] = [] result["summary"] = {"success_rows": 0, "success_gate_pass": False} return result benchmark_map = { (int(r["donor_prompt_index"]), int(r["recipient_prompt_index"]), int(r["decision_pos"])): r for r in benchmark_result.get("accepted_rows", []) } candidate_rows = [] for row in use_site_result.get("per_row", []): if row.get("readout_class") != "sentence_reset" or row.get("readout_text") != ".": continue p = float(row["baseline_actual_token_prob"]) if p < 0.3 or p > 0.8: continue delta = float(row["actual_token_prob_shift_target"] - row["actual_token_prob_shift_reverse"]) coeff = float(row["signed_direction_coeff_shift_target"]) score = ( int(delta > 0), delta, p * (1.0 - p), coeff, ) key = (int(row["donor_prompt_index"]), int(row["recipient_prompt_index"]), int(row["decision_pos"])) merged = dict(row) merged["benchmark_row"] = benchmark_map.get(key, {}) merged["_score"] = score candidate_rows.append(merged) candidate_rows.sort(key=lambda r: r["_score"], reverse=True) selected_rows = candidate_rows[:n_rows] result["selected_rows"] = [ { "donor_prompt_index": int(r["donor_prompt_index"]), "recipient_prompt_index": int(r["recipient_prompt_index"]), "decision_pos": int(r["decision_pos"]), "baseline_actual_token_prob": float(r["baseline_actual_token_prob"]), "prior_target_minus_reverse_prob": float(r["actual_token_prob_shift_target"] - r["actual_token_prob_shift_reverse"]), "prior_signed_direction_coeff_shift": float(r["signed_direction_coeff_shift_target"]), } for r in selected_rows ] print( "[factor-trace-dot-direction] " f"selected_rows={len(selected_rows)} exact_dot_only=True", flush=True, ) sequence_map = {idx: seq.to(device) for idx, seq in enumerate(corpus_sequences)} tracked_features = [62, 105, 121, 9, 137] per_row = [] direction = direction.to(device=device, dtype=torch.float32) t0 = time.perf_counter() for row_idx, row in enumerate(selected_rows, start=1): recipient_idx = int(row["recipient_prompt_index"]) seq_full = sequence_map[recipient_idx] decision_pos = int(row["decision_pos"]) intervention_pos = int(row.get("intervention_pos", max(0, decision_pos - 1))) wrong_time_pos = max(0, intervention_pos - 10) seq_end = min(int(seq_full.shape[1]), decision_pos + 3) seq = seq_full[:, :seq_end] baseline = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, ) baseline_logits = baseline["logits_by_pos"][decision_pos] baseline_prob = _token_prob_from_logits(baseline_logits, target_token_id) token_ids = seq_full[0].detach().cpu().tolist() ctx_start = max(0, decision_pos - 24) ctx_end = min(len(token_ids), decision_pos + 2) context_text = tokenizer.decode( token_ids[ctx_start:ctx_end], skip_special_tokens=False, clean_up_tokenization_spaces=False, ) alpha_rows = [] for alpha in alphas: scaled = direction * float(alpha) inject = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_state_delta=scaled, intervention_pos=intervention_pos, ) reverse = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_state_delta=-scaled, intervention_pos=intervention_pos, ) wrong_time = teacher_forced_feature_rollout( model=model, input_ids=seq, layer_idx=layer_idx, head_idx=head_idx, sae=sae, sae_type=sae_type, tracked_features=tracked_features, intervention_state_delta=scaled, intervention_pos=wrong_time_pos, ) inject_logits = inject["logits_by_pos"][decision_pos] reverse_logits = reverse["logits_by_pos"][decision_pos] wrong_time_logits = wrong_time["logits_by_pos"][decision_pos] inject_prob = _token_prob_from_logits(inject_logits, target_token_id) reverse_prob = _token_prob_from_logits(reverse_logits, target_token_id) wrong_time_prob = _token_prob_from_logits(wrong_time_logits, target_token_id) alpha_rows.append( { "alpha": float(alpha), "baseline_prob": baseline_prob, "inject_prob": inject_prob, "reverse_prob": reverse_prob, "wrong_time_prob": wrong_time_prob, "inject_minus_reverse": float(inject_prob - reverse_prob), "inject_minus_baseline": float(inject_prob - baseline_prob), "reverse_minus_baseline": float(reverse_prob - baseline_prob), "wrong_time_minus_baseline": float(wrong_time_prob - baseline_prob), } ) best = max(alpha_rows, key=lambda item: item["inject_minus_reverse"]) success = bool( best["inject_prob"] > best["baseline_prob"] and best["reverse_prob"] < best["baseline_prob"] and best["inject_minus_reverse"] > 0.05 ) per_row.append( { "donor_prompt_index": int(row["donor_prompt_index"]), "recipient_prompt_index": recipient_idx, "decision_pos": decision_pos, "intervention_pos": intervention_pos, "wrong_time_pos": wrong_time_pos, "context_text": context_text, "baseline_prob": baseline_prob, "prior_target_minus_reverse_prob": float(row["actual_token_prob_shift_target"] - row["actual_token_prob_shift_reverse"]), "prior_signed_direction_coeff_shift": float(row["signed_direction_coeff_shift_target"]), "alphas": alpha_rows, "best_alpha": float(best["alpha"]), "best_inject_prob": float(best["inject_prob"]), "best_reverse_prob": float(best["reverse_prob"]), "best_wrong_time_prob": float(best["wrong_time_prob"]), "best_inject_minus_reverse": float(best["inject_minus_reverse"]), "success": success, } ) print( "[factor-trace-dot-direction] " f"row {row_idx}/{len(selected_rows)} donor={row['donor_prompt_index']}→recipient={recipient_idx}@{decision_pos} " f"base={baseline_prob:.3f} best_alpha={best['alpha']:.1f} " f"inject={best['inject_prob']:.3f} reverse={best['reverse_prob']:.3f} " f"wrong_time={best['wrong_time_prob']:.3f} delta={best['inject_minus_reverse']:+.4f} " f"success={success}", flush=True, ) success_rows = sum(int(r["success"]) for r in per_row) result["status"] = "completed" result["rows"] = per_row result["summary"] = { "success_rows": int(success_rows), "n_rows": int(len(per_row)), "success_gate_pass": bool(success_rows >= 3), "max_inject_minus_reverse": float(max((r["best_inject_minus_reverse"] for r in per_row), default=0.0)), "mean_best_inject_minus_reverse": float(np.mean([r["best_inject_minus_reverse"] for r in per_row])) if per_row else 0.0, "total_time_s": float(time.perf_counter() - t0), } return result