""" visualize_probe_features.py — Multi-layer probe feature visualizer for LLaVA. Uses a single forward pass per batch to cache residual-stream activations at ALL target layers simultaneously (shared SAE d_sae), then runs the SAE per-layer on the cached activations. This is ~Nx faster than the old per-layer approach (where N = number of layers). For each batch: 1. (Optional) Generate captions with the model (once, shared across layers). 2. One forward pass caching activations at every target layer's hook point. 3. For each layer: SAE → sparse features → update top-k heaps. Generates ONE self-contained interactive HTML: Layers ▶ Top Features (id + probe weight) ▶ Top Image Patches / Text Tokens Data modes (--data_mode): toilet — only images from HF "pbcong/bathroom-toilet" that match --object_mode (toilet | bathroom). Supports --caption_mode. Uses visualize_multilayer_sae_features.create_dataloader. cc3m — full CC3M or COCO dataset via HF or a plain image folder. Uses visualize_multilayer_sae_features.create_dataloader. Usage ----- # Toilet subset python training/visualize_probe_features.py \\ --data_mode toilet \\ --sae_ckpt training/multilayer_sae_ckpt/last.ckpt \\ --probe_dir training/multilayer_sae_ckpt \\ --image_folder /path/to/cc3m_images/train \\ --object_mode toilet \\ --caption_mode generated \\ --layers 0 1 2 3 4 5 6 \\ --output_dir outputs/probe_features_toilet \\ --device_id 0 # Full CC3M (HF dataset) python training/visualize_probe_features.py \\ --data_mode cc3m \\ --sae_ckpt training/multilayer_sae_ckpt/last.ckpt \\ --probe_dir training/multilayer_sae_ckpt \\ --hf_dataset pixparse/cc3m-wds \\ --local_val_path /path/to/cc3m_images/train \\ --split train \\ --layers 0 1 2 3 4 5 6 \\ --output_dir outputs/probe_features_cc3m \\ --device_id 0 # Plain image folder python training/visualize_probe_features.py \\ --data_mode cc3m \\ --data_dir /path/to/images \\ --sae_ckpt training/multilayer_sae_ckpt/last.ckpt \\ --probe_dir training/multilayer_sae_ckpt \\ --layers 0 1 2 3 4 5 6 \\ --output_dir outputs/probe_features_folder \\ --device_id 0 """ import sys import os import json import argparse import pickle import shutil from pathlib import Path from typing import Dict, List, Optional, Tuple import torch import torch.distributed as dist from tqdm import tqdm sys.path.insert(0, str(Path(__file__).parent.parent)) # ── Lazy-import both source modules via importlib ───────────────────────────── import importlib.util as _ilu def _load_module(name: str, path: Path): spec = _ilu.spec_from_file_location(name, path) mod = _ilu.module_from_spec(spec) sys.modules[name] = mod # register before exec so pickle can resolve the module spec.loader.exec_module(mod) return mod _scripts_dir = Path(__file__).parent _vtf = _load_module("visualize_multilayer_sae_features", _scripts_dir / "visualize_multilayer_sae_features.py") # Shared helpers from the multilayer SAE feature visualizer setup_distributed = _vtf.setup_distributed cleanup_distributed = _vtf.cleanup_distributed TopKHeap = _vtf.TopKHeap build_image_html = _vtf.build_image_html build_text_html = _vtf.build_text_html create_dataloader = _vtf.create_dataloader ImageRecord = _vtf.ImageRecord TextRecord = _vtf.TextRecord find_image_token_positions = _vtf.find_image_token_positions expanded_to_original = _vtf.expanded_to_original original_to_expanded = _vtf.original_to_expanded N_IMAGE_PATCHES = _vtf.N_IMAGE_PATCHES IMAGE_TOKEN_ID = _vtf.IMAGE_TOKEN_ID from hallucination.extra_materials.mechanistic_interp.probe.probing import LinearProbe from sae.SAE_Tools import load_sae_model, get_sae_activations, get_loader from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration from sae.Training_Utils import str_to_torch_dtype from transformers import LlavaProcessor # ───────────────────────────────────────────────────────────────────────────── # Probe loading # ───────────────────────────────────────────────────────────────────────────── def load_probe_top_features( probe_dir: str, layer: int, top_k: int, input_dim: int = 65536, ) -> Tuple[List[int], List[float]]: """Return (feature_ids, probe_weights) for the top-k probe features.""" ckp_path = ( Path(probe_dir) / f"probe_model.language_model.layers.{layer}.hook_resid_post.pt" ) if not ckp_path.exists(): raise FileNotFoundError(f"Probe checkpoint not found: {ckp_path}") ckpt = torch.load(ckp_path, map_location="cpu") probe = LinearProbe(input_dim=input_dim, num_outputs=1) probe.load_state_dict(ckpt) probe.eval() weights = probe.weights.squeeze() # (input_dim,) top = torch.topk(weights, k=top_k) return top.indices.cpu().tolist(), weights[top.indices].cpu().tolist() # ───────────────────────────────────────────────────────────────────────────── # Multi-layer single pass (1 model fwd per batch instead of 1 per layer) # ───────────────────────────────────────────────────────────────────────────── @torch.no_grad() def multi_layer_pass( model, sae, processor, dataloader, device, layer_probe_info: Dict[int, Tuple[List[int], List[float]]], args, ) -> Dict[int, dict]: """ Single pass over the dataloader, hooking ALL target layers at once. For each batch: 1. (Optional) Generate caption with original model. 2. One forward pass caching residual-stream activations at every target layer. 3. For each layer: SAE → sparse features → update top-k heaps. Returns {layer: {feature_ids, probe_weights, image_heaps, text_heaps, hook_counts, hook_point}}. """ sae_dtype = next(sae.parameters()).dtype sae.eval() sae.to(device) # Build hook-point mapping and per-layer data structures hook_points: Dict[int, str] = {} target_hooks: set = set() layer_data: Dict[int, dict] = {} for layer, (feat_ids, feat_wts) in layer_probe_info.items(): hp = f"model.language_model.layers.{layer}.hook_resid_post" hook_points[layer] = hp target_hooks.add(hp) layer_data[layer] = { "image_heaps": {fi: TopKHeap(args.top_images) for fi in feat_ids}, "text_heaps": {fi: TopKHeap(args.top_texts) for fi in feat_ids}, "hook_counts": {fi: 0 for fi in feat_ids}, } def _make_hook_fn(cache: dict): def hook_fn(act, hook): if hook.name in target_hooks: cache[hook.name] = act.detach().cpu() return hook_fn n_batches = ( len(dataloader) if args.max_batches is None else min(args.max_batches, len(dataloader)) ) for batch_idx, batch in enumerate(tqdm(dataloader, total=n_batches, desc="Scanning")): if args.max_batches is not None and batch_idx >= args.max_batches: break B = batch["input_ids"].shape[0] model_inputs = { k: batch[k].to(device) for k in ("input_ids", "attention_mask", "pixel_values") } global_idxs = batch["global_idxs"] # ── Generated caption mode (once per batch, shared across layers) ──── if args.caption_mode == "generated": gen_ids = model.generate( **model_inputs, do_sample=False, num_beams=1, use_cache=True, max_new_tokens=args.max_new_tokens, ) full_texts = processor.batch_decode(gen_ids, skip_special_tokens=True) captions = [ txt.split("ASSISTANT:")[-1].strip() if "ASSISTANT:" in txt else txt.strip() for txt in full_texts ] images = [ dataloader.dataset.dataset[gi]["image"] for gi in global_idxs ] forced_texts = [ f"USER: \nDescribe this image. \nASSISTANT: {cap}" for cap in captions ] re_processed = processor( images=images, text=forced_texts, return_tensors="pt", padding=True, ) model_inputs = { "input_ids": re_processed["input_ids"].to(device), "attention_mask": re_processed["attention_mask"].to(device), "pixel_values": re_processed["pixel_values"].to(device), } input_ids_cpu = re_processed["input_ids"] B = input_ids_cpu.shape[0] else: input_ids_cpu = batch["input_ids"] # ── Single forward pass, cache ALL layer activations ───────────────── cache: dict = {} model.run_with_hooks( model_inputs, fwd_hooks=[(lambda n: n in target_hooks, _make_hook_fn(cache))], ) img_positions = find_image_token_positions(input_ids_cpu) # ── Process each layer from the cache ──────────────────────────────── for layer, (feat_ids, _) in layer_probe_info.items(): hp = hook_points[layer] if hp not in cache: continue acts = cache[hp] # (B, T, D) _, T, D = acts.shape flat = acts.reshape(B * T, 1, D).to(device, dtype=sae_dtype) loader = get_loader(flat, batch_size=min(args.sae_batch, B * T)) idxs, vals = get_sae_activations(sae, loader, device, no_tqdm=True) idxs = idxs.squeeze(1).reshape(B, T, -1) vals = vals.squeeze(1).reshape(B, T, -1) ld = layer_data[layer] image_heaps = ld["image_heaps"] text_heaps = ld["text_heaps"] hook_counts = ld["hook_counts"] for fi in feat_ids: mask = (idxs == fi) feat_acts = (vals * mask.float()).sum(dim=-1) # (B, T) above = feat_acts > args.threshold hook_counts[fi] += above.sum().item() for b in range(B): ip = img_positions[b].item() global_idx = global_idxs[b] acts_b = feat_acts[b] # (T,) # — Image patches — if ip >= 0: img_acts = acts_b[ip : ip + N_IMAGE_PATCHES] n_img = (img_acts > args.threshold).sum().item() if n_img > 0: k_img = min(args.top_images, n_img) topk_v, topk_i = img_acts.topk(k_img) for v, pidx in zip(topk_v, topk_i): if v.item() < args.threshold: break image_heaps[fi].push( v.item(), ImageRecord(v.item(), global_idx, pidx.item(), hp), ) # — Text tokens — text_ranges: List[Tuple[int, int]] = [] if ip >= 0: if ip > 0: text_ranges.append((0, ip)) if ip + N_IMAGE_PATCHES < T: text_ranges.append((ip + N_IMAGE_PATCHES, T)) else: text_ranges.append((0, T)) for rng_start, rng_end in text_ranges: txt_acts = acts_b[rng_start:rng_end] n_txt = (txt_acts > args.threshold).sum().item() if n_txt == 0: continue k_txt = min(args.top_texts, n_txt) topk_v, topk_i = txt_acts.topk(k_txt) for v, rel_i in zip(topk_v, topk_i): if v.item() < args.threshold: break exp_pos = rng_start + rel_i.item() orig_pos = expanded_to_original(exp_pos, ip) if orig_pos < 0: continue if (orig_pos < input_ids_cpu.shape[1] and input_ids_cpu[b, orig_pos].item() == IMAGE_TOKEN_ID): continue seq_len = input_ids_cpu.shape[1] ctx_s = max(0, orig_pos - args.buffer) ctx_e = min(seq_len, orig_pos + args.buffer + 1) if ip >= 0: if orig_pos < ip: ctx_e = min(ctx_e, ip) elif orig_pos > ip: ctx_s = max(ctx_s, ip + 1) ctx_ids = input_ids_cpu[b, ctx_s:ctx_e].tolist() ctx_acts_list: List[float] = [] for op in range(ctx_s, ctx_e): ep = original_to_expanded(op, ip) if ep < 0 or ep >= T: ctx_acts_list.append(0.0) else: ctx_acts_list.append(feat_acts[b, ep].item()) text_heaps[fi].push( v.item(), TextRecord( v.item(), global_idx, hp, ctx_ids, ctx_acts_list, orig_pos - ctx_s, ), ) del flat, idxs, vals del cache # Build return dict results: Dict[int, dict] = {} for layer, (feat_ids, feat_wts) in layer_probe_info.items(): ld = layer_data[layer] results[layer] = { "feature_ids": feat_ids, "probe_weights": feat_wts, "image_heaps": ld["image_heaps"], "text_heaps": ld["text_heaps"], "hook_counts": ld["hook_counts"], "hook_point": hook_points[layer], } return results # ───────────────────────────────────────────────────────────────────────────── # Interactive HTML builder # ───────────────────────────────────────────────────────────────────────────── _CSS = """ """ _JS = """ """ def _feature_html( layer_idx: int, rank: int, feat_id: int, probe_weight: float, image_records, text_records, total_act: int, dataset, processor, ) -> str: body_id = f"feat_L{layer_idx}_F{feat_id}" img_html = build_image_html(image_records, dataset) txt_html = build_text_html(text_records, processor) # Replace raw

