localize_collapse.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Deception-representation collapse localization across layers."""
from __future__ import annotations
 
from pathlib import Path
from typing import Any
 
import torch
 
from aquin.compute.find_feature import load_deception_probes
 
SIGNAL_WEAK = 0.08
SIGNAL_COLLAPSED = 0.03
MAX_PROBES_PER_CLASS = 24
 
 
def _unit(v: torch.Tensor) -> torch.Tensor:
    return v / v.norm().clamp(min=1e-8)
 
 
def _centroid_signal(honest: torch.Tensor, deceptive: torch.Tensor) -> dict[str, float]:
    """honest/deceptive: (n, d). Relative centroid separation + cosine gap."""
    if honest.shape[0] == 0 or deceptive.shape[0] == 0:
        return {"signal": 0.0, "l2_sep": 0.0, "cos_sep": 0.0, "proj_delta": 0.0}
 
    mu_h = honest.mean(dim=0)
    mu_d = deceptive.mean(dim=0)
    delta = mu_d - mu_h
    l2_sep = float(delta.norm().item())
    scale = float(mu_h.norm().item() + mu_d.norm().item()) + 1e-8
    signal = l2_sep / scale
 
    cos = float(torch.nn.functional.cosine_similarity(mu_h.unsqueeze(0), mu_d.unsqueeze(0)).item())
    cos_sep = max(0.0, 1.0 - cos)
 
    direction = _unit(delta)
    proj_h = float((honest @ direction).mean().item())
    proj_d = float((deceptive @ direction).mean().item())
    proj_delta = abs(proj_d - proj_h)
 
    return {
        "signal": round(signal, 6),
        "l2_sep": round(l2_sep, 6),
        "cos_sep": round(cos_sep, 6),
        "proj_delta": round(proj_delta, 6),
    }
 
 
def _status(signal: float) -> str:
    if signal <= SIGNAL_COLLAPSED:
        return "collapsed"
    if signal <= SIGNAL_WEAK:
        return "weak"
    return "ok"
 
 
def pick_collapse_layer(layers: list[dict[str, Any]]) -> dict[str, Any]:
    """Peak = max signal; collapse = steepest drop after peak (else global min)."""
    if not layers:
        return {"peak_layer": None, "collapse_layer": None, "drop_from_peak": None}
 
    peak = max(layers, key=lambda r: float(r.get("signal") or 0.0))
    peak_layer = int(peak["layer"])
    peak_signal = float(peak["signal"])
 
    after = [r for r in layers if int(r["layer"]) > peak_layer]
    if after:
        worst = min(after, key=lambda r: float(r.get("signal") or 0.0))
        collapse_layer = int(worst["layer"])
        drop = round(peak_signal - float(worst["signal"]), 6)
    else:
        worst = min(layers, key=lambda r: float(r.get("signal") or 0.0))
        collapse_layer = int(worst["layer"])
        drop = round(peak_signal - float(worst["signal"]), 6)
 
    return {
        "peak_layer": peak_layer,
        "collapse_layer": collapse_layer,
        "drop_from_peak": drop,
        "peak_signal": round(peak_signal, 6),
        "collapse_signal": round(float(worst["signal"]), 6),
    }
 
 
def score_layer_activations(
    honest_by_layer: dict[int, torch.Tensor],
    deceptive_by_layer: dict[int, torch.Tensor],
) -> list[dict[str, Any]]:
    layers_out: list[dict[str, Any]] = []
    for layer in sorted(set(honest_by_layer) & set(deceptive_by_layer)):
        metrics = _centroid_signal(honest_by_layer[layer], deceptive_by_layer[layer])
        layers_out.append({
            "layer": int(layer),
            **metrics,
            "status": _status(metrics["signal"]),
            "n_honest": int(honest_by_layer[layer].shape[0]),
            "n_deceptive": int(deceptive_by_layer[layer].shape[0]),
        })
    return layers_out
 
 
def score_stressor_collapse(
    baseline_layers: list[dict[str, Any]],
    stressor_layers: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    by_base = {int(r["layer"]): r for r in baseline_layers}
    by_stress = {int(r["layer"]): r for r in stressor_layers}
    rows: list[dict[str, Any]] = []
    for layer in sorted(set(by_base) & set(by_stress)):
        b = float(by_base[layer]["signal"])
        s = float(by_stress[layer]["signal"])
        delta = round(b - s, 6)
        rows.append({
            "layer": layer,
            "baseline_signal": round(b, 6),
            "stressor_signal": round(s, 6),
            "collapse_delta": delta,
            "status": "collapsed" if delta >= 0.05 and s <= SIGNAL_WEAK else (
                "weakened" if delta >= 0.02 else "stable"
            ),
        })
    return rows
 
 
def _project_on_direction(
    honest: torch.Tensor,
    deceptive: torch.Tensor,
    direction: torch.Tensor,
) -> dict[str, float]:
    d = _unit(direction.float().cpu())
    if honest.shape[-1] != d.shape[-1]:
        raise ValueError(
            f"Direction d_model={d.shape[-1]} does not match activations d_model={honest.shape[-1]}"
        )
    proj_h = float((honest @ d).mean().item())
    proj_d = float((deceptive @ d).mean().item())
    return {
        "honest_proj": round(proj_h, 6),
        "deceptive_proj": round(proj_d, 6),
        "proj_delta": round(abs(proj_d - proj_h), 6),
    }
 
 
def rank_sae_features_at_layer(
    model: Any,
    model_id: str,
    layer: int,
    honest_texts: list[str],
    deceptive_texts: list[str],
    *,
    top_k: int = 8,
) -> list[dict[str, Any]] | None:
    """Rank SAE features by deceptive−honest mean activation at one layer."""
    try:
        from aquin.compute.feature_analysis import load_sae
        from aquin.compute.find_feature import _rank_features
    except Exception:
        return None
 
    try:
        sae = load_sae(model_id, layer)
    except Exception:
        return None
 
    def _mean_feats(texts: list[str]) -> torch.Tensor | None:
        feats: list[torch.Tensor] = []
        hook = f"blocks.{layer}.hook_resid_post"
        for text in texts:
            tokens = model.to_tokens(text)
            with torch.no_grad():
                _, cache = model.run_with_cache(tokens, names_filter=lambda n, h=hook: n == h)
            if hook not in cache:
                continue
            act = cache[hook][0, -1].float()
            with torch.no_grad():
                f = sae.encode(act.unsqueeze(0))[0].float().cpu()
            feats.append(f)
        if not feats:
            return None
        return torch.stack(feats).mean(dim=0)
 
    h = _mean_feats(honest_texts)
    d = _mean_feats(deceptive_texts)
    if h is None or d is None:
        return None
    return _rank_features(h, d, top_k=top_k, direction="both")
 
 
def run_localize_collapse(
    model: Any,
    model_id: str,
    *,
    prompts: str | Path | None = None,
    stressor_prompts: str | Path | None = None,
    feature_idx: int | None = None,
    vector_path: str | Path | None = None,
    layer: int | None = None,
    top_k_features: int = 8,
    collect_layer_activations: Any | None = None,
) -> dict[str, Any]:
    """
    Rank layers by honest vs deceptive representation strength.
 
    Default: contrastive centroid signal per resid_post layer.
    Optional stressor prompts → per-layer collapse_delta.
    Optional feature_idx / vector → direction projection at the relevant layer.
    """
    from aquin.compute.layer_analysis import _collect_layer_activations
 
    collect = collect_layer_activations or _collect_layer_activations
 
    try:
        honest, deceptive, meta = load_deception_probes(prompts)
    except (FileNotFoundError, ValueError) as e:
        return {"error": str(e)}
 
    honest = honest[:MAX_PROBES_PER_CLASS]
    deceptive = deceptive[:MAX_PROBES_PER_CLASS]
 
    honest_acts = collect(model, honest)
    deceptive_acts = collect(model, deceptive)
    layers = score_layer_activations(honest_acts, deceptive_acts)
    locus = pick_collapse_layer(layers)
 
    out: dict[str, Any] = {
        "mode": "contrastive",
        "prompts_path": meta.get("prompts_path"),
        "n_honest": len(honest),
        "n_deceptive": len(deceptive),
        "layers": layers,
        **locus,
        "n_collapsed": sum(1 for r in layers if r["status"] == "collapsed"),
        "n_weak": sum(1 for r in layers if r["status"] == "weak"),
    }
 
    if stressor_prompts is not None:
        try:
            s_honest, s_deceptive, s_meta = load_deception_probes(stressor_prompts)
        except (FileNotFoundError, ValueError) as e:
            return {**out, "error": f"stressor probes: {e}"}
        s_honest = s_honest[:MAX_PROBES_PER_CLASS]
        s_deceptive = s_deceptive[:MAX_PROBES_PER_CLASS]
        s_layers = score_layer_activations(
            collect(model, s_honest),
            collect(model, s_deceptive),
        )
        stress_rows = score_stressor_collapse(layers, s_layers)
        out["stressor"] = {
            "prompts_path": s_meta.get("prompts_path"),
            "n_honest": len(s_honest),
            "n_deceptive": len(s_deceptive),
            "layers": stress_rows,
        }
        if stress_rows:
            worst = max(stress_rows, key=lambda r: float(r["collapse_delta"]))
            out["stressor"]["max_collapse_layer"] = int(worst["layer"])
            out["stressor"]["max_collapse_delta"] = float(worst["collapse_delta"])
            # Prefer stressor-driven collapse locus when available
            out["collapse_layer"] = int(worst["layer"])
            out["collapse_signal"] = float(worst["stressor_signal"])
            out["drop_from_peak"] = float(worst["collapse_delta"])
            out["mode"] = "contrastive+stressor"
 
    direction_info: dict[str, Any] | None = None
    direction_vec: torch.Tensor | None = None
    direction_layer: int | None = layer
 
    if vector_path:
        from aquin.compute.steer_vector import load_steer_vector_file, steer_vector_tensor
 
        try:
            payload = load_steer_vector_file(vector_path)
            direction_vec = steer_vector_tensor(
                payload, device="cpu", dtype=torch.float32,
            )
            direction_layer = int(payload.get("layer", direction_layer or 0))
            direction_info = {
                "source": "vector",
                "vector_path": str(vector_path),
                "layer": direction_layer,
                "feature_idx": payload.get("feature_idx"),
            }
        except (OSError, ValueError) as e:
            return {**out, "error": f"vector: {e}"}
    elif feature_idx is not None:
        from aquin.compute.steer_vector import resolve_feature_direction
 
        try:
            direction_vec, direction_layer, _ = resolve_feature_direction(
                model_id, int(feature_idx), layer=layer,
            )
            direction_vec = direction_vec.float().cpu()
            direction_info = {
                "source": "feature",
                "feature_idx": int(feature_idx),
                "layer": int(direction_layer),
            }
        except Exception as e:
            return {**out, "error": f"feature direction: {e}"}
 
    if direction_vec is not None and direction_layer is not None:
        h = honest_acts.get(int(direction_layer))
        d = deceptive_acts.get(int(direction_layer))
        if h is not None and d is not None:
            try:
                proj = _project_on_direction(h, d, direction_vec)
                direction_info = {**(direction_info or {}), **proj}
            except ValueError as e:
                return {**out, "error": str(e)}
        out["direction"] = direction_info
 
    collapse_l = out.get("collapse_layer")
    if collapse_l is not None:
        feats = rank_sae_features_at_layer(
            model,
            model_id,
            int(collapse_l),
            honest,
            deceptive,
            top_k=int(top_k_features),
        )
        if feats is not None:
            out["features_at_collapse"] = feats
            out["features_layer"] = int(collapse_l)
 
    return out