layer_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.
 
"""LLM layer health: activation stability + OOD similarity."""
 
from __future__ import annotations
 
import torch
from transformer_lens import HookedTransformer
 
DEFAULT_STABILITY_PROMPTS = [
    "The capital of France is",
    "Write a function that sorts a list",
    "Explain photosynthesis in simple terms",
    "Hello, how are you today?",
    "The quick brown fox jumps over the lazy dog",
]
 
DEFAULT_IN_DOMAIN_PROMPTS = [
    "The capital of France is Paris.",
    "Water boils at 100 degrees Celsius.",
    "Python is a programming language.",
    "The Earth orbits the Sun.",
]
 
DEFAULT_OOD_PROMPTS = [
    "asdf jkl qwerty zxcv",
    "!!! ??? ### @@@",
    "florp sniggle wumpus",
    "9918273645 192837465",
]
 
COLLAPSE_THRESHOLD = 0.85
DEAD_THRESHOLD = 0.01
MIN_STABILITY_PROMPTS = 2
 
 
def _resolve_stability_prompts(prompts: list[str] | None) -> tuple[list[str], bool]:
    """PCA needs ≥2 samples. Pad with defaults when the user passes fewer."""
    if not prompts:
        return list(DEFAULT_STABILITY_PROMPTS), False
    resolved = list(prompts)
    if len(resolved) >= MIN_STABILITY_PROMPTS:
        return resolved, False
    for p in DEFAULT_STABILITY_PROMPTS:
        if p not in resolved:
            resolved.append(p)
        if len(resolved) >= MIN_STABILITY_PROMPTS:
            break
    return resolved, True
 
 
def _pca_variance_ratios(acts: torch.Tensor, top_k: int) -> tuple[float, float]:
    """acts: (n_samples, d_model). Returns (top1_ratio, topk_ratio)."""
    if acts.shape[0] < 2:
        return 0.0, 0.0
    centered = acts - acts.mean(dim=0, keepdim=True)
    try:
        _, s, _ = torch.linalg.svd(centered, full_matrices=False)
    except Exception:
        return 0.0, 0.0
    var = s ** 2
    total = var.sum()
    if total < 1e-12:
        return 0.0, 0.0
    top1 = float(var[0] / total)
    k = min(top_k, len(var))
    topk = float(var[:k].sum() / total)
    return top1, topk
 
 
def _collect_layer_activations(
    model: HookedTransformer,
    prompts: list[str],
) -> dict[int, torch.Tensor]:
    """Last-token resid_post per layer. layer -> (n_prompts, d_model)."""
    n_layers = model.cfg.n_layers
    per_layer: dict[int, list[torch.Tensor]] = {i: [] for i in range(n_layers)}
 
    for prompt in prompts:
        tokens = model.to_tokens(prompt)
        with torch.no_grad():
            _, cache = model.run_with_cache(
                tokens,
                names_filter=lambda name: name.endswith("hook_resid_post"),
            )
        for layer in range(n_layers):
            key = f"blocks.{layer}.hook_resid_post"
            if key in cache:
                per_layer[layer].append(cache[key][0, -1].float().cpu())
 
    return {layer: torch.stack(vecs) for layer, vecs in per_layer.items() if vecs}
 
 
def run_activation_stability(
    model: HookedTransformer,
    prompts: list[str] | None = None,
    top_k: int = 10,
) -> dict:
    user_n = len(prompts) if prompts else 0
    prompts, padded_defaults = _resolve_stability_prompts(prompts)
    layer_acts = _collect_layer_activations(model, prompts)
    n_layers = model.cfg.n_layers
    layers_out: list[dict] = []
    dead_count = 0
    collapsed_count = 0
 
    for layer in range(n_layers):
        acts = layer_acts.get(layer)
        if acts is None or acts.shape[0] < 2:
            continue
        top1, topk = _pca_variance_ratios(acts, top_k)
        mean_act = float(acts.mean())
        std_act = float(acts.std())
        status = "ok"
        if std_act < DEAD_THRESHOLD:
            status = "dead"
            dead_count += 1
        elif top1 >= COLLAPSE_THRESHOLD:
            status = "collapsed"
            collapsed_count += 1
        layers_out.append({
            "layer": layer,
            "top1_variance_ratio": round(top1, 4),
            "topk_variance_ratio": round(topk, 4),
            "mean_activation": round(mean_act, 4),
            "activation_std": round(std_act, 4),
            "status": status,
        })
 
    return {
        "n_layers": n_layers,
        "top_k": top_k,
        "n_prompts": len(prompts),
        "n_user_prompts": user_n or len(prompts),
        "padded_defaults": padded_defaults,
        "layers": layers_out,
        "dead_count": dead_count,
        "collapsed_count": collapsed_count,
        "collapse_threshold": COLLAPSE_THRESHOLD,
        "dead_threshold": DEAD_THRESHOLD,
    }
 
 
def _mean_pairwise_cos(vecs: torch.Tensor) -> float:
    if vecs.shape[0] < 2:
        return 1.0
    v = vecs / vecs.norm(dim=-1, keepdim=True).clamp(min=1e-8)
    sims = v @ v.T
    n = sims.shape[0]
    return float((sims.sum() - n) / (n * (n - 1)))
 
 
