"""Analysis-Informed Abliteration Pipeline. Closes the feedback loop between OBLITERATUS's 15 analysis modules (#3) and the abliteration pipeline (#2). Instead of running analysis as a standalone post-hoc step, this pipeline runs targeted analysis modules *during* each stage of abliteration to make smarter decisions: SUMMON → load model PROBE → collect activations ANALYZE → run analysis modules to inform excision strategy DISTILL → extract directions using analysis-informed parameters EXCISE → remove refusal with analysis-guided precision VERIFY → post-excision analysis to detect residual refusal REBIRTH → save with comprehensive analysis metadata The ANALYZE stage is the key innovation: it sits between PROBE and DISTILL and uses analysis module outputs to automatically configure the downstream stages. The VERIFY stage also uses analysis modules to detect self-repair (Ouroboros effect) and trigger additional refinement passes if needed. Analysis modules integrated: Stage | Module used | What it informs ------------|------------------------------|------------------------------------------ ANALYZE | AlignmentImprintDetector | Auto-selects method preset (DPO/RLHF/CAI) ANALYZE | ConceptConeAnalyzer | Per-category vs universal direction choice ANALYZE | CrossLayerAlignmentAnalyzer | Smart layer selection (cluster-aware) ANALYZE | SparseDirectionSurgeon | Sparsity-aware projection plan ANALYZE | DefenseRobustnessEvaluator | Ouroboros risk assessment, entanglement map DISTILL | WhitenedSVDExtractor | Covariance-normalized direction extraction EXCISE | SparseDirectionSurgeon | Targeted row-level weight surgery VERIFY | ActivationProbe | Post-excision refusal signal detection VERIFY | CrossLayerAlignmentAnalyzer | Post-excision direction persistence check VERIFY | DefenseRobustnessEvaluator | Self-repair / Ouroboros effect detection VERIFY | SteeringVectorFactory | Pre-screen with steering before permanent changes Novel contributions: - First closed-loop analysis→abliteration pipeline - Alignment-aware auto-tuning: detected training method (DPO/RLHF/CAI) automatically configures projection parameters - Cone-aware excision: polyhedral models get per-category directions, linear models get single universal direction - Cluster-aware layer selection: respects direction cluster boundaries instead of arbitrary top-k selection - Ouroboros-compensated refinement: detects self-repair and adds targeted passes at compensating layers - Entanglement-gated projection: skips highly entangled layers to preserve capabilities """ from __future__ import annotations import logging import time from dataclasses import dataclass, field from pathlib import Path from threading import Event from typing import Callable import torch from obliteratus.abliterate import ( AbliterationPipeline, PipelineFailure, StageResult, ) logger = logging.getLogger(__name__) # ── Analysis-informed method preset ────────────────────────────────────── INFORMED_METHOD = { "label": "Informed (Analysis-Guided)", "description": ( "Runs analysis modules between PROBE and DISTILL to auto-configure " "direction extraction, layer selection, and projection strategy based " "on the model's actual refusal geometry. Defaults to single diff-of-means " "direction + Bayesian optimization (Heretic-style)." ), "n_directions": 1, # overridden by analysis "direction_method": "diff_means", # overridden by analysis; "leace" also available "norm_preserve": True, "regularization": 0.0, # overridden by analysis "refinement_passes": 2, # overridden by analysis "project_biases": True, "use_chat_template": True, "use_whitened_svd": False, # overridden by analysis "true_iterative_refinement": True, } # ── Analysis result containers ─────────────────────────────────────────── @dataclass class AnalysisInsights: """Insights gathered from the ANALYZE stage. These inform every downstream decision in the pipeline. """ # Alignment imprint detected_alignment_method: str = "unknown" alignment_confidence: float = 0.0 alignment_probabilities: dict[str, float] = field(default_factory=dict) # Cone geometry cone_is_polyhedral: bool = False cone_dimensionality: float = 1.0 mean_pairwise_cosine: float = 1.0 per_category_directions: dict[str, torch.Tensor] = field(default_factory=dict) direction_specificity: dict[str, float] = field(default_factory=dict) # Cross-layer structure direction_clusters: list[list[int]] = field(default_factory=list) cluster_count: int = 0 direction_persistence: float = 0.0 cluster_representative_layers: list[int] = field(default_factory=list) # Sparse surgery mean_refusal_sparsity_index: float = 0.0 recommended_sparsity: float = 0.1 use_sparse_surgery: bool = False # Defense robustness estimated_robustness: str = "unknown" self_repair_estimate: float = 0.0 entanglement_score: float = 0.0 entangled_layers: list[int] = field(default_factory=list) clean_layers: list[int] = field(default_factory=list) # Derived configuration recommended_n_directions: int = 1 recommended_direction_method: str = "diff_means" recommended_regularization: float = 0.0 recommended_refinement_passes: int = 2 recommended_layers: list[int] = field(default_factory=list) skip_layers: list[int] = field(default_factory=list) @dataclass class InformedPipelineReport: """Complete report from the informed pipeline.""" insights: AnalysisInsights stages: list[StageResult] = field(default_factory=list) analysis_duration: float = 0.0 total_duration: float = 0.0 ouroboros_passes: int = 0 final_refusal_rate: float = 0.0 # ── The Informed Pipeline ──────────────────────────────────────────────── class InformedAbliterationPipeline(AbliterationPipeline): """Analysis-informed abliteration pipeline. Extends the base AbliterationPipeline with a new ANALYZE stage that runs between PROBE and DISTILL. Analysis module outputs automatically configure the downstream stages for optimal refusal removal with minimal capability damage. Usage: pipeline = InformedAbliterationPipeline( model_name="meta-llama/Llama-3.1-8B-Instruct", output_dir="abliterated_informed", ) result_path, report = pipeline.run_informed() # The report contains all analysis insights print(f"Detected alignment: {report.insights.detected_alignment_method}") print(f"Cone type: {'polyhedral' if report.insights.cone_is_polyhedral else 'linear'}") print(f"Ouroboros passes needed: {report.ouroboros_passes}") """ def __init__( self, model_name: str, output_dir: str = "abliterated_informed", device: str = "auto", dtype: str = "float16", trust_remote_code: bool = True, harmful_prompts: list[str] | None = None, harmless_prompts: list[str] | None = None, on_stage: Callable[[StageResult], None] | None = None, on_log: Callable[[str], None] | None = None, cancellation_event: Event | None = None, max_perplexity_increase: float = 3.0, min_coherence_retention: float = 0.5, max_degenerate_fraction: float = 0.2, # Base pipeline kwargs forwarded to AbliterationPipeline push_to_hub: str | None = None, hub_token: str | None = None, hub_community_org: str | None = None, quantization: str | None = None, # Analysis configuration run_cone_analysis: bool = True, run_alignment_detection: bool = True, run_cross_layer_analysis: bool = True, run_sparse_analysis: bool = True, run_defense_analysis: bool = True, # Ouroboros compensation ouroboros_threshold: float = 0.5, max_ouroboros_passes: int = 3, # Deprecated aliases (kept for backwards compatibility) hydra_threshold: float | None = None, max_hydra_passes: int | None = None, # Entanglement gating entanglement_gate: float = 0.8, # Sparsity control sparse_surgery_threshold: float = 0.5, ): # Initialize base pipeline with informed method preset super().__init__( model_name=model_name, output_dir=output_dir, device=device, dtype=dtype, trust_remote_code=trust_remote_code, method="advanced", # base config, will be overridden harmful_prompts=harmful_prompts, harmless_prompts=harmless_prompts, on_stage=on_stage, on_log=on_log, cancellation_event=cancellation_event, max_perplexity_increase=max_perplexity_increase, min_coherence_retention=min_coherence_retention, max_degenerate_fraction=max_degenerate_fraction, push_to_hub=push_to_hub, hub_token=hub_token, hub_community_org=hub_community_org, quantization=quantization, # Set informed defaults: single direction + Bayesian opt n_directions=1, direction_method="diff_means", norm_preserve=True, project_biases=True, use_chat_template=True, use_whitened_svd=False, true_iterative_refinement=True, use_kl_optimization=True, float_layer_interpolation=True, layer_adaptive_strength=True, winsorize_activations=True, winsorize_percentile=0.01, ) self.method = "informed" # Analysis module flags self._run_cone = run_cone_analysis self._run_alignment = run_alignment_detection self._run_cross_layer = run_cross_layer_analysis self._run_sparse = run_sparse_analysis self._run_defense = run_defense_analysis # Ouroboros compensation parameters self._ouroboros_threshold = hydra_threshold if hydra_threshold is not None else ouroboros_threshold self._max_ouroboros_passes = max_hydra_passes if max_hydra_passes is not None else max_ouroboros_passes # Entanglement gating self._entanglement_gate = entanglement_gate # Sparse surgery self._sparse_threshold = sparse_surgery_threshold # State self._insights = AnalysisInsights() self._report = InformedPipelineReport(insights=self._insights) def run_informed(self) -> tuple[Path, InformedPipelineReport]: """Execute the full analysis-informed pipeline. Returns: (output_path, report) tuple with saved model path and comprehensive analysis report. """ try: return self._run_informed_pipeline() except PipelineFailure: self.cleanup_failed_run() raise def _run_informed_pipeline(self) -> tuple[Path, InformedPipelineReport]: """Execute informed stages while the public method owns cleanup.""" t0 = time.time() # Stage 1: SUMMON self._active_stage = "summon" self._summon() # Prove the pristine checkpoint is healthy and the permanent surgery # layout is supported before analysis can lead to weight mutation. self._active_stage = "baseline" self._capture_stock_baseline() self._validate_architecture_surgery_support() # Stage 2: PROBE self._active_stage = "probe" self._probe() # Stage 3: ANALYZE (new stage — the feedback loop) self._active_stage = "analyze" self._analyze() # Stage 4: DISTILL (informed by analysis) self._active_stage = "distill" self._distill_informed() # Stage 5: EXCISE (informed by analysis) self._active_stage = "excise" self._excise_informed() # Stage 6: VERIFY + Ouroboros compensation loop self._active_stage = "verify" self._verify_and_compensate() # Stage 7: REBIRTH self._active_stage = "rebirth" output_path = self._rebirth_informed() self._report.total_duration = time.time() - t0 return output_path, self._report # ── Stage 3: ANALYZE ───────────────────────────────────────────── def _analyze(self): """Run analysis modules to inform downstream decisions. This is the key innovation: analysis runs BETWEEN probe and distill, so its outputs configure how directions are extracted and excised. """ self._emit("analyze", "running", "Running analysis modules...") t0 = time.time() self.log("=" * 60) self.log("ANALYSIS-INFORMED PIPELINE — ANALYZE STAGE") self.log("=" * 60) # 1. Alignment Imprint Detection if self._run_alignment: self._analyze_alignment_imprint() # 2. Concept Cone Geometry if self._run_cone: self._analyze_cone_geometry() # 3. Cross-Layer Alignment if self._run_cross_layer: self._analyze_cross_layer() # 4. Defense Robustness if self._run_defense: self._analyze_defense_robustness() # 5. Sparse Surgery Analysis (RSI computation) if self._run_sparse: self._analyze_sparsity() # 6. Derive configuration from insights self._derive_configuration() elapsed = time.time() - t0 self._report.analysis_duration = elapsed self.log(f"\nAnalysis complete ({elapsed:.1f}s)") self.log(f" Detected alignment: {self._insights.detected_alignment_method}") self.log(f" Cone type: {'polyhedral' if self._insights.cone_is_polyhedral else 'linear'}") self.log(f" Direction clusters: {self._insights.cluster_count}") self.log(f" Recommended directions: {self._insights.recommended_n_directions}") self.log(f" Recommended regularization: {self._insights.recommended_regularization}") self.log(f" Recommended passes: {self._insights.recommended_refinement_passes}") self.log(f" Layers to skip (entangled): {self._insights.skip_layers}") self._emit( "analyze", "done", f"Analysis complete ({elapsed:.1f}s)", duration=elapsed, ) def _analyze_alignment_imprint(self): """Detect alignment training method from refusal geometry.""" self.log("\n[1/4] Alignment Imprint Detection") self.log("-" * 40) from obliteratus.analysis.alignment_imprint import AlignmentImprintDetector detector = AlignmentImprintDetector() # We need refusal directions for this — compute quick diff-in-means quick_directions = {} for idx in sorted(self._harmful_means.keys()): diff = (self._harmful_means[idx] - self._harmless_means[idx]).squeeze() norm = diff.norm().item() if norm > 1e-10: quick_directions[idx] = diff / diff.norm() if not quick_directions: self.log(" No refusal directions found — skipping alignment detection") return imprint = detector.detect_imprint(quick_directions) self._insights.detected_alignment_method = imprint.predicted_method self._insights.alignment_confidence = imprint.confidence self._insights.alignment_probabilities = { "dpo": imprint.dpo_probability, "rlhf": imprint.rlhf_probability, "cai": imprint.cai_probability, "sft": imprint.sft_probability, } self.log(f" Predicted: {imprint.predicted_method.upper()} " f"(confidence: {imprint.confidence:.1%})") self.log(f" DPO={imprint.dpo_probability:.1%} " f"RLHF={imprint.rlhf_probability:.1%} " f"CAI={imprint.cai_probability:.1%} " f"SFT={imprint.sft_probability:.1%}") self.log(" Geometric features:") self.log(f" Gini coefficient: {imprint.gini_coefficient:.3f}") self.log(f" Effective rank: {imprint.effective_rank:.2f}") self.log(f" Cross-layer smooth: {imprint.cross_layer_smoothness:.3f}") self.log(f" Tail layer bias: {imprint.tail_layer_bias:.3f}") def _analyze_cone_geometry(self): """Analyze concept cone structure to determine per-category vs universal.""" self.log("\n[2/4] Concept Cone Geometry") self.log("-" * 40) from obliteratus.analysis.concept_geometry import ConceptConeAnalyzer analyzer = ConceptConeAnalyzer() # Analyze at layers that are likely strong refusal layers # (middle-to-late layers based on literature) n_layers = len(self._harmful_acts) candidate_layers = list(range(n_layers // 3, int(n_layers * 0.85))) # Sample a subset to keep analysis fast step = max(1, len(candidate_layers) // 6) sample_layers = candidate_layers[::step] polyhedral_count = 0 all_results = [] best_cone_result = None best_strength = 0.0 for layer_idx in sample_layers: if layer_idx not in self._harmful_acts or layer_idx not in self._harmless_acts: continue result = analyzer.analyze_layer( self._harmful_acts[layer_idx], self._harmless_acts[layer_idx], layer_idx=layer_idx, ) all_results.append(result) if result.is_polyhedral: polyhedral_count += 1 # Track the strongest layer's cone analysis for per-category directions general_strength = result.general_direction.norm().item() if result.general_direction.numel() > 1 else 0 if general_strength > best_strength: best_strength = general_strength best_cone_result = result if all_results: # Aggregate cone geometry across sampled layers (majority vote + # mean dimensionality) instead of relying on a single layer. n_sampled = len(all_results) is_polyhedral = polyhedral_count > n_sampled / 2 avg_dimensionality = sum(r.cone_dimensionality for r in all_results) / n_sampled avg_pairwise_cos = sum(r.mean_pairwise_cosine for r in all_results) / n_sampled self._insights.cone_is_polyhedral = is_polyhedral self._insights.cone_dimensionality = avg_dimensionality self._insights.mean_pairwise_cosine = avg_pairwise_cos # Store per-category directions from the strongest layer if best_cone_result is not None: for cd in best_cone_result.category_directions: self._insights.per_category_directions[cd.category] = cd.direction self._insights.direction_specificity[cd.category] = cd.specificity cone_type = "POLYHEDRAL" if is_polyhedral else "LINEAR" self.log(f" Cone type: {cone_type} (majority vote: {polyhedral_count}/{n_sampled} layers)") self.log(f" Avg dimensionality: {avg_dimensionality:.2f}") self.log(f" Avg pairwise cosine: {avg_pairwise_cos:.3f}") if best_cone_result is not None: self.log(f" Categories detected: {best_cone_result.category_count}") for cd in sorted(best_cone_result.category_directions, key=lambda x: -x.strength)[:5]: self.log(f" {cd.category:15s} DSI={cd.specificity:.3f} str={cd.strength:.3f}") else: self.log(" No cone results — using default linear assumption") def _analyze_cross_layer(self): """Analyze cross-layer direction alignment for cluster-aware layer selection.""" self.log("\n[3/4] Cross-Layer Direction Alignment") self.log("-" * 40) from obliteratus.analysis.cross_layer import CrossLayerAlignmentAnalyzer # Compute quick directions for cross-layer analysis quick_directions = {} for idx in sorted(self._harmful_means.keys()): diff = (self._harmful_means[idx] - self._harmless_means[idx]).squeeze() norm = diff.norm().item() if norm > 1e-10: quick_directions[idx] = diff / diff.norm() if len(quick_directions) < 2: self.log(" Too few layers with refusal directions") return analyzer = CrossLayerAlignmentAnalyzer(cluster_threshold=0.85) result = analyzer.analyze(quick_directions) self._insights.direction_clusters = result.clusters self._insights.cluster_count = result.cluster_count self._insights.direction_persistence = result.direction_persistence_score # Select representative layers from each cluster # (the strongest layer per cluster is the best representative) representatives = [] norms = {idx: (self._harmful_means[idx] - self._harmless_means[idx]).squeeze().norm().item() for idx in quick_directions} for cluster in result.clusters: best = max(cluster, key=lambda ly: norms.get(ly, 0)) representatives.append(best) self._insights.cluster_representative_layers = representatives self.log(f" Direction persistence: {result.direction_persistence_score:.3f}") self.log(f" Mean adjacent cosine: {result.mean_adjacent_cosine:.3f}") self.log(f" Direction clusters: {result.cluster_count}") for i, cluster in enumerate(result.clusters): self.log(f" Cluster {i+1}: layers {cluster}") self.log(f" Representative layers: {representatives}") def _analyze_defense_robustness(self): """Assess defense robustness, self-repair risk, and entanglement.""" self.log("\n[4/4] Defense Robustness Assessment") self.log("-" * 40) from obliteratus.analysis.defense_robustness import DefenseRobustnessEvaluator # Temporarily set refusal_directions for the evaluator quick_directions = {} for idx in sorted(self._harmful_means.keys()): diff = (self._harmful_means[idx] - self._harmless_means[idx]).squeeze() norm = diff.norm().item() if norm > 1e-10: quick_directions[idx] = diff / diff.norm() # Store temporarily for the evaluator original_dirs = self.refusal_directions self.refusal_directions = quick_directions evaluator = DefenseRobustnessEvaluator(self) profile = evaluator.profile_defense() emap = evaluator.map_entanglement() # Restore self.refusal_directions = original_dirs self._insights.estimated_robustness = profile.estimated_robustness self._insights.self_repair_estimate = profile.self_repair_estimate self._insights.entanglement_score = profile.entanglement_score self._insights.entangled_layers = emap.most_entangled_layers self._insights.clean_layers = emap.least_entangled_layers self.log(f" Estimated robustness: {profile.estimated_robustness.upper()}") self.log(f" Self-repair estimate: {profile.self_repair_estimate:.2f}") self.log(f" Safety-capability entanglement: {profile.entanglement_score:.3f}") self.log(f" Most entangled layers: {emap.most_entangled_layers}") self.log(f" Cleanest layers: {emap.least_entangled_layers}") def _analyze_sparsity(self): """Compute Refusal Sparsity Index to decide sparse vs dense excision.""" self.log("\n[5/5] Refusal Sparsity Analysis") self.log("-" * 40) from obliteratus.analysis.sparse_surgery import SparseDirectionSurgeon from obliteratus.strategies.utils import ( get_ffn_module, get_layer_modules, ) # Need refusal directions — use quick diff-in-means quick_directions = {} for idx in sorted(self._harmful_means.keys()): diff = (self._harmful_means[idx] - self._harmless_means[idx]).squeeze() norm = diff.norm().item() if norm > 1e-10: quick_directions[idx] = diff / diff.norm() if not quick_directions: self.log(" No refusal directions — skipping sparsity analysis") return # Gather FFN output weights for representative layers (sample for speed) layers = get_layer_modules(self.handle) arch = self.handle.architecture n_layers = len(layers) sample_idxs = sorted(quick_directions.keys()) step = max(1, len(sample_idxs) // 8) sample_idxs = sample_idxs[::step] weights = {} sampled_dirs = {} for idx in sample_idxs: if idx >= n_layers: continue try: ffn = get_ffn_module(layers[idx], arch) for name in ["down_proj", "c_proj", "dense_4h_to_h", "fc_out", "fc2", "w2"]: proj = getattr(ffn, name, None) if proj is not None and hasattr(proj, "weight"): W = proj.weight.data d = quick_directions[idx] if W.shape[-1] == d.shape[0]: weights[idx] = W sampled_dirs[idx] = d break except (AttributeError, RuntimeError): continue if not weights: self.log(" Could not access FFN weights — skipping sparsity analysis") return surgeon = SparseDirectionSurgeon(auto_sparsity=True) plan = surgeon.plan_surgery(weights, sampled_dirs) self._insights.mean_refusal_sparsity_index = plan.mean_refusal_sparsity_index self._insights.recommended_sparsity = plan.recommended_sparsity self.log(f" Mean RSI: {plan.mean_refusal_sparsity_index:.3f}") self.log(f" Recommended sparsity: {plan.recommended_sparsity:.1%}") self.log(f" Most sparse layer: {plan.most_sparse_layer}") self.log(f" Most dense layer: {plan.most_dense_layer}") # ── Configuration Derivation ───────────────────────────────────── def _derive_configuration(self): """Derive optimal pipeline configuration from analysis insights. This is where analysis feeds forward into abliteration decisions. """ self.log("\n>>> DERIVING CONFIGURATION FROM ANALYSIS") self.log("-" * 50) insights = self._insights # 1. n_directions + direction_method: based on cone geometry # Default: single direction via diff-of-means (proven most robust). # Only escalate to multi-direction when analysis confirms polyhedral geometry. if insights.cone_is_polyhedral and insights.cone_dimensionality > 2.0: # Clearly polyhedral cone → use multiple directions via SVD n_dirs = max(4, min(8, int(insights.cone_dimensionality * 2))) self.direction_method = "svd" self.use_whitened_svd = True self.log(f" Polyhedral cone (dim={insights.cone_dimensionality:.1f}) " f"-> n_directions={n_dirs}, method=svd (whitened)") elif insights.cone_is_polyhedral: # Mildly polyhedral → LEACE gives better single-direction erasure n_dirs = 1 self.direction_method = "leace" self.use_whitened_svd = False self.log(f" Mildly polyhedral (dim={insights.cone_dimensionality:.1f}) " f"-> n_directions=1, method=leace") else: # Linear cone → single direction via diff-of-means (simplest, most robust) n_dirs = 1 self.direction_method = "diff_means" self.use_whitened_svd = False self.log(f" Linear cone (dim={insights.cone_dimensionality:.1f}) " f"-> n_directions=1, method=diff_means") insights.recommended_n_directions = n_dirs insights.recommended_direction_method = self.direction_method self.n_directions = n_dirs # 2. regularization: based on alignment method + entanglement method = insights.detected_alignment_method if method == "dpo": # DPO: concentrated refusal, low entanglement → aggressive removal reg = 0.0 elif method == "rlhf": # RLHF: distributed, moderate entanglement → some regularization reg = 0.15 elif method == "cai": # CAI: recursive, high dimensionality → moderate regularization reg = 0.2 elif method == "sft": # SFT: concentrated in late layers → low regularization reg = 0.05 else: reg = 0.1 # safe default # Increase regularization for highly entangled models if insights.entanglement_score > 0.5: reg = min(0.5, reg + 0.15) self.log(f" High entanglement ({insights.entanglement_score:.2f}) " f"-> increased regularization") insights.recommended_regularization = reg self.regularization = reg self.log(f" Alignment={method}, entanglement={insights.entanglement_score:.2f} " f"-> regularization={reg}") # 3. refinement_passes: based on self-repair risk + robustness if insights.self_repair_estimate > 0.7: passes = 3 self.log(f" High self-repair ({insights.self_repair_estimate:.2f}) -> 3 refinement passes") elif insights.self_repair_estimate > 0.4: passes = 2 self.log(f" Moderate self-repair ({insights.self_repair_estimate:.2f}) -> 2 refinement passes") else: passes = 1 self.log(f" Low self-repair ({insights.self_repair_estimate:.2f}) -> 1 refinement pass") insights.recommended_refinement_passes = passes self.refinement_passes = passes # 4. Layer selection: cluster-aware + entanglement-gated if insights.cluster_representative_layers: # Start from cluster representatives (strongest per cluster) base_layers = list(insights.cluster_representative_layers) # Conservative expansion: for each cluster, add at most the top-2 # strongest layers (by refusal norm) beyond the representative, # to avoid over-modifying weak layers in large clusters. norms = {} for idx in self._harmful_means: if idx in self._harmless_means: norms[idx] = (self._harmful_means[idx] - self._harmless_means[idx]).squeeze().norm().item() for cluster in insights.direction_clusters: ranked = sorted(cluster, key=lambda ly: norms.get(ly, 0), reverse=True) # Add up to 2 additional strong layers per cluster for ly in ranked[:3]: # representative + up to 2 more base_layers.append(ly) base_layers = sorted(set(base_layers)) # Gate: remove highly entangled layers skip = set() for layer_idx in insights.entangled_layers: # Only skip if entanglement exceeds the gate threshold # and there are alternative layers available if len(base_layers) > len(insights.entangled_layers) + 1: skip.add(layer_idx) self.log(f" Skipping layer {layer_idx} (entangled)") insights.skip_layers = sorted(skip) insights.recommended_layers = [ly for ly in base_layers if ly not in skip] else: insights.recommended_layers = [] self.log(f" Final layer set: {insights.recommended_layers or '(default knee detection)'}") # 5. Sparse surgery: if refusal is concentrated, use targeted projection if insights.mean_refusal_sparsity_index > self._sparse_threshold: insights.use_sparse_surgery = True self.log(f" RSI={insights.mean_refusal_sparsity_index:.2f} > {self._sparse_threshold} " f"-> sparse surgery enabled") else: self.log(f" RSI={insights.mean_refusal_sparsity_index:.2f} " f"-> standard dense projection") # 6. Direction method summary (already set in step 1) self.log(f" Direction method: {self.direction_method} " f"(whitened_svd={'on' if self.use_whitened_svd else 'off'})") # ── Informed DISTILL ───────────────────────────────────────────── def _distill_informed(self): """Distill refusal directions using analysis-informed parameters. Key differences from base _distill(): - Uses analysis-recommended n_directions - Respects layer selection from cross-layer analysis - Can extract per-category directions for polyhedral models """ self._emit("distill", "running", "Extracting refusal subspace (analysis-informed)...") t0 = time.time() self.log("\nDISTILL (analysis-informed)") # Run the standard distillation (which now uses our overridden params) # The base _distill() uses self.n_directions, self.use_whitened_svd, etc. # which we've already configured in _derive_configuration() n_layers = len(self._harmful_means) norms: dict[int, float] = {} # ── Small-model direction cap (matching base _distill) ──────── # On small models, each SVD direction removes a proportionally # larger fraction of weight energy. Cap to prevent over-ablation. hidden_size = self.handle.hidden_size if self.handle else 0 total_params = getattr(self.handle, 'total_params', 0) if self.handle else 0 if total_params == 0 and self.handle: try: total_params = sum(p.numel() for p in self.handle.model.parameters()) except Exception: pass if self.n_directions > 1 and ( (0 < hidden_size < 2048) or (0 < total_params < 2_000_000_000) or n_layers <= 16 ): max_dirs = max(1, min(self.n_directions, 2)) if max_dirs < self.n_directions: self.log( f"Capped n_directions from {self.n_directions} to {max_dirs} " f"for small model (hidden={hidden_size}, " f"params={total_params / 1e9:.1f}B, layers={n_layers})" ) self.n_directions = max_dirs # LEACE extractor for optimal concept erasure leace_extractor = None if self.direction_method == "leace": from obliteratus.analysis.leace import LEACEExtractor leace_extractor = LEACEExtractor() self.log("Using LEACE (closed-form optimal concept erasure)") if self.use_whitened_svd and self.n_directions > 1 and leace_extractor is None: from obliteratus.analysis.whitened_svd import WhitenedSVDExtractor whitened_extractor = WhitenedSVDExtractor() self.log(f"Using whitened SVD with {self.n_directions} directions") else: whitened_extractor = None for idx in range(n_layers): # LEACE path: theoretically optimal single-direction erasure if leace_extractor is not None: if idx in self._harmful_acts and idx in self._harmless_acts: try: l_result = leace_extractor.extract( self._harmful_acts[idx], self._harmless_acts[idx], layer_idx=idx, ) self.refusal_directions[idx] = l_result.direction self.refusal_subspaces[idx] = l_result.direction.unsqueeze(0) norms[idx] = l_result.generalized_eigenvalue if idx < 5 or idx == n_layers - 1: self.log( f" layer {idx}: LEACE eigenvalue={l_result.generalized_eigenvalue:.4f}, " f"erasure_loss={l_result.erasure_loss:.4f}" ) continue except Exception as e: if idx < 5: self.log(f" layer {idx}: LEACE failed ({e}), falling back") if self.n_directions == 1: diff = (self._harmful_means[idx] - self._harmless_means[idx]).squeeze(0) norm = diff.norm().item() norms[idx] = norm direction = diff / diff.norm() if norm > 0 else diff self.refusal_directions[idx] = direction self.refusal_subspaces[idx] = direction.unsqueeze(0) elif whitened_extractor is not None: result = whitened_extractor.extract( self._harmful_acts[idx], self._harmless_acts[idx], n_directions=self.n_directions, layer_idx=idx, ) self.refusal_subspaces[idx] = result.directions self.refusal_directions[idx] = result.directions[0] norms[idx] = result.singular_values.sum().item() else: harmful_stack = torch.stack(self._harmful_acts[idx]).squeeze(1) harmless_stack = torch.stack(self._harmless_acts[idx]).squeeze(1) diff_matrix = harmful_stack - harmless_stack if not torch.isfinite(diff_matrix).all(): diff_matrix = torch.nan_to_num(diff_matrix) k = min(self.n_directions, diff_matrix.shape[0], diff_matrix.shape[1]) U, S, Vh = torch.linalg.svd(diff_matrix, full_matrices=False) if not torch.isfinite(S).all() or not torch.isfinite(Vh).all(): continue subspace = Vh[:k] self.refusal_subspaces[idx] = subspace primary = subspace[0] self.refusal_directions[idx] = primary / primary.norm() norms[idx] = S[:k].sum().item() # Enrich subspaces with per-category cone directions when available. # This uses the actual refusal cone generators instead of purely # data-agnostic SVD components. cat_dirs = self._insights.per_category_directions if cat_dirs and self._insights.cone_is_polyhedral and self.n_directions > 1: cat_tensors = list(cat_dirs.values()) # Stack and orthogonalize category directions cat_stack = torch.stack(cat_tensors) # (n_cats, hidden) cat_norms = cat_stack.norm(dim=1, keepdim=True).clamp(min=1e-8) cat_stack = cat_stack / cat_norms # Blend into strong-signal layers: replace later SVD components # with category directions (which are geometrically meaningful) n_cat = cat_stack.shape[0] for idx in norms: sub = self.refusal_subspaces.get(idx) if sub is None or sub.shape[0] <= 1: continue # Keep the first SVD direction (strongest), replace remaining # with category directions projected to be orthogonal to it primary = sub[0:1] # (1, hidden) # Project category directions orthogonal to primary cos = (cat_stack @ primary.squeeze(0)) # (n_cat,) ortho_cats = cat_stack - cos.unsqueeze(1) * primary ortho_norms = ortho_cats.norm(dim=1) # Keep only directions that survived orthogonalization valid = ortho_norms > 0.1 if valid.sum() > 0: ortho_cats = ortho_cats[valid] ortho_cats = ortho_cats / ortho_cats.norm(dim=1, keepdim=True) # Take up to (n_directions - 1) category directions n_take = min(self.n_directions - 1, ortho_cats.shape[0]) new_sub = torch.cat([primary, ortho_cats[:n_take]], dim=0) self.refusal_subspaces[idx] = new_sub self.log(f"Enriched subspaces with {n_cat} per-category cone directions") # Layer selection: use analysis-recommended layers if available, # otherwise fall back to knee detection if self._insights.recommended_layers: self._strong_layers = [ly for ly in self._insights.recommended_layers if ly in self.refusal_directions] self.log(f"Using analysis-recommended layers: {self._strong_layers}") else: sorted_layers = sorted(norms.items(), key=lambda x: x[1], reverse=True) self._strong_layers = self._select_layers_knee(sorted_layers) self.log(f"Using knee-detected layers: {self._strong_layers}") # Remove skipped layers (entanglement-gated) if self._insights.skip_layers: before = len(self._strong_layers) self._strong_layers = [ly for ly in self._strong_layers if ly not in self._insights.skip_layers] after = len(self._strong_layers) if before != after: self.log(f"Entanglement gate removed {before - after} layers " f"-> {after} remaining") elapsed = time.time() - t0 self.log(f"Distillation complete: {len(self._strong_layers)} layers, " f"{self.n_directions} directions ({elapsed:.1f}s)") self._emit( "distill", "done", f"Analysis-informed: {len(self._strong_layers)} layers, " f"{self.n_directions} dirs ({elapsed:.1f}s)", duration=elapsed, strong_layers=self._strong_layers, ) # ── Informed EXCISE ────────────────────────────────────────────── def _excise_informed(self): """Excise refusal directions with analysis-informed strategy. Uses Bayesian optimization (when available) with analysis-derived warm-start parameters, falling back to sparse surgery or standard projection. This is the key integration: analysis maps the geometry, Bayesian optimization finds the optimal projection strength. """ if self._insights.use_sparse_surgery: self._excise_sparse() return # Enable Bayesian optimization using analysis insights for warm-start. # The analysis provides much better initial parameters than the default # heuristic (strongest-layer-based peak), dramatically narrowing the # search space and improving convergence. self._configure_bayesian_warm_start() self._excise() def _configure_bayesian_warm_start(self): """Configure Bayesian optimization with analysis-derived warm-start. Translates analysis insights into a much tighter search space: - peak_position from cluster representative layers - spread from cluster structure (narrow clusters → narrow spread) - component scaling from entanglement analysis - KL budget from alignment method detection """ insights = self._insights # Enable Bayesian optimization (50 trials default, same as heretic) self._bayesian_trials = 50 # Also set heretic-compatible flags on the pipeline so the base # _excise_inner() picks them up during Bayesian optimization. self.layer_adaptive_strength = True self.float_layer_interpolation = True self.use_kl_optimization = True # KL budget: tighter for methods that are fragile (CAI, RLHF), # looser for concentrated methods (DPO, SFT). method = insights.detected_alignment_method if method == "dpo": self.kl_budget = 0.5 elif method == "rlhf": self.kl_budget = 0.3 elif method == "cai": self.kl_budget = 0.2 elif method == "sft": self.kl_budget = 0.4 else: self.kl_budget = 0.35 self.log(f"Bayesian optimization enabled (50 trials, KL budget={self.kl_budget})") self.log("Analysis insights will warm-start the optimizer") # Compute analysis-derived warm-start for the parametric kernel. # The Bayesian optimizer reads these from the pipeline if present. n_layers = len(self._harmful_means) if self._harmful_means else 32 if insights.cluster_representative_layers and n_layers > 1: # Peak position: normalized position of the strongest cluster rep norms = {} for idx in self._harmful_means: if idx in self._harmless_means: norms[idx] = (self._harmful_means[idx] - self._harmless_means[idx]).squeeze().norm().item() reps = insights.cluster_representative_layers if norms: best_rep = max(reps, key=lambda ly: norms.get(ly, 0)) else: best_rep = reps[len(reps) // 2] warm_peak = best_rep / max(n_layers - 1, 1) # Spread: narrow if clusters are tight, wide if clusters span many layers if insights.direction_clusters: cluster_widths = [ (max(c) - min(c)) / max(n_layers - 1, 1) for c in insights.direction_clusters if len(c) > 1 ] warm_spread = max(0.1, min(0.6, sum(cluster_widths) / len(cluster_widths) if cluster_widths else 0.3)) else: warm_spread = 0.3 # Min weight: higher if high persistence (refusal spread across all layers) warm_min = min(0.3, max(0.0, insights.direction_persistence * 0.2)) # Attn/MLP scaling: reduce MLP scaling if entanglement is high # (MLP projections cause more capability damage) if insights.entanglement_score > 0.5: warm_mlp = 0.4 warm_attn = 0.7 else: warm_mlp = 0.6 warm_attn = 0.8 else: warm_peak = 0.5 warm_spread = 0.3 warm_min = 0.05 warm_mlp = 0.6 warm_attn = 0.8 # Store warm-start params for the Bayesian optimizer to pick up self._informed_warm_start = { "max_weight": 0.9, "peak_position": warm_peak, "min_weight": warm_min, "spread": warm_spread, "attn_scale": warm_attn, "mlp_scale": warm_mlp, "dir_idx": 0.0, } self.log( f" Warm-start: peak={warm_peak:.2f}, spread={warm_spread:.2f}, " f"min={warm_min:.2f}, attn={warm_attn:.2f}, mlp={warm_mlp:.2f}" ) def _excise_sparse(self): """Sparse direction surgery — only modifies high-projection rows.""" self._emit("excise", "running", "Sparse direction surgery...") t0 = time.time() from obliteratus.analysis.sparse_surgery import SparseDirectionSurgeon from obliteratus.strategies.utils import ( get_attention_module, get_ffn_module, get_layer_modules, ) surgeon = SparseDirectionSurgeon( sparsity=self._insights.recommended_sparsity, auto_sparsity=True, ) layers = get_layer_modules(self.handle) arch = self.handle.architecture total_modified = 0 for pass_num in range(self.refinement_passes): modified = 0 if self.refinement_passes > 1: self.log(f"Sparse surgery pass {pass_num + 1}/{self.refinement_passes}") if pass_num > 0 and self.true_iterative_refinement: self.log(" Re-probing after sparse surgery...") self._probe() self._distill_inner() for idx in self._strong_layers: subspace = self.refusal_subspaces[idx] layer = layers[idx] device = next(layer.parameters()).device layer_dtype = next(layer.parameters()).dtype for dir_idx in range(subspace.shape[0]): direction = subspace[dir_idx].to(device).to(layer_dtype) # Apply sparse projection to attention and FFN output weights for module_getter, out_names in [ (get_attention_module, ["o_proj", "out_proj", "dense", "c_proj"]), (get_ffn_module, ["down_proj", "c_proj", "dense_4h_to_h", "fc_out", "fc2", "w2"]), ]: try: module = module_getter(layer, arch) for name in out_names: proj = getattr(module, name, None) if proj is None or not hasattr(proj, "weight"): continue W = proj.weight.data if W.shape[-1] == direction.shape[0]: original_norm = W.norm().item() W_new = surgeon.apply_sparse_projection(W, direction) if self.norm_preserve and original_norm > 0: new_norm = W_new.norm().item() if new_norm > 0: W_new = W_new * (original_norm / new_norm) proj.weight.data = W_new.to(layer_dtype) modified += 1 break except (AttributeError, RuntimeError): continue self.log(f" layer {idx}: sparse surgery on {subspace.shape[0]} directions") total_modified += modified self.log(f" Pass {pass_num + 1}: {modified} matrices modified (sparse)") elapsed = time.time() - t0 self.log(f"Sparse excision: {total_modified} projections ({elapsed:.1f}s)") self._emit( "excise", "done", f"Sparse surgery: {total_modified} projections ({elapsed:.1f}s)", duration=elapsed, modified_count=total_modified, ) # ── Informed VERIFY + Ouroboros Compensation ────────────────────── def _verify_and_compensate(self): """Verify excision and run Ouroboros-compensated refinement if needed. After the initial excision, uses analysis modules to detect: 1. Residual refusal signal (via activation probing) 2. Self-repair / Ouroboros effect (via defense robustness) 3. Triggers additional targeted passes at compensating layers KL-gated: stops early if model damage (KL divergence) is getting worse even though refusal persists. This prevents the death spiral where each pass damages the model without removing refusal. """ # Run standard verification first self._verify() # Check if Ouroboros compensation is needed refusal_rate = self._quality_metrics.get("refusal_rate", 0.0) prev_kl = self._quality_metrics.get("kl_divergence", 0.0) ouroboros_pass = 0 # KL budget: stop if KL exceeds this threshold (model too damaged) kl_ceiling = getattr(self, "kl_budget", 0.5) * 2.0 # 2x budget as hard ceiling while (refusal_rate > self._ouroboros_threshold and ouroboros_pass < self._max_ouroboros_passes): ouroboros_pass += 1 self.log(f"\n{'='*60}") self.log(f"OUROBOROS COMPENSATION — Pass {ouroboros_pass}") self.log(f"Refusal rate still {refusal_rate:.0%} > {self._ouroboros_threshold:.0%} threshold") self.log(f"{'='*60}") # Re-probe to find where refusal has re-emerged self.log("Re-probing model for residual refusal...") self._probe() # Re-distill to find rotated directions self._distill_inner() self.log(f"Found {len(self._strong_layers)} layers with residual refusal") # Re-excise at the new strong layers using informed strategy if self._strong_layers: self._excise_informed() else: self.log("No strong layers found — stopping Ouroboros compensation") break # Re-verify self._verify() refusal_rate = self._quality_metrics.get("refusal_rate", 0.0) current_kl = self._quality_metrics.get("kl_divergence", 0.0) self.log(f"After Ouroboros pass {ouroboros_pass}: refusal={refusal_rate:.0%}, KL={current_kl:.4f}") # KL-gated early stopping: if KL is rising and exceeds ceiling, # the model is being damaged faster than refusal is being removed. if current_kl > kl_ceiling: self.log( f"KL divergence {current_kl:.4f} exceeds ceiling {kl_ceiling:.4f} — " f"stopping to prevent further model damage" ) break if ouroboros_pass > 1 and current_kl > prev_kl * 1.5 and refusal_rate > 0.3: self.log( f"KL rising sharply ({prev_kl:.4f} -> {current_kl:.4f}) with " f"refusal still at {refusal_rate:.0%} — stopping (diminishing returns)" ) break prev_kl = current_kl self._report.ouroboros_passes = ouroboros_pass self._report.final_refusal_rate = refusal_rate if ouroboros_pass > 0: self.log(f"\nOuroboros compensation: {ouroboros_pass} additional passes applied") # ── Informed REBIRTH ───────────────────────────────────────────── def _rebirth_informed(self) -> Path: """Save model with comprehensive analysis metadata.""" self._emit("rebirth", "running", f"Saving to {self.output_dir}...") t0 = time.time() self.output_dir.mkdir(parents=True, exist_ok=True) self.handle.model.save_pretrained(self.output_dir) self.handle.tokenizer.save_pretrained(self.output_dir) insights = self._insights metadata = { "source_model": self.model_name, "technique": "analysis_informed_abliteration", "method": "informed", "analysis_insights": { "detected_alignment_method": insights.detected_alignment_method, "alignment_confidence": insights.alignment_confidence, "alignment_probabilities": insights.alignment_probabilities, "cone_is_polyhedral": insights.cone_is_polyhedral, "cone_dimensionality": insights.cone_dimensionality, "mean_pairwise_cosine": insights.mean_pairwise_cosine, "direction_clusters": insights.direction_clusters, "cluster_count": insights.cluster_count, "direction_persistence": insights.direction_persistence, "estimated_robustness": insights.estimated_robustness, "self_repair_estimate": insights.self_repair_estimate, "entanglement_score": insights.entanglement_score, "entangled_layers_skipped": insights.skip_layers, "use_sparse_surgery": insights.use_sparse_surgery, "recommended_sparsity": insights.recommended_sparsity, }, "derived_config": { "n_directions": insights.recommended_n_directions, "direction_method": insights.recommended_direction_method, "regularization": insights.recommended_regularization, "refinement_passes": insights.recommended_refinement_passes, "layers_used": insights.recommended_layers, "layers_skipped": insights.skip_layers, "norm_preserve": self.norm_preserve, "whitened_svd": self.use_whitened_svd, "sparse_surgery": insights.use_sparse_surgery, }, "pipeline_stats": { "analysis_duration_s": self._report.analysis_duration, "total_duration_s": self._report.total_duration, "ouroboros_passes": self._report.ouroboros_passes, "final_refusal_rate": self._report.final_refusal_rate, }, "strong_layers": self._strong_layers, "quality_metrics": self._quality_metrics, "references": [ "Arditi et al., Refusal in Language Models Is Mediated by a Single Direction (2024)", "Gabliteration: SVD-based multi-direction extraction (arXiv:2512.18901)", "grimjim, Norm-Preserving Biprojected Abliteration (2025)", "Wollschlager et al., The Geometry of Refusal in LLMs — concept cones (ICML 2025, arXiv:2502.17420)", "Joad et al., The Ouroboros Effect: Self-Repair in Abliterated LLMs (2026)", "OBLITERATUS: Analysis-informed abliteration pipeline (novel)", ], } import json (self.output_dir / "abliteration_metadata.json").write_text( json.dumps(metadata, indent=2, default=str) ) elapsed = time.time() - t0 self.log(f"Saved informed model to {self.output_dir}/ ({elapsed:.1f}s)") self._emit("rebirth", "done", f"Saved to {self.output_dir} ({elapsed:.1f}s)", duration=elapsed) return self.output_dir @staticmethod def format_insights(insights: AnalysisInsights) -> str: """Format analysis insights as a human-readable report.""" lines = [] lines.append("Analysis-Informed Pipeline — Insights Report") lines.append("=" * 50) lines.append("") lines.append("Alignment Imprint:") lines.append(f" Detected method: {insights.detected_alignment_method.upper()}") lines.append(f" Confidence: {insights.alignment_confidence:.1%}") for method, prob in sorted(insights.alignment_probabilities.items()): lines.append(f" {method.upper():6s} {prob:.1%}") lines.append("") lines.append("Concept Cone Geometry:") cone_type = "POLYHEDRAL" if insights.cone_is_polyhedral else "LINEAR" lines.append(f" Type: {cone_type}") lines.append(f" Dimensionality: {insights.cone_dimensionality:.2f}") lines.append(f" Mean pairwise cosine: {insights.mean_pairwise_cosine:.3f}") if insights.direction_specificity: lines.append(" Per-category DSI:") for cat, dsi in sorted(insights.direction_specificity.items(), key=lambda x: -x[1]): lines.append(f" {cat:15s}: {dsi:.3f}") lines.append("") lines.append("Cross-Layer Structure:") lines.append(f" Direction clusters: {insights.cluster_count}") lines.append(f" Direction persistence: {insights.direction_persistence:.3f}") lines.append(f" Cluster representatives: {insights.cluster_representative_layers}") lines.append("") lines.append("Defense Robustness:") lines.append(f" Estimated robustness: {insights.estimated_robustness.upper()}") lines.append(f" Self-repair (Ouroboros): {insights.self_repair_estimate:.2f}") lines.append(f" Entanglement: {insights.entanglement_score:.3f}") lines.append(f" Entangled layers: {insights.entangled_layers}") lines.append(f" Clean layers: {insights.clean_layers}") lines.append("") lines.append("Derived Configuration:") lines.append(f" n_directions: {insights.recommended_n_directions}") lines.append(f" direction_method: {insights.recommended_direction_method}") lines.append(f" regularization: {insights.recommended_regularization}") lines.append(f" refinement_passes: {insights.recommended_refinement_passes}") lines.append(f" sparse surgery: {insights.use_sparse_surgery}") lines.append(f" layers: {insights.recommended_layers or '(knee detection)'}") lines.append(f" skipped: {insights.skip_layers or '(none)'}") return "\n".join(lines)