OBLITERATUS/obliteratus/analysis/wasserstein_optimal.py

"""Wasserstein-optimal refusal direction extraction.

Standard abliteration selects r to maximize the harmful-vs-harmless mean
shift (r^T d)^2. But this ignores the distributional cost: projecting out
a direction that has high variance in the harmless distribution causes
large distortion even for harmless inputs.

The Wasserstein-optimal direction minimizes the ratio of distributional
cost to refusal removal effectiveness:

    r* = argmin_{||r||=1} [W_2^2(mu_harmless, mu_projected)] / [(r^T d)^2]

where W_2^2 decomposes into a mean-shift term and a Bures divergence term
(Theorem A.5 in the paper, Appendix A.2).

This reduces to a generalized eigenvalue problem:

    r* = argmin_{||r||=1} [(r^T m)^2 + r^T Sigma r] / [(r^T d)^2]

where m is the harmless mean, Sigma is the harmless covariance, and d is
the harmful-harmless mean difference.

The solution is the eigenvector corresponding to the smallest eigenvalue of:
    (m m^T + Sigma) r = lambda (d d^T) r

In practice, since d d^T is rank-1, we use a Rayleigh quotient approach.

Comparison with other methods:
    - Difference-in-means: maximizes (r^T d)^2 only
    - Whitened SVD (Fisher): maximizes (r^T d)^2 / (r^T Sigma r)
    - Wasserstein-optimal: minimizes [(r^T m)^2 + r^T Sigma r] / (r^T d)^2
      (accounts for both mean shift AND covariance distortion)

The Wasserstein direction should produce lower KL divergence on harmless
prompts than Fisher-optimal, at the cost of slightly weaker refusal removal.

References:
    - Dowson & Landau (1982): The Frechet distance between multivariate normals
    - Givens & Shortt (1984): A class of Wasserstein metrics
    - OBLITERATUS paper Appendix A.2, Corollary A.2
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

import torch

logger = logging.getLogger(__name__)


@dataclass
class WassersteinDirectionResult:
    """Result of Wasserstein-optimal direction extraction for one layer."""

    layer_idx: int
    direction: torch.Tensor             # (hidden_dim,) optimal direction
    wasserstein_cost: float             # W_2^2 cost for this direction
    mean_shift_component: float         # (r^T m)^2 portion
    bures_component: float              # r^T Sigma r portion (exact when r is eigenvector of Sigma, lower bound otherwise)
    refusal_projection: float           # (r^T d)^2
    cost_effectiveness_ratio: float     # W_2^2 / (r^T d)^2


@dataclass
class WassersteinComparisonResult:
    """Comparison of Wasserstein-optimal vs other directions."""

    layer_idx: int
    wasserstein_direction: torch.Tensor
    fisher_direction: torch.Tensor | None
    dim_direction: torch.Tensor | None    # difference-in-means

    wasserstein_cost_ratio: float
    fisher_cost_ratio: float | None
    dim_cost_ratio: float | None

    cosine_wasserstein_fisher: float | None
    cosine_wasserstein_dim: float | None

    improvement_over_fisher: float | None  # % reduction in cost ratio
    improvement_over_dim: float | None


@dataclass
class MultiLayerWassersteinResult:
    """Aggregated Wasserstein-optimal results across layers."""

    per_layer: dict[int, WassersteinDirectionResult]
    best_layer: int
    mean_cost_ratio: float
    comparison: dict[int, WassersteinComparisonResult] | None


class WassersteinOptimalExtractor:
    """Extract Wasserstein-optimal refusal directions.

    Solves the generalized eigenvalue problem that minimizes the 2-Wasserstein
    cost of abliteration on harmless inputs per unit of refusal removed.
    """

    def __init__(
        self,
        regularization_eps: float = 1e-4,
        n_candidates: int = 100,
    ):
        """
        Args:
            regularization_eps: Regularization for covariance matrix.
            n_candidates: Number of candidate directions to evaluate when
                the generalized eigenvalue problem is ill-conditioned.
        """
        self.regularization_eps = regularization_eps
        self.n_candidates = n_candidates

    def extract(
        self,
        harmful_activations: list[torch.Tensor],
        harmless_activations: list[torch.Tensor],
        layer_idx: int = 0,
    ) -> WassersteinDirectionResult:
        """Extract the Wasserstein-optimal refusal direction for one layer.

        Args:
            harmful_activations: List of (hidden_dim,) tensors from harmful prompts.
            harmless_activations: List of (hidden_dim,) tensors from harmless prompts.
            layer_idx: Index of the layer.

        Returns:
            WassersteinDirectionResult with the optimal direction and cost analysis.
        """
        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_b, d = B.shape

        # Compute statistics
        mu_h = H.mean(dim=0)       # harmful mean
        mu_b = B.mean(dim=0)       # harmless mean (m in the formulation)
        diff = mu_h - mu_b          # d in the formulation

        # Harmless covariance
        B_centered = B - mu_b.unsqueeze(0)
        Sigma = (B_centered.T @ B_centered) / max(n_b - 1, 1)
        Sigma = Sigma + self.regularization_eps * torch.eye(d, device=Sigma.device)

        # Cost matrix: C = m m^T + Sigma
        # This is the numerator of our objective
        cost_matrix = mu_b.unsqueeze(1) @ mu_b.unsqueeze(0) + Sigma  # (d, d)

        # Effectiveness matrix: E = d d^T (rank-1)
        # This is the denominator

        # The generalized eigenvalue problem: C r = lambda E r
        # Since E = d d^T is rank-1, we can solve this analytically.
        #
        # For any r, the Rayleigh quotient is:
        #   Q(r) = (r^T C r) / (r^T d)^2
        #
        # The minimum over all r with r^T d != 0 is achieved by:
        #   r* = C^{-1} d / ||C^{-1} d||
        #
        # (This is the standard result for rank-1 denominator GEP)

        # Solve: C^{-1} d
        try:
            C_inv_d = torch.linalg.solve(cost_matrix, diff)
        except RuntimeError:
            # Fallback: use pseudoinverse
            logger.warning("Cost matrix singular, using pseudoinverse at layer %d", layer_idx)
            C_inv_d = torch.linalg.lstsq(cost_matrix, diff.unsqueeze(1)).solution.squeeze(1)

        # Normalize to unit vector
        r_opt = C_inv_d / C_inv_d.norm().clamp(min=1e-10)

        # Compute cost components
        mean_shift = (r_opt @ mu_b).item() ** 2
        bures = (r_opt @ Sigma @ r_opt).item()
        wasserstein_cost = mean_shift + bures
        refusal_proj = (r_opt @ diff).item() ** 2
        cost_ratio = wasserstein_cost / max(refusal_proj, 1e-12)

        return WassersteinDirectionResult(
            layer_idx=layer_idx,
            direction=r_opt,
            wasserstein_cost=wasserstein_cost,
            mean_shift_component=mean_shift,
            bures_component=bures,
            refusal_projection=refusal_proj,
            cost_effectiveness_ratio=cost_ratio,
        )

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

        Args:
            harmful_acts: {layer_idx: [activations]} from harmful prompts.
            harmless_acts: {layer_idx: [activations]} from harmless prompts.

        Returns:
            MultiLayerWassersteinResult with per-layer results.
        """
        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,
            )

        if not results:
            return MultiLayerWassersteinResult(
                per_layer={}, best_layer=0, mean_cost_ratio=0.0, comparison=None,
            )

        best = min(results.items(), key=lambda x: x[1].cost_effectiveness_ratio)
        mean_ratio = sum(r.cost_effectiveness_ratio for r in results.values()) / len(results)

        return MultiLayerWassersteinResult(
            per_layer=results,
            best_layer=best[0],
            mean_cost_ratio=mean_ratio,
            comparison=None,
        )

    def compare_with_alternatives(
        self,
        wasserstein_result: WassersteinDirectionResult,
        harmful_activations: list[torch.Tensor],
        harmless_activations: list[torch.Tensor],
        fisher_direction: torch.Tensor | None = None,
        dim_direction: torch.Tensor | None = None,
    ) -> WassersteinComparisonResult:
        """Compare Wasserstein-optimal direction with Fisher and diff-in-means.

        Args:
            wasserstein_result: Result from extract().
            harmful_activations: Harmful prompt activations.
            harmless_activations: Harmless prompt activations.
            fisher_direction: Direction from whitened SVD (Fisher-optimal).
            dim_direction: Direction from difference-in-means.

        Returns:
            WassersteinComparisonResult with head-to-head comparison.
        """
        H = torch.stack(harmful_activations).float()
        B = torch.stack(harmless_activations).float()
        if H.dim() == 3:
            H = H.squeeze(1)
        if B.dim() == 3:
            B = B.squeeze(1)

        mu_b = B.mean(dim=0)
        mu_h = H.mean(dim=0)
        diff = mu_h - mu_b
        n_b = B.shape[0]
        B_c = B - mu_b.unsqueeze(0)
        Sigma = (B_c.T @ B_c) / max(n_b - 1, 1) + self.regularization_eps * torch.eye(B.shape[1])

        w_dir = wasserstein_result.direction

        def cost_ratio(r):
            r = r.float().squeeze()
            r = r / r.norm().clamp(min=1e-10)
            ms = (r @ mu_b).item() ** 2
            bur = (r @ Sigma @ r).item()
            rp = (r @ diff).item() ** 2
            return (ms + bur) / max(rp, 1e-12)

        w_ratio = wasserstein_result.cost_effectiveness_ratio

        fisher_ratio = None
        cos_wf = None
        imp_fisher = None
        if fisher_direction is not None:
            f = fisher_direction.float().squeeze()
            f = f / f.norm().clamp(min=1e-10)
            fisher_ratio = cost_ratio(f)
            cos_wf = abs((w_dir @ f).item())
            if fisher_ratio > 0:
                imp_fisher = (fisher_ratio - w_ratio) / fisher_ratio * 100

        dim_ratio = None
        cos_wd = None
        imp_dim = None
        if dim_direction is not None:
            dm = dim_direction.float().squeeze()
            dm = dm / dm.norm().clamp(min=1e-10)
            dim_ratio = cost_ratio(dm)
            cos_wd = abs((w_dir @ dm).item())
            if dim_ratio > 0:
                imp_dim = (dim_ratio - w_ratio) / dim_ratio * 100

        return WassersteinComparisonResult(
            layer_idx=wasserstein_result.layer_idx,
            wasserstein_direction=w_dir,
            fisher_direction=fisher_direction,
            dim_direction=dim_direction,
            wasserstein_cost_ratio=w_ratio,
            fisher_cost_ratio=fisher_ratio,
            dim_cost_ratio=dim_ratio,
            cosine_wasserstein_fisher=cos_wf,
            cosine_wasserstein_dim=cos_wd,
            improvement_over_fisher=imp_fisher,
            improvement_over_dim=imp_dim,
        )

    @staticmethod
    def format_report(result: MultiLayerWassersteinResult) -> str:
        """Format Wasserstein-optimal extraction results."""
        lines = []
        lines.append("Wasserstein-Optimal Refusal Direction Extraction")
        lines.append("=" * 50)
        lines.append("")

        if not result.per_layer:
            lines.append("No layers analyzed.")
            return "\n".join(lines)

        lines.append(f"Best layer (lowest cost ratio): {result.best_layer}")
        lines.append(f"Mean cost-effectiveness ratio: {result.mean_cost_ratio:.6f}")
        lines.append("")

        for idx in sorted(result.per_layer.keys()):
            r = result.per_layer[idx]
            lines.append(f"Layer {idx}:")
            lines.append(f"  W2 cost: {r.wasserstein_cost:.6f}")
            lines.append(f"    Mean shift: {r.mean_shift_component:.6f}")
            lines.append(f"    Bures:      {r.bures_component:.6f}")
            lines.append(f"  Refusal projection: {r.refusal_projection:.6f}")
            lines.append(f"  Cost ratio: {r.cost_effectiveness_ratio:.6f}")
            lines.append("")

        return "\n".join(lines)