feature_analysis.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""
Ingested from inspection-backend/feature_analysis.py.
Import adaptation only: model_config -> aquin.compute.model_loader, sae -> aquin.compute.sae.
"""
from __future__ import annotations
 
import hashlib
 
import torch
from transformer_lens import HookedTransformer
 
from aquin.compute.sae import SparseAutoencoder
from aquin.compute.model_loader import (
    get_config,
    get_config as _get_config,
    resolve_model_id,
)
from aquin.compute.device import resolve_compute_device, synchronize_device
 
TOP_K_FEATURES = 10
DEVICE = resolve_compute_device()
 
_sae_cache: dict[tuple, SparseAutoencoder] = {}
_norm_cache: dict[tuple, dict | None] = {}
_session_label_cache = {}
 
_sae = None
_norm = None
 
_kernel_feature_acts = None
_kernel_resid = None
_kernel_top_features = []
 
 
def _get_sae_path_for_layer(model_id: str, layer: int | None = None):
    from aquin.compute.model_loader import resolve_sae_path
 
    return resolve_sae_path(model_id, layer)
 
 
def _get_norm_path_for_layer(model_id: str, layer: int | None = None):
    from pathlib import Path
 
    from aquin.compute.model_loader import resolve_model_id, resolve_sae_checkpoint_path
 
    cfg = get_config(model_id)
    short = resolve_model_id(model_id)
    resolved_layer = int(layer if layer is not None else cfg["sae_layer"])
 
    sae_path = resolve_sae_checkpoint_path(short, resolved_layer)
    if sae_path is not None:
        parent = sae_path.parent
        for candidate in (
            parent / f"norm_layer{resolved_layer}.pt",
            parent / "norm.pt",
            parent / f"_acts_layer{resolved_layer}" / "norm.pt",
        ):
            if candidate.is_file():
                return candidate
 
    base = Path.home() / ".aquin" / "sae" / short
    if "norm_layers" in cfg:
        rel = cfg["norm_layers"].get(resolved_layer)
        if rel:
            p = base / Path(rel).name
            if p.exists():
                return p
    p_rel = cfg.get("norm_path")
    if p_rel:
        p = base / Path(p_rel).name
        return p if p.exists() else None
    return None
 
 
def _load_sae_native(model_id: str, layer: int | None = None) -> SparseAutoencoder:
    sae_path = _get_sae_path_for_layer(model_id, layer)
    short = resolve_model_id(model_id)
    cfg = get_config(short)
    resolved_layer = int(layer if layer is not None else cfg["sae_layer"])
    if not sae_path.exists():
        raise FileNotFoundError(
            f"SAE not found at {sae_path}. Run: aquin load sae {short}-l{resolved_layer}"
        )
    from aquin.compute.model_loader import corrupt_sae_checkpoint_hint
 
    corrupt = corrupt_sae_checkpoint_hint(short, resolved_layer)
    if corrupt:
        raise ValueError(corrupt)
    return SparseAutoencoder.load(sae_path, device=DEVICE)
 
 
def load_sae(model_id: str = "llama-3.2-1b", layer: int | None = None) -> SparseAutoencoder:
    global _sae
    short = resolve_model_id(model_id)
    cfg = get_config(short)
    resolved_layer = int(layer if layer is not None else cfg["sae_layer"])
    key = (short, resolved_layer)
    if key not in _sae_cache:
        from aquin.compute.model_loader import load_sae_from_disk
 
        _sae_cache[key] = load_sae_from_disk(short, resolved_layer, device=DEVICE)
        if short == "llama-3.2-1b" and layer is None:
            _sae = _sae_cache[key]
    return _sae_cache[key]
 
 
def load_norm(model_id: str = "llama-3.2-1b", layer: int | None = None) -> dict | None:
    global _norm
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    key = (model_id, resolved_layer)
    if key not in _norm_cache:
        norm_path = _get_norm_path_for_layer(model_id, resolved_layer)
        if norm_path and norm_path.exists():
            from aquin.compute.torch_io import load_norm_stats
 
            loaded = load_norm_stats(norm_path, map_location=DEVICE)
            _norm_cache[key] = loaded
        else:
            _norm_cache[key] = None
        if model_id == "llama-3.2-1b" and layer is None:
            _norm = _norm_cache[key]
    return _norm_cache[key]
 
 
def normalize(x: torch.Tensor, model_id: str = "llama-3.2-1b", layer: int | None = None) -> torch.Tensor:
    n = load_norm(model_id, layer)
    if n is None:
        return x
    mean, std = n["mean"], n["std"]
    if mean.shape[-1] != x.shape[-1]:
        return x
    return (x - mean) / std
 
 
def _sae_feature_ablate_hook(
    value: torch.Tensor,
    pos: int,
    feature_idx: int,
    sae: SparseAutoencoder,
    model_id: str,
    sae_layer: int,
) -> torch.Tensor:
    """Zero one SAE feature at a residual position (same logic as label_feature_causally)."""
    resid_here = value[0, pos]
    feat_acts = sae.encode(normalize(resid_here.unsqueeze(0), model_id, sae_layer)).squeeze(0)
    feat_ablated = feat_acts.clone()
    feat_ablated[feature_idx] = 0.0
    delta = sae.decode(feat_ablated) - sae.decode(feat_acts)
    if delta.dim() > 1:
        delta = delta.squeeze(0)
    value[0, pos] = resid_here + delta
    return value
 
 
def label_feature_causally(fi: int, prompt: str, model: HookedTransformer, sae: SparseAutoencoder, client, top_k: int = 5, model_id: str = "llama-3.2-1b", layer: int | None = None) -> str:
    cfg = get_config(model_id)
    sae_layer = layer if layer is not None else cfg["sae_layer"]
 
    tokens = model.to_tokens(prompt[:512])
    if tokens.shape[1] < 4:
        return f"feature_{fi}"
 
    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=f"blocks.{sae_layer}.hook_resid_post",
            return_type=None,
        )
    resid = cache[f"blocks.{sae_layer}.hook_resid_post"][0]
    acts = sae.encode(normalize(resid, model_id, sae_layer))
 
    firing_pos = (acts[:, fi] > 0.5).nonzero(as_tuple=True)[0].tolist()
    if not firing_pos:
        top_pos = int(acts[:, fi].argmax().item())
        if acts[top_pos, fi].item() < 0.01:
            return f"feature_{fi}"
        firing_pos = [top_pos]
 
    causal_examples = []
 
    for pos in firing_pos[:3]:
        act_val = acts[pos, fi].item()
 
        with torch.no_grad():
            baseline_logits = model(tokens)[0, -1]
        baseline_probs = torch.softmax(baseline_logits, dim=-1)
 
        def ablate_feature(value, hook, pos=pos):
            return _sae_feature_ablate_hook(value, pos, fi, sae, model_id, sae_layer)
 
        with torch.no_grad():
            ablated_logits = model.run_with_hooks(
                tokens,
                fwd_hooks=[(f"blocks.{sae_layer}.hook_resid_post", ablate_feature)]
            )[0, -1]
        ablated_probs = torch.softmax(ablated_logits, dim=-1)
 
        delta = baseline_probs - ablated_probs
        topk_boosted = delta.topk(top_k)
        topk_suppressed = (-delta).topk(top_k)
 
        boosted = [
            (model.tokenizer.decode([idx.item()]).strip(), round(val.item(), 4))
            for idx, val in zip(topk_boosted.indices, topk_boosted.values)
            if val > 0.001
        ]
        suppressed = [
            (model.tokenizer.decode([idx.item()]).strip(), round(val.item(), 4))
            for idx, val in zip(topk_suppressed.indices, topk_suppressed.values)
            if val > 0.001
        ]
 
        context = model.tokenizer.decode(tokens[0, max(0, pos - 4):pos + 5].tolist())
        pivot = model.tokenizer.decode([tokens[0, pos].item()]).strip()
 
        causal_examples.append({
            "activation": round(act_val, 3),
            "context": context,
            "pivot": pivot,
            "boosts": boosted,
            "suppresses": suppressed,
        })
 
    if not causal_examples:
        return f"feature_{fi}"
 
    causal_examples.sort(key=lambda x: x["activation"], reverse=True)
 
    formatted = "\n".join(
        f'  Context: "...{ex["context"]}..." (token: "{ex["pivot"]}", activation: {ex["activation"]})\n'
        f'  Causally boosts: {ex["boosts"]}\n'
        f'  Causally suppresses: {ex["suppresses"]}'
        for ex in causal_examples
    )
 
    prompt_text = f"""A sparse autoencoder feature causally influences a language model's predictions as follows:
 
