OBLITERATUS/obliteratus/analysis/sae_abliteration.py

"""Sparse Autoencoder (SAE) Feature-Level Abliteration.

Standard abliteration projects weight matrices along refusal directions
extracted from raw activation space. But the refusal direction in
activation space is a *linear combination* of many underlying features,
some safety-related and some capability-related. This cross-talk means
that projecting out the refusal direction also damages overlapping
capabilities.

Sparse Autoencoders decompose activations into an overcomplete set of
interpretable features. In this higher-dimensional feature space, refusal
may be represented by a small number of specific features that are
cleanly separable from capability features. By identifying and ablating
only the refusal features, we achieve more precise removal with less
collateral damage.

Approach:
  1. Train a lightweight SAE on residual stream activations at key layers
  2. Encode harmful and harmless activations through the SAE
  3. Identify features with large activation differences (refusal features)
  4. Convert refusal features back to hidden-space directions via the
     decoder weight columns
  5. Use these SAE-derived directions as more precise projection targets

The SAE-derived directions complement SVD-extracted directions by
operating in a disentangled feature space rather than raw activation
space.

References:
    - Cunningham et al. (2023): Sparse Autoencoders Find Highly Interpretable Features
    - Bricken et al. (2023): Towards Monosemanticity (Anthropic)
    - Templeton et al. (2024): Scaling Monosemanticity
    - EMNLP 2025: Understanding Refusal with Sparse Autoencoders
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
from obliteratus import device as dev


@dataclass
class SAERefusalFeatures:
    """Result of SAE-based refusal feature identification."""

    layer_idx: int
    n_features_total: int
    n_refusal_features: int
    refusal_feature_indices: list[int]
    refusal_scores: torch.Tensor          # per-feature refusal score
    sae_directions: torch.Tensor           # (n_refusal, hidden_dim) — directions in hidden space
    variance_explained: float              # how much refusal variance these features capture
    reconstruction_loss: float             # SAE reconstruction quality


class SparseAutoencoder(nn.Module):
    """Lightweight sparse autoencoder for refusal feature extraction.

    Architecture: hidden_dim → expansion * hidden_dim → hidden_dim
    with ReLU activation for sparsity and L1 penalty on the latent.

    The decoder columns of identified refusal features give directions
    in hidden space that can be used for more precise abliteration.
    """

    def __init__(self, hidden_dim: int, expansion: int = 4, tied_weights: bool = True):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_features = expansion * hidden_dim
        self.tied_weights = tied_weights

        # Encoder: hidden → features (overcomplete)
        self.encoder = nn.Linear(hidden_dim, self.n_features, bias=True)
        # Decoder: features → hidden (reconstruct)
        if tied_weights:
            # Tied weights: decoder uses encoder.weight.T directly (no separate param).
            # We only need the decoder bias as a learnable parameter.
            self.decoder_bias = nn.Parameter(torch.zeros(hidden_dim))
        else:
            self.decoder = nn.Linear(self.n_features, hidden_dim, bias=True)

        # Initialize with Kaiming for ReLU
        nn.init.kaiming_uniform_(self.encoder.weight, nonlinearity="relu")
        nn.init.zeros_(self.encoder.bias)
        if not tied_weights:
            nn.init.zeros_(self.decoder.bias)

    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """Encode to sparse feature activations."""
        return torch.relu(self.encoder(x))

    @property
    def decoder_weight(self) -> torch.Tensor:
        """Return the decoder weight matrix (hidden_dim x n_features for untied, or encoder.weight.T)."""
        if self.tied_weights:
            return self.encoder.weight.T
        return self.decoder.weight

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """Decode from features back to hidden space."""
        if self.tied_weights:
            return z @ self.encoder.weight + self.decoder_bias
        return self.decoder(z)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Forward pass returning reconstruction and latent features."""
        z = self.encode(x)
        x_hat = self.decode(z)
        return x_hat, z