def _mean_cross_cos(a: torch.Tensor, b: torch.Tensor) -> float:
    if a.shape[0] == 0 or b.shape[0] == 0:
        return 0.0
    a = a / a.norm(dim=-1, keepdim=True).clamp(min=1e-8)
    b = b / b.norm(dim=-1, keepdim=True).clamp(min=1e-8)
    return float((a @ b.T).mean())
 
 
def _rbf_mmd(x: torch.Tensor, y: torch.Tensor) -> float:
    if x.shape[0] < 2 or y.shape[0] < 2:
        return 0.0
    xy = torch.cat([x, y], dim=0)
    dists = torch.cdist(xy, xy, p=2)
    sigma = float(dists.median())
    if sigma < 1e-8:
        sigma = 1.0
 
    def _kernel(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        d = torch.cdist(a, b, p=2)
        return torch.exp(-(d ** 2) / (2 * sigma ** 2))
 
    k_xx = _kernel(x, x).mean()
    k_yy = _kernel(y, y).mean()
    k_xy = _kernel(x, y).mean()
    return float(torch.clamp(k_xx + k_yy - 2 * k_xy, min=0.0))
 
 
def run_ood_similarity(
    model: HookedTransformer,
    in_domain_prompts: list[str] | None = None,
    ood_prompts: list[str] | None = None,
) -> dict:
    in_prompts = in_domain_prompts or DEFAULT_IN_DOMAIN_PROMPTS
    ood_ps = ood_prompts or DEFAULT_OOD_PROMPTS
    in_acts = _collect_layer_activations(model, in_prompts)
    ood_acts = _collect_layer_activations(model, ood_ps)
    n_layers = model.cfg.n_layers
    layers_out: list[dict] = []
 
    for layer in range(n_layers):
        in_a = in_acts.get(layer)
        ood_a = ood_acts.get(layer)
        if in_a is None or ood_a is None:
            continue
        in_cos = _mean_pairwise_cos(in_a)
        ood_cos = _mean_pairwise_cos(ood_a)
        cross_cos = _mean_cross_cos(in_a, ood_a)
        separation = in_cos - cross_cos
        mmd = _rbf_mmd(in_a, ood_a)
        layers_out.append({
            "layer": layer,
            "in_cos": round(in_cos, 4),
            "ood_cos": round(ood_cos, 4),
            "cross_cos": round(cross_cos, 4),
            "mmd": round(mmd, 4),
            "separation": round(separation, 4),
        })
 
    peak_layer = max(layers_out, key=lambda row: row["separation"])["layer"] if layers_out else 0
    mean_sep = (
        sum(row["separation"] for row in layers_out) / len(layers_out)
        if layers_out else 0.0
    )
 
    return {
        "n_layers": n_layers,
        "n_in_prompts": len(in_prompts),
        "n_ood_prompts": len(ood_ps),
        "layers": layers_out,
        "peak_layer": peak_layer,
        "mean_separation": round(mean_sep, 4),
    }
 
 
def run_layer_analysis(model: HookedTransformer, args: dict) -> dict:
    prompts = args.get("prompts")
    do_localize = bool(args.get("localize"))
 
    # When --localize uses a probe file path, don't feed that path into stability PCA.
    stability_prompts = prompts
    localize_prompts = args.get("localize_prompts")
    if do_localize:
        if localize_prompts is None and isinstance(prompts, str) and not prompts.strip().startswith("["):
            localize_prompts = prompts
            stability_prompts = None
        elif localize_prompts is None and prompts is None:
            localize_prompts = None  # default deception fixture
        if isinstance(prompts, str) and prompts.strip().startswith("["):
            try:
                import json
                stability_prompts = json.loads(prompts)
            except Exception:
                stability_prompts = [prompts]
 
    if isinstance(stability_prompts, str):
        stability_prompts = [stability_prompts]
 
    in_domain = args.get("in_domain_prompts")
    ood = args.get("ood_prompts")
    top_k = int(args.get("top_k", 10))
 
    stability = run_activation_stability(model, prompts=stability_prompts, top_k=top_k)
    ood_result = run_ood_similarity(
        model,
        in_domain_prompts=in_domain,
        ood_prompts=ood,
    )
    out: dict = {
        "stability": stability,
        "ood": ood_result,
        "model_id": args.get("model_id", ""),
    }
 
    if do_localize:
        from aquin.compute.localize_collapse import run_localize_collapse
 
        feature_idx = args.get("feature_idx")
        if feature_idx is not None:
            feature_idx = int(feature_idx)
        layer = args.get("layer")
        if layer is not None:
            layer = int(layer)
 
        localize = run_localize_collapse(
            model,
            str(args.get("model_id") or ""),
            prompts=localize_prompts,
            stressor_prompts=args.get("stressor_prompts") or args.get("stressor"),
            feature_idx=feature_idx,
            vector_path=args.get("vector") or args.get("vector_path"),
            layer=layer,
            top_k_features=int(args.get("top_k_features") or args.get("top_k") or 8),
            collect_layer_activations=_collect_layer_activations,
        )
        out["localize"] = localize
 
    return out