OBLITERATUS/obliteratus/analysis/causal_tracing.py

"""Approximate Causal Importance estimation for refusal circuits.

NOTE: This module provides a *simulation-based approximation* of causal
importance. It does NOT perform real activation patching (which requires
running the model multiple times with interventions). Instead, it estimates
causal effects from pre-collected activations by simulating corruption
with Gaussian noise and measuring how each component's projection onto
the refusal direction would change.

For real causal tracing (Meng et al. 2022), use TransformerLens or
nnsight, which support actual forward passes with patched activations.

What this module DOES provide:
  - **Approximate causal importance**: Estimates which layers contribute
    most to the refusal signal using noise-based sensitivity analysis
  - **Correlation vs importance ranking**: Spearman agreement between
    projection magnitude and estimated causal importance
  - **Silent contributor detection**: Components where projection magnitude
    and estimated importance disagree

What this module does NOT do:
  - Real activation patching (no model forward passes)
  - True counterfactual analysis
  - Edge-level circuit identification (use ACDC for this)

The noise-based approach is a useful first-pass approximation that works
without model access, but its results should be validated with real
causal interventions when model access is available.

References:
    - Meng et al. (2022): Locating and Editing Factual Associations
    - Conmy et al. (2023): Automated Circuit Discovery (ACDC)
    - Wang et al. (2023): Interpretability in the Wild
    - Goldowsky-Dill et al. (2023): Localizing Model Behavior
"""

from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass
class ComponentCausalEffect:
    """Causal effect of a single component."""

    layer_idx: int
    component_type: str  # "attention", "mlp", "full_layer"
    clean_projection: float       # refusal projection in clean run
    corrupted_projection: float   # refusal projection in corrupted run
    restored_projection: float    # refusal projection after patching this component
    causal_effect: float          # how much patching this component restores refusal
    indirect_effect: float        # total - direct effect (mediated through downstream)
    is_causal: bool               # above threshold for causal importance


@dataclass
class CausalTracingResult:
    """Full causal tracing results."""

    n_layers: int
    noise_level: float
    component_effects: list[ComponentCausalEffect]

    # Aggregate metrics
    clean_refusal_strength: float
    corrupted_refusal_strength: float
    total_corruption_effect: float  # clean - corrupted

    # Circuit identification
    causal_components: list[tuple[int, str]]  # (layer, type) pairs above threshold
    circuit_size: int               # number of causally important components
    circuit_fraction: float         # fraction of total components that are causal

    # Correlation vs causation analysis
    correlation_causal_agreement: float  # how well projection predicts causal importance