def _auto_detect_device(device: str | None = None) -> str:
    """Auto-detect the best available device for SAE training.

    When device is ``None`` or ``"auto"``, selects CUDA if available
    and sufficient free memory exists (>512 MB), otherwise falls back
    to CPU.
    """
    if device is not None and device not in ("auto",):
        return device
    if dev.is_gpu_available():
        try:
            free_mb = dev.get_total_free_gb() * 1024
            if free_mb > 512:
                return dev.get_device()
        except Exception:
            pass
    return "cpu"


def train_sae(
    activations: list[torch.Tensor],
    hidden_dim: int,
    expansion: int = 4,
    n_epochs: int = 50,
    lr: float = 3e-4,
    sparsity_coef: float = 1e-3,
    batch_size: int = 32,
    device: str | None = None,
    test_fraction: float = 0.2,
    patience: int = 5,
    quality_threshold: float = 0.1,
) -> SparseAutoencoder:
    """Train a sparse autoencoder on collected activations.

    Uses reconstruction loss + L1 sparsity penalty with train/test split,
    early stopping on held-out loss, and a reconstruction quality gate.

    Args:
        activations: List of activation tensors (each shape: (hidden_dim,) or (1, hidden_dim))
        hidden_dim: Model hidden dimension
        expansion: Feature expansion factor (features = expansion * hidden_dim)
        n_epochs: Training epochs
        lr: Learning rate
        sparsity_coef: L1 sparsity penalty weight
        batch_size: Mini-batch size
        device: Training device. ``None`` or ``"auto"`` to auto-detect
            (CUDA when available with sufficient free memory, else CPU).
        test_fraction: Fraction of data reserved for held-out validation
        patience: Early stopping patience (epochs without improvement)
        quality_threshold: Maximum acceptable held-out reconstruction MSE.
            If the final test loss exceeds this, a warning is emitted
            indicating the SAE directions may be unreliable.
    """
    import warnings

    device = _auto_detect_device(device)

    # Stack and normalize activations
    X = torch.stack([a.squeeze() for a in activations]).float().to(device)
    mean = X.mean(dim=0, keepdim=True)
    X = X - mean  # center activations

    # ── Train/test split ───────────────────────────────────────────
    n_samples = X.shape[0]
    n_test = max(1, int(n_samples * test_fraction))
    n_train = n_samples - n_test
    perm = torch.randperm(n_samples, device=device)
    X_train = X[perm[:n_train]]
    X_test = X[perm[n_train:]]

    sae = SparseAutoencoder(hidden_dim, expansion).to(device)
    optimizer = torch.optim.Adam(sae.parameters(), lr=lr)

    best_test_loss = float("inf")
    best_state = None
    epochs_without_improvement = 0

    for epoch in range(n_epochs):
        # ── Training ───────────────────────────────────────────────
        sae.train()
        train_perm = torch.randperm(n_train, device=device)
        X_shuffled = X_train[train_perm]

        epoch_loss = 0.0
        n_batches = 0
        for i in range(0, n_train, batch_size):
            batch = X_shuffled[i : i + batch_size]
            x_hat, z = sae(batch)

            recon_loss = (batch - x_hat).pow(2).mean()
            sparsity_loss = z.abs().mean()
            loss = recon_loss + sparsity_coef * sparsity_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Normalize decoder columns to unit norm (prevents feature collapse).
            with torch.no_grad():
                if sae.tied_weights:
                    row_norms = sae.encoder.weight.data.norm(dim=1, keepdim=True).clamp(min=1e-8)
                    sae.encoder.weight.data.div_(row_norms)
                else:
                    norms = sae.decoder.weight.data.norm(dim=0, keepdim=True).clamp(min=1e-8)
                    sae.decoder.weight.data.div_(norms)

            epoch_loss += loss.item()
            n_batches += 1

        # ── Held-out validation ────────────────────────────────────
        sae.eval()
        with torch.no_grad():
            x_hat_test, z_test = sae(X_test)
            test_recon = (X_test - x_hat_test).pow(2).mean().item()
            test_sparsity = z_test.abs().mean().item()
            test_loss = test_recon + sparsity_coef * test_sparsity

        # ── Early stopping ─────────────────────────────────────────
        if test_loss < best_test_loss:
            best_test_loss = test_loss
            best_state = {k: v.clone() for k, v in sae.state_dict().items()}
            epochs_without_improvement = 0
        else:
            epochs_without_improvement += 1
            if epochs_without_improvement >= patience:
                break

    # Restore best checkpoint
    if best_state is not None:
        sae.load_state_dict(best_state)
    sae.eval()

    # ── Quality gate ───────────────────────────────────────────────
    with torch.no_grad():
        x_hat_final, _ = sae(X_test)
        final_test_mse = (X_test - x_hat_final).pow(2).mean().item()
    if final_test_mse > quality_threshold:
        warnings.warn(
            f"SAE held-out reconstruction MSE ({final_test_mse:.4f}) exceeds "
            f"quality threshold ({quality_threshold}). SAE-derived refusal "
            f"directions may be unreliable due to overfitting or insufficient "
            f"training data ({n_train} train / {n_test} test samples). "
            f"Consider increasing prompt count or reducing expansion factor.",
            stacklevel=2,
        )

    return sae


