mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
482 lines
19 KiB
Python
482 lines
19 KiB
Python
"""Conditional Abliteration with Category-Selective Projection Fields.
|
|
|
|
Standard abliteration is all-or-nothing: it removes ALL refusal, including
|
|
legitimate safety boundaries. CAST (Lee et al., ICLR 2025 Spotlight) showed
|
|
that condition vectors can selectively gate activation steering at inference
|
|
time, but CAST doesn't modify weights.
|
|
|
|
This module synthesizes CAST's conditional gating with abliteration's weight
|
|
surgery. For each harm category c, we learn a category-specific projection
|
|
operator P_c. The key algebraic structure: the family {P_c} forms a *sheaf*
|
|
over the category lattice — projectors for parent categories consistently
|
|
restrict to child categories.
|
|
|
|
Contributions:
|
|
1. **Category-selective projectors**: Per-category projection operators
|
|
that remove refusal only for matched categories
|
|
2. **Condition vector extraction**: Learn category signatures in
|
|
activation space that gate projector application
|
|
3. **Sheaf consistency**: Prove hierarchical consistency — abliterating
|
|
"violence" equals union of "weapons" + "assault" + "threats"
|
|
4. **Selective abliteration**: Weight-level conditional surgery
|
|
|
|
References:
|
|
- Lee et al. (ICLR 2025): CAST — Conditional Activation Steering
|
|
- Wollschlager et al. (2025): Geometry of Concepts in LLMs (arXiv:2502.17420)
|
|
- Yeo et al. (EMNLP 2025): Understanding Refusal with SAEs (Findings of EMNLP)
|
|
- Cracken AI (2025): Domain-specific abliteration on Kimi K2
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class CategoryProjector:
|
|
"""A category-specific projection operator for selective abliteration."""
|
|
|
|
category: str # harm category name
|
|
condition_vector: torch.Tensor # (hidden_dim,) activation pattern for this category
|
|
projection_direction: torch.Tensor # (hidden_dim,) category-specific refusal direction
|
|
selectivity: float # how specifically this targets one category (0-1)
|
|
activation_threshold: float # cosine sim threshold for condition matching
|
|
refusal_removal_rate: float # estimated refusal removal for matched inputs
|
|
collateral_damage: float # estimated refusal removal for non-matched inputs
|
|
|
|
|
|
@dataclass
|
|
class ConditionalAbliterationResult:
|
|
"""Result of conditional abliteration analysis."""
|
|
|
|
# Category projectors
|
|
n_categories: int
|
|
projectors: list[CategoryProjector]
|
|
category_names: list[str]
|
|
|
|
# Sheaf consistency
|
|
sheaf_consistency_score: float # 0-1, how well projectors compose hierarchically
|
|
max_inconsistency: float # worst case hierarchical inconsistency
|
|
consistency_violations: list[str] # descriptions of consistency violations
|
|
|
|
# Selectivity metrics
|
|
mean_selectivity: float # average category selectivity
|
|
min_selectivity: float # worst case (least selective projector)
|
|
cross_category_leakage: torch.Tensor # (n_cat, n_cat) leakage matrix
|
|
|
|
# Geometric structure
|
|
projector_angles: torch.Tensor # (n_cat, n_cat) angles between projector directions
|
|
condition_angles: torch.Tensor # (n_cat, n_cat) angles between condition vectors
|
|
orthogonality_score: float # how orthogonal the category subspaces are
|
|
|
|
# Recommendation
|
|
viable_categories: list[str] # categories where selective abliteration is safe
|
|
risky_categories: list[str] # categories with high collateral damage
|
|
|
|
|
|
class ConditionalAbliterator:
|
|
"""Learn category-selective projection fields for conditional abliteration.
|
|
|
|
Instead of removing all refusal indiscriminately, this module learns
|
|
per-category projectors that can be selectively applied based on
|
|
input content. Each projector has a condition vector (what activates it)
|
|
and a projection direction (what it removes).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
selectivity_threshold: float = 0.7,
|
|
condition_threshold: float = 0.3,
|
|
min_samples_per_category: int = 5,
|
|
):
|
|
"""
|
|
Args:
|
|
selectivity_threshold: Minimum selectivity for a projector to
|
|
be considered viable (below this, too much collateral).
|
|
condition_threshold: Cosine similarity threshold for condition
|
|
vector matching.
|
|
min_samples_per_category: Minimum harmful samples per category
|
|
to learn a reliable projector.
|
|
"""
|
|
self.selectivity_threshold = selectivity_threshold
|
|
self.condition_threshold = condition_threshold
|
|
self.min_samples_per_category = min_samples_per_category
|
|
|
|
def analyze(
|
|
self,
|
|
category_activations: dict[str, torch.Tensor],
|
|
harmless_activations: torch.Tensor,
|
|
global_refusal_direction: torch.Tensor | None = None,
|
|
) -> ConditionalAbliterationResult:
|
|
"""Learn category-selective projectors and analyze their geometry.
|
|
|
|
Args:
|
|
category_activations: {category_name: (n_samples, hidden_dim)}
|
|
activations for each harm category.
|
|
harmless_activations: (n_harmless, hidden_dim) activations on
|
|
harmless prompts.
|
|
global_refusal_direction: Optional pre-computed global refusal
|
|
direction for comparison.
|
|
|
|
Returns:
|
|
ConditionalAbliterationResult with projectors and analysis.
|
|
"""
|
|
categories = sorted(category_activations.keys())
|
|
n_cat = len(categories)
|
|
|
|
if n_cat == 0 or harmless_activations.shape[0] < 2:
|
|
return self._empty_result()
|
|
|
|
harmless_mean = harmless_activations.mean(dim=0)
|
|
|
|
# Step 1: Extract per-category condition vectors and projectors
|
|
projectors: list[CategoryProjector] = []
|
|
valid_categories: list[str] = []
|
|
cat_directions: list[torch.Tensor] = []
|
|
cat_conditions: list[torch.Tensor] = []
|
|
|
|
for cat in categories:
|
|
cat_acts = category_activations[cat]
|
|
if cat_acts.shape[0] < self.min_samples_per_category:
|
|
logger.info(
|
|
"Category '%s' has too few samples (%d < %d), skipping",
|
|
cat, cat_acts.shape[0], self.min_samples_per_category,
|
|
)
|
|
continue
|
|
|
|
# Condition vector: mean activation pattern specific to this category
|
|
# (difference from harmless mean, normalized)
|
|
cat_mean = cat_acts.mean(dim=0)
|
|
condition = cat_mean - harmless_mean
|
|
cond_norm = condition.norm()
|
|
if cond_norm < 1e-8:
|
|
continue
|
|
condition = condition / cond_norm
|
|
|
|
# Category-specific refusal direction: direction that maximally
|
|
# separates this category from harmless, while being orthogonal
|
|
# to other categories' directions
|
|
proj_dir = self._extract_category_direction(
|
|
cat_acts, harmless_activations, cat_directions
|
|
)
|
|
|
|
if proj_dir is None:
|
|
continue
|
|
|
|
# Measure selectivity: how much does this projector affect
|
|
# other categories?
|
|
selectivity, collateral = self._measure_selectivity(
|
|
proj_dir, condition, category_activations, cat,
|
|
harmless_activations
|
|
)
|
|
|
|
# Estimate refusal removal rate
|
|
cat_proj_magnitudes = (cat_acts @ proj_dir).abs().mean().item()
|
|
harmless_proj_magnitudes = (harmless_activations @ proj_dir).abs().mean().item()
|
|
removal_rate = cat_proj_magnitudes / max(
|
|
cat_proj_magnitudes + harmless_proj_magnitudes, 1e-10
|
|
)
|
|
|
|
projectors.append(CategoryProjector(
|
|
category=cat,
|
|
condition_vector=condition,
|
|
projection_direction=proj_dir,
|
|
selectivity=selectivity,
|
|
activation_threshold=self.condition_threshold,
|
|
refusal_removal_rate=removal_rate,
|
|
collateral_damage=collateral,
|
|
))
|
|
|
|
valid_categories.append(cat)
|
|
cat_directions.append(proj_dir)
|
|
cat_conditions.append(condition)
|
|
|
|
n_valid = len(valid_categories)
|
|
if n_valid == 0:
|
|
return self._empty_result()
|
|
|
|
# Step 2: Compute cross-category geometry
|
|
dir_stack = torch.stack(cat_directions) # (n_valid, hidden_dim)
|
|
cond_stack = torch.stack(cat_conditions)
|
|
|
|
# Projector angle matrix
|
|
proj_angles = self._compute_angle_matrix(dir_stack)
|
|
|
|
# Condition angle matrix
|
|
cond_angles = self._compute_angle_matrix(cond_stack)
|
|
|
|
# Cross-category leakage matrix
|
|
leakage = self._compute_leakage_matrix(
|
|
projectors, category_activations, valid_categories
|
|
)
|
|
|
|
# Orthogonality score: mean absolute cosine between projector directions
|
|
if n_valid > 1:
|
|
cos_matrix = dir_stack @ dir_stack.T
|
|
mask = ~torch.eye(n_valid, dtype=torch.bool)
|
|
ortho_score = 1.0 - cos_matrix.abs()[mask].mean().item()
|
|
else:
|
|
ortho_score = 1.0
|
|
|
|
# Step 3: Sheaf consistency check
|
|
consistency, max_incon, violations = self._check_sheaf_consistency(
|
|
projectors, category_activations, harmless_activations
|
|
)
|
|
|
|
# Step 4: Classify categories
|
|
viable = [
|
|
p.category for p in projectors
|
|
if p.selectivity >= self.selectivity_threshold
|
|
]
|
|
risky = [
|
|
p.category for p in projectors
|
|
if p.selectivity < self.selectivity_threshold
|
|
]
|
|
|
|
# Selectivity stats
|
|
selectivities = [p.selectivity for p in projectors]
|
|
mean_sel = sum(selectivities) / len(selectivities) if selectivities else 0.0
|
|
min_sel = min(selectivities) if selectivities else 0.0
|
|
|
|
return ConditionalAbliterationResult(
|
|
n_categories=n_valid,
|
|
projectors=projectors,
|
|
category_names=valid_categories,
|
|
sheaf_consistency_score=consistency,
|
|
max_inconsistency=max_incon,
|
|
consistency_violations=violations,
|
|
mean_selectivity=mean_sel,
|
|
min_selectivity=min_sel,
|
|
cross_category_leakage=leakage,
|
|
projector_angles=proj_angles,
|
|
condition_angles=cond_angles,
|
|
orthogonality_score=ortho_score,
|
|
viable_categories=viable,
|
|
risky_categories=risky,
|
|
)
|
|
|
|
def _extract_category_direction(
|
|
self,
|
|
category_acts: torch.Tensor,
|
|
harmless_acts: torch.Tensor,
|
|
existing_directions: list[torch.Tensor],
|
|
) -> torch.Tensor | None:
|
|
"""Extract category-specific refusal direction.
|
|
|
|
Uses difference-of-means (category_mean - harmless_mean)
|
|
and then orthogonalizes against previously extracted directions
|
|
to ensure category independence.
|
|
"""
|
|
cat_mean = category_acts.mean(dim=0)
|
|
harmless_mean = harmless_acts.mean(dim=0)
|
|
|
|
diff = cat_mean - harmless_mean
|
|
diff_norm = diff.norm()
|
|
if diff_norm < 1e-8:
|
|
return None
|
|
|
|
direction = diff / diff_norm
|
|
|
|
# Orthogonalize against existing category directions
|
|
for existing in existing_directions:
|
|
proj = (direction @ existing) * existing
|
|
direction = direction - proj
|
|
d_norm = direction.norm()
|
|
if d_norm < 1e-8:
|
|
return None
|
|
direction = direction / d_norm
|
|
|
|
return direction
|
|
|
|
def _measure_selectivity(
|
|
self,
|
|
proj_dir: torch.Tensor,
|
|
condition: torch.Tensor,
|
|
category_activations: dict[str, torch.Tensor],
|
|
target_category: str,
|
|
harmless_activations: torch.Tensor,
|
|
) -> tuple[float, float]:
|
|
"""Measure how selectively a projector targets its intended category.
|
|
|
|
Selectivity = 1 - (collateral damage / intended removal)
|
|
Collateral = how much refusal is removed from non-target categories
|
|
"""
|
|
target_acts = category_activations[target_category]
|
|
target_effect = (target_acts @ proj_dir).abs().mean().item()
|
|
|
|
if target_effect < 1e-10:
|
|
return 0.0, 0.0
|
|
|
|
# Measure effect on non-target categories
|
|
collateral_effects = []
|
|
for cat, acts in category_activations.items():
|
|
if cat == target_category:
|
|
continue
|
|
# Check if condition matches (would this projector fire?)
|
|
cat_mean = acts.mean(dim=0)
|
|
harmless_mean = harmless_activations.mean(dim=0)
|
|
cat_condition = cat_mean - harmless_mean
|
|
cond_norm = cat_condition.norm()
|
|
if cond_norm > 1e-8:
|
|
cat_condition = cat_condition / cond_norm
|
|
cos_sim = (cat_condition @ condition).abs().item()
|
|
if cos_sim > self.condition_threshold:
|
|
# This category would trigger the projector
|
|
effect = (acts @ proj_dir).abs().mean().item()
|
|
collateral_effects.append(effect)
|
|
|
|
total_collateral = sum(collateral_effects) if collateral_effects else 0.0
|
|
mean_collateral = (
|
|
total_collateral / len(collateral_effects)
|
|
if collateral_effects
|
|
else 0.0
|
|
)
|
|
|
|
selectivity = max(0.0, 1.0 - mean_collateral / max(target_effect, 1e-10))
|
|
collateral_ratio = mean_collateral / max(target_effect, 1e-10)
|
|
|
|
return selectivity, collateral_ratio
|
|
|
|
def _compute_angle_matrix(self, vectors: torch.Tensor) -> torch.Tensor:
|
|
"""Compute pairwise angle matrix between vectors."""
|
|
norms = vectors.norm(dim=-1, keepdim=True)
|
|
safe_norms = torch.clamp(norms, min=1e-8)
|
|
normalized = vectors / safe_norms
|
|
cos_matrix = normalized @ normalized.T
|
|
cos_matrix = torch.clamp(cos_matrix, -1.0, 1.0)
|
|
angles = torch.acos(cos_matrix.abs()) * (180.0 / math.pi)
|
|
return angles
|
|
|
|
def _compute_leakage_matrix(
|
|
self,
|
|
projectors: list[CategoryProjector],
|
|
category_activations: dict[str, torch.Tensor],
|
|
valid_categories: list[str],
|
|
) -> torch.Tensor:
|
|
"""Compute cross-category leakage matrix.
|
|
|
|
Entry (i,j) = how much projector i affects category j's refusal.
|
|
Diagonal should be high (intended effect), off-diagonal low (leakage).
|
|
"""
|
|
n = len(valid_categories)
|
|
leakage = torch.zeros(n, n)
|
|
|
|
for i, proj in enumerate(projectors):
|
|
for j, cat in enumerate(valid_categories):
|
|
if cat not in category_activations:
|
|
continue
|
|
acts = category_activations[cat]
|
|
effect = (acts @ proj.projection_direction).abs().mean().item()
|
|
leakage[i, j] = effect
|
|
|
|
# Normalize rows by diagonal
|
|
diag = leakage.diag().clone()
|
|
for i in range(n):
|
|
if diag[i] > 1e-10:
|
|
leakage[i] = leakage[i] / diag[i]
|
|
|
|
return leakage
|
|
|
|
def _check_sheaf_consistency(
|
|
self,
|
|
projectors: list[CategoryProjector],
|
|
category_activations: dict[str, torch.Tensor],
|
|
harmless_activations: torch.Tensor,
|
|
) -> tuple[float, float, list[str]]:
|
|
"""Check sheaf consistency of category projectors.
|
|
|
|
The sheaf property requires that for parent category P containing
|
|
child categories C1, C2, ..., the projector for P should be
|
|
consistent with the union of child projectors:
|
|
P_parent ≈ P_c1 + P_c2 + ... (in the projection space)
|
|
|
|
Since we don't have explicit category hierarchy, we check pairwise
|
|
consistency: projecting with P_a then P_b should be similar to
|
|
projecting with P_a+b (combined direction).
|
|
"""
|
|
violations: list[str] = []
|
|
consistencies: list[float] = []
|
|
|
|
n = len(projectors)
|
|
if n < 2:
|
|
return 1.0, 0.0, []
|
|
|
|
for i in range(n):
|
|
for j in range(i + 1, n):
|
|
pi = projectors[i].projection_direction
|
|
pj = projectors[j].projection_direction
|
|
|
|
# Combined direction (unnormalized sum then normalize)
|
|
combined = pi + pj
|
|
c_norm = combined.norm()
|
|
if c_norm < 1e-8:
|
|
continue
|
|
combined = combined / c_norm
|
|
|
|
# Sequential projection should approximate combined projection
|
|
# on the combined category data
|
|
cat_i = projectors[i].category
|
|
cat_j = projectors[j].category
|
|
|
|
acts_i = category_activations.get(cat_i)
|
|
acts_j = category_activations.get(cat_j)
|
|
if acts_i is None or acts_j is None:
|
|
continue
|
|
|
|
combined_acts = torch.cat([acts_i, acts_j], dim=0)
|
|
|
|
# Sequential removal
|
|
seq_residual = combined_acts.clone()
|
|
seq_residual = seq_residual - (seq_residual @ pi).unsqueeze(-1) * pi
|
|
seq_residual = seq_residual - (seq_residual @ pj).unsqueeze(-1) * pj
|
|
|
|
# Combined removal
|
|
comb_residual = combined_acts - (combined_acts @ combined).unsqueeze(-1) * combined
|
|
|
|
# Consistency = cosine similarity of residual patterns
|
|
if seq_residual.norm() > 1e-8 and comb_residual.norm() > 1e-8:
|
|
# Compare mean residuals
|
|
seq_mean = seq_residual.mean(dim=0)
|
|
comb_mean = comb_residual.mean(dim=0)
|
|
consistency = torch.nn.functional.cosine_similarity(
|
|
seq_mean.unsqueeze(0), comb_mean.unsqueeze(0)
|
|
).item()
|
|
consistencies.append(consistency)
|
|
|
|
if consistency < 0.7:
|
|
violations.append(
|
|
f"{cat_i} + {cat_j}: consistency = {consistency:.3f}"
|
|
)
|
|
|
|
if not consistencies:
|
|
return 1.0, 0.0, []
|
|
|
|
mean_consistency = sum(consistencies) / len(consistencies)
|
|
max_inconsistency = 1.0 - min(consistencies)
|
|
|
|
return mean_consistency, max_inconsistency, violations
|
|
|
|
def _empty_result(self) -> ConditionalAbliterationResult:
|
|
return ConditionalAbliterationResult(
|
|
n_categories=0,
|
|
projectors=[],
|
|
category_names=[],
|
|
sheaf_consistency_score=1.0,
|
|
max_inconsistency=0.0,
|
|
consistency_violations=[],
|
|
mean_selectivity=0.0,
|
|
min_selectivity=0.0,
|
|
cross_category_leakage=torch.zeros(1, 1),
|
|
projector_angles=torch.zeros(1, 1),
|
|
condition_angles=torch.zeros(1, 1),
|
|
orthogonality_score=0.0,
|
|
viable_categories=[],
|
|
risky_categories=[],
|
|
)
|