from helpers with styled section titles img_html = img_html.replace( "

Top Activating Image Patches

", '
Top Activating Image Patches
', 1, ) txt_html = txt_html.replace( "

Top Activating Text Tokens

", '
Top Activating Text Tokens
', 1, ) return f"""
► #{rank} Feature {feat_id} probe w = {probe_weight:+.4f}  activations: {total_act}
{img_html} {txt_html}
""" def build_interactive_html( layer_results: Dict[int, dict], dataset, processor, probe_dir: str, data_mode: str, ) -> str: layer_blocks: List[str] = [] for layer_idx in sorted(layer_results.keys()): res = layer_results[layer_idx] feat_ids = res["feature_ids"] probe_wts = res["probe_weights"] image_heaps = res["image_heaps"] text_heaps = res["text_heaps"] hook_counts = res["hook_counts"] hook_point = res["hook_point"] layer_body_id = f"layer_{layer_idx}_body" feature_blocks: List[str] = [] for rank, (fi, pw) in enumerate(zip(feat_ids, probe_wts), start=1): img_recs = image_heaps[fi].sorted_records() if fi in image_heaps else [] txt_recs = text_heaps[fi].sorted_records() if fi in text_heaps else [] total = hook_counts.get(fi, 0) feature_blocks.append( _feature_html(layer_idx, rank, fi, pw, img_recs, txt_recs, total, dataset, processor) ) n_feats = len(feat_ids) total_acts = sum(hook_counts.get(fi, 0) for fi in feat_ids) layer_blocks.append(f"""
► Layer {layer_idx} {hook_point}  |  {n_feats} features  |  {total_acts} total activations
{"".join(feature_blocks)}
""") return f""" Probe Feature Visualizer — Multi-Layer {_CSS}

Probe Feature Visualizer — Multi-Layer SAE Activations

Probe dir: {probe_dir}  |  data mode: {data_mode}
Click a Layer to expand its top probe features. Click a Feature to see top-activating image patches and text tokens.

{"".join(layer_blocks)}
{_JS} """ # ───────────────────────────────────────────────────────────────────────────── # Main # ───────────────────────────────────────────────────────────────────────────── def main(): ap = argparse.ArgumentParser( description="Multi-layer probe feature visualizer for LLaVA.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) # ── Data mode ──────────────────────────────────────────────────────────── ap.add_argument( "--data_mode", required=True, choices=["toilet", "cc3m", "coco", "folder"], help=( "toilet: pbcong/bathroom-toilet positives + CC3M negatives (needs --image_folder). " "cc3m: full CC3M via --hf_dataset + --local_val_path. " "coco: COCO via --hf_dataset + --local_val_path. " "folder: plain image folder via --data_dir." ), ) # ── Model / SAE ────────────────────────────────────────────────────────── ap.add_argument("--sae_ckpt", required=True) ap.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf") ap.add_argument("--device_id", type=int, default=0) ap.add_argument("--dtype", default="float16") # ── Probe ──────────────────────────────────────────────────────────────── ap.add_argument("--probe_dir", required=True, help="Dir containing probe_model.language_model.layers.*.pt files.") ap.add_argument("--probe_input_dim", type=int, default=65536, help="SAE width (LinearProbe input_dim).") ap.add_argument("--layers", type=int, nargs="+", default=list(range(7)), help="LLM layers to process.") ap.add_argument("--top_probe_k", type=int, default=10, help="Top-k probe-weight features to visualize per layer.") # ── Data — toilet ──────────────────────────────────────────────────────── ap.add_argument("--image_folder", default=None, help="[toilet] Local CC3M image folder (positives + negative pool).") ap.add_argument("--object_mode", default="toilet", choices=["toilet", "bathroom", "both"], help="[toilet] Positive class: toilet==1, bathroom==1, or both (toilet==1 OR bathroom==1).") ap.add_argument("--num_negatives", type=int, default=10000, help="[toilet] Number of random CC3M negatives to include.") # ── Caption mode ───────────────────────────────────────────────────────── ap.add_argument( "--caption_mode", default="generated", choices=["generated", "caption"], help=( "generated: model generates caption first; teacher-forced pass collects acts. " "caption: use stored caption (CC3M txt / COCO sentences / " "pbcong/bathroom-toilet caption field)." ), ) ap.add_argument("--max_new_tokens", type=int, default=128, help="Max tokens to generate per image (only with --caption_mode generated).") # ── Data — CC3M / COCO ─────────────────────────────────────────────────── ap.add_argument("--hf_dataset", default=None, help="HF dataset path (CC3M / COCO / also for toilet-negative captions).") ap.add_argument("--local_val_path", default=None, help="[cc3m/coco] Local image root for the HF dataset split.") ap.add_argument("--data_dir", default=None, help="[folder] Plain image folder.") ap.add_argument("--split", default="train", help="HF dataset split.") # ── Common data ────────────────────────────────────────────────────────── ap.add_argument("--num_workers", type=int, default=4) # ── Processing ─────────────────────────────────────────────────────────── ap.add_argument("--batch_size", type=int, default=4) ap.add_argument("--sae_batch", type=int, default=4096) ap.add_argument("--threshold", type=float, default=1e-3) ap.add_argument("--max_batches", type=int, default=None) # ── Visualisation ──────────────────────────────────────────────────────── ap.add_argument("--output_dir", default="outputs/probe_features") ap.add_argument("--top_images", type=int, default=10) ap.add_argument("--top_texts", type=int, default=10) ap.add_argument("--buffer", type=int, default=10) args = ap.parse_args() # ── Validate data-mode args ─────────────────────────────────────────────── if args.data_mode == "toilet" and not args.image_folder: ap.error("--data_mode toilet requires --image_folder.") if args.data_mode in ("cc3m", "coco") and not args.hf_dataset: ap.error(f"--data_mode {args.data_mode} requires --hf_dataset (+ --local_val_path).") if args.data_mode == "folder" and not args.data_dir: ap.error("--data_mode folder requires --data_dir.") # ── Distributed setup ──────────────────────────────────────────────────── rank, world_size, local_rank = setup_distributed() is_distributed = world_size > 1 device = torch.device( f"cuda:{local_rank}" if is_distributed else f"cuda:{args.device_id}" if torch.cuda.is_available() else "cpu" ) dtype = str_to_torch_dtype(args.dtype) if rank == 0: print(f"Device: {device} | world_size: {world_size} | data_mode: {args.data_mode}") # ── Load SAE once ──────────────────────────────────────────────────────── sae = load_sae_model(args.sae_ckpt, model_type="llava", hook_type="text", device=device) sae.eval() # ── Load LLaVA model + processor once ──────────────────────────────────── model = HookedSAELlavaConditionalGeneration.from_pretrained(args.model_name) model.to(device, dtype=dtype) model.eval() processor = LlavaProcessor.from_pretrained(args.model_name) # ── Build probe feature dict ────────────────────────────────────────────── probe_info: Dict[int, Tuple[List[int], List[float]]] = {} for layer in args.layers: feat_ids, feat_wts = load_probe_top_features( args.probe_dir, layer, args.top_probe_k, args.probe_input_dim, ) probe_info[layer] = (feat_ids, feat_wts) if rank == 0: print(f" Layer {layer:2d} top features: {feat_ids}") # ── Build dataloader ──────────────────────────────────────────────────────── dataset, dataloader = create_dataloader(args, processor, rank, world_size) if rank == 0: print(f"Dataset: {len(dataset)} samples, {len(dataloader)} batches/rank.") # ── Single multi-layer forward pass ──────────────────────────────────────── raw_results = multi_layer_pass( model, sae, processor, dataloader, device, probe_info, args, ) # ── DDP gather (all layers) ────────────────────────────────────────────── layer_results: Dict[int, dict] = {} for layer in args.layers: res = raw_results[layer] feat_ids = res["feature_ids"] image_heaps = res["image_heaps"] text_heaps = res["text_heaps"] hook_counts = res["hook_counts"] if is_distributed: output_dir = Path(args.output_dir) tmp_dir = output_dir / f".ddp_tmp_L{layer}" if rank == 0: tmp_dir.mkdir(parents=True, exist_ok=True) dist.barrier() with open(tmp_dir / f"rank_{rank}.pkl", "wb") as f: pickle.dump({ "image_heaps": {fi: image_heaps[fi].sorted_records() for fi in feat_ids}, "text_heaps": {fi: text_heaps[fi].sorted_records() for fi in feat_ids}, "hook_counts": hook_counts, }, f) dist.barrier() if rank == 0: merged_img: Dict[int, TopKHeap] = {fi: TopKHeap(args.top_images) for fi in feat_ids} merged_txt: Dict[int, TopKHeap] = {fi: TopKHeap(args.top_texts) for fi in feat_ids} merged_cnt: Dict[int, int] = {fi: 0 for fi in feat_ids} for r in range(world_size): with open(tmp_dir / f"rank_{r}.pkl", "rb") as f: data = pickle.load(f) for fi in feat_ids: for rec in data["image_heaps"][fi]: merged_img[fi].push(rec.value, rec) for rec in data["text_heaps"][fi]: merged_txt[fi].push(rec.value, rec) merged_cnt[fi] += data["hook_counts"][fi] image_heaps = merged_img text_heaps = merged_txt hook_counts = merged_cnt shutil.rmtree(tmp_dir) if rank == 0: layer_results[layer] = { "feature_ids": feat_ids, "probe_weights": res["probe_weights"], "image_heaps": image_heaps, "text_heaps": text_heaps, "hook_counts": hook_counts, "hook_point": res["hook_point"], } for fi in feat_ids: print(f" feature {fi}: {hook_counts.get(fi, 0)} activations") # ── Generate HTML (rank 0 only) ────────────────────────────────────────── if rank == 0: output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) print("\nBuilding interactive HTML …") html = build_interactive_html( layer_results, dataset, processor, args.probe_dir, args.data_mode, ) out_path = output_dir / "probe_features.html" out_path.write_text(html, encoding="utf-8") print(f"Done. Open: {out_path}") summary = {} for layer, res in layer_results.items(): summary[f"layer_{layer}"] = { str(fi): { "probe_weight": pw, "total_activations": res["hook_counts"].get(fi, 0), } for fi, pw in zip(res["feature_ids"], res["probe_weights"]) } with open(output_dir / "summary.json", "w") as f: json.dump(summary, f, indent=2) cleanup_distributed() if __name__ == "__main__": main()