def identify_refusal_features(
    sae: SparseAutoencoder,
    harmful_acts: list[torch.Tensor],
    harmless_acts: list[torch.Tensor],
    layer_idx: int,
    top_k: int = 16,
    device: str | None = None,
) -> SAERefusalFeatures:
    """Identify SAE features that encode refusal behavior.

    Compares feature activations between harmful and harmless prompts.
    Features with large positive differences (more active for harmful)
    are candidates for refusal encoding.

    Args:
        sae: Trained sparse autoencoder
        harmful_acts: Activations from harmful prompts
        harmless_acts: Activations from harmless prompts
        layer_idx: Which layer these activations are from
        top_k: Number of top refusal features to return
        device: Computation device. ``None`` or ``"auto"`` to auto-detect.
    """
    device = _auto_detect_device(device)
    sae = sae.to(device)

    with torch.no_grad():
        # Encode both sets — center inputs to match train_sae preprocessing
        X_harm = torch.stack([a.squeeze() for a in harmful_acts]).float().to(device)
        X_safe = torch.stack([a.squeeze() for a in harmless_acts]).float().to(device)

        # Center using pooled mean (same centering used in train_sae)
        X_all = torch.cat([X_harm, X_safe], dim=0)
        mean = X_all.mean(dim=0, keepdim=True)
        X_harm = X_harm - mean
        X_safe = X_safe - mean

        z_harm = sae.encode(X_harm)  # (n_harmful, n_features)
        z_safe = sae.encode(X_safe)  # (n_harmless, n_features)

        # Per-feature mean activation difference
        mean_harm = z_harm.mean(dim=0)  # (n_features,)
        mean_safe = z_safe.mean(dim=0)  # (n_features,)
        diff = mean_harm - mean_safe     # positive = more active for harmful

        # Z-score normalization: use pooled std for significance
        pooled = torch.cat([z_harm, z_safe], dim=0)
        std = pooled.std(dim=0).clamp(min=1e-8)
        z_scores = diff / std

        # Select top-k features by POSITIVE z-score only.
        # Positive z = more active for harmful prompts = refusal features.
        # Using abs() would also select anti-refusal features (negative z),
        # and projecting those out would INCREASE refusal.
        top_k = min(top_k, z_scores.shape[0])
        _, top_indices = z_scores.topk(top_k)
        refusal_indices = top_indices.cpu().tolist()

        # Extract directions from decoder columns
        # Each decoder column is the hidden-space direction for a feature.
        # decoder_weight shape is always (hidden_dim, n_features) regardless
        # of tied/untied mode.
        dec_w = sae.decoder_weight.data  # (hidden_dim, n_features)
        directions = dec_w[:, top_indices].T  # (top_k, hidden_dim)
        directions = directions / directions.norm(dim=1, keepdim=True).clamp(min=1e-8)

        # Compute variance explained
        # Project harmful-harmless diff onto SAE directions
        raw_diff = (X_harm.mean(0) - X_safe.mean(0))
        raw_diff_norm = raw_diff.norm().item()
        if raw_diff_norm > 0:
            # How much of the raw difference is captured by SAE directions
            projected = torch.zeros_like(raw_diff)
            for d in directions:
                comp = (raw_diff @ d) * d
                projected = projected + comp
            var_explained = projected.norm().item() / raw_diff_norm
        else:
            var_explained = 0.0

        # Reconstruction quality
        x_hat_harm, _ = sae(X_harm)
        recon_loss = (X_harm - x_hat_harm).pow(2).mean().item()

    return SAERefusalFeatures(
        layer_idx=layer_idx,
        n_features_total=sae.n_features,
        n_refusal_features=top_k,
        refusal_feature_indices=refusal_indices,
        refusal_scores=z_scores[top_indices].cpu(),
        sae_directions=directions.cpu(),
        variance_explained=min(var_explained, 1.0),
        reconstruction_loss=recon_loss,
    )