{formatted}
 
Based on what this feature causally promotes and suppresses in this context, give a concise 2-6 word label for its functional role. Good label examples: "promotes plural nouns", "suppresses hedging language", "boosts location names after prepositions".
 
Reply with ONLY the label."""
 
    if client is None:
        return f"feature_{fi}"
    try:
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt_text}],
            max_tokens=20,
            temperature=0.0,
        )
        return resp.choices[0].message.content.strip().strip('"')
    except Exception as e:
        print(f"[label] error on feature {fi}: {e}", flush=True)
        return f"feature_{fi}"
 
 
def get_causal_label(fi: int, prompt: str, model: HookedTransformer, sae: SparseAutoencoder, client, model_id: str = "llama-3.2-1b", layer: int | None = None) -> str:
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    prompt_hash = hashlib.md5(prompt.encode()).hexdigest()[:8]
    key = (fi, prompt_hash, model_id, resolved_layer)
    if key in _session_label_cache:
        return _session_label_cache[key]
 
    label = label_feature_causally(fi, prompt, model, sae, client, model_id=model_id, layer=resolved_layer)
 
    _session_label_cache[key] = label
    return label
 
 
def format_feature_ref(feature_idx: int, label: str | None = None) -> str:
    """Display form: index plus causal label (never index alone when label is known)."""
    if label:
        return f"{feature_idx} · {label}"
    return str(feature_idx)
 
 
def prompt_for_labeling(ctx: dict | None = None, args: dict | None = None) -> str:
    args = args or {}
    ctx = ctx or {}
    mem = ctx.get("state", {}).get("memory", {})
    return args.get("prompt") or mem.get("lastPrompt") or "Hello"
 
 
def _attach_label_to_feature_dict(
    f: dict,
    *,
    prompt: str,
    model: HookedTransformer,
    sae: SparseAutoencoder,
    client,
    model_id: str,
    layer: int | None,
    seen: set[int] | None = None,
) -> None:
    fi = int(f["feature_idx"])
    if seen is not None and fi in seen and f.get("label"):
        f["feature_ref"] = format_feature_ref(fi, f["label"])
        return
    if not f.get("label"):
        f["label"] = get_causal_label(fi, prompt, model, sae, client, model_id=model_id, layer=layer)
    f["feature_ref"] = format_feature_ref(fi, f["label"])
    if seen is not None:
        seen.add(fi)
 
 
def label_inspection_features(
    feat_result: dict,
    *,
    prompt: str,
    model: HookedTransformer,
    client,
    model_id: str = "llama-3.2-1b",
    layer: int | None = None,
) -> dict:
    """Attach causal labels to inspection feature lists (top + attribution)."""
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    sae = load_sae(model_id, resolved_layer)
    seen: set[int] = set()
 
    for f in feat_result.get("top_response_features", []):
        _attach_label_to_feature_dict(
            f, prompt=prompt, model=model, sae=sae, client=client,
            model_id=model_id, layer=resolved_layer, seen=seen,
        )
    for attr in feat_result.get("attribution", []):
        for f in attr.get("driven_by_features", []):
            _attach_label_to_feature_dict(
                f, prompt=prompt, model=model, sae=sae, client=client,
                model_id=model_id, layer=resolved_layer, seen=seen,
            )
    return feat_result
 
 
def enrich_feature_tool_result(
    result: dict,
    *,
    prompt: str,
    model: HookedTransformer,
    client,
    model_id: str = "llama-3.2-1b",
    layer: int | None = None,
    label_neighbors: bool = False,
) -> dict:
    """Add label + feature_ref to a feature-logits or feature-neighbors payload."""
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    sae = load_sae(model_id, resolved_layer)
 
    fi = result.get("feature_idx")
    if fi is not None:
        label = get_causal_label(
            int(fi), prompt, model, sae, client, model_id=model_id, layer=resolved_layer,
        )
        result["label"] = label
        result["feature_ref"] = format_feature_ref(int(fi), label)
 
    if label_neighbors:
        for n in result.get("neighbors", []):
            nfi = int(n["feature_idx"])
            nlabel = get_causal_label(
                nfi, prompt, model, sae, client, model_id=model_id, layer=resolved_layer,
            )
            n["label"] = nlabel
            n["feature_ref"] = format_feature_ref(nfi, nlabel)
 
    return result
 
 
def resolve_feature_label(
    feature_idx: int,
    *,
    ctx: dict,
    args: dict | None = None,
    layer: int | None = None,
) -> str:
    """Resolve a causal label using session context (for steer / UI tools)."""
    from aquin.compute.model_loader import (
        get_active_model_id,
        get_loaded_model,
        load_model,
        resolve_model_id,
    )
 
    args = args or {}
    model_id = (
        args.get("model_id")
        or ctx.get("state", {}).get("activeModelId")
        or get_active_model_id()
        or "llama-3.2-1b"
    )
    model_id = resolve_model_id(model_id)
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    prompt = prompt_for_labeling(ctx, args)
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    sae = load_sae(model_id, resolved_layer)
    from aquin.compute.openai_client import get_openai_client
    return get_causal_label(
        int(feature_idx), prompt, model, sae, get_openai_client(ctx),
        model_id=model_id, layer=resolved_layer,
    )
 
 
def _run_sae_pass(prompt: str, response: str, model: HookedTransformer, top_k: int = TOP_K_FEATURES, model_id: str = "llama-3.2-1b", layer: int | None = None) -> dict:
    from aquin.compute.model_loader import require_sae_layer
 
    cfg = get_config(model_id)
    if layer is not None:
        sae_layer = require_sae_layer(model_id, int(layer), command="trace")
    else:
        sae_layer = int(cfg["sae_layer"])
    sae = load_sae(model_id, sae_layer)
 
    full_ctx = f"{prompt}\n{response}"
    tokens = model.to_tokens(full_ctx)
 
    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=f"blocks.{sae_layer}.hook_resid_post",
            return_type=None,
        )
    resid = cache[f"blocks.{sae_layer}.hook_resid_post"][0]
 
    with torch.no_grad():
        feature_acts = sae.encode(normalize(resid, model_id, sae_layer))
 
    seq_len = resid.shape[0]
    prompt_ctx = f"{prompt}\n"
    prompt_ctx_len = model.to_tokens(prompt_ctx, prepend_bos=True).shape[1]
 
    all_strs = [model.to_string([tokens[0, i].item()]) for i in range(seq_len)]
    prompt_strs = all_strs[1:prompt_ctx_len]
    response_strs = all_strs[prompt_ctx_len:]
    prompt_idxs = list(range(1, prompt_ctx_len))
    response_idxs = list(range(prompt_ctx_len, seq_len))
 
    def feats_for_unlabeled(positions):
        out = []
        for pos in positions:
            acts = feature_acts[pos]
            topk = acts.topk(top_k)
            out.append([
                {
                    "feature_idx": int(i),
                    "activation": round(float(v), 3),
                    "label": f"feature_{int(i)}",
                }
                for i, v in zip(topk.indices, topk.values) if v > 0.001
            ])
        return out
 
    prompt_features = feats_for_unlabeled(prompt_idxs)
    response_features = feats_for_unlabeled(response_idxs)
 
    resp_acts = feature_acts[prompt_ctx_len:]
    resp_top = resp_acts.max(0).values.topk(20)
    top_response_features = []
    for idx, val in zip(resp_top.indices.tolist(), resp_top.values.tolist()):
        if val < 0.001:
            continue
        best_pos = int(resp_acts[:, idx].argmax().item())
        top_response_features.append({
            "feature_idx": idx,
            "activation": round(val, 3),
            "label": f"feature_{idx}",
            "token": response_strs[best_pos].strip() if best_pos < len(response_strs) else "",
            "token_idx": best_pos,
        })
 
    attribution = []
    for ri, rpos in enumerate(response_idxs):
        resp_tok = response_strs[ri].strip() if ri < len(response_strs) else ""
        if not resp_tok or resp_tok in ("the", "a", "an", "is", "of", ".", ","):
            continue
        resp_feats = feature_acts[rpos]
        top_resp = resp_feats.topk(top_k)
        driven_by = []
        for fidx, fval in zip(top_resp.indices.tolist(), top_resp.values.tolist()):
            if fval < 0.001:
                continue
            prompt_feat_acts = feature_acts[prompt_idxs, fidx]
            active = (prompt_feat_acts > 0.001).nonzero(as_tuple=True)[0].tolist()
            if not active:
                continue
            driven_by.append({
                "feature_idx": fidx,
                "label": f"feature_{fidx}",
                "activation": round(fval, 3),
                "also_in_prompt_positions": active,
                "also_in_prompt_tokens": [prompt_strs[p].strip() for p in active if p < len(prompt_strs)],
            })
        if driven_by:
            attribution.append({
                "response_token": resp_tok,
                "response_ti": ri,
                "driven_by_features": sorted(driven_by, key=lambda x: x["activation"], reverse=True)[:5],
            })
 
    synchronize_device()
 
    global _kernel_feature_acts, _kernel_resid, _kernel_top_features
    _kernel_feature_acts = feature_acts.detach().cpu()
    _kernel_resid = resid.detach().cpu()
    _kernel_top_features = top_response_features
 
    return {
        "prompt_tokens": [t.strip() for t in prompt_strs],
        "response_tokens": [t.strip() for t in response_strs],
        "prompt_features": prompt_features,
        "response_features": response_features,
        "top_response_features": top_response_features,
        "attribution": attribution,
        "sae_layer": sae_layer,
    }
 
 
def run_feature_analysis_unlabeled(prompt: str, response: str, model: HookedTransformer, model_id: str = "llama-3.2-1b", layer: int | None = None) -> dict:
    return _run_sae_pass(prompt, response, model, model_id=model_id, layer=layer)
 
 
def run_feature_analysis(prompt: str, response: str, model: HookedTransformer, client, top_k: int = TOP_K_FEATURES, model_id: str = "llama-3.2-1b", layer: int | None = None) -> dict:
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    sae = load_sae(model_id, resolved_layer)
    result = _run_sae_pass(prompt, response, model, top_k, model_id=model_id, layer=resolved_layer)
 
    def fill_labels(features_list):
        for pos_feats in features_list:
            for f in pos_feats:
                fi = int(f["feature_idx"])
                f["label"] = get_causal_label(fi, prompt, model, sae, client, model_id=model_id, layer=resolved_layer)
                f["feature_ref"] = format_feature_ref(fi, f["label"])
 
    fill_labels(result["prompt_features"])
    fill_labels(result["response_features"])
 
    for f in result["top_response_features"]:
        fi = int(f["feature_idx"])
        f["label"] = get_causal_label(fi, prompt, model, sae, client, model_id=model_id, layer=resolved_layer)
        f["feature_ref"] = format_feature_ref(fi, f["label"])
 
    for attr in result["attribution"]:
        for f in attr["driven_by_features"]:
            fi = int(f["feature_idx"])
            f["label"] = get_causal_label(fi, prompt, model, sae, client, model_id=model_id, layer=resolved_layer)
            f["feature_ref"] = format_feature_ref(fi, f["label"])
 
    return result
 
 
def get_feature_logits(
    feature_idx: int,
    model: HookedTransformer,
    model_id: str = "llama-3.2-1b",
    layer: int | None = None,
    top_k: int = 10,
) -> dict:
    """Top vocab tokens boosted/suppressed by an SAE decoder direction (W_dec @ W_U)."""
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    sae = load_sae(model_id, resolved_layer)
 
    device = model.W_U.device
    dtype = model.W_U.dtype
    steer_vec = sae.W_dec[feature_idx].to(device=device, dtype=dtype)
 
    from aquin.compute.hf_llm_shim import project_residual_to_logits
 
    with torch.no_grad():
        logits = project_residual_to_logits(model, steer_vec)
 
    top_pos = logits.topk(top_k)
    top_neg = logits.topk(top_k, largest=False)
 
    def _row(idx: int, val: float) -> dict:
        token = model.to_string([int(idx)]).strip()
        return {"token": token, "logit": round(float(val), 4)}
 
    boosts = [_row(int(i), float(v)) for i, v in zip(top_pos.indices, top_pos.values)]
    suppresses = [_row(int(i), float(v)) for i, v in zip(top_neg.indices, top_neg.values)]
 
    return {
        "feature_idx": feature_idx,
        "layer": resolved_layer,
        "boosts": boosts,
        "suppresses": suppresses,
        "top": boosts,
        "bottom": suppresses,
    }
 
 
def get_feature_neighbors(
    feature_idx: int,
    model_id: str = "llama-3.2-1b",
    layer: int | None = None,
    top_k: int = 8,
) -> dict:
    """Cosine-nearest SAE features in decoder weight space."""
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    sae = load_sae(model_id, resolved_layer)
 
    with torch.no_grad():
        W = sae.W_dec.float()
        W = W / W.norm(dim=-1, keepdim=True).clamp(min=1e-8)
        query = W[feature_idx]
        sims = W @ query
        sims[feature_idx] = -1.0
        top = sims.topk(top_k)
 
    neighbors = [
        {"feature_idx": int(i), "similarity": round(float(s), 4)}
        for i, s in zip(top.indices.tolist(), top.values.tolist())
    ]
    return {"feature_idx": feature_idx, "layer": resolved_layer, "neighbors": neighbors}