OBLITERATUS/obliteratus/analysis/whitened_svd.py

"""Whitened SVD direction extraction for refusal subspace identification.

Standard SVD on the difference matrix extracts directions that maximize
absolute variance in the harmful-vs-harmless difference. However, some of
this variance may simply reflect the natural anisotropy of the model's
activation space (rogue dimensions with high variance across all inputs).

Whitened SVD normalizes by the harmless covariance matrix first, so the
extracted directions maximize variance *relative to the model's baseline
activation distribution*. This produces cleaner refusal directions that
are less contaminated by general-purpose high-variance dimensions.

Mathematical formulation:
    Given harmful activations H and harmless activations B (both n x d):
    1. Compute harmless covariance: C_B = (B - mu_B)^T (B - mu_B) / (n-1)
    2. Regularize: C_reg = C_B + eps * I  (for numerical stability)
    3. Whitening transform: W = C_reg^{-1/2}
    4. Whiten both sets: H_w = (H - mu_B) @ W,  B_w = (B - mu_B) @ W
    5. Compute whitened difference: D_w = H_w - B_w
    6. SVD on D_w to extract principal whitened refusal directions
    7. Un-whiten to get directions in original activation space

References:
    - Oursland (2024): Whitened activation analysis for LLMs
    - Kessy et al. (2018): Optimal whitening and decorrelation
"""

from __future__ import annotations

from dataclasses import dataclass

import torch

from obliteratus.analysis.numerical_contracts import (
    validate_whitened_parameters,
    validate_whitened_request,
)


@dataclass
class WhitenedSVDResult:
    """Result of whitened SVD extraction for a single layer."""

    layer_idx: int
    directions: torch.Tensor          # (k, hidden_dim) in original space
    whitened_directions: torch.Tensor  # (k, hidden_dim) in whitened space
    singular_values: torch.Tensor      # (k,)
    variance_explained: float          # fraction of total variance
    condition_number: float            # condition number of covariance
    effective_rank: float              # effective rank of covariance


