"""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, }