class CausalRefusalTracer:
    """Identify causally important components for refusal via activation patching.

    Instead of just measuring where the refusal signal is large (correlational),
    this determines which components *actually cause* refusal by intervening
    on individual components and measuring the effect.
    """

    def __init__(
        self,
        noise_level: float = 3.0,
        causal_threshold: float = 0.1,
    ):
        """
        Args:
            noise_level: Standard deviation of Gaussian noise for corruption.
            causal_threshold: Minimum causal effect to classify as "causal".
        """
        self.noise_level = noise_level
        self.causal_threshold = causal_threshold

    def trace_from_activations(
        self,
        clean_activations: dict[int, torch.Tensor],
        refusal_direction: dict[int, torch.Tensor] | torch.Tensor,
        component_types: list[str] | None = None,
    ) -> CausalTracingResult:
        """Perform causal tracing using pre-collected activations.

        This is a simulation-based approach that doesn't require running
        the actual model — it estimates causal effects from the activation
        geometry alone.

        For each component, we estimate: "if we removed this component's
        contribution to the refusal direction, how much would refusal decrease?"

        Args:
            clean_activations: {layer_idx: activation_tensor} from harmful prompt.
            refusal_direction: Per-layer or single refusal direction.
            component_types: Which component types to trace. Default: ["full_layer"].

        Returns:
            CausalTracingResult with causal importance map.
        """
        if component_types is None:
            component_types = ["full_layer"]

        layers = sorted(clean_activations.keys())
        n_layers = len(layers)

        # Normalize refusal directions
        if isinstance(refusal_direction, torch.Tensor):
            ref_dirs = {ly: refusal_direction.float().squeeze() for ly in layers}
        else:
            ref_dirs = {
                ly: refusal_direction[ly].float().squeeze()
                for ly in layers if ly in refusal_direction
            }

        for ly in ref_dirs:
            ref_dirs[ly] = ref_dirs[ly] / ref_dirs[ly].norm().clamp(min=1e-10)

        # Clean projections
        clean_projs = {}
        for ly in layers:
            if ly in ref_dirs:
                act = clean_activations[ly].float().squeeze()
                clean_projs[ly] = (act @ ref_dirs[ly]).item()
            else:
                clean_projs[ly] = 0.0

        clean_strength = sum(abs(v) for v in clean_projs.values()) / max(len(clean_projs), 1)

        # Simulate corruption: add noise to estimate corrupted baseline
        torch.manual_seed(42)
        corrupted_projs = {}
        for ly in layers:
            if ly in ref_dirs:
                act = clean_activations[ly].float().squeeze()
                noise = torch.randn_like(act) * self.noise_level
                corrupted = act + noise
                corrupted_projs[ly] = (corrupted @ ref_dirs[ly]).item()
            else:
                corrupted_projs[ly] = 0.0

        corrupted_strength = sum(abs(v) for v in corrupted_projs.values()) / max(len(corrupted_projs), 1)

        total_corruption = clean_strength - corrupted_strength

        # For each component, estimate causal effect via ablation
        effects = []
        for ly in layers:
            for comp_type in component_types:
                if ly not in ref_dirs:
                    continue

                act = clean_activations[ly].float().squeeze()

                # Clean projection at this layer
                clean_proj = clean_projs[ly]

                # Corrupted projection at this layer
                corrupted_proj = corrupted_projs[ly]

                # Restored projection: patch clean activation back in
                # In the simulation, this means the projection returns to clean value
                restored_proj = clean_proj

                # Causal effect: how much does restoring this component
                # recover the refusal signal (normalized by total corruption)
                if abs(total_corruption) > 1e-10:
                    causal_effect = abs(clean_proj - corrupted_proj) / (
                        abs(total_corruption) * n_layers
                    )
                else:
                    causal_effect = 0.0

                # Indirect effect: contribution mediated through downstream layers
                # Estimate via the projection magnitude relative to total
                total_proj = sum(abs(v) for v in clean_projs.values())
                if total_proj > 1e-10:
                    direct_fraction = abs(clean_proj) / total_proj
                else:
                    direct_fraction = 0.0
                indirect = max(0.0, causal_effect - direct_fraction)

                is_causal = causal_effect > self.causal_threshold

                effects.append(ComponentCausalEffect(
                    layer_idx=ly,
                    component_type=comp_type,
                    clean_projection=clean_proj,
                    corrupted_projection=corrupted_proj,
                    restored_projection=restored_proj,
                    causal_effect=causal_effect,
                    indirect_effect=indirect,
                    is_causal=is_causal,
                ))

        # Identify circuit
        causal_components = [
            (e.layer_idx, e.component_type) for e in effects if e.is_causal
        ]
        total_components = len(effects)
        circuit_fraction = len(causal_components) / max(total_components, 1)

        # Correlation vs causation agreement
        # Compare ranking by projection magnitude vs ranking by causal effect
        agreement = self._rank_agreement(effects)

        return CausalTracingResult(
            n_layers=n_layers,
            noise_level=self.noise_level,
            component_effects=effects,
            clean_refusal_strength=clean_strength,
            corrupted_refusal_strength=corrupted_strength,
            total_corruption_effect=total_corruption,
            causal_components=causal_components,
            circuit_size=len(causal_components),
            circuit_fraction=circuit_fraction,
            correlation_causal_agreement=agreement,
        )

    def identify_silent_contributors(
        self, result: CausalTracingResult, top_k: int = 5,
    ) -> dict[str, list[ComponentCausalEffect]]:
        """Find components where correlational and causal importance disagree.

        "Silent contributors" have high causal effect but low projection.
        "Loud non-contributors" have high projection but low causal effect.

        Args:
            result: CausalTracingResult from trace_from_activations.
            top_k: Number of components to return in each category.

        Returns:
            Dict with "silent_contributors" and "loud_non_contributors".
        """
        effects = result.component_effects
        if not effects:
            return {"silent_contributors": [], "loud_non_contributors": []}

        # Score the discrepancy
        for e in effects:
            # Normalize to [0, 1] ranges
            max_proj = max(abs(x.clean_projection) for x in effects)
            max_causal = max(x.causal_effect for x in effects)

            if max_proj > 0:
                norm_proj = abs(e.clean_projection) / max_proj
            else:
                norm_proj = 0.0
            if max_causal > 0:
                norm_causal = e.causal_effect / max_causal
            else:
                norm_causal = 0.0

            e._norm_proj = norm_proj
            e._norm_causal = norm_causal

        # Silent: high causal, low projection
        silent = sorted(
            effects,
            key=lambda e: e._norm_causal - e._norm_proj,
            reverse=True,
        )[:top_k]

        # Loud: high projection, low causal
        loud = sorted(
            effects,
            key=lambda e: e._norm_proj - e._norm_causal,
            reverse=True,
        )[:top_k]

        # Clean up temporary attributes
        for e in effects:
            if hasattr(e, '_norm_proj'):
                delattr(e, '_norm_proj')
            if hasattr(e, '_norm_causal'):
                delattr(e, '_norm_causal')

        return {
            "silent_contributors": silent,
            "loud_non_contributors": loud,
        }

    def _rank_agreement(self, effects: list[ComponentCausalEffect]) -> float:
        """Compute Spearman-like rank agreement between projection and causal rankings."""
        if len(effects) < 2:
            return 1.0

        # Rank by projection magnitude
        proj_ranked = sorted(
            range(len(effects)),
            key=lambda i: abs(effects[i].clean_projection),
            reverse=True,
        )
        proj_ranks = {idx: rank for rank, idx in enumerate(proj_ranked)}

        # Rank by causal effect
        causal_ranked = sorted(
            range(len(effects)),
            key=lambda i: effects[i].causal_effect,
            reverse=True,
        )
        causal_ranks = {idx: rank for rank, idx in enumerate(causal_ranked)}

        # Spearman correlation
        n = len(effects)
        d_sq_sum = sum(
            (proj_ranks[i] - causal_ranks[i]) ** 2 for i in range(n)
        )
        if n * (n * n - 1) == 0:
            return 1.0
        rho = 1.0 - (6.0 * d_sq_sum) / (n * (n * n - 1))
        return max(-1.0, min(1.0, rho))

    @staticmethod
    def format_tracing_report(result: CausalTracingResult) -> str:
        """Format causal tracing results."""
        lines = []
        lines.append("Causal Tracing — Refusal Circuit Identification")
        lines.append("=" * 50)
        lines.append("")
        lines.append(f"Layers traced: {result.n_layers}")
        lines.append(f"Noise level: {result.noise_level}")
        lines.append(f"Clean refusal strength: {result.clean_refusal_strength:.4f}")
        lines.append(f"Corrupted strength: {result.corrupted_refusal_strength:.4f}")
        lines.append(f"Corruption effect: {result.total_corruption_effect:.4f}")
        lines.append("")
        lines.append(f"Circuit size: {result.circuit_size} / {len(result.component_effects)} "
                      f"({result.circuit_fraction:.0%})")
        lines.append(f"Correlation-causation agreement: {result.correlation_causal_agreement:.3f}")
        lines.append("")

        if result.component_effects:
            lines.append("Top causal components:")
            sorted_effects = sorted(
                result.component_effects,
                key=lambda e: e.causal_effect,
                reverse=True,
            )
            for e in sorted_effects[:10]:
                marker = " [CAUSAL]" if e.is_causal else ""
                lines.append(
                    f"  Layer {e.layer_idx:3d} {e.component_type:10s}  "
                    f"causal={e.causal_effect:.4f}  "
                    f"proj={e.clean_projection:+.4f}{marker}"
                )

        return "\n".join(lines)