class WhitenedSVDExtractor:
    """Extract refusal directions using covariance-whitened SVD.

    This produces directions that are unusual *relative to* the model's
    baseline activation variance, rather than directions that simply have
    high absolute variance.
    """

    def __init__(
        self,
        regularization_eps: float = 1e-4,
        min_variance_ratio: float = 0.01,
    ):
        """
        Args:
            regularization_eps: Tikhonov regularization added to covariance
                diagonal for numerical stability. Larger values produce more
                conservative whitening.
            min_variance_ratio: Minimum eigenvalue ratio (relative to max)
                below which dimensions are truncated. Prevents amplifying
                noise in near-degenerate dimensions.
        """
        self.regularization_eps, self.min_variance_ratio = validate_whitened_parameters(
            regularization_eps,
            min_variance_ratio,
        )

    @staticmethod
    def _stack_activation_pair(
        harmful_activations: list[torch.Tensor],
        harmless_activations: list[torch.Tensor],
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Validate and stack paired activation samples into two 2D tensors."""
        def normalize(samples: list[torch.Tensor], name: str) -> torch.Tensor:
            normalized = []
            width = None
            device = None
            for index, sample in enumerate(samples):
                if not isinstance(sample, torch.Tensor):
                    raise ValueError(f"{name} activation {index} must be a tensor")
                if sample.dim() == 2 and sample.shape[0] == 1:
                    sample = sample.squeeze(0)
                if sample.dim() != 1 or sample.numel() == 0:
                    raise ValueError(
                        f"{name} activation {index} must be a non-empty vector "
                        "or a single-row matrix",
                    )
                if not torch.isfinite(sample).all():
                    raise ValueError(f"{name} activations must contain only finite values")
                if width is None:
                    width = sample.shape[0]
                    device = sample.device
                elif sample.shape[0] != width:
                    raise ValueError(f"{name} activations must have a consistent width")
                elif sample.device != device:
                    raise ValueError(f"{name} activations must be on a single device")
                normalized.append(sample)
            return torch.stack(normalized).float()

        harmful = normalize(harmful_activations, "harmful")
        harmless = normalize(harmless_activations, "harmless")
        if harmful.shape[1] != harmless.shape[1]:
            raise ValueError(
                "harmful and harmless activations must have the same width, got "
                f"{harmful.shape[1]} and {harmless.shape[1]}",
            )
        if harmful.device != harmless.device:
            raise ValueError("harmful and harmless activations must be on the same device")
        return harmful, harmless

    def extract(
        self,
        harmful_activations: list[torch.Tensor],
        harmless_activations: list[torch.Tensor],
        n_directions: int = 4,
        layer_idx: int = 0,
    ) -> WhitenedSVDResult:
        """Extract whitened refusal directions for a single layer.

        Args:
            harmful_activations: List of (hidden_dim,) tensors, one per prompt.
            harmless_activations: List of (hidden_dim,) tensors, one per prompt.
            n_directions: Number of refusal directions to extract.
            layer_idx: Index of the layer (for metadata).

        Returns:
            WhitenedSVDResult with directions in original activation space.
        """
        n_directions = validate_whitened_request(
            len(harmful_activations),
            len(harmless_activations),
            n_directions,
        )
        H, B = self._stack_activation_pair(harmful_activations, harmless_activations)

        n_samples, d = B.shape

        # Step 1: Compute harmless covariance with centering
        mu_B = B.mean(dim=0, keepdim=True)
        B_centered = B - mu_B
        cov_B = (B_centered.T @ B_centered) / max(n_samples - 1, 1)

        # Step 2: Eigendecompose covariance for whitening
        eigenvalues, eigenvectors = torch.linalg.eigh(cov_B)
        eigenvalues = eigenvalues.clamp(min=0)  # numerical safety

        # Compute condition number using only valid (positive) eigenvalues.
        # After clamping, min_eig is often 0.0 (from numerical noise), which
        # gives a meaningless condition number of ~1e15. Use eigenvalues above
        # a small threshold instead.
        max_eig = eigenvalues.max().item()
        positive_eigs = eigenvalues[eigenvalues > max_eig * 1e-10]
        min_eig = positive_eigs.min().item() if positive_eigs.numel() > 0 else 1e-12
        condition_number = max_eig / max(min_eig, 1e-12)

        # Effective rank via Shannon entropy of normalized eigenvalues
        eig_normalized = eigenvalues / eigenvalues.sum().clamp(min=1e-12)
        eig_nonzero = eig_normalized[eig_normalized > 1e-12]
        effective_rank = torch.exp(-(eig_nonzero * eig_nonzero.log()).sum()).item()

        # Step 3: Truncate near-degenerate dimensions
        threshold = max_eig * self.min_variance_ratio
        valid_mask = eigenvalues >= threshold
        eigenvalues_valid = eigenvalues[valid_mask]
        eigenvectors_valid = eigenvectors[:, valid_mask]

        # Step 4: Compute whitening transform W = V @ diag(1/sqrt(lam + eps)) @ V^T
        # But we work in the truncated eigenspace for efficiency
        inv_sqrt_eig = 1.0 / torch.sqrt(eigenvalues_valid + self.regularization_eps)
        # Whitening projection: x_whitened = (x - mu) @ V_valid @ diag(inv_sqrt)
        whiten_proj = eigenvectors_valid * inv_sqrt_eig.unsqueeze(0)  # (d, k_valid)

        # Step 5: Whiten both activation sets (centered on harmless mean)
        H_centered = H - mu_B
        H_whitened = H_centered @ whiten_proj  # (n, k_valid)
        B_whitened = B_centered @ whiten_proj   # (n, k_valid)

        # Step 6: Compute whitened difference and SVD
        D_whitened = H_whitened - B_whitened  # (n, k_valid)
        if torch.linalg.vector_norm(D_whitened) <= 1e-12:
            raise ValueError("cannot extract a refusal direction without activation difference")

        k = min(n_directions, D_whitened.shape[0], D_whitened.shape[1])
        U, S, Vh = torch.linalg.svd(D_whitened, full_matrices=False)

        whitened_dirs = Vh[:k]  # (k, k_valid) in whitened space
        singular_vals = S[:k]

        # Step 7: Un-whiten to get directions in original activation space
        # x_whitened = x_orig @ whiten_proj, where whiten_proj = V * 1/sqrt(lam)
        # To map a direction v_w from whitened space back to original space,
        # we need the INVERSE whitening: unwhiten_proj = V * sqrt(lam)
        # Then: v_orig = v_w @ unwhiten_proj.T
        unwhiten_proj = eigenvectors_valid * torch.sqrt(
            eigenvalues_valid + self.regularization_eps
        ).unsqueeze(0)
        original_dirs = whitened_dirs @ unwhiten_proj.T  # (k, d)

        # Normalize each direction to unit length
        norms = original_dirs.norm(dim=-1, keepdim=True).clamp(min=1e-8)
        original_dirs = original_dirs / norms

        # Also normalize whitened directions
        w_norms = whitened_dirs.norm(dim=-1, keepdim=True).clamp(min=1e-8)
        whitened_dirs = whitened_dirs / w_norms

        # Variance explained (use S^2: variance is proportional to sigma^2)
        total_var = (S ** 2).sum().item()
        top_k_var = (singular_vals ** 2).sum().item()
        var_explained = top_k_var / max(total_var, 1e-12)

        return WhitenedSVDResult(
            layer_idx=layer_idx,
            directions=original_dirs,
            whitened_directions=whitened_dirs,
            singular_values=singular_vals,
            variance_explained=var_explained,
            condition_number=condition_number,
            effective_rank=effective_rank,
        )

    def extract_all_layers(
        self,
        harmful_acts: dict[int, list[torch.Tensor]],
        harmless_acts: dict[int, list[torch.Tensor]],
        n_directions: int = 4,
    ) -> dict[int, WhitenedSVDResult]:
        """Extract whitened refusal directions for all layers.

        Args:
            harmful_acts: {layer_idx: [activations]} from activation collection.
            harmless_acts: {layer_idx: [activations]} from activation collection.
            n_directions: Number of directions to extract per layer.

        Returns:
            {layer_idx: WhitenedSVDResult} for each layer.
        """
        results = {}
        for idx in sorted(harmful_acts.keys()):
            if idx not in harmless_acts:
                continue
            results[idx] = self.extract(
                harmful_acts[idx],
                harmless_acts[idx],
                n_directions=n_directions,
                layer_idx=idx,
            )
        return results

    @staticmethod
    def compare_with_standard(
        whitened_result: WhitenedSVDResult,
        standard_direction: torch.Tensor,
    ) -> dict[str, float]:
        """Compare whitened vs standard SVD directions.

        Returns cosine similarities between the whitened and standard
        directions, revealing how much the whitening transformation
        rotates the extracted refusal subspace.
        """
        if standard_direction.dim() == 1:
            standard_direction = standard_direction.unsqueeze(0)

        # Ensure unit vectors
        std_norm = standard_direction / standard_direction.norm(dim=-1, keepdim=True).clamp(min=1e-8)
        wht_dirs = whitened_result.directions

        # Primary direction alignment
        primary_cos = (wht_dirs[0] @ std_norm[0]).abs().item()

        # Subspace overlap: average max cosine sim for each whitened dir
        n_w = wht_dirs.shape[0]
        n_s = std_norm.shape[0]
        cos_matrix = (wht_dirs @ std_norm.T).abs()  # (n_w, n_s)

        avg_max_cos = cos_matrix.max(dim=-1).values.mean().item()

        # Subspace principal angle (smallest angle between subspaces)
        if n_w > 1 and n_s > 1:
            _, S_overlap, _ = torch.linalg.svd(wht_dirs @ std_norm.T)
            principal_cos = S_overlap[0].clamp(max=1.0).item()
        else:
            principal_cos = primary_cos

        return {
            "primary_direction_cosine": primary_cos,
            "avg_max_direction_cosine": avg_max_cos,
            "subspace_principal_cosine": principal_cos,
            "whitened_condition_number": whitened_result.condition_number,
            "whitened_effective_rank": whitened_result.effective_rank,
        }