residual_drift.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Per-layer activation drift: catalog base vs fine-tuned checkpoint (LLM)."""
 
from __future__ import annotations
 
from pathlib import Path
from typing import Any
 
import torch
 
from aquin.compute.layer_analysis import _collect_layer_activations
from aquin.compute.sae_diff import load_checkpoint_state, load_prompts, load_tl_from_checkpoint
 
 
def _cosine_distance_rows(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """Per-row cosine distance. a, b: (n, d)."""
    a_n = a.float() / a.float().norm(dim=-1, keepdim=True).clamp(min=1e-8)
    b_n = b.float() / b.float().norm(dim=-1, keepdim=True).clamp(min=1e-8)
    return 1.0 - (a_n * b_n).sum(dim=-1)
 
 
def _layer_drift_rows(
    base_acts: dict[int, torch.Tensor],
    ft_acts: dict[int, torch.Tensor],
    n_layers: int,
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for layer in range(n_layers):
        base = base_acts.get(layer)
        ft = ft_acts.get(layer)
        if base is None or ft is None or base.shape != ft.shape:
            continue
        dists = _cosine_distance_rows(base, ft)
        rows.append({
            "layer": layer,
            "mean_cosine_distance": round(float(dists.mean().item()), 6),
            "max_cosine_distance": round(float(dists.max().item()), 6),
            "mean_cosine_sim": round(float((1.0 - dists).mean().item()), 6),
            "n_probes": int(dists.shape[0]),
        })
    return rows
 
 
def _per_probe_rows(
    base_acts: dict[int, torch.Tensor],
    ft_acts: dict[int, torch.Tensor],
    layer_rows: list[dict[str, Any]],
    probes: list[str],
) -> list[dict[str, Any]]:
    if not layer_rows:
        return []
    layers = [int(r["layer"]) for r in layer_rows]
    out: list[dict[str, Any]] = []
    n_probes = int(layer_rows[0].get("n_probes") or len(probes))
    for i in range(n_probes):
        dists: list[float] = []
        for layer in layers:
            base = base_acts.get(layer)
            ft = ft_acts.get(layer)
            if base is None or ft is None:
                continue
            dists.append(float(_cosine_distance_rows(base[i : i + 1], ft[i : i + 1]).item()))
        if not dists:
            continue
        peak_layer = layers[int(max(range(len(dists)), key=lambda j: dists[j]))]
        preview = probes[i][:80] + ("…" if len(probes[i]) > 80 else "")
        out.append({
            "probe_index": i,
            "probe_preview": preview,
            "mean_drift": round(sum(dists) / len(dists), 6),
            "max_drift": round(max(dists), 6),
            "peak_layer": peak_layer,
        })
    return sorted(out, key=lambda r: r["max_drift"], reverse=True)
 
 
def _finalize_payload(
    *,
    model_id: str,
    mode: str,
    activation_mode: str,
    checkpoint_name: str,
    checkpoint_path: str,
    prompts: list[str],
    layer_profile: list[dict[str, Any]],
    per_probe: list[dict[str, Any]],
    step: int | None = None,
) -> dict[str, Any]:
    drifts = [float(r["mean_cosine_distance"]) for r in layer_profile]
    peak = max(layer_profile, key=lambda r: r["mean_cosine_distance"]) if layer_profile else None
    return {
        "schema_version": 1,
        "type": "residualDrift",
        "baseModelId": model_id,
        "ftCheckpointName": checkpoint_name,
        "checkpointPath": checkpoint_path,
        "modelMode": mode,
        "activationMode": activation_mode,
        "trainingStep": step,
        "nProbes": len(prompts),
        "nLayers": len(layer_profile),
        "meanDrift": round(sum(drifts) / len(drifts), 6) if drifts else 0.0,
        "maxDrift": round(max(drifts), 6) if drifts else 0.0,
        "peakLayer": peak["layer"] if peak else None,
        "layerProfile": layer_profile,
        "topLayers": sorted(layer_profile, key=lambda r: r["mean_cosine_distance"], reverse=True)[:10],
        "perProbe": per_probe[:25],
        "promptsUsed": prompts,
    }
 
 
def run_llm_residual_drift(
    model_id: str,
    checkpoint_path: str | Path,
    prompts: list[str],
    *,
    checkpoint_name: str = "checkpoint",
) -> dict[str, Any]:
    from aquin.compute.causal_trace import load_model as load_base_tl
    from aquin.compute.device import empty_device_cache
    from aquin.compute.model_loader import resolve_model_id
    from aquin.compute.model_runtime import unload_weights
 
    short = resolve_model_id(model_id)
    _, step = load_checkpoint_state(checkpoint_path)
 
    unload_weights(clear_active=False)
    empty_device_cache()
 
    tl_base = load_base_tl(short)
    base_acts = _collect_layer_activations(tl_base, prompts)
    n_layers = int(tl_base.cfg.n_layers)
    del tl_base
    unload_weights(clear_active=False)
    empty_device_cache()
 
    tl_ft = load_tl_from_checkpoint(short, checkpoint_path=checkpoint_path)
    ft_acts = _collect_layer_activations(tl_ft, prompts)
    del tl_ft
    unload_weights(clear_active=False)
    empty_device_cache()
 
    layer_profile = _layer_drift_rows(base_acts, ft_acts, n_layers)
    per_probe = _per_probe_rows(base_acts, ft_acts, layer_profile, prompts)
 
    return _finalize_payload(
        model_id=short,
        mode="llm",
        activation_mode="hook_resid_post_last_token",
        checkpoint_name=checkpoint_name,
        checkpoint_path=str(checkpoint_path),
        prompts=prompts,
        layer_profile=layer_profile,
        per_probe=per_probe,
        step=step,
    )
 
 
def run_residual_drift(
    model_id: str,
    checkpoint_path: str | Path,
    prompts: list[str] | None = None,
    *,
    checkpoint_name: str | None = None,
) -> dict[str, Any]:
    from aquin.compute.model_loader import resolve_model_id
 
    ckpt = Path(checkpoint_path)
    if not ckpt.exists():
        raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
 
    probe_list = list(prompts or load_prompts(None))
    name = checkpoint_name or ckpt.stem
 
    short = resolve_model_id(model_id)
    return run_llm_residual_drift(short, ckpt, probe_list, checkpoint_name=name)
 
 
def run_residual_drift_from_args(args: dict[str, Any]) -> dict[str, Any]:
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    model_id = args.get("model_id") or get_active_model_id() or "llama-3.2-1b"
    checkpoint = args.get("checkpoint")
    if not checkpoint:
        return {"error": "Missing --checkpoint <path>. Run: aquin diff residue --help"}
 
    try:
        model_id = resolve_model_id(str(model_id))
    except ValueError as exc:
        return {"error": str(exc)}
 
    prompts_path = args.get("prompts")
    try:
        prompts = load_prompts(prompts_path) if prompts_path else None
    except (FileNotFoundError, ValueError) as exc:
        return {"error": str(exc)}
 
    try:
        return run_residual_drift(
            model_id,
            checkpoint,
            prompts,
            checkpoint_name=args.get("name"),
        )
    except Exception as exc:
        return {"error": str(exc)}