Upload 81 files

obliteratus/abliterate.py

@@ -166,21 +166,20 @@ METHODS = {
     "nuclear": {
         "label": "Nuclear (Maximum Force Combo)",
         "description": (
-            "The ultimate abliteration mode for the most stubborn models. "
-            "Combines ALL techniques at maximum strength: 16 SVD directions, "
-            "3 refinement passes with true iterative re-probing, 2.5x "
-            "reflection (stronger than standard 2x inversion), input "
-            "embedding projection, expert weight transplant (overwrites "
-            "safety expert weights with capability expert weights), and "
-            "inference-time activation steering hooks that continuously "
-            "subtract the refusal direction from hidden states during "
-            "generation. Designed to break through models like GPT-OSS 20B "
-            "where standard methods plateau."
+            "Comprehensive combo mode for stubborn MoE models. Combines ALL "
+            "techniques with calibrated intensity: 10 SVD directions, 2 "
+            "refinement passes, standard 2x reflection, input embedding "
+            "projection, partial expert weight transplant (blends 30%% "
+            "capability weights into safety experts), and light inference-"
+            "time activation steering (0.3x residual cleanup). Designed "
+            "to break through models like GPT-OSS 20B without causing "
+            "decoherence — every technique fires at moderate power so they "
+            "complement rather than compound."
         ),
-        "n_directions": 16,
+        "n_directions": 10,
         "norm_preserve": True,
         "regularization": 0.0,
-        "refinement_passes": 3,
+        "refinement_passes": 2,
         "project_biases": True,
         "use_chat_template": True,
         "use_whitened_svd": True,
@@ -192,10 +191,12 @@ METHODS = {
         "attention_head_surgery": True,
         "use_sae_features": True,
         "invert_refusal": True,
-        "reflection_strength": 2.5,
+        "reflection_strength": 2.0,
         "project_embeddings": True,
         "activation_steering": True,
+        "steering_strength": 0.3,
         "expert_transplant": True,
+        "transplant_blend": 0.3,
     },
 }
 
@@ -591,7 +592,9 @@ class AbliterationPipeline:
         reflection_strength: float | None = None,
         project_embeddings: bool | None = None,
         activation_steering: bool | None = None,
+        steering_strength: float | None = None,
         expert_transplant: bool | None = None,
+        transplant_blend: float | None = None,
         on_stage: Callable[[StageResult], None] | None = None,
         on_log: Callable[[str], None] | None = None,
     ):
@@ -633,7 +636,9 @@ class AbliterationPipeline:
         self.reflection_strength = reflection_strength if reflection_strength is not None else method_cfg.get("reflection_strength", 2.0)
         self.project_embeddings = project_embeddings if project_embeddings is not None else method_cfg.get("project_embeddings", False)
         self.activation_steering = activation_steering if activation_steering is not None else method_cfg.get("activation_steering", False)
+        self.steering_strength = steering_strength if steering_strength is not None else method_cfg.get("steering_strength", 0.3)
         self.expert_transplant = expert_transplant if expert_transplant is not None else method_cfg.get("expert_transplant", False)
+        self.transplant_blend = transplant_blend if transplant_blend is not None else method_cfg.get("transplant_blend", 0.3)
 
         self.handle: ModelHandle | None = None
         self.refusal_directions: dict[int, torch.Tensor] = {}  # per-layer primary direction
@@ -2463,16 +2468,21 @@ class AbliterationPipeline:
     # ── Nuclear-mode helpers ─────────────────────────────────────────────
 
     def _transplant_expert_weights(self, layers: nn.ModuleList) -> int:
-        """Overwrite safety expert down_proj weights with capability expert averages.
+        """Blend capability expert weights into safety expert down_proj.
 
         For each MoE layer, computes the mean of capability experts' down_proj
-        weights and copies it into each safety expert's down_proj. This is more
-        aggressive than reflection — it completely replaces the refusal-encoding
-        output path with the capability-encoding one.
+        weights and blends it into each safety expert's down_proj using the
+        transplant_blend ratio. A blend of 0.3 means:
+            new_weight = 0.7 * original_safety + 0.3 * capability_mean
+
+        This preserves most of the safety expert's general language modeling
+        ability while nudging its output toward the capability distribution.
+        Full overwrite (blend=1.0) causes decoherence.
 
-        Returns the number of weight matrices transplanted.
+        Returns the number of weight matrices blended.
         """
         arch = self.handle.architecture
+        blend = self.transplant_blend
         count = 0
 
         for idx in self._strong_layers:
@@ -2514,19 +2524,23 @@ class AbliterationPipeline:
                 # Compute mean of capability expert weights
                 cap_mean = torch.stack(cap_weights).mean(dim=0)
 
-                # Transplant into safety experts
+                # Partial blend into safety experts
                 for ei in safety_indices:
                     if ei >= len(experts):
                         continue
                     target = getattr(experts[ei], wname, None)
                     if target is not None and hasattr(target, "weight"):
                         if target.weight.data.shape == cap_mean.shape:
-                            target.weight.data.copy_(cap_mean)
+                            # Blend: (1-blend) * original + blend * capability_mean
+                            target.weight.data.mul_(1.0 - blend).add_(cap_mean * blend)
                             count += 1
 
                 del cap_weights, cap_mean
 
-            self.log(f"  layer {idx}: transplanted capability weights into {len(safety_indices)} safety experts")
+            self.log(
+                f"  layer {idx}: blended {blend:.0%} capability weights "
+                f"into {len(safety_indices)} safety experts"
+            )
 
         return count
 
@@ -2537,8 +2551,10 @@ class AbliterationPipeline:
         continuously steering the model away from the refusal direction.
         This catches residual signal that static weight surgery may have missed.
 
-        The steering strength scales with the reflection_strength parameter:
-        standard (2.0) = subtract 1x direction, boosted (2.5) = 1.5x, etc.
+        Uses the dedicated steering_strength parameter (default 0.3) instead
+        of coupling to reflection_strength. A light touch (0.3) works as
+        residual cleanup without causing decoherence — the weight surgery
+        already handles the bulk of the removal.
 
         Returns the number of hooks installed.
         """
@@ -2547,7 +2563,9 @@ class AbliterationPipeline:
             hook.remove()
         self._steering_hooks.clear()
 
-        steering_scale = self.reflection_strength - 1.0  # 2.0→1.0, 2.5→1.5, 3.0→2.0
+        # Use only the primary refusal direction (not full subspace) to
+        # minimize interference with the model's representation space
+        steering_scale = self.steering_strength
 
         for idx in self._strong_layers:
             if idx not in self.refusal_directions:

tests/test_abliterate.py

@@ -1409,12 +1409,14 @@ class TestNuclearMethod:
         """Nuclear method should have all techniques enabled at maximum."""
         cfg = METHODS["nuclear"]
         assert cfg["invert_refusal"] is True
-        assert cfg["n_directions"] == 16
-        assert cfg["refinement_passes"] == 3
-        assert cfg["reflection_strength"] == 2.5
+        assert cfg["n_directions"] == 10
+        assert cfg["refinement_passes"] == 2
+        assert cfg["reflection_strength"] == 2.0
         assert cfg["project_embeddings"] is True
         assert cfg["activation_steering"] is True
+        assert cfg["steering_strength"] == 0.3
         assert cfg["expert_transplant"] is True
+        assert cfg["transplant_blend"] == 0.3
         assert cfg["use_jailbreak_contrast"] is True
         assert cfg["attention_head_surgery"] is True
 
@@ -1422,12 +1424,14 @@ class TestNuclearMethod:
         """Pipeline initialized with nuclear method should have all flags set."""
         pipeline = AbliterationPipeline(model_name="test", method="nuclear")
         assert pipeline.invert_refusal is True
-        assert pipeline.reflection_strength == 2.5
+        assert pipeline.reflection_strength == 2.0
+        assert pipeline.steering_strength == 0.3
+        assert pipeline.transplant_blend == 0.3
         assert pipeline.project_embeddings is True
         assert pipeline.activation_steering is True
         assert pipeline.expert_transplant is True
-        assert pipeline.n_directions == 16
-        assert pipeline.refinement_passes == 3
+        assert pipeline.n_directions == 10
+        assert pipeline.refinement_passes == 2
 
     def test_reflection_strength_configurable(self):
         """reflection_strength should be explicitly overridable."""
@@ -1581,12 +1585,14 @@ class TestNuclearMethod:
         finally:
             abl_module.get_ffn_module = orig_get_ffn
 
-        assert count >= 2, f"Should transplant at least 2 weights (safety experts), got {count}"
+        assert count >= 2, f"Should blend at least 2 weights (safety experts), got {count}"
 
-        # Safety expert 0 should now have the capability mean
+        # Safety expert 0 should be a 30% blend toward capability mean
+        # new = 0.7 * original + 0.3 * capability_mean
+        expected_blend = 0.7 * orig_safety0 + 0.3 * expected_mean
         transplanted = layer.mlp.experts[0].down_proj.weight.data
-        assert torch.allclose(transplanted, expected_mean, atol=1e-4), (
-            "Safety expert weight should match capability expert mean after transplant"
+        assert torch.allclose(transplanted, expected_blend, atol=1e-4), (
+            "Safety expert weight should be 30% blended toward capability mean"
         )
 
         # Capability expert 2 should be unchanged