OBLITERATUS/obliteratus/analysis/leace.py

"""LEACE-inspired direction extraction for refusal concept erasure.

This module implements Fisher's Linear Discriminant (FLD) direction for
concept erasure, inspired by LEACE (Belrose et al. 2023).

IMPORTANT: This is NOT a faithful implementation of LEACE as described in
the paper. Key difference:

  - **True LEACE** uses the *total* covariance Sigma_X for whitening:
        P* = I - W^{-1} P_{W Sigma_XZ} W  where W = Sigma_X^{-1/2}
    For binary concepts, this yields: v = Sigma_X^{-1} delta

  - **This implementation** uses *within-class* covariance S_w:
        v = S_w^{-1} delta
    This is Fisher's Linear Discriminant direction, which maximizes
    class separability relative to within-class spread.

For binary concepts, Sigma_X = S_w + p(1-p) * delta @ delta^T,
so the two directions differ when the between-class scatter is
non-negligible relative to within-class scatter. In high-dimensional
settings (d >> 1) with moderate class separation, the difference
is typically small but non-zero.

The FLD direction is still a strong choice for refusal erasure — it
handles rogue dimensions (high-variance but non-discriminative) better
than plain diff-of-means, and is a closed-form solution with no
iterative optimization.

Advantages over SVD:
    - Within-class normalization prevents high-variance but
      non-discriminative dimensions from dominating
    - No hyperparameters beyond regularization epsilon
    - Closed-form solution (no iterative optimization)

References:
    - Belrose et al. (2023): LEACE: Perfect linear concept erasure in
      closed form. NeurIPS 2023.
    - Ravfogel et al. (2022): RLACE: Adversarial concept erasure
      (iterative precursor to LEACE).
    - Fisher (1936): The use of multiple measurements in taxonomic
      problems. Annals of Eugenics.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass
class LEACEResult:
    """Result of LEACE direction extraction for a single layer."""

    layer_idx: int
    direction: torch.Tensor           # (hidden_dim,) unit vector
    generalized_eigenvalue: float     # lambda from GEP (discriminability)
    within_class_condition: float     # condition number of S_w
    mean_diff_norm: float             # ||mu_1 - mu_0||
    erasure_loss: float               # expected squared distortion from erasure


class LEACEExtractor:
    """Extract refusal directions via Fisher's Linear Discriminant.

    Finds the direction that maximally separates harmful from harmless
    activations relative to within-class variance (v = S_w^{-1} delta).
    See module docstring for how this relates to true LEACE.
    """

    def __init__(
        self,
        regularization_eps: float = 1e-4,
        shrinkage: float = 0.0,
    ):
        """
        Args:
            regularization_eps: Tikhonov regularization for S_w inversion.
                Larger values produce more conservative (but stable) results.
            shrinkage: Ledoit-Wolf shrinkage toward identity (0..1).
                0 = no shrinkage, 1 = full shrinkage to scaled identity.
                Useful when n_samples < hidden_dim.
        """
        self.regularization_eps = regularization_eps
        self.shrinkage = shrinkage

    def extract(
        self,
        harmful_activations: list[torch.Tensor],
        harmless_activations: list[torch.Tensor],
        layer_idx: int = 0,
    ) -> LEACEResult:
        """Extract the LEACE direction for a single layer.

        Args:
            harmful_activations: List of (hidden_dim,) tensors from harmful prompts.
            harmless_activations: List of (hidden_dim,) tensors from harmless prompts.
            layer_idx: Layer index (for metadata).

        Returns:
            LEACEResult with the optimal erasure direction.
        """
        H = torch.stack(harmful_activations).float()   # (n_h, d)
        B = torch.stack(harmless_activations).float()   # (n_b, d)

        if H.dim() == 3:
            H = H.squeeze(1)
        if B.dim() == 3:
            B = B.squeeze(1)

        n_h, d = H.shape
        n_b = B.shape[0]

        # Class-conditional means
        mu_h = H.mean(dim=0)  # (d,)
        mu_b = B.mean(dim=0)  # (d,)

        # Mean difference (between-class direction)
        delta = mu_h - mu_b  # (d,)
        delta_norm = delta.norm().item()

        # Within-class covariance: S_w = (S_h + S_b) / 2
        # where S_h = (H - mu_h)^T (H - mu_h) / (n_h - 1) etc.
        H_centered = H - mu_h.unsqueeze(0)
        B_centered = B - mu_b.unsqueeze(0)

        S_h = (H_centered.T @ H_centered) / max(n_h - 1, 1)
        S_b = (B_centered.T @ B_centered) / max(n_b - 1, 1)
        S_w = (S_h + S_b) / 2.0  # (d, d)

        # Apply Ledoit-Wolf shrinkage if requested
        if self.shrinkage > 0:
            trace_S_w = S_w.trace().item()
            S_w = (1 - self.shrinkage) * S_w + self.shrinkage * (trace_S_w / d) * torch.eye(d, device=S_w.device)

        # Regularize S_w for numerical stability
        S_w_reg = S_w + self.regularization_eps * torch.eye(d, device=S_w.device)

        # Condition number of S_w (for diagnostics)
        try:
            eigs_w = torch.linalg.eigvalsh(S_w_reg)
            eigs_w = eigs_w.clamp(min=0)
            pos_eigs = eigs_w[eigs_w > eigs_w.max() * 1e-10]
            condition = (pos_eigs.max() / pos_eigs.min()).item() if pos_eigs.numel() > 0 else float('inf')
        except Exception:
            condition = float('inf')

        # LEACE direction via S_w^{-1} @ delta
        # The generalized eigenvector for rank-1 S_between = delta @ delta^T
        # reduces to: v = S_w^{-1} @ delta (up to normalization)
        try:
            # Use solve for numerical stability (avoids explicit inverse)
            v = torch.linalg.solve(S_w_reg, delta)  # (d,)
        except torch.linalg.LinAlgError:
            # Fallback: pseudoinverse
            v = torch.linalg.lstsq(S_w_reg, delta.unsqueeze(1)).solution.squeeze(1)

        # Normalize to unit length
        v_norm = v.norm()
        if v_norm > 1e-8:
            direction = v / v_norm
        else:
            # Degenerate case: fall back to normalized mean difference
            direction = delta / max(delta_norm, 1e-8)

        # Generalized eigenvalue: lambda = delta^T @ S_w^{-1} @ delta
        # This measures how discriminable the classes are after whitening
        gen_eigenvalue = (delta @ v).item()

        # Erasure loss: expected squared distortion E[||x - x'||^2]
        # For rank-1 projection: loss = v^T @ S_total @ v where S_total
        # is the total (pooled) covariance
        all_acts = torch.cat([H, B], dim=0)
        mu_total = all_acts.mean(dim=0)
        centered_total = all_acts - mu_total.unsqueeze(0)
        S_total = (centered_total.T @ centered_total) / max(all_acts.shape[0] - 1, 1)
        erasure_loss = (direction @ S_total @ direction).item()

        return LEACEResult(
            layer_idx=layer_idx,
            direction=direction,
            generalized_eigenvalue=gen_eigenvalue,
            within_class_condition=condition,
            mean_diff_norm=delta_norm,
            erasure_loss=erasure_loss,
        )

    def extract_all_layers(
        self,
        harmful_acts: dict[int, list[torch.Tensor]],
        harmless_acts: dict[int, list[torch.Tensor]],
    ) -> dict[int, LEACEResult]:
        """Extract LEACE directions for all layers.

        Args:
            harmful_acts: {layer_idx: [activations]} from activation collection.
            harmless_acts: {layer_idx: [activations]} from activation collection.

        Returns:
            {layer_idx: LEACEResult} 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],
                layer_idx=idx,
            )
        return results

    @staticmethod
    def compare_with_diff_of_means(
        leace_result: LEACEResult,
        harmful_mean: torch.Tensor,
        harmless_mean: torch.Tensor,
    ) -> dict[str, float]:
        """Compare LEACE direction with simple diff-of-means.

        Returns cosine similarity and diagnostic metrics showing how much
        the within-class normalization rotates the direction.
        """
        diff = harmful_mean.squeeze() - harmless_mean.squeeze()
        diff_norm = diff.norm()
        if diff_norm > 1e-8:
            diff_normalized = diff / diff_norm
        else:
            diff_normalized = diff

        cosine_sim = (leace_result.direction @ diff_normalized).abs().item()

        return {
            "cosine_similarity": cosine_sim,
            "leace_eigenvalue": leace_result.generalized_eigenvalue,
            "leace_erasure_loss": leace_result.erasure_loss,
            "within_class_condition": leace_result.within_class_condition,
            "mean_diff_norm": leace_result.mean_diff_norm,
        }