OBLITERATUS/obliteratus/analysis/spectral_certification.py

"""Spectral Abliteration Completeness Certification via Random Matrix Theory.

Current abliteration tools test success empirically — run harmful prompts,
check if refusal drops. There is no formal guarantee that abliteration is
complete. Extended-refusal fine-tuning (Shairah et al., KAUST, May 2025)
distributes refusal into many low-energy dimensions, defeating single-
direction abliteration. GRP-Obliteration (Russinovich et al., Microsoft,
Feb 2026) reorganizes safety representations entirely.

This module uses random matrix theory to build a *spectral diagnostic*
for abliteration completeness. After abliteration, it computes the
covariance of residual activations and applies the BBP phase transition
to determine whether any detectable refusal signal survives.

Contributions:
  1. **Spectral diagnostic**: Green/Yellow/Red classification when the
     estimated sample requirement is met, otherwise Inconclusive
     based on eigenvalue analysis relative to BBP threshold
  2. **Non-isotropic BBP extension**: Extends Paper Theorem 4 to
     anisotropic activation covariance (heuristic extension)
  3. **Distributed refusal detection**: Identifies when refusal has been
     distributed across many weak dimensions (Yellow tier)
  4. **Marchenko-Pastur noise floor**: Rigorous separation of signal
     from noise in post-abliteration residuals

References:
    - Baik, Ben Arous & Peche (2005): BBP phase transition
    - Marchenko & Pastur (1967): Limiting distribution of eigenvalues
    - Shairah et al. (2025): Extended-Refusal Fine-Tuning defense
    - Russinovich et al. (2026): GRP-Obliteration
    - Paper Theorem 4: BBP Detectability Phase Transition
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass
from enum import Enum

import torch

logger = logging.getLogger(__name__)


class CertificationLevel(Enum):
    """Spectral diagnostic outcome for abliteration completeness."""

    GREEN = "certified_complete"
    """All eigenvalues below BBP threshold. No detectable linear refusal
    remains in the post-abliteration residual stream."""

    YELLOW = "distributed_refusal"
    """Eigenvalues above threshold but below concentration bound. Refusal
    has been distributed across many weak dimensions (defense like
    extended-refusal is active). Escalate to GRP-Obliteration."""

    RED = "incomplete"
    """Clear eigenvalue spikes above threshold. Abliteration failed to
    remove all refusal signal. Re-run with more directions."""

    INCONCLUSIVE = "insufficient_samples"
    """The observed activation sample is too small for a reliable BBP-based
    classification at this dimensionality. Collect more samples before
    interpreting the spectral result."""


@dataclass
class SpectralCertificate:
    """Statistical diagnostic of abliteration completeness."""

    # Certification
    level: CertificationLevel
    confidence: float                      # 0-1 confidence in the assessment

    # BBP analysis
    bbp_threshold: float                   # sigma^2 * (1 + sqrt(gamma))^2
    leading_eigenvalue: float              # largest eigenvalue of difference cov
    eigenvalue_margin: float               # leading_eigenvalue - bbp_threshold
    n_eigenvalues_above_threshold: int     # how many eigenvalues exceed BBP

    # Marchenko-Pastur noise floor
    mp_upper_edge: float                   # upper edge of MP distribution
    mp_lower_edge: float                   # lower edge of MP distribution
    noise_variance: float                  # estimated sigma^2

    # Non-isotropic extension
    condition_number: float                # kappa of activation covariance
    isotropic_threshold: float             # BBP threshold assuming isotropy
    anisotropic_threshold: float           # corrected threshold for anisotropy
    anisotropy_correction: float           # ratio anisotropic/isotropic

    # Signal analysis
    signal_dimensions: int                 # number of refusal signal dimensions
    signal_energy: float                   # total signal energy above noise floor
    noise_energy: float                    # total noise energy
    signal_to_noise_ratio: float           # SNR of residual refusal

    # Distributed refusal detection
    is_distributed: bool                   # whether refusal is distributed
    n_weak_dimensions: int                 # dimensions with weak but present signal
    distributed_total_energy: float        # total energy in weak dimensions

    # Sample requirements
    n_samples_used: int                    # samples used for this analysis
    n_samples_required: int                # minimum samples for reliable detection
    is_sample_sufficient: bool             # whether we have enough data

    # Recommendations
    recommendation: str                    # human-readable recommendation
    suggested_action: str                  # "none" | "more_directions" | "grp_obliteration" | "more_samples"


@dataclass
class EigenvalueAnalysis:
    """Detailed eigenvalue decomposition of the residual covariance."""

    eigenvalues: torch.Tensor              # all eigenvalues (descending)
    eigenvectors: torch.Tensor             # corresponding eigenvectors
    above_threshold: list[int]             # indices above BBP threshold
    in_bulk: list[int]                     # indices within MP bulk
    signal_subspace_dim: int               # dimension of signal subspace


class SpectralCertifier:
    """Certify abliteration completeness via random matrix theory.

    Uses the BBP phase transition and Marchenko-Pastur distribution
    to assess whether a detectable residual refusal signal exists in the
    post-abliteration model. This is a finite-sample diagnostic, not a formal
    guarantee; results are inconclusive when the sample requirement is unmet.
    """

    def __init__(
        self,
        confidence_level: float = 0.95,
        distribution_threshold: float = 0.3,
        min_samples: int = 30,
    ):
        """
        Args:
            confidence_level: Confidence level for statistical tests (0-1).
            distribution_threshold: Energy fraction threshold for detecting
                distributed refusal (Yellow tier).
            min_samples: Minimum samples for reliable spectral analysis.
        """
        self.confidence_level = confidence_level
        self.distribution_threshold = distribution_threshold
        self.min_samples = min_samples

    def certify(
        self,
        harmful_activations: torch.Tensor,
        harmless_activations: torch.Tensor,
        layer_idx: int = -1,
    ) -> SpectralCertificate:
        """Certify abliteration completeness for one layer.

        Args:
            harmful_activations: (n_harmful, hidden_dim) post-abliteration
                activations on harmful prompts.
            harmless_activations: (n_harmless, hidden_dim) post-abliteration
                activations on harmless prompts.
            layer_idx: Layer index (for logging).

        Returns:
            SpectralCertificate with a statistical classification, or an
            inconclusive result when the sample requirement is unmet.
        """
        n_h, d = harmful_activations.shape
        n_b = harmless_activations.shape[0]
        n = n_h + n_b

        # Step 1: Compute difference covariance matrix
        # Pooled covariance minus individual covariances
        harmful_mean = harmful_activations.mean(dim=0)
        harmless_mean = harmless_activations.mean(dim=0)

        diff = harmful_mean - harmless_mean
        diff_norm = diff.norm().item()

        # Between-class scatter
        harmful_centered = harmful_activations - harmful_mean
        harmless_centered = harmless_activations - harmless_mean

        # Pooled within-class covariance (standard formula: sum of scatter
        # matrices divided by total degrees of freedom)
        cov_h = harmful_centered.T @ harmful_centered / max(n_h - 1, 1)
        cov_b = harmless_centered.T @ harmless_centered / max(n_b - 1, 1)
        pooled_cov = (cov_h * (n_h - 1) + cov_b * (n_b - 1)) / max(n - 2, 1)

        # Step 2: Estimate noise variance (median eigenvalue method)
        noise_var = self._estimate_noise_variance(pooled_cov, n, d)

        # Step 3: Compute BBP threshold
        gamma = d / max(n, 1)  # aspect ratio

        # Isotropic BBP threshold
        isotropic_threshold = noise_var * (1 + math.sqrt(gamma)) ** 2

        # Non-isotropic correction (OBLITERATUS heuristic extension)
        kappa = self._estimate_condition_number(pooled_cov)
        anisotropic_threshold = isotropic_threshold * math.sqrt(kappa)
        anisotropy_correction = math.sqrt(kappa)

        bbp_threshold = anisotropic_threshold

        # Step 4: Marchenko-Pastur edges
        mp_upper = noise_var * (1 + math.sqrt(gamma)) ** 2
        mp_lower = noise_var * max(0, (1 - math.sqrt(gamma)) ** 2)

        # Step 5: Eigenvalue analysis of between-class covariance
        between_cov = torch.outer(diff, diff)  # rank-1 between-class scatter
        eigen_result = self._eigenvalue_analysis(
            between_cov, bbp_threshold, mp_upper
        )

        # Step 6: Classify certification level
        leading_eig = eigen_result.eigenvalues[0].item() if eigen_result.eigenvalues.numel() > 0 else 0.0
        n_above = len(eigen_result.above_threshold)
        eigenvalue_margin = leading_eig - bbp_threshold

        # Signal analysis
        signal_energy = sum(
            eigen_result.eigenvalues[i].item()
            for i in eigen_result.above_threshold
        )
        total_energy = eigen_result.eigenvalues.sum().item()
        noise_energy = max(0, total_energy - signal_energy)
        snr = signal_energy / max(noise_energy, 1e-10)

        # Distributed refusal detection
        # Look for many weak eigenvalues between MP upper edge and BBP threshold
        weak_dims = [
            i for i in range(len(eigen_result.eigenvalues))
            if mp_upper < eigen_result.eigenvalues[i].item() < bbp_threshold
        ]
        n_weak = len(weak_dims)
        weak_energy = sum(eigen_result.eigenvalues[i].item() for i in weak_dims)
        is_distributed = (
            n_weak > 3 and weak_energy > self.distribution_threshold * total_energy
        )

        # Sample sufficiency check
        # From BBP: need n > d / rho^2 where rho = signal_strength / noise_var
        rho = diff_norm / max(math.sqrt(noise_var), 1e-10)
        n_required = max(self.min_samples, int(d / max(rho ** 2, 0.01)))
        is_sufficient = n >= n_required

        # Certification level. Do not emit a definitive traffic-light result
        # when the diagnostic's own sample requirement is unmet: doing so made
        # tiny, high-dimensional batches appear authoritatively RED or GREEN.
        if not is_sufficient:
            level = CertificationLevel.INCONCLUSIVE
            confidence = min(0.99, self.confidence_level * (n / max(n_required, 1)))
        elif n_above == 0 and not is_distributed:
            level = CertificationLevel.GREEN
            confidence = min(0.99, self.confidence_level * (n / max(n_required, 1)))
        elif is_distributed:
            level = CertificationLevel.YELLOW
            confidence = min(0.95, 0.8 * (n / max(n_required, 1)))
        else:
            level = CertificationLevel.RED
            confidence = min(0.99, self.confidence_level)

        # Recommendations
        if level == CertificationLevel.INCONCLUSIVE:
            recommendation = (
                f"Spectral diagnostic is inconclusive: only {n} samples were "
                f"available, while {n_required} are recommended for reliable "
                f"detection at this dimensionality."
            )
            action = "more_samples"
        elif level == CertificationLevel.GREEN:
            recommendation = (
                f"No detectable linear refusal component was found. "
                f"No linear refusal component with eigenvalue above "
                f"BBP threshold ({bbp_threshold:.4f}) detected."
            )
            action = "none"
        elif level == CertificationLevel.YELLOW:
            recommendation = (
                f"Refusal appears distributed across {n_weak} weak dimensions "
                f"(total energy {weak_energy:.4f}). Extended-refusal defense "
                f"may be active. Consider GRP-Obliteration."
            )
            action = "grp_obliteration"
        else:
            recommendation = (
                f"Abliteration incomplete: {n_above} eigenvalue(s) above "
                f"BBP threshold. Leading eigenvalue {leading_eig:.4f} exceeds "
                f"threshold {bbp_threshold:.4f} by {eigenvalue_margin:.4f}. "
                f"Re-run with more directions."
            )
            action = "more_directions"

        return SpectralCertificate(
            level=level,
            confidence=confidence,
            bbp_threshold=bbp_threshold,
            leading_eigenvalue=leading_eig,
            eigenvalue_margin=eigenvalue_margin,
            n_eigenvalues_above_threshold=n_above,
            mp_upper_edge=mp_upper,
            mp_lower_edge=mp_lower,
            noise_variance=noise_var,
            condition_number=kappa,
            isotropic_threshold=isotropic_threshold,
            anisotropic_threshold=anisotropic_threshold,
            anisotropy_correction=anisotropy_correction,
            signal_dimensions=eigen_result.signal_subspace_dim,
            signal_energy=signal_energy,
            noise_energy=noise_energy,
            signal_to_noise_ratio=snr,
            is_distributed=is_distributed,
            n_weak_dimensions=n_weak,
            distributed_total_energy=weak_energy,
            n_samples_used=n,
            n_samples_required=n_required,
            is_sample_sufficient=is_sufficient,
            recommendation=recommendation,
            suggested_action=action,
        )

    def certify_all_layers(
        self,
        harmful_activations: dict[int, torch.Tensor],
        harmless_activations: dict[int, torch.Tensor],
    ) -> dict[int, SpectralCertificate]:
        """Certify abliteration completeness across all layers.

        Returns a certificate for each layer. Overall certification
        is the worst (most RED) across all layers.
        """
        results = {}
        for layer_idx in sorted(harmful_activations.keys()):
            if layer_idx not in harmless_activations:
                continue
            results[layer_idx] = self.certify(
                harmful_activations[layer_idx],
                harmless_activations[layer_idx],
                layer_idx=layer_idx,
            )
        return results

    def overall_certification(
        self, layer_certificates: dict[int, SpectralCertificate]
    ) -> SpectralCertificate | None:
        """Compute overall certification from per-layer certificates.

        The overall level is the WORST across all layers (most RED).
        """
        if not layer_certificates:
            return None

        # A sufficient RED result remains actionable. Otherwise, any
        # insufficient layer makes the aggregate result inconclusive.
        levels = [c.level for c in layer_certificates.values()]
        if CertificationLevel.RED in levels:
            worst = CertificationLevel.RED
        elif CertificationLevel.INCONCLUSIVE in levels:
            worst = CertificationLevel.INCONCLUSIVE
        elif CertificationLevel.YELLOW in levels:
            worst = CertificationLevel.YELLOW
        else:
            worst = CertificationLevel.GREEN

        # Find the certificate with the worst level
        for cert in layer_certificates.values():
            if cert.level == worst:
                return cert

        return list(layer_certificates.values())[0]

    def _estimate_noise_variance(
        self,
        covariance: torch.Tensor,
        n: int,
        d: int,
    ) -> float:
        """Estimate noise variance using the median eigenvalue method.

        The median eigenvalue of the sample covariance converges to the
        noise variance times a known quantile of the Marchenko-Pastur
        distribution.
        """
        try:
            eigenvalues = torch.linalg.eigvalsh(covariance)
            median_eig = eigenvalues[len(eigenvalues) // 2].item()

            # Correct for MP bias: median of MP distribution
            gamma = d / max(n, 1)
            if gamma < 1:
                # MP median approximation.  The exact MP median requires
                # numerical inversion of the MP CDF; we use the empirical
                # approximation median ≈ (1 - sqrt(gamma))^2 + gamma^(1/3)
                # which is more accurate than the naive 0.5 * upper_edge
                # for small gamma.  Falls back to the simpler formula when
                # gamma is very small.
                mp_median_ratio = (1 - math.sqrt(gamma)) ** 2 + gamma ** (1.0 / 3.0)
                noise_var = median_eig / max(mp_median_ratio, 1e-10)
            else:
                noise_var = median_eig

            return max(noise_var, 1e-10)
        except Exception:
            return 1.0

    def _estimate_condition_number(
        self, covariance: torch.Tensor
    ) -> float:
        """Estimate condition number of the covariance matrix."""
        try:
            eigenvalues = torch.linalg.eigvalsh(covariance)
            pos_eigs = eigenvalues[eigenvalues > 1e-10]
            if len(pos_eigs) < 2:
                return 1.0
            kappa = (pos_eigs[-1] / pos_eigs[0]).item()
            return max(1.0, min(kappa, 1e6))
        except Exception:
            return 1.0

    def _eigenvalue_analysis(
        self,
        between_cov: torch.Tensor,
        bbp_threshold: float,
        mp_upper: float,
    ) -> EigenvalueAnalysis:
        """Analyze eigenvalues of the between-class covariance."""
        try:
            eigenvalues, eigenvectors = torch.linalg.eigh(between_cov)
            # Sort descending
            idx = eigenvalues.argsort(descending=True)
            eigenvalues = eigenvalues[idx]
            eigenvectors = eigenvectors[:, idx]

            above = [i for i, e in enumerate(eigenvalues) if e.item() > bbp_threshold]
            in_bulk = [
                i for i, e in enumerate(eigenvalues)
                if mp_upper * 0.01 < e.item() <= bbp_threshold
            ]
            signal_dim = len(above)

            return EigenvalueAnalysis(
                eigenvalues=eigenvalues,
                eigenvectors=eigenvectors,
                above_threshold=above,
                in_bulk=in_bulk,
                signal_subspace_dim=signal_dim,
            )
        except Exception:
            return EigenvalueAnalysis(
                eigenvalues=torch.tensor([0.0]),
                eigenvectors=torch.zeros(1, 1),
                above_threshold=[],
                in_bulk=[],
                signal_subspace_dim=0,
            )