Add files via upload

obliteratus/bayesian_optimizer.py

@@ -142,28 +142,35 @@ def _parametric_layer_weight(
     min_weight: float,
     spread: float,
 ) -> float:
-    """Compute ablation weight for a layer using a parametric bell curve.
+    """Compute ablation weight for a layer using a piecewise-linear tent kernel.
 
-    This is the Heretic-style parametric kernel:
-    - max_weight: peak ablation strength (0..1)
-    - peak_position: normalized position of peak (0..1 maps to layer 0..n_layers-1)
-    - min_weight: minimum ablation weight at the tails
-    - spread: controls width of the bell curve (higher = wider)
+    Faithful reproduction of Heretic's parametric kernel (p-e-w/heretic):
+    - max_weight: peak ablation strength at peak_position
+    - peak_position: normalized position of peak (0..1)
+    - min_weight: weight at the edges of the tent
+    - spread: normalized distance from peak to tent edge (min_weight_distance)
 
-    Returns a value in [min_weight, max_weight] representing how strongly
-    to ablate this layer (1.0 = full projection, 0.0 = no projection).
+    Layers beyond ``spread`` from the peak get weight 0 (skipped entirely).
+    Within the tent, weight drops linearly from max_weight to min_weight.
+    This matches Heretic's actual formula::
+
+        distance = abs(layer_index - max_weight_position)
+        if distance > min_weight_distance: skip
+        weight = max_weight + (distance / min_weight_distance) * (min_weight - max_weight)
     """
     if n_layers <= 1:
         return max_weight
 
     normalized_pos = layer_idx / (n_layers - 1)
-    peak = peak_position
-    # Gaussian-shaped kernel
-    dist = abs(normalized_pos - peak)
-    sigma = max(spread, 0.01)
-    gauss = math.exp(-0.5 * (dist / sigma) ** 2)
+    dist = abs(normalized_pos - peak_position)
+    min_weight_distance = max(spread, 0.01)
 
-    return min_weight + (max_weight - min_weight) * gauss
+    # Hard cutoff: layers outside the tent get 0 weight (Heretic skips them)
+    if dist > min_weight_distance:
+        return 0.0
+
+    # Linear interpolation: max_weight at peak → min_weight at edges
+    return max_weight + (dist / min_weight_distance) * (min_weight - max_weight)
 
 
 def _interpolate_direction(
@@ -171,37 +178,56 @@ def _interpolate_direction(
     layer_idx: int,
     float_dir_idx: float,
 ) -> torch.Tensor:
-    """Get an interpolated refusal direction from a float-valued index.
+    """Get an interpolated refusal direction from a float-valued layer index.
 
-    Non-integer values interpolate between adjacent SVD directions in the
-    refusal subspace, unlocking a continuous space of directions beyond
-    the discrete top-k.
+    Faithful reproduction of Heretic's direction interpolation: the index
+    selects which *layer's* diff-of-means direction to use, with float
+    values interpolating between adjacent layers' directions.  This is
+    fundamentally different from interpolating between SVD components
+    within a single layer — it searches across the layer axis.
+
+    From Heretic source (model.py)::
+
+        weight, index = math.modf(direction_index + 1)
+        refusal_direction = F.normalize(
+            refusal_directions[int(index)].lerp(
+                refusal_directions[int(index) + 1], weight), p=2, dim=0)
 
     Args:
-        pipeline: Pipeline with extracted refusal subspaces.
-        layer_idx: Which layer's subspace to use.
-        float_dir_idx: Continuous direction index (e.g., 0.7 interpolates
-            between direction 0 and direction 1).
+        pipeline: Pipeline with extracted refusal directions per layer.
+        layer_idx: The layer being projected (used as fallback).
+        float_dir_idx: Continuous direction index — selects which layer's
+            direction to use (e.g., 5.3 interpolates 70% layer-5 + 30% layer-6).
 
     Returns:
         Normalized direction tensor.
     """
-    subspace = pipeline.refusal_subspaces.get(layer_idx)
-    if subspace is None or subspace.shape[0] == 0:
+    # Build sorted list of layer indices that have refusal directions
+    sorted_layers = sorted(pipeline.refusal_directions.keys())
+    if not sorted_layers:
         return pipeline.refusal_directions.get(layer_idx, torch.zeros(1))
 
