From fd203318fa7e7ff74c94a87a196cfef3a080935a Mon Sep 17 00:00:00 2001 From: rick Date: Fri, 5 Jun 2026 20:42:30 -0600 Subject: [PATCH] add SAE feature coactivation analysis --- experiments/coactivation.py | 275 ++++++++++++++++++++++++++++++++++++ 1 file changed, 275 insertions(+) create mode 100644 experiments/coactivation.py diff --git a/experiments/coactivation.py b/experiments/coactivation.py new file mode 100644 index 0000000..0734417 --- /dev/null +++ b/experiments/coactivation.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python3 +""" +SAE feature coactivation analysis for VLA models. + +Finds which SAE features consistently fire together. If concept_id already +told you *which* features matter for a concept, this tells you *what they +travel with* -- revealing feature circuits inside the VLA. + +Works on the sparse codes produced by encoding activations through a trained +SAE. For each concept feature, computes Jaccard similarity against every other +feature across all task episodes, then ranks the top co-firing partners. + +Run concept_id.py first to get concept_results, then point this at the same +SAE dir and activations. + +Examples: + python experiments/coactivation.py \ + --sae-dir outputs/xvla_saes/libero_object \ + --activations-dir outputs/xvla_experiments/baseline_libero_object/activations \ + --concept-results results/concept_id/libero_object/all_layers.json \ + --suite libero_object + + python experiments/coactivation.py \ + --sae-dir outputs/groot_saes/libero_goal \ + --activations-dir outputs/groot_experiments/baseline_libero_goal/activations \ + --concept-results results/concept_id/libero_goal/all_layers.json \ + --suite libero_goal --layers eagle_L04 eagle_L08 --top-k 15 +""" + +import gc +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional + +import numpy as np +import torch +import tyro + +import sys +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from experiments.sae_hooks import TopKSAE +from experiments.concept_identification import get_concept_task_mapping + + +@dataclass +class CoactivationConfig: + + sae_dir: str = "" + activations_dir: str = "" + + concept_results: str = "" + # Path to concept_id all_layers.json -- used to pull the concept features + # we actually care about so we're not computing jaccard over 8k x 8k + + suite: str = "libero_object" + layers: Optional[List[str]] = None + output_dir: Optional[str] = None + + top_k: int = 10 + # top co-firing partners to report per concept feature + + top_concept_features: int = 5 + # how many features per concept to trace (ranked by concept_id score) + + max_samples: int = 200_000 + # cap total samples to keep memory sane + + +def load_sae_checkpoint(sae_path: Path, device: str): + data = torch.load(sae_path, map_location="cpu", weights_only=True) + cfg = data["config"] + sae = TopKSAE(cfg["input_dim"], cfg["hidden_dim"], k=cfg.get("k", 64)) + sae.load_state_dict(data["sae_state_dict"]) + sae.eval().to(device) + act_mean = data.get("activation_mean", data.get("act_mean", torch.zeros(cfg["input_dim"]))) + act_std = data.get("activation_std", data.get("act_std", torch.ones(cfg["input_dim"]))) + return sae, act_mean.to(device), act_std.to(device) + + +def load_all_activations(act_dir: Path, layer_name: str, max_samples: int) -> Optional[torch.Tensor]: + # grab everything, don't care about task splits here -- we just want the + # full distribution of co-firing patterns across the whole suite + all_acts = [] + total = 0 + + for task_dir in sorted(act_dir.glob("task*")): + if not task_dir.is_dir(): + continue + for ep_dir in sorted(task_dir.iterdir()): + if not ep_dir.is_dir(): + continue + pt = ep_dir / f"{layer_name}.pt" + if not pt.exists(): + continue + try: + t = torch.load(pt, map_location="cpu", weights_only=True).float() + if t.dim() > 2: + t = t.reshape(-1, t.shape[-1]) + all_acts.append(t) + total += t.shape[0] + except Exception: + continue + if total >= max_samples: + break + + # flat file fallback (pre-concatenated layout) + if not all_acts: + for pt in sorted(act_dir.glob(f"{layer_name}_task*.pt")): + try: + t = torch.load(pt, map_location="cpu", weights_only=True).float() + if t.dim() > 2: + t = t.reshape(-1, t.shape[-1]) + all_acts.append(t) + except Exception: + continue + + if not all_acts: + return None + + combined = torch.cat(all_acts, dim=0) + if combined.shape[0] > max_samples: + idx = torch.randperm(combined.shape[0])[:max_samples] + combined = combined[idx] + return combined + + +def jaccard_for_feature(binary_codes: np.ndarray, feature_idx: int) -> np.ndarray: + # binary_codes is [N, hidden_dim] bool array + # returns jaccard similarity of feature_idx against every other feature + a = binary_codes[:, feature_idx] # [N] + # intersection: both active + inter = (binary_codes & a[:, None]).sum(axis=0).astype(np.float32) + # union: either active + union = (binary_codes | a[:, None]).sum(axis=0).astype(np.float32) + jac = np.where(union > 0, inter / union, 0.0) + jac[feature_idx] = 0.0 # don't report self + return jac + + +def pull_concept_features(concept_results: dict, layer_name: str, top_n: int) -> Dict[str, List[int]]: + # returns {concept_name: [feature_idx, ...]} for the layer + layer_data = concept_results.get(layer_name, {}) + out = {} + for cat, concepts in layer_data.items(): + for concept_name, info in concepts.items(): + feats = info.get("top_features", [])[:top_n] + if feats: + key = f"{cat}/{concept_name}" + out[key] = [f["feature_idx"] for f in feats] + return out + + +def analyze_layer(sae, act_mean, act_std, acts: torch.Tensor, + concept_features: Dict[str, List[int]], top_k: int, device: str) -> dict: + + # encode activations -> sparse codes + print(f" encoding {acts.shape[0]:,} samples through SAE...") + batch_size = 4096 + codes_list = [] + with torch.no_grad(): + for i in range(0, acts.shape[0], batch_size): + batch = acts[i:i+batch_size].to(device) + batch_norm = (batch - act_mean) / (act_std + 1e-8) + z = sae.encode(batch_norm) + codes_list.append(z.cpu()) + codes = torch.cat(codes_list, dim=0) # [N, hidden_dim] + + # binarize -- a feature is "active" if it's nonzero (TopK so it's clean) + binary = (codes.abs() > 0).numpy() + + results = {} + for concept_key, feat_indices in concept_features.items(): + concept_result = {} + for feat_idx in feat_indices: + jac = jaccard_for_feature(binary, feat_idx) + top_idx = np.argsort(jac)[::-1][:top_k] + partners = [] + for partner_idx in top_idx: + if jac[partner_idx] < 1e-4: + break + partners.append({ + "feature_idx": int(partner_idx), + "jaccard": float(jac[partner_idx]), + "co_fire_rate": float(binary[:, partner_idx].mean()), + }) + concept_result[str(feat_idx)] = { + "fire_rate": float(binary[:, feat_idx].mean()), + "top_coactive": partners, + } + results[concept_key] = concept_result + + return results + + +def main(cfg: CoactivationConfig): + device = "cuda" if torch.cuda.is_available() else "cpu" + sae_dir = Path(cfg.sae_dir) + act_dir = Path(cfg.activations_dir) + + if not cfg.concept_results: + raise ValueError("--concept-results is required (point at concept_id all_layers.json)") + + with open(cfg.concept_results) as f: + concept_results = json.load(f) + + output_dir = Path(cfg.output_dir) if cfg.output_dir else Path(f"results/coactivation/{cfg.suite}") + output_dir.mkdir(parents=True, exist_ok=True) + + if cfg.layers: + layer_names = cfg.layers + else: + layer_names = sorted([d.name for d in sae_dir.iterdir() + if d.is_dir() and (d / "sae_best.pt").exists()]) + + print(f"Coactivation analysis: {cfg.suite}") + print(f" {len(layer_names)} layers, top_k={cfg.top_k}, concept_features={cfg.top_concept_features}") + + all_results = {} + + for layer_name in layer_names: + print(f"\n--- {layer_name} ---") + + sae_path = sae_dir / layer_name / "sae_best.pt" + if not sae_path.exists(): + print(" [skip] no SAE") + continue + + concept_feats = pull_concept_features(concept_results, layer_name, cfg.top_concept_features) + if not concept_feats: + print(" [skip] no concept features for this layer in concept_results") + continue + + n_feats_total = sum(len(v) for v in concept_feats.values()) + print(f" tracing {n_feats_total} concept features across {len(concept_feats)} concepts") + + acts = load_all_activations(act_dir, layer_name, cfg.max_samples) + if acts is None: + print(" [skip] no activations found") + continue + print(f" loaded {acts.shape[0]:,} samples") + + sae, act_mean, act_std = load_sae_checkpoint(sae_path, device) + + layer_out = analyze_layer(sae, act_mean, act_std, acts, concept_feats, + cfg.top_k, device) + all_results[layer_name] = layer_out + + # quick summary print + for concept_key, feat_data in layer_out.items(): + for feat_idx, info in feat_data.items(): + if info["top_coactive"]: + top = info["top_coactive"][0] + print(f" {concept_key} feat {feat_idx}: " + f"top partner={top['feature_idx']} jaccard={top['jaccard']:.3f}") + + out_path = output_dir / f"{layer_name}.json" + with open(out_path, "w") as f: + json.dump(layer_out, f, indent=2) + print(f" saved: {out_path}") + + del sae, acts + gc.collect() + if device == "cuda": + torch.cuda.empty_cache() + + with open(output_dir / "all_layers.json", "w") as f: + json.dump(all_results, f, indent=2) + print(f"\ndone. results -> {output_dir}") + + +if __name__ == "__main__": + cfg = tyro.cli(CoactivationConfig) + main(cfg)