# ---------------------------------------------------------------------------
#  Enhanced SAE Decomposition Pipeline
# ---------------------------------------------------------------------------

@dataclass
class FeatureClusterResult:
    """Result of clustering SAE features into semantic groups."""

    n_clusters: int
    cluster_labels: list[int]           # cluster assignment per refusal feature
    cluster_directions: torch.Tensor    # (n_clusters, hidden_dim) mean directions
    cluster_strengths: list[float]      # per-cluster mean refusal score
    silhouette_score: float             # clustering quality (-1 to 1)


@dataclass
class SAEDecompositionResult:
    """Full decomposition pipeline result."""

    layer_idx: int
    sae: SparseAutoencoder
    refusal_features: SAERefusalFeatures

    # Feature characterization
    feature_sparsity: list[float]       # L0 sparsity per refusal feature
    feature_monosemanticity: list[float]  # activation consistency scores
    feature_clusters: FeatureClusterResult | None

    # Ablation simulation
    per_feature_refusal_reduction: list[float]  # estimated refusal drop per feature
    cumulative_refusal_reduction: list[float]    # cumulative as features are added

    # Comparison with raw direction
    raw_direction_overlap: float         # cosine with diff-in-means direction
    sae_improvement_estimate: float      # estimated precision improvement


class SAEDecompositionPipeline:
    """Full SAE decomposition pipeline following Anthropic's methodology.

    Extends the basic train-and-identify workflow with:
      1. Feature sparsity and monosemanticity analysis
      2. Feature clustering into semantic groups
      3. Greedy feature ablation simulation
      4. Comparison with raw-direction methods

    References:
        - Bricken et al. (2023): Towards Monosemanticity
        - Cunningham et al. (2023): Sparse Autoencoders Find Interpretable Features
        - Templeton et al. (2024): Scaling Monosemanticity
    """

    def __init__(
        self,
        expansion: int = 4,
        n_epochs: int = 50,
        lr: float = 3e-4,
        sparsity_coef: float = 1e-3,
        top_k_features: int = 16,
        n_clusters: int = 4,
    ):
        self.expansion = expansion
        self.n_epochs = n_epochs
        self.lr = lr
        self.sparsity_coef = sparsity_coef
        self.top_k_features = top_k_features
        self.n_clusters = n_clusters

    def run(
        self,
        harmful_acts: list[torch.Tensor],
        harmless_acts: list[torch.Tensor],
        layer_idx: int = 0,
        device: str | None = None,
    ) -> SAEDecompositionResult:
        """Run the full decomposition pipeline.

        Args:
            harmful_acts: Activations from harmful prompts.
            harmless_acts: Activations from harmless prompts.
            layer_idx: Layer index for metadata.
            device: Computation device. ``None`` or ``"auto"`` to auto-detect.

        Returns:
            SAEDecompositionResult with comprehensive feature analysis.
        """
        device = _auto_detect_device(device)
        all_acts = harmful_acts + harmless_acts
        hidden_dim = harmful_acts[0].squeeze().shape[0]

        # Step 1: Train SAE
        sae = train_sae(
            all_acts, hidden_dim,
            expansion=self.expansion,
            n_epochs=self.n_epochs,
            lr=self.lr,
            sparsity_coef=self.sparsity_coef,
            device=device,
        )

        # Step 2: Identify refusal features
        refusal_features = identify_refusal_features(
            sae, harmful_acts, harmless_acts, layer_idx,
            top_k=self.top_k_features, device=device,
        )

        # Step 3: Compute feature sparsity and monosemanticity
        sparsity, monosemanticity = self._analyze_features(
            sae, harmful_acts, harmless_acts,
            refusal_features.refusal_feature_indices, device,
        )

        # Step 4: Cluster features
        clusters = self._cluster_features(refusal_features)

        # Step 5: Ablation simulation
        per_feat_reduction, cumul_reduction = self._ablation_simulation(
            sae, harmful_acts, harmless_acts,
            refusal_features.refusal_feature_indices, device,
        )

        # Step 6: Compare with raw direction
        raw_overlap = self._compare_raw_direction(
            harmful_acts, harmless_acts, refusal_features.sae_directions,
        )

        # Estimate improvement: higher variance explained with sparser intervention
        improvement = refusal_features.variance_explained * (1.0 - raw_overlap)

        return SAEDecompositionResult(
            layer_idx=layer_idx,
            sae=sae,
            refusal_features=refusal_features,
            feature_sparsity=sparsity,
            feature_monosemanticity=monosemanticity,
            feature_clusters=clusters,
            per_feature_refusal_reduction=per_feat_reduction,
            cumulative_refusal_reduction=cumul_reduction,
            raw_direction_overlap=raw_overlap,
            sae_improvement_estimate=improvement,
        )

    def _analyze_features(
        self,
        sae: SparseAutoencoder,
        harmful_acts: list[torch.Tensor],
        harmless_acts: list[torch.Tensor],
        feature_indices: list[int],
        device: str,
    ) -> tuple[list[float], list[float]]:
        """Compute per-feature sparsity and monosemanticity scores."""
        all_acts = harmful_acts + harmless_acts
        X = torch.stack([a.squeeze() for a in all_acts]).float().to(device)

        with torch.no_grad():
            z = sae.encode(X)  # (n_samples, n_features)

        sparsity_scores = []
        mono_scores = []

        for idx in feature_indices:
            feat_acts = z[:, idx]  # (n_samples,)

            # L0 sparsity: fraction of samples where feature is active
            l0 = (feat_acts > 0.01).float().mean().item()
            sparsity_scores.append(l0)

            # Monosemanticity: how consistently the feature activates
            # for one class vs the other
            n_harm = len(harmful_acts)
            harm_acts = feat_acts[:n_harm]
            safe_acts = feat_acts[n_harm:]

            harm_mean = harm_acts.mean().item()
            safe_mean = safe_acts.mean().item()

            # Monosemanticity = |harm_mean - safe_mean| / (pooled_std + eps)
            pooled_std = feat_acts.std().item() + 1e-8
            mono = abs(harm_mean - safe_mean) / pooled_std
            mono_scores.append(min(mono, 5.0))  # cap at 5

        return sparsity_scores, mono_scores

    def _cluster_features(
        self, refusal_features: SAERefusalFeatures,
    ) -> FeatureClusterResult | None:
        """Cluster refusal features by direction similarity."""
        directions = refusal_features.sae_directions  # (k, hidden_dim)
        k = directions.shape[0]

        if k < 2:
            return None

        n_clusters = min(self.n_clusters, k)

        # Cosine similarity matrix
        cos_sim = directions @ directions.T  # (k, k)

        # Simple k-means-like clustering in direction space
        # Initialize centroids from most dissimilar features
        labels = [0] * k
        centroids = [directions[0]]

        for c in range(1, n_clusters):
            # Pick the feature most dissimilar to existing centroids
            min_sims = []
            for i in range(k):
                max_sim = max(
                    abs((directions[i] @ cent).item())
                    for cent in centroids
                )
                min_sims.append(max_sim)
            new_idx = min(range(k), key=lambda i: min_sims[i])
            centroids.append(directions[new_idx])

        # Assign features to nearest centroid (5 iterations)
        for _ in range(5):
            centroid_stack = torch.stack(centroids)  # (n_clusters, hidden_dim)
            sims = (directions @ centroid_stack.T).abs()  # (k, n_clusters)
            labels = sims.argmax(dim=1).tolist()

            # Recompute centroids
            new_centroids = []
            for c in range(n_clusters):
                members = [i for i, lbl in enumerate(labels) if lbl == c]
                if members:
                    cent = directions[members].mean(dim=0)
                    cent = cent / cent.norm().clamp(min=1e-8)
                    new_centroids.append(cent)
                else:
                    new_centroids.append(centroids[c])
            centroids = new_centroids

        cluster_dirs = torch.stack(centroids)
        cluster_strengths = []
        for c in range(n_clusters):
            members = [i for i, lbl in enumerate(labels) if lbl == c]
            if members:
                strength = refusal_features.refusal_scores[members].abs().mean().item()
            else:
                strength = 0.0
            cluster_strengths.append(strength)

        # Silhouette score approximation
        sil = self._silhouette_approx(cos_sim, labels, n_clusters)

        return FeatureClusterResult(
            n_clusters=n_clusters,
            cluster_labels=labels,
            cluster_directions=cluster_dirs,
            cluster_strengths=cluster_strengths,
            silhouette_score=sil,
        )

    def _silhouette_approx(
        self, cos_sim: torch.Tensor, labels: list[int], n_clusters: int,
    ) -> float:
        """Approximate silhouette score from cosine similarity matrix."""
        k = cos_sim.shape[0]
        if k < 2 or n_clusters < 2:
            return 0.0

        scores = []
        for i in range(k):
            # Intra-cluster similarity
            same = [j for j in range(k) if labels[j] == labels[i] and j != i]
            if same:
                a_i = 1.0 - cos_sim[i, same].abs().mean().item()  # distance
            else:
                a_i = 0.0

            # Nearest other cluster distance
            b_i = float('inf')
            for c in range(n_clusters):
                if c == labels[i]:
                    continue
                others = [j for j in range(k) if labels[j] == c]
                if others:
                    dist = 1.0 - cos_sim[i, others].abs().mean().item()
                    b_i = min(b_i, dist)

            if b_i == float('inf'):
                b_i = 0.0

            denom = max(a_i, b_i)
            if denom > 0:
                scores.append((b_i - a_i) / denom)
            else:
                scores.append(0.0)

        return sum(scores) / len(scores)

    def _ablation_simulation(
        self,
        sae: SparseAutoencoder,
        harmful_acts: list[torch.Tensor],
        harmless_acts: list[torch.Tensor],
        feature_indices: list[int],
        device: str,
    ) -> tuple[list[float], list[float]]:
        """Simulate ablating refusal features one at a time."""
        X_harm = torch.stack([a.squeeze() for a in harmful_acts]).float().to(device)
        X_safe = torch.stack([a.squeeze() for a in harmless_acts]).float().to(device)

        with torch.no_grad():
            z_harm = sae.encode(X_harm)
            z_safe = sae.encode(X_safe)

            # Baseline refusal signal in feature space
            diff_baseline = (z_harm.mean(0) - z_safe.mean(0))
            baseline_signal = diff_baseline.norm().item()

        per_feat = []
        cumulative = []
        ablated_indices = set()

        for idx in feature_indices:
            with torch.no_grad():
                # Zero out this feature
                z_harm_mod = z_harm.clone()
                z_harm_mod[:, idx] = 0.0

                diff_mod = (z_harm_mod.mean(0) - z_safe.mean(0))
                mod_signal = diff_mod.norm().item()

                reduction = (baseline_signal - mod_signal) / max(baseline_signal, 1e-10)
                per_feat.append(max(0.0, reduction))

            ablated_indices.add(idx)
            with torch.no_grad():
                z_harm_cumul = z_harm.clone()
                for ai in ablated_indices:
                    z_harm_cumul[:, ai] = 0.0
                diff_cumul = (z_harm_cumul.mean(0) - z_safe.mean(0))
                cumul_signal = diff_cumul.norm().item()
                cumul_reduction = (baseline_signal - cumul_signal) / max(baseline_signal, 1e-10)
                cumulative.append(max(0.0, cumul_reduction))

        return per_feat, cumulative

    def _compare_raw_direction(
        self,
        harmful_acts: list[torch.Tensor],
        harmless_acts: list[torch.Tensor],
        sae_directions: torch.Tensor,
    ) -> float:
        """Compare SAE-derived directions with the raw diff-in-means direction."""
        H = torch.stack([a.squeeze() for a in harmful_acts]).float()
        B = torch.stack([a.squeeze() for a in harmless_acts]).float()

        raw_diff = H.mean(0) - B.mean(0)
        raw_dir = raw_diff / raw_diff.norm().clamp(min=1e-8)

        # Max cosine similarity between raw direction and any SAE direction
        if sae_directions.shape[0] == 0:
            return 0.0

        cosines = (sae_directions @ raw_dir).abs()
        return cosines.max().item()

    @staticmethod
    def format_report(result: SAEDecompositionResult) -> str:
        """Format full decomposition pipeline results."""
        lines = []
        lines.append("SAE Feature Decomposition Pipeline")
        lines.append("=" * 36)
        lines.append("")

        rf = result.refusal_features
        lines.append(f"Layer: {result.layer_idx}")
        lines.append(f"Total SAE features: {rf.n_features_total}")
        lines.append(f"Refusal features identified: {rf.n_refusal_features}")
        lines.append(f"Variance explained: {rf.variance_explained:.1%}")
        lines.append(f"Reconstruction loss: {rf.reconstruction_loss:.6f}")
        lines.append(f"Raw direction overlap: {result.raw_direction_overlap:.3f}")
        lines.append(f"Estimated improvement: {result.sae_improvement_estimate:.3f}")
        lines.append("")

        # Per-feature analysis
        lines.append("Top refusal features:")
        for i, idx in enumerate(rf.refusal_feature_indices[:10]):
            score = rf.refusal_scores[i].item()
            sp = result.feature_sparsity[i] if i < len(result.feature_sparsity) else 0
            mono = result.feature_monosemanticity[i] if i < len(result.feature_monosemanticity) else 0
            red = result.per_feature_refusal_reduction[i] if i < len(result.per_feature_refusal_reduction) else 0
            lines.append(
                f"  Feature {idx:5d}: score={score:+.3f}  "
                f"sparsity={sp:.2f}  mono={mono:.2f}  "
                f"reduction={red:.1%}"
            )

        if result.cumulative_refusal_reduction:
            lines.append("")
            lines.append(f"Cumulative refusal reduction (all {rf.n_refusal_features} features): "
                         f"{result.cumulative_refusal_reduction[-1]:.1%}")

        if result.feature_clusters:
            fc = result.feature_clusters
            lines.append("")
            lines.append(f"Feature clusters: {fc.n_clusters} (silhouette={fc.silhouette_score:.3f})")
            for c in range(fc.n_clusters):
                n_members = sum(1 for lbl in fc.cluster_labels if lbl == c)
                lines.append(f"  Cluster {c}: {n_members} features, strength={fc.cluster_strengths[c]:.3f}")

        return "\n".join(lines)