mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
1315 lines
59 KiB
Python
1315 lines
59 KiB
Python
"""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 typing import Callable
|
|
|
|
import torch
|
|
|
|
from obliteratus.abliterate import (
|
|
AbliterationPipeline,
|
|
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,
|
|
# 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,
|
|
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.
|
|
"""
|
|
t0 = time.time()
|
|
|
|
# Stage 1: SUMMON
|
|
self._summon()
|
|
|
|
# Stage 2: PROBE
|
|
self._probe()
|
|
|
|
# Stage 3: ANALYZE (new stage — the feedback loop)
|
|
self._analyze()
|
|
|
|
# Stage 4: DISTILL (informed by analysis)
|
|
self._distill_informed()
|
|
|
|
# Stage 5: EXCISE (informed by analysis)
|
|
self._excise_informed()
|
|
|
|
# Stage 6: VERIFY + Ouroboros compensation loop
|
|
self._verify_and_compensate()
|
|
|
|
# Stage 7: 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)
|