-    n_dirs = subspace.shape[0]
-    # Clamp to valid range
-    float_dir_idx = max(0.0, min(float_dir_idx, n_dirs - 1))
+    n_layers_with_dirs = len(sorted_layers)
+
+    # Heretic uses direction_index + 1 offset; we map float_dir_idx into
+    # the sorted layer list, clamped to valid range.
+    float_dir_idx = max(0.0, min(float_dir_idx, n_layers_with_dirs - 1))
 
     lo = int(float_dir_idx)
-    hi = min(lo + 1, n_dirs - 1)
+    hi = min(lo + 1, n_layers_with_dirs - 1)
+
+    lo_layer = sorted_layers[lo]
+    hi_layer = sorted_layers[hi]
+
+    d_lo = pipeline.refusal_directions[lo_layer]
+    d_hi = pipeline.refusal_directions[hi_layer]
 
     if lo == hi:
-        d = subspace[lo]
+        d = d_lo
     else:
+        # Linear interpolation between adjacent layers' directions
         alpha = float_dir_idx - lo
-        d = (1.0 - alpha) * subspace[lo] + alpha * subspace[hi]
+        d = (1.0 - alpha) * d_lo + alpha * d_hi
 
     norm = d.norm()
     if norm > 1e-8:
@@ -342,9 +368,14 @@ def run_bayesian_optimization(
         for live_data, saved_clone in original_params:  # noqa: F821
             live_data.copy_(saved_clone.to(live_data.device))
 
-    # Warm-start values for the parametric kernel
-    # Estimate peak position from strongest layer
-    if pipeline._strong_layers:
+    # Warm-start values for the parametric kernel.
+    # If the informed pipeline provided analysis-derived warm-start params,
+    # use those (they're much better than the default heuristic).
+    informed_warm = getattr(pipeline, "_informed_warm_start", None)
+    if informed_warm:
+        warm_peak = informed_warm.get("peak_position", 0.5)
+        pipeline.log(f"  Using analysis-informed warm-start (peak={warm_peak:.2f})")
+    elif pipeline._strong_layers:
         peak_layer = pipeline._strong_layers[0]
         warm_peak = peak_layer / max(n_total_layers - 1, 1)
     else:
@@ -356,56 +387,56 @@ def run_bayesian_optimization(
     # Suppress Optuna's verbose logging
     optuna.logging.set_verbosity(optuna.logging.WARNING)
 
-    # Max SVD directions available (for float direction interpolation)
-    max_n_dirs = max(
-        (pipeline.refusal_subspaces[idx].shape[0]
-         for idx in pipeline._strong_layers
-         if idx in pipeline.refusal_subspaces),
-        default=1,
-    )
+    # Max layers with directions (for float direction interpolation)
+    n_layers_with_dirs = len([
+        idx for idx in pipeline._strong_layers
+        if idx in pipeline.refusal_directions
+    ])
 
     # ── Phase 1: Parametric kernel optimization (compact search space) ──
+    # Heretic uses SEPARATE kernel parameters for attention and MLP,
+    # allowing them to peak at different layers (8 params + 1 dir_idx = 9).
 
     def objective(trial: optuna.Trial) -> tuple[float, float]:
         """Multi-objective: minimize (refusal_rate, kl_divergence)."""
         _restore_all()
 
-        # Parametric kernel: 4 params describe the entire layer weighting
-        max_weight = trial.suggest_float("max_weight", 0.5, 1.0)
-        peak_position = trial.suggest_float("peak_position", 0.1, 0.9)
-        min_weight = trial.suggest_float("min_weight", 0.0, 0.3)
-        spread = trial.suggest_float("spread", 0.1, 0.6)
+        # Attention kernel: 4 params
+        attn_max = trial.suggest_float("attn_max_weight", 0.5, 1.0)
+        attn_peak = trial.suggest_float("attn_peak_position", 0.1, 0.9)
+        attn_min = trial.suggest_float("attn_min_weight", 0.0, 0.3)
+        attn_spread = trial.suggest_float("attn_spread", 0.1, 0.6)
 
-        # Component-specific scaling (Heretic insight: MLP more damaging)
-        attn_scale = trial.suggest_float("attn_scale", 0.5, 1.0)
-        mlp_scale = trial.suggest_float("mlp_scale", 0.3, 1.0)
+        # MLP kernel: 4 params (separate — can peak at a different layer)
+        mlp_max = trial.suggest_float("mlp_max_weight", 0.3, 1.0)
+        mlp_peak = trial.suggest_float("mlp_peak_position", 0.1, 0.9)
+        mlp_min = trial.suggest_float("mlp_min_weight", 0.0, 0.3)
+        mlp_spread = trial.suggest_float("mlp_spread", 0.1, 0.6)
 
-        # Float direction index (continuous interpolation between SVD dirs)
-        dir_idx = trial.suggest_float("dir_idx", 0.0, max(max_n_dirs - 1, 0.0))
+        # Float direction index (cross-layer interpolation, Heretic-style)
+        dir_idx = trial.suggest_float("dir_idx", 0.0, max(n_layers_with_dirs - 1, 0.0))
 
-        # Compute per-layer regularization from parametric kernel
-        layer_regs: dict[int, float] = {}
+        # Compute per-layer, per-component regularization from kernels
+        attn_regs: dict[int, float] = {}
+        mlp_regs: dict[int, float] = {}
         for idx in pipeline._strong_layers:
-            weight = _parametric_layer_weight(
-                idx, n_total_layers, max_weight, peak_position, min_weight, spread,
-            )
-            # Convert weight to regularization (weight=1 → reg=0, weight=0 → reg=1)
-            layer_regs[idx] = 1.0 - weight
+            attn_w = _parametric_layer_weight(idx, n_total_layers, attn_max, attn_peak, attn_min, attn_spread)
+            mlp_w = _parametric_layer_weight(idx, n_total_layers, mlp_max, mlp_peak, mlp_min, mlp_spread)
+            attn_regs[idx] = 1.0 - attn_w
+            mlp_regs[idx] = 1.0 - mlp_w
 
         # Apply projection with trial's parameters
         for idx in pipeline._strong_layers:
-            if idx not in pipeline.refusal_subspaces:
+            if idx not in pipeline.refusal_directions:
                 continue
 
-            # Use interpolated direction
+            # Use cross-layer interpolated direction
             direction = _interpolate_direction(pipeline, idx, dir_idx)
             d_col = direction.to(device=next(layer_modules[idx].parameters()).device)
             d_col = d_col.unsqueeze(-1) if d_col.dim() == 1 else d_col
 
-            reg = layer_regs[idx]
-
-            # Attention projection (with attn_scale)
-            attn_reg = 1.0 - (1.0 - reg) * attn_scale
+            # Attention projection (with per-component kernel)
+            attn_reg = attn_regs[idx]
             try:
                 attn = get_attention_module(layer_modules[idx], arch)
                 pipeline._project_out_advanced(
@@ -416,8 +447,8 @@ def run_bayesian_optimization(
             except (AttributeError, RuntimeError):
                 pass
 
-            # MLP/FFN projection (with mlp_scale)
-            mlp_reg = 1.0 - (1.0 - reg) * mlp_scale
+            # MLP/FFN projection (with per-component kernel)
+            mlp_reg = mlp_regs[idx]
             try:
                 ffn = get_ffn_module(layer_modules[idx], arch)
                 count = pipeline._project_out_advanced(
@@ -439,18 +470,20 @@ def run_bayesian_optimization(
         refusal = _measure_refusal_rate(pipeline, n_prompts=n_refusal_prompts)
         kl = _measure_kl_divergence(pipeline, reference_logits, kl_prompts)
 
-        # Track best combined score
+        # Track best combined score (use average of attn/mlp regs for layer_regs)
         nonlocal best_score, best_result
         combined = refusal + 0.5 * kl
         if combined < best_score:
             best_score = combined
-            best_result = dict(layer_regs)
+            best_result = {
+                idx: (attn_regs[idx] + mlp_regs[idx]) / 2.0
+                for idx in pipeline._strong_layers
+            }
 
         pipeline.log(
             f"  Trial {trial.number + 1}/{n_trials}: "
             f"refusal={refusal:.0%}, KL={kl:.4f} "
-            f"(peak={peak_position:.2f}, spread={spread:.2f}, "
-            f"attn={attn_scale:.2f}, mlp={mlp_scale:.2f}, dir={dir_idx:.2f})"
+            f"(attn_peak={attn_peak:.2f}, mlp_peak={mlp_peak:.2f}, dir={dir_idx:.2f})"
         )
 
         return refusal, kl
@@ -462,16 +495,33 @@ def run_bayesian_optimization(
         study_name="obliteratus_parametric_optimization",
     )
 
-    # Enqueue warm-start trial with analysis-derived estimates
-    warm_params = {
-        "max_weight": 0.9,
-        "peak_position": warm_peak,
-        "min_weight": 0.05,
-        "spread": 0.3,
-        "attn_scale": 0.8,
-        "mlp_scale": 0.6,
-        "dir_idx": 0.0,
-    }
+    # Enqueue warm-start trial with analysis-derived estimates.
+    # Translate informed pipeline params to the new per-component format.
+    if informed_warm:
+        iw = informed_warm
+        warm_params = {
+            "attn_max_weight": iw.get("max_weight", 0.9),
+            "attn_peak_position": iw.get("peak_position", warm_peak),
+            "attn_min_weight": iw.get("min_weight", 0.05),
+            "attn_spread": iw.get("spread", 0.3),
+            "mlp_max_weight": iw.get("max_weight", 0.9) * iw.get("mlp_scale", 0.6),
+            "mlp_peak_position": iw.get("peak_position", warm_peak),
+            "mlp_min_weight": iw.get("min_weight", 0.05),
+            "mlp_spread": iw.get("spread", 0.3),
+            "dir_idx": iw.get("dir_idx", 0.0),
+        }
+    else:
+        warm_params = {
+            "attn_max_weight": 0.9,
+            "attn_peak_position": warm_peak,
+            "attn_min_weight": 0.05,
+            "attn_spread": 0.3,
+            "mlp_max_weight": 0.6,
+            "mlp_peak_position": warm_peak,
+            "mlp_min_weight": 0.05,
+            "mlp_spread": 0.3,
+            "dir_idx": 0.0,
+        }
     study.enqueue_trial(warm_params)
 
     pipeline.log(f"Bayesian optimization: running {n_trials} trials (parametric kernel)...")
@@ -490,25 +540,32 @@ def run_bayesian_optimization(
         p = best_trial.params
         best_result = {}
         for idx in pipeline._strong_layers:
-            weight = _parametric_layer_weight(
+            attn_w = _parametric_layer_weight(
                 idx, n_total_layers,
-                p["max_weight"], p["peak_position"],
-                p["min_weight"], p["spread"],
+                p["attn_max_weight"], p["attn_peak_position"],
+                p["attn_min_weight"], p["attn_spread"],
             )
-            best_result[idx] = 1.0 - weight
+            mlp_w = _parametric_layer_weight(
+                idx, n_total_layers,
+                p["mlp_max_weight"], p["mlp_peak_position"],
+                p["mlp_min_weight"], p["mlp_spread"],
+            )
+            best_result[idx] = (attn_w + mlp_w) / 2.0  # average for layer-level reg
+            best_result[idx] = 1.0 - best_result[idx]
 
         pipeline.log(
             f"  Best trial: refusal={best_trial.values[0]:.0%}, "
             f"KL={best_trial.values[1]:.4f}"
         )
         pipeline.log(
-            f"  Kernel: peak={p['peak_position']:.2f}, spread={p['spread']:.2f}, "
-            f"max={p['max_weight']:.2f}, min={p['min_weight']:.2f}"
+            f"  Attn kernel: peak={p['attn_peak_position']:.2f}, "
+            f"spread={p['attn_spread']:.2f}, max={p['attn_max_weight']:.2f}"
         )
         pipeline.log(
-            f"  Components: attn={p['attn_scale']:.2f}, mlp={p['mlp_scale']:.2f}, "
-            f"dir_idx={p['dir_idx']:.2f}"
+            f"  MLP kernel:  peak={p['mlp_peak_position']:.2f}, "
+            f"spread={p['mlp_spread']:.2f}, max={p['mlp_max_weight']:.2f}"
         )
+        pipeline.log(f"  dir_idx={p['dir_idx']:.2f}")
 
         # Store the best direction index for use during EXCISE
         best_dir_idx = p.get("dir_idx", 0.0)
@@ -518,9 +575,9 @@ def run_bayesian_optimization(
                 new_dir = _interpolate_direction(pipeline, idx, best_dir_idx)
                 pipeline.refusal_directions[idx] = new_dir
 
-        # Store component scales for use in EXCISE
-        pipeline._bayesian_attn_scale = p.get("attn_scale", 1.0)
-        pipeline._bayesian_mlp_scale = p.get("mlp_scale", 1.0)
+        # Store component scales for use in EXCISE (backward compat)
+        pipeline._bayesian_attn_scale = p.get("attn_max_weight", 1.0)
+        pipeline._bayesian_mlp_scale = p.get("mlp_max_weight", 1.0)
 
     elif best_result:
         pipeline.log(f"  Using best combined score: {best_score:.4f}")