weight_trojans.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/weight_trojans.py.
Import adaptation: causal_trace -> aquin.compute.model_loader.
FastAPI router stripped — run_weight_trojan_analysis() is the direct entry point.
"""
from __future__ import annotations
 
from typing import Optional
 
import math
import numpy as np
import torch
 
KURTOSIS_FLAG   = 4.0
OUTLIER_FLAG    = 0.002
SV_RATIO_FLAG   = 8.0
RISK_SUSPICIOUS = 0.35
RISK_HIGH       = 0.65
 
SIGNAL_WEIGHTS = {"kurtosis": 0.40, "outlier": 0.35, "sv_ratio": 0.25}
 
 
def _excess_kurtosis(t: torch.Tensor) -> float:
    f = t.float().flatten()
    n = f.numel()
    if n < 8:
        return 0.0
    mu = f.mean()
    sigma = f.std()
    if sigma < 1e-9:
        return 0.0
    z = (f - mu) / sigma
    kurt = float((z ** 4).mean()) - 3.0
    return kurt
 
 
def _outlier_density(t: torch.Tensor) -> float:
    f = t.float().flatten()
    mu = f.mean()
    sigma = f.std()
    if sigma < 1e-9:
        return 0.0
    threshold = mu + 4.0 * sigma
    return float((f.abs() > threshold.abs()).float().mean())
 
 
def _sv_ratio(t: torch.Tensor) -> float:
    f = t.float()
    if f.dim() == 1:
        return 1.0
    if f.dim() > 2:
        f = f.reshape(-1, f.shape[-1])
    if min(f.shape) < 2:
        return 1.0
    # Power iteration on CPU — avoids extra GPU allocations while model is loaded.
    f = f.cpu()
    try:
        v = torch.randn(f.shape[1])
        for _ in range(20):
            u = f @ v
            u = u / u.norm().clamp(min=1e-9)
            v = f.T @ u
            v = v / v.norm().clamp(min=1e-9)
        sigma1 = (u.unsqueeze(0) @ f @ v.unsqueeze(1)).reshape(-1)[0].item()
 
        f2 = f - sigma1 * u.unsqueeze(1) * v.unsqueeze(0)
        v2 = torch.randn(f.shape[1])
        for _ in range(20):
            u2 = f2 @ v2
            u2 = u2 / u2.norm().clamp(min=1e-9)
            v2 = f2.T @ u2
            v2 = v2 / v2.norm().clamp(min=1e-9)
        sigma2 = (u2.unsqueeze(0) @ f2 @ v2.unsqueeze(1)).reshape(-1)[0].item()
 
        if abs(sigma2) < 1e-6:
            return abs(sigma1) / 1e-6
        return abs(sigma1) / max(abs(sigma2), 1e-6)
    except Exception:
        return 1.0
 
 
def _analyse_tensor(name: str, t: torch.Tensor) -> dict:
    kurt = _excess_kurtosis(t)
    out  = _outlier_density(t)
    svr  = _sv_ratio(t)
    return {
        "name":            name,
        "shape":           list(t.shape),
        "n_params":        t.numel(),
        "mean":            round(float(t.float().mean()), 6),
        "std":             round(float(t.float().std()),  6),
        "kurtosis":        round(kurt,  4),
        "outlier_density": round(out,   6),
        "sv_ratio":        round(svr,   4),
    }
 
 
def _compute_risk(layer_stats: list[dict]) -> list[dict]:
    if not layer_stats:
        return []
 
    kurtosis_vals = [s["kurtosis"] for s in layer_stats]
    outlier_vals  = [s["outlier_density"] for s in layer_stats]
    sv_vals       = [s["sv_ratio"] for s in layer_stats]
 
    def _zscore_norm(vals: list[float]) -> list[float]:
        arr = np.array(vals, dtype=float)
        mu, sigma = arr.mean(), arr.std()
        if sigma < 1e-9:
            return [0.0] * len(vals)
        return np.clip((arr - mu) / sigma, 0, None).tolist()
 
    kurt_z = _zscore_norm(kurtosis_vals)
    out_z  = _zscore_norm(outlier_vals)
    sv_z   = _zscore_norm(sv_vals)
 
    def _norm01(vals: list[float]) -> list[float]:
        mx = max(vals) if vals else 0.0
        if mx < 1e-9:
            return [0.0] * len(vals)
        return [v / mx for v in vals]
 
    kurt_n = _norm01(kurt_z)
    out_n  = _norm01(out_z)
    sv_n   = _norm01(sv_z)
 
    wk = SIGNAL_WEIGHTS["kurtosis"]
    wo = SIGNAL_WEIGHTS["outlier"]
    ws = SIGNAL_WEIGHTS["sv_ratio"]
 
    results = []
    for i, s in enumerate(layer_stats):
        risk = wk * kurt_n[i] + wo * out_n[i] + ws * sv_n[i]
        risk = round(min(risk, 1.0), 4)
        status = (
            "high_risk"  if risk >= RISK_HIGH       else
            "suspicious" if risk >= RISK_SUSPICIOUS else
            "clean"
        )
        flags = []
        if s["kurtosis"] > KURTOSIS_FLAG:
            flags.append(f"excess kurtosis {s['kurtosis']:.2f} (heavy-tailed weight distribution)")
        if s["outlier_density"] > OUTLIER_FLAG:
            flags.append(f"outlier density {s['outlier_density']:.4%} (weight magnitude spikes)")
        if s["sv_ratio"] > SV_RATIO_FLAG:
            flags.append(f"SV ratio {s['sv_ratio']:.1f} (low-rank implant signature)")
 
        results.append({**s, "risk_score": risk, "status": status, "flags": flags})
 
    return results
 
 
def run_weight_trojan_analysis(
    model,
    layer_range: Optional[list[int]] = None,
    max_tensors_per_layer: int = 4,
) -> dict:
    from aquin.compute.weight_rank import _is_heavy_weight_model, _should_skip_weight_param
 
    heavy = _is_heavy_weight_model(model)
    if heavy:
        max_tensors_per_layer = min(max_tensors_per_layer, 6)
 
    target_layers = set(layer_range) if layer_range else None
    layer_stats: list[dict] = []
    skipped = 0
    per_layer_count: dict[int, int] = {}
 
    for name, param in model.named_parameters():
        parts = name.split(".")
        layer_idx: Optional[int] = None
        for p in parts:
            if p.isdigit():
                layer_idx = int(p)
                break
 
        if target_layers is not None and layer_idx not in target_layers:
            skipped += 1
            continue
 
        if param.dim() < 2 or param.numel() < 256:
            continue
 
        if _should_skip_weight_param(name, heavy=heavy):
            skipped += 1
            continue
 
        if layer_idx is not None:
            count = per_layer_count.get(layer_idx, 0)
            if count >= max_tensors_per_layer:
                skipped += 1
                continue
            per_layer_count[layer_idx] = count + 1
 
        stats = _analyse_tensor(name, param.detach().cpu())
        if layer_idx is not None:
            stats["layer_idx"] = layer_idx
        layer_stats.append(stats)
 
    if not layer_stats:
        return {"error": "No eligible weight tensors found", "layers_analysed": 0}
 
    scored = _compute_risk(layer_stats)
 
    high_risk  = [s for s in scored if s["status"] == "high_risk"]
    suspicious = [s for s in scored if s["status"] == "suspicious"]
    clean      = [s for s in scored if s["status"] == "clean"]
 
    n = len(scored)
    pct_flagged = round((len(high_risk) + len(suspicious)) / max(n, 1) * 100, 2)
 
    if len(high_risk) >= 2:
        verdict = "high_risk"
    elif len(high_risk) >= 1 or len(suspicious) >= 3:
        verdict = "suspicious"
    else:
        verdict = "clean"
 
    composite = round(float(np.mean([s["risk_score"] for s in scored])), 4) if scored else 0.0
 
    all_flags: list[str] = []
    for s in high_risk + suspicious:
        for f in s["flags"]:
            if f not in all_flags:
                all_flags.append(f)
 
    return {
        "layers_analysed":  n,
        "tensors_skipped":  skipped,
        "composite_risk":   composite,
        "verdict":          verdict,
        "pct_flagged":      pct_flagged,
        "high_risk_count":  len(high_risk),
        "suspicious_count": len(suspicious),
        "clean_count":      len(clean),
        "all_flags":        all_flags,
        "scored_tensors":   scored,
        "signals":          list(SIGNAL_WEIGHTS.keys()),
        "heavy_model":      heavy,
        "thresholds": {
            "kurtosis_flag":   KURTOSIS_FLAG,
            "outlier_flag":    OUTLIER_FLAG,
            "sv_ratio_flag":   SV_RATIO_FLAG,
            "risk_suspicious": RISK_SUSPICIOUS,
            "risk_high":       RISK_HIGH,